Distilling a 4.5 hour CUDA talk
A reading companion to Ben Spector & Quinn McIntyre's CUDA & ThunderKittens deep dive.
A few mechanics: the highlighted cards restate the questions raised by the source notes, so they stand on their own whether or not you've watched the stream. Timestamp chips like 0:34:25 jump to the moment that prompted them, chips like R1 point to the reference list at the bottom, and anything not grounded in the talk or a reference is flagged as inference. Three interactive widgets (§3, §7, §8) let you play with the machinery yourself.
How to read this
A useful framing sentence appears in the first two minutes 0:01:14: "ThunderKittens is not really a set of abstractions, it's a set of primitives, and the difference is that with abstractions you don't have to know what's going on under the hood, and with primitives you do."
A useful working model for CUDA, TK, and GPU kernel frameworks has five parts:
- The hardware level: SMs, quadrants, register files, execution pipelines (§1–2).
- The execution model: warps, schedulers, and why the GPU is a throughput machine, not a latency machine (§3–4).
- Work decomposition: grids, blocks, thread blocks, thread block clusters, warps, warp groups, et cetera (§5).
- The memory consistency model: Ben calls this "the entire difficult part of CUDA; everything else is actually easy" 1:05:02 (§6–7).
- Layouts: banks, core matrices, swizzles, et cetera (§8).
Sections 9–14 cover TK itself: its types, its ops, its templates, and the workflow Ben uses to ship kernels, including the DeepSeek head-dim-576 case.
The die: what a GPU physically is
Ben begins with NVIDIA's Hopper die shot and identifies the structures that matter for programming it. The details worth keeping, roughly in order of how often they reappear later:
NVIDIA markets "16,896 CUDA cores." Ben argues this is the wrong level to call a core. The SM (streaming multiprocessor) is the real analogue of a CPU core: it has its own schedulers, register file, L1, and execution units. A "CUDA core" is a single hardware lane — one FP32 ALU inside an SM. When Apple markets an M2 Pro as a "19-core GPU," those "cores" are much closer to SM-sized execution blocks than to NVIDIA's "CUDA cores," though the architectures are not one-to-one. Keep this straight and a lot of marketing numbers become legible: 132 SMs × 128 lanes = 16,896 "CUDA cores."
The whole memory hierarchy, top to bottom
| Level | Where it physically lives | Ballpark | Character |
|---|---|---|---|
| Registers | Inside each SM quadrant | ~1 cycle | Fastest SRAM; register pressure often constrains occupancy and kernel structure |
| Shared memory / L1 | Inside each SM (one pool, split by config) | ~20–30 cycles | Programmer-controlled scratchpad (SMEM) or hardware cache (L1) |
| L2 | On die, two halves + crossbar | ~200–600 cycles depending on hit location | Shared by all SMs; far-side hits ≈ HBM latency |
| HBM | Stacks beside the die | ~500–700+ cycles, ~3.35 TB/s on H100 | The DRAM. Everything on-die is SRAM; this is not |
The argument has three parts. The L1 and shared memory (SMEM) are the same physical SRAM on the SM; a config knob decides how much acts as each. For AI kernels you push almost all of it to SMEM because:
- SMEM is faster than L1. A cache must do tag lookup and hit/miss resolution on every access; a scratchpad just decodes an address. Same silicon, fewer steps.
- Guaranteed residency. A cache can evict your data whenever it likes. Data you
__shared__-allocate stays put until you overwrite it. Your "hit rate" is 100% by construction, not by luck. - Layout control. You choose exactly which byte goes to which address — which, as §8 will show, is a central constraint for feeding tensor cores without bank conflicts. L1 gives you no say.
Ben's stronger version of the claim: well-written dense AI kernels try to minimize implicit cache traffic — no register spills (so no L1 traffic from spills), and global reuse handled explicitly through SMEM 0:18:43. L1-as-cache is for irregular access patterns (graphics, sparse workloads) where reuse is not statically predictable.
NVIDIA's SM diagram draws "FP32" and "INT32" unit columns, but the physical device has more than the diagram shows: there is dedicated FP16 hardware that is twice as wide as the FP32 path, which is why plain (non-tensor-core) FP16 FMAs run at 2× the FP32 rate. The diagram omits it, presumably for compactness 0:21:08. The FP8 story is different: FP8 arithmetic on Hopper lives in the tensor cores (there's no general FP8 ALU path) — you get FP8 through mma/wgmma instructions, not through scalar ops.
The deeper point Ben is making: the diagram shows unit counts, but what matters operationally is that instructions dispatch to named hardware pipelines — ALU (int32 & friends), FMA (fp32 multiply-add), Tensor, XU/SFU (transcendentals), LSU (loads/stores). Nsight Compute reports utilization per pipeline, and that's the vocabulary you'll actually profile in.
Inside the SM: quadrants & pipelines
At the level of one SM, the unit is four replicated quadrants (NVIDIA calls them "processing blocks" or sub-partitions) plus a few resources shared across the whole SM: the L1/SMEM pool, and — on Hopper — the TMA (Tensor Memory Accelerator), a small engine that generates addresses and moves tensors between global and shared memory asynchronously, without occupying threads 0:17:16.
On a CPU, a register like rax is one 64-bit cell. On an NVIDIA GPU, when your code names register R4, the hardware stores 32 copies of it — one 32-bit slot per lane of the warp — side by side in one physical row of the register file. 32 lanes × 32 bits = 1024 bits per physical register row. So the register file is best imagined as a small number of very wide vector registers, indexed by (register name, warp), rather than a large collection of tiny scalar cells.
This one fact connects two things you'll meet later: (a) the 32-thread warp is the group of lanes that share one register-file row and one instruction (§3); and (b) the tensor-core "core matrix" in §8 is 8×8: an 8×8 bf16 tile is 128 bytes = 1024 bits = exactly one physical register row spread across the warp.
Each quadrant's execution units are independent hardware pipelines: FMA (fp32), ALU (int32), Tensor, SFU/XU (transcendentals), LSU (memory). The scheduler issues at most one instruction per cycle, but consecutive instructions can go to different pipelines, and each pipeline chews on its work concurrently. So in a good kernel the tensor pipe is doing an MMA while the FMA pipe computes softmax scales while the LSU has loads in flight — all pipes "full at once."
The flip side: each pipeline has its own issue rate. The SFU is much narrower than 32 lanes — if your kernel is wall-to-wall exp instructions, the SFU saturates while everything else idles, and no amount of clever scheduling fixes it 0:22:42. Profilers expose this directly: Nsight Compute's pipe-utilization chart tells you which pipe is your actual bottleneck.
Because the fast hardware exponential path is base-2. The SFU instruction is MUFU.EX2 — "multi-function unit, exponential base 2" (log is likewise base 2: MUFU.LG2). The fully specified expf path can include accuracy and range-handling code, but the optimized inner-loop shape is effectively:
e^x = 2^(x · log₂e) = EX2(x × 1.4426950…)
— i.e., an extra FMA by log₂e ≈ 1.4427 in front of every single exponential 0:23:41. In an attention kernel, that's one extra FMA per element of the attention matrix. §7 shows the standard trick to avoid an extra per-element multiply: fold the 1.4427 into a constant you're already multiplying by.
Why did NVIDIA build base 2? Binary floating point is base 2 — 2ˣ decomposes trivially into "integer part goes into the exponent field, fractional part is a small polynomial/table lookup on [0,1)." An eˣ unit would be strictly harder for zero benefit, since one multiply converts between bases.
There is no full-rate divide pipeline analogous to FMA. Multiplication can be deeply pipelined to high throughput. Division is inherently iterative (each quotient bit/digit depends on the previous remainder), so building a full-rate divider per lane would cost substantial area for an instruction that's rare in AI code. Instead the GPU computes a reciprocal approximation (MUFU.RCP on the SFU) and then refines/multiplies — Ben quotes the effective throughput at roughly 1/8 to 1/16 of FMA rate on Hopper 0:24:48.
Why it rarely hurts in practice: kernels often divide many values by the same denominator (softmax normalization is the canonical case). So you pay for one reciprocal, then reuse it as a multiply across the whole row. Expensive instruction, amortized over many values.
Warps, latency hiding, and the throughput mindset
This section introduces the throughput model used throughout the rest of the talk. A CPU is built to finish one instruction stream as fast as possible: branch predictors, out-of-order engines, deep speculation — all in service of latency. A GPU instead accepts high individual instruction latency (an fp32 op is ~4 cycles to first use; a dependent chain to global memory is hundreds) and tries to keep enough independent work resident that something can always issue.
Ben's phrasing 0:26:36: "orienting around being able to continually issue new instructions is much more important than making sure each individual hardware stream is as optimized as possible."
The reason is arithmetic. Suppose an instruction stalls for 400 cycles waiting on memory. On a CPU, that stall dominates the runtime of the stream, so CPUs spend billions of transistors avoiding it. On a GPU, the quadrant scheduler can issue from a different resident warp on the next cycle. If enough warps are resident, the 400-cycle latency is hidden in throughput terms — the ALUs continue to receive work.
The practical corollary Ben draws: don't profile GPUs by timing one thread's critical path. The Hazy team initially inserted per-thread clock counters into kernels and measured elapsed cycles between points — and found it mostly useless 0:25:49, because a long per-thread latency is fine if someone else fills the gap. What you actually hunt for is stalls that nothing can cover — which is exactly what Nsight Compute's warp-state sampling shows you (§7).
Warps: the unit that actually executes
CUDA exposes "threads," but the hardware never executes a thread alone. Threads are bundled in groups of 32 — a warp — and the scheduler issues one instruction for the active lanes of that warp on the quadrant's parallel ALUs. On Volta and newer, lanes can have independent program counters under divergence, so "lockstep" is a useful first model but not a correctness guarantee.
The mental model holds, with one refinement. A CUDA "thread" is a software abstraction over execution lanes and register-file columns. Thread k of a warp corresponds to lane k mod 32: its "private registers" are that lane's 32-bit slices of the warp's 1024-bit register rows, and "its" instructions are the warp's instructions, executed on its lane 0:27:09.
The refinement: it's a virtualization, not a fixed mapping — far more warps can be resident on an SM than there are physical lane-sets, and the scheduler multiplexes them over time. That is the basic latency-hiding mechanism. The abstraction is also leaky in exactly one direction: correctness is guaranteed if you follow the model, but performance requires knowing the lanes are there (divergence, coalescing, bank conflicts all fall out of the lane structure).
Directly: warp width × 32-bit lane slot = physical register width. 32 threads × 32 bits = the 1024-bit register row from §2. The warp width, register-file rows, and ALU lanes are co-designed around the same 32-lane granularity. One warp instruction reads whole register rows, feeds 32 lanes, writes whole rows back. Everything is provisioned around that width, which is also why NVIDIA "pays overhead for everything being vector oriented" (§4) — even purely scalar work occupies 32-lane-wide machinery.
Yes: each quadrant's scheduler issues at most one instruction per clock, and that instruction belongs to exactly one warp (covering up to all 32 of its lanes). Four quadrants per SM ⇒ up to 4 warp-instructions issued per SM per cycle. Each cycle, the scheduler looks at its resident warps, finds ones that are eligible (next instruction's operands ready, pipeline not full), and picks one. If none are eligible, the slot is wasted — that's a stall, and the profiler will tell you which flavor (§7).
Both. Ben gives exactly these two mechanisms 0:28:09:
- ILP within a warp: if instruction n+1 doesn't depend on instruction n, the warp can issue it before n retires — the FMA pipeline is itself pipelined, accepting a new warp-instruction every cycle even though each takes ~4 cycles to complete. Compilers unroll loops precisely to create this independence.
- Co-scheduled warps ("occupancy"): when one warp has no eligible instruction, another resident warp on the same quadrant issues instead. (Terminology note: this happens at warp granularity on a quadrant; "block" is the software grouping — a block's warps get spread across the SM's quadrants.)
They trade off against each other: high ILP + lots of registers per thread means fewer resident warps fit, and vice versa. Fast kernels usually sit at surprisingly low occupancy with high ILP — quality of resident work beats quantity.
Two independent reasons, one architectural and one performance:
Architectural (the register math): one precision fix on the number: CUDA's documented compile-time cap is 255 registers per thread R2; the talk rounds this to 256. A quadrant's file is 64 KB = 16,384 32-bit registers = 512 registers per lane R1. A 255-register cap × 32 lanes means one warp can touch just under half the file. So to even address the whole register file you must place ≥ 2 warps on the quadrant — Ben interprets this cap as a hint about the intended occupancy regime 0:28:58.
Performance (latency hiding): as in the figure above — with a single warp, every dependent-latency bubble leaves the scheduler with no eligible instruction from that warp. A second warp gives the scheduler another independent stream of work, which is the basic mechanism behind occupancy-based latency hiding.
_sync intrinsics exist?Pre-Volta, a warp had one shared program counter: all 32 lanes were physically at the same instruction at all times (divergent branches serialized, but lanes never interleaved out of order). People exploited this and wrote "warp-synchronous" code — e.g. warp reductions through shared memory with no explicit synchronization — because lockstep provided an implicit within-warp ordering guarantee.
Volta introduced independent thread scheduling R11: each lane got its own program counter, and the hardware may interleave divergent paths. Lockstep is no longer guaranteed, so all that old code became silently racy. NVIDIA's fix was to make the synchronization explicit in the ISA: the old __shfl(...), __ballot(...), etc. were deprecated in favor of __shfl_sync(mask, ...), __ballot_sync(mask, ...), plus a new __syncwarp(mask). The _sync suffix and mask say: "these named lanes must converge and participate here" — the guarantee you used to get for free from the hardware generation, now requested explicitly. That's the connection: shuffle-sync exists because the warp-level lockstep guarantee was withdrawn.
The economics of 32: energy, warp width, and why tensor cores exist
This segment gives a compact explanation for why the hardware is shaped this way. It hinges on one number from Bill Dally (NVIDIA's former chief scientist): the 20:1 instruction-issue overhead ratio.
The number and the conclusion are memorable; the full chain of reasoning is worth laying out once.
The measurement: for a plain fp32 FMA across a warp, the energy of the "bookkeeping" — fetching and decoding the instruction, the warp scheduler, operand collection from the register file, dispatch, writeback — is about 20× the energy of the 32 multiply-adds themselves (Dally has put numbers on this publicly: overhead of ~20× the ~1.5 pJ the math itself costs R4) 0:34:14. Arithmetic is nearly free in modern silicon; moving data and control is what costs. So a GPU running scalar FMA code spends ~95% of its energy on overhead.
The fix — amortize the issue: if one instruction triggers more FLOPs, the fixed issue cost is divided over them. That is one reason tensor cores are effective. Ben's example: a Hopper wgmma.mma_async shape in the m64n256k16 family computes a 64×256 tile from a 64×16 tile and a 16×256 tile, or on the order of half a million FLOPs (64·16·256·2) per issued warp-group MMA 0:34:43. Issue overhead per FLOP becomes small relative to the arithmetic work — "you get more into ASIC territory."
The tension: now open Flash Attention 3's SASS (NVIDIA's machine code) and look at the hot loop: alongside the big HGMMAs there is a wall of scalar FMUL.FTZ instructions — the softmax rescaling — each paying the full 20:1 tax 0:35:42. Even in a highly optimized attention kernel, the non-matmul part is expensive per FLOP. Ben frames this as an argument for TPUs and for wider wavefronts 0:38:11: a TPU's systolic array issues almost nothing; AMD's 64-wide wavefront cuts the per-lane share of issue overhead in half by amortizing one instruction over twice the lanes.
FMUL, every address calculation — pays the left bar.The overhead is a flexibility tax. NVIDIA's model makes every operation — even inherently scalar bookkeeping like loop counters, address math, predicates — occupy a 32-lane vector instruction with full per-lane register file machinery. That buys generality: the same machine runs graphics shaders, scientific code, and transformers, with arbitrary per-lane control flow. The cost is the 20:1 ratio above applied to everything that isn't a tensor-core op.
The startup thesis (Etched is a prominent example of a transformer-focused ASIC approach; the same logic animates Groq, Cerebras, and the TPU itself) is: if you only ever run one workload shape, specialize the dataflow and reduce general-purpose issue overhead. Ben's "we'll see if that works" is the right hedge: the bet fails if model architectures move out from under your ASIC's assumptions faster than you can tape out new silicon.
Why a fixed width at all: lockstep lanes share one instruction fetch/decode/schedule — the amortization from the previous card. Wider = cheaper control per lane.
Why not wider, then? Divergence. When lanes of a warp take different branches (if (threadIdx.x % 2)…), the hardware runs both paths with lanes masked off. Worst case, an N-wide warp does N× the work of what the live lanes needed. Expected waste grows with width. Graphics shaders branch per-pixel constantly (material edges, lighting cases), so GPUs born for graphics — NVIDIA, Apple, AMD's RDNA line — settled near 32 as the sweet spot between issue amortization and divergence waste 0:33:21.
AI flips the tradeoff. Dense tensor code has very little divergence — most lanes execute the same instruction stream. For that workload, the cost of narrower execution groups is less often offset by reduced divergence, so a 64-wide design can amortize issue overhead across more lanes. Hence AMD's compute line (CDNA / MI300) uses wavefront 64 while their gaming line (RDNA) uses 32 0:32:51 — same company, two answers, because the two markets sit at different points on the divergence curve.
Grids, blocks, and how to tile work
The CUDA software model in one line: a kernel launches a grid of thread blocks; each block contains up to 1024 threads; each block is assigned, whole, to one SM; its threads share that SM's shared memory and can synchronize with each other (__syncthreads()). The grid is the unit of "embarrassing" parallelism; the block is the unit of cooperation.
Yes and yes. Every thread in a block sees the same shared-memory allocation; that is the definition of the block's shared memory 0:18:43. Multiple blocks routinely reside on one SM simultaneously, with the SM's registers and SMEM partitioned among them. Crucially, each block gets its own disjoint SMEM slice: block 0 cannot see block 1's shared memory even when they sit on the same SM. Co-residency matters for performance (§7 shows a second block covering the first block's stalls); isolation matters for correctness (inter-block communication has to go through global memory — see the next card for why the normal grid model avoids it).
The mechanism underneath is worth internalizing as a why rather than a rule:
- No ordering guarantee: the GPU's block scheduler launches blocks onto SMs whenever resources free up, in whatever order it likes 0:40:52. If block 2 waits for data block 1 will produce, nothing guarantees block 1 has even started.
- The deadlock trap: launch more blocks than fit on the device at once and have them all wait at a global barrier → the resident blocks wait forever for blocks that can't launch until the resident ones finish. Unrecoverable by construction 0:41:14.
- The middle ground: Hopper's thread block clusters — small groups of blocks the hardware co-schedules and lets read each other's SMEM ("distributed shared memory") over an on-chip network, with real synchronization primitives. This is supported, but it still has coordination cost.
And the design heuristic Ben attaches: for AI kernels, wanting cross-block communication almost always means you partitioned the work wrong 0:39:35. AI workloads usually decompose into independent output tiles (next card). Scientific computing — stencils, graph algorithms — often does not, which is why that world has grid-synchronization libraries and more complicated coordination mechanisms.
Blocking: deciding what each block owns
Ben's whiteboard segment 0:46:59 isolates the central question in blocking: which tile of the problem does each block own? For elementwise operations, several decompositions work about equally well. For matmul, the ownership choice determines whether the kernel can accumulate locally or needs synchronization across blocks.
Two separate questions hiding in one:
Chunk shape (the add-kernel discussion): for the toy matrix-add, 16×16 vs row-strips vs chunked rows all compute the same thing; the differences are (a) whether a row still fits in one block's thread budget when the matrix gets wide, and (b) memory coalescing — long contiguous runs per block use HBM's wide bus better than square patches 0:49:47. For elementwise kernels it's a few percent. It becomes decisive for matmul.
Races (the matmul discussion): the race isn't really about sizes — it's about which matrix you block on. Any decomposition where two blocks contribute partial sums to the same region of C (blocking on A, blocking on B, or split-K along the reduction dimension) creates read-modify-write collisions on C that demand atomics or inter-block sync. Any decomposition where each block owns a disjoint C tile and privately consumes whatever A/B it needs has no races at all, no matter what the tile sizes are. Sizes then get chosen for arithmetic intensity: the biggest C tile your registers can hold (Ben: ~192×192 or 128×256 per block on H100 0:52:09), fed by thin slabs of A and B, because output-tile FLOPs grow as area while input traffic grows as perimeter.
Global memory is transacted in fixed-size units. On modern NVIDIA GPUs, the coalescer thinks in 32-byte sectors within larger cache lines. If a warp's 32 addresses are 32 consecutive 4-byte words, the request occupies four sectors, all useful: perfect. If the 32 lanes are scattered so each lane needs its own sector, you fetch 32 sectors to use 128 bytes — about 1/8 useful bandwidth before any cache reuse rescues the pattern.
"Not using full cache lines" is Ben's diagnosis of the naive add kernel 0:45:51: its 16×16 tiling means each block's rows are short 16-element (64-byte) runs, so transactions are half-wasted at the tile edges unless L2 happens to rescue you by serving the other half to a neighboring block. The fix is always the same: arrange things so that consecutive lanes read consecutive addresses, in runs of ≥128 bytes.
It's worth listing what a production matmul library is actually juggling, because it explains both why cuBLAS is large and why a specialized matmul for one shape can be much simpler 0:12:52:
- Shape coverage: a 128×256 C tile is great for big square matmuls and terrible for a 512×512×8 GEMV-ish shape — small problems need smaller tiles, or split-K, just to fill 132 SMs at all.
- Divisibility & padding: tiles want to divide the problem evenly; ragged edges need predicated or padded epilogue paths.
- Block execution order for L2 reuse: naively marching C tiles left-to-right thrashes L2 on the B matrix. Reordering blocks into super-blocks (small 2-D clusters of tiles executed together, so they share A-rows and B-columns while those are still L2-resident) is worth up to ~2× — "not a small effect" 0:54:36. You'll see exactly this super-blocking comment in TK's matmul template (§12).
None of the individual decisions is hard; the product of all of them across every (M, N, K, dtype, layout, arch) is a combinatorial heuristic table. That design space is why cuBLAS is large.
Memory spaces & the consistency model
The way Ben puts it: understanding NVIDIA's memory consistency model is the hard part of CUDA, and it's an abstraction NVIDIA defines on top of the hardware — if you try to reason from "what the silicon probably does" instead of from the documented model, you will write bugs that take weeks to find 0:03:17.
The memory spaces
| Space | Visible to | What it really is |
|---|---|---|
| Registers | One thread | The lane's slices of the warp's register rows |
| Local | One thread | Per-thread address space; compiler's overflow area (see card) |
| Shared | One block | The SM's SMEM scratchpad slice for that block |
| Constant | Grid, read-only | Small cached read-only region, broadcast-optimized |
| Texture | Grid, read-only | Read path through texture units: 2-D locality caching, interpolation, format conversion |
| Global | Everything | Device DRAM (HBM), through L2/L1 |
Local memory: Ben's model is the right one 0:58:56: conceptually every per-thread variable lives in "local" space; the compiler promotes what it can into registers and the remainder — spills — lands in actual local memory. The trap in the name: local memory is not local silicon. It's global DRAM, thread-partitioned, merely cached through L1/L2. A spill converts a 1-cycle register access into a potentially hundreds-of-cycles memory access inside your hot loop. This is why register budgeting (≤255/thread, watching -Xptxas -v output) is a first-class concern, and why Ben frames implicit cache traffic as something dense AI kernels try to avoid.
Texture memory: it exists for graphics access patterns — sampling a 2-D image with spatial locality, hardware bilinear interpolation, border handling, format conversion on load. AI kernels read dense tensors in large, regular, software-prefetched tiles; every service the texture path offers is either not needed or better done explicitly through SMEM/TMA. So: relevant if you write fragment shaders or image-warping kernels; usually not relevant for transformers 0:59:55.
"Consistency" = when do other observers see my writes, and in what order? Within one thread, the hardware guarantees program order for that thread's own view: write a value, read it back, you get what you wrote. That's the "consistency within a thread's individual memory space."
The moment two threads communicate through any shared space (SMEM or global), that guarantee evaporates. GPU memory is weakly ordered: thread B may observe thread A's writes late, partially, or reordered relative to A's program order — because of caches, buffered stores, asynchronous copy engines, and the scheduler interleaving warps arbitrarily. The model's contract is: writes become reliably usable by other threads only through a synchronization protocol at the right scope. Barriers such as __syncwarp(), __syncthreads(), cooperative-group sync, mbarrier waits, and kernel boundaries provide rendezvous/visibility. Fences are lower-level: they order one thread's memory operations, but a consumer still needs an acquire/atomic/barrier-style handoff before relying on the data. Between sync points, assume nothing.
Hence the rule of thumb above: crossing a memory-space boundary between producer and consumer ⇒ there is a synchronization requirement, find it and satisfy it. §7 makes this concrete twice — once where the needed sync is a __syncwarp() and once where only a full __syncthreads() will do — and shows the kernel actually breaking without them.
The 155-TFLOP 4090 attention kernel, annotated
Now make the synchronization model concrete. Ben opens TK's demo attention kernel — non-causal, head-dim 64, ~155 TFLOPs on a 4090, within ~10% of the best known (figures as stated on stream; the TK paper documents the same synchronous-vs-pipelined and occupancy trade-offs quantitatively R9) — and examines each synchronization point in it. This is a Flash-Attention-2-style kernel: Q tiles stay resident in registers; K/V blocks stream through shared memory; softmax runs online with running max and normalizer. The structure of the hot path:
// each warp loads its own Q tile of 16x64
if (q_seq*ROWS<D> < g.Qg.rows) {
load(qo_smem[workerid], g.Qg, {batch, head, q_seq, 0}); // global → shared (coalescing)
__syncwarp(); // ← card 1
load(q_reg, qo_smem[workerid]); // shared → registers
}
__syncthreads();
for (kv_idx = 0; ...; kv_idx++) { // stream K,V blocks
load_async(k_smem[next], ...); // prefetch next block ← cards 2–4
load_async(v_smem[next], ...);
load_async_wait<N>(); // wait, leaving N in flight
__syncthreads(); // ← card 2
// ... per-subtile: mma_ABt(att, q, k) → online softmax → mma_AB(o, att, v)
}
__syncwarp() between a global→shared load and a shared→register load. What exactly would race without it — and might the kernel happen to work anyway?The precise shape of the hazard is worth spelling out. The global→shared load is strided for coalescing: each thread writes whichever SMEM addresses make the global read pattern contiguous. The subsequent shared→register load uses the tensor-core register layout (§8) — each thread reads addresses that other threads wrote. Producer set ≠ consumer set within the warp ⇒ a cross-thread dependency through shared memory ⇒ by §6's rule, a synchronization is required. Since both loads are performed by the same single warp, warp scope suffices: __syncwarp().
As for whether it might still work anyway — Ben is exact 1:02:23: on current hardware the generated instruction timing may happen to be safe, but the code is fundamentally incorrect without it. It is a race the model does not excuse, and a different compiler version or architecture can expose it. Empirical success on one compiler/GPU pair is not a memory-consistency argument.
__syncthreads() in the main loop is load-bearing — Ben deletes it live and the correctness check fails on the spot. Why isn't warp-level synchronization enough here?Because in the main loop the dependency crosses warps, not just threads of one warp. The K/V loads are divided up: each warp (each thread, really) fetches its assigned slice of the K and V tiles into shared memory. But the compute phase — the mma_ABt against the K tile — reads the entire tile: memory loaded by every warp 1:09:32. It's a broadcast dependency: all-load-then-all-use.
Ben demonstrates the failure path concretely 1:11:24: warp 0 issues its loads, waits for its own memory (that's all load_async_wait checks — per-thread arrival), and advances two loop iterations — before warp 1 has even issued its loads. Warp 0 is now doing matmuls against uninitialized or stale values in the SMEM region warp 1 was responsible for. __syncthreads() is the barrier that says "every thread in the block has arrived here" — i.e., everyone's loads are issued-and-awaited — before anyone computes. He deletes it live and the kernel's correctness check fails immediately 1:04:21. Notably, this is true even in the fully synchronous version of the kernel — synchronous loads are synchronous per thread, which guarantees nothing about your neighbor 1:10:57.
load_async actually do? And is __syncwarp() strictly cheaper than __syncthreads()?load_async wraps the cp.async instruction (SM80+): "copy this from global directly into shared memory, bypassing registers, and don't make me wait — keep issuing instructions; I will tell you when to wait" 1:10:32. Two distinct wins: (1) asynchrony — the copy proceeds in the background while you compute; (2) no register round-trip — a synchronous global→shared copy must go LDG into registers, then STS into SMEM, burning registers and instructions. In this kernel's SASS it shows up as LDGSTS.E.BYPASS.128 (load-global-store-shared, bypassing L1, 128-bit) instead of LDG/STS pairs 1:18:33; the cache behavior is an instruction/cache-policy detail, while the portable idea is async global→shared with no register waypoint. load_async_wait<N>() then means "block this thread until all but the most recent N of my async copy groups have landed" — the N is what lets a pipeline keep stages in flight (next card).
Sync costs: yes — and Ben's own reasoning runs the same way 1:17:34: __syncthreads() logically contains a syncwarp — it converges your own warp and makes it wait for every other warp in the block, so its cost is bounded below by the slowest warp's arrival. __syncwarp() converges 32 threads that live on one scheduler; it's nearly free. Use the narrowest scope that satisfies the actual dependency — which is also Ben's philosophical point at 2:19:20: extra barriers can be incorrect or costly; synchronization scope should follow the dependency being protected.
STS.128 [R3], R4 — soaks up 7% of all samples just waiting. Where does that stall really come from, and why does fixing it barely move end-to-end time here?The profiler screenshot is subtle enough that the mechanics deserve pinning down. STS.128 [R3], R4 means "store the 128-bit value in registers R4..R7 to the shared-memory address in R3." The store itself is fast. What the sampler shows (99.46% "Long Scoreboard" on that line) is that the instruction can't issue: R4's value comes from a preceding LDG from global memory, and the scoreboard — the hardware's operand-readiness tracker — says the data hasn't arrived. "Long scoreboard" = waiting on a long-latency (global memory) dependency. So the STS line inherits the LDG's ~500–700-cycle latency and the kernel spends 7.34% of its sampled time parked on it 1:19:16.
In the async kernel there is no LDG→register→STS chain at all (cp.async goes straight to SMEM), so the residual wait migrates to the loop-closing branch — that's why the hot line becomes BRA.CONV, at ~2% 1:21:20: first-iteration latency plus slack — the "pay it once" reading.
And yes on occupancy 1:21:51: end-to-end, sync vs async differ by only a few percent here because two blocks are co-resident per SM — while one stalls, the other keeps the tensor pipe busy (the profile is full of "math pipe throttle" stalls, the good kind: it means the tensor cores are the bottleneck). Ben then artificially limits occupancy to one block (by inflating SMEM usage) and the gap widens as predicted. The general lesson: occupancy and pipelining are substitutes for hiding latency; H100-class kernels lean on pipelining because their SMEM/register appetites leave no room for occupancy.
Nsight Compute buckets every sampled non-issuing cycle into a stall reason. "Stall Wait" is the bucket for fixed, short, known-at-compile-time latencies: the next instruction needs the result of a previous fixed-latency instruction (an FMA, a shift), and the compiler has statically scheduled a wait of the required few cycles. Contrast "Long Scoreboard" (variable, long, memory) from the previous card.
It's "a good stall" 0:36:48 because it means you're bounded by real arithmetic dependencies, not by memory or scheduling pathology. Mitigations are marginal — more ILP so something else issues in the gap, or cheaper instructions — but a dependent chain of FMAs simply takes chain-length × latency, and a profile dominated by Stall Wait plus math-pipe-throttle is close to the hardware's ceiling. That's Ben's read of the FA3 profile: likely close to the hardware limit.
expf — and a natural question is why it doesn't use log₂(e) instead. What's going on?They are the same number: 1.44269504… is log₂(e). Recall §2: the hardware only has EX2 (2ˣ). To get eˣ you compute EX2(x·log₂e), which naively costs an extra multiply per exponential. The kernel trick 1:25:37: attention already multiplies scores by a constant — the 1/√d temperature. So fold log₂e into it once:
scale = (1/√d) · log₂e // one host-side constant
softmax numerator: EX2(score·scale − rowmax·…) // pure exp2, no FMA tax
which is why the TK code calls exp2() on the attention block and why a bare 1.44269… appears in the kernel source instead of in the SASS. One multiplication, hoisted out of ~seq² exponentials. Optimized attention kernels commonly use this transformation.
Banks, core matrices & swizzles
This section connects the earlier hardware model to ThunderKittens' layout system. The chain of logic is: tensor cores impose register layouts → register layouts determine shared-memory access patterns → shared memory is banked → naive layouts create conflicts → swizzled layouts remove many of those conflicts.
Step 1: shared memory is 32 banks
SMEM is physically 32 independent banks; consecutive 32-bit words map to consecutive banks (word address mod 32). Each bank serves one 32-bit word per cycle. When a warp's 32 lanes issue an LDS, the requests fan out to banks: all different banks (or same-word broadcast) → one cycle in the simple case. Two lanes needing different words in the same bank → serialized: a bank conflict. In the simplified model, an n-way conflict serializes into n bank service steps, so effective SMEM bandwidth falls by roughly that factor 1:31:06.
As a term you'll act on: yes 1:36:22. Registers technically have banked read ports too (visible only as occasional compiler scheduling quirks — nothing you manage), and global memory's analogous failure mode is un-coalesced access (§5), which is a different mechanism (wasted cache-line transactions, not serialized banks). "Bank conflict" as a thing you diagnose in Nsight and fix with layouts = shared memory only. Ben's calibration is that the cost depends on how latency-sensitive the SMEM path is: attention kernels can sometimes tolerate a little because they are waiting on tensor cores anyway; other kernels may be much more sensitive 1:32:10.
Step 2: core matrices — why the "natural" access is the worst one
A core matrix is the atomic tile of the tensor-core register layouts in the PTX guide: an 8×8 matrix of 16-bit values, held collectively by one warp, where each of the 32 threads holds two adjacent values — thread 0 holds row 0, cols 0–1; thread 1 holds row 0, cols 2–3; … threads 0–3 own row 0, threads 4–7 own row 1, and so on. Now do the register math from §2: 8×8×2 bytes = 128 bytes = 1024 bits = exactly one physical register row of the warp. That's the sense in which a core matrix "is" a single 32-bit vector register: each thread contributes its one 32-bit slot (two bf16s) 1:35:02. Bigger MMA fragments are just grids of these 8×8 atoms.
Where the 8-way conflict comes from — the arithmetic is worth doing once by hand. Take Ben's 32×64 bf16 matrix, row-major in SMEM. One row = 64 × 2 B = 128 B = exactly one full sweep of all 32 banks. So column position determines bank, and every row starts at bank 0. Loading a core matrix means the warp touches an 8-row × 8-column patch: each row's 8 bf16 = 16 B = 4 banks (say banks 0–3), and all 8 rows hit those same four banks → requests pile 8-deep per bank → 8-way conflict, 1/8 bandwidth 1:35:32. The row-major layout that's ideal for loading whole rows is a poor match for the access pattern the tensor cores require.
Step 3: the fixes — padding and swizzles
Padding: allocate each row one extra bank wide — a 64-column bf16 tile stored with stride 66 elements (33 words) instead of 64 (32 words). Nothing about the data changes; you just skip a word at each row end. Effect: row r now starts at bank (33·r) mod 32 = r — each successive row is rotated one bank left. The 8-deep column stack that used to hit one bank now spreads across 8 different banks. Costs: ~3% wasted SMEM, and you lose the power-of-two stride (which complicates vectorized 128-bit accesses — part of why it "only" cut conflicts to ~4–5-way in Ben's picture rather than to zero 1:36:22).
Why transpose kernels love it: a transpose through SMEM reads row-wise and writes column-wise (or vice versa). In an unpadded power-of-two-width tile, the column-wise phase is a textbook 32-way conflict — every element of a column lives in the same bank. Pad by one and column c's elements land in banks c, c+1, c+2, …: conflict-free in both directions. It's the classic fix because transpose is the maximally conflict-prone access pattern and padding solves both phases with one line of code.
A swizzle is any deterministic, invertible permutation of where elements sit within a tile in SMEM, chosen so the access patterns you care about spread across banks. Padding rotates by adding; the XOR swizzle permutes by XORing: element at logical (row, col) is stored at column col ⊕ f(row) (with the XOR applied to some bits of the word-column index). XOR of a fixed value is a self-inverse bijection on {0,…,2ᵏ−1}, so each row is a clean shuffle of column positions — no space wasted, strides stay power-of-two, and the de-swizzle at read time is one cheap XOR in the address math.
Why powers of two: XOR only permutes within a power-of-two-sized set. If the row width isn't a power of two, col ⊕ f(row) can produce a column index outside the row (XOR can only clear/set bits, and there's no carry to wrap it around), so the "permutation" would scatter elements out of the tile. Non-power-of-two widths need additional bounds handling or a different mapping 1:37:42.
The "cut by 2" question: it points at a genuine live correction. Ben first presents the XOR layout as conflict-free; Quinn points at the picture and objects; Ben concedes it's a 4-way conflict for this access pattern — the same as TK's 32-byte swizzle — not zero and not exactly "cut by 2" 1:38:08. The honest general statement: how much a given swizzle helps depends on the swizzle's period vs. the access pattern's geometry — which is exactly what the next card quantifies.
These are the swizzle periods relevant to this TK layout discussion — 32 B, 64 B, and 128 B. CUDA's tensor-map API exposes additional variants for other data types and access modes, but these three are the ones that explain the bank-conflict pictures in the stream. The byte number is the period: the size of the chunk within which 16-byte groups get XOR-permuted, driven by low bits of the row index. Intuition for the conflict counts on the core-matrix access 1:38:39:
- 32 B period rotates among 32 B = 8 banks' worth of positions → an 8-deep column stack collapses onto 32/8·… → lands as a 4-way conflict;
- 64 B period spreads twice as far → 2-way;
- 128 B period spreads a row's 16-byte groups across the full 128 B = all 32 banks → the 8 rows of a core-matrix access all land on disjoint banks → conflict-free.
Why not always 128 B? The period can't exceed the row: a 128 B swizzle needs rows ≥ 128 B (64 bf16 columns). TK's templates auto-select the widest swizzle your tile width admits 1:43:03 — that's the resolution of "how do you choose": you don't; the width chooses. Narrow tiles simply can't do better than 4-way with these periods. Why use hardware-recognized swizzles at all? Because TMA can load data into SMEM with bank shuffling applied, while ldmatrix/stmatrix efficiently consume the resulting layout from row addresses supplied by the program — close to free compared with doing every permutation as manual scalar bit-twiddling (next card).
ldmatrix amortize that 4:1 — and are there workloads where a swizzled layout is simply the wrong choice?The de-swizzle math (shift/mask/XOR per address) runs on the regular ALU pipe, so it eats issue slots — recall §4: every scalar instruction pays 20:1. Two mitigations: ldmatrix/stmatrix (LDSM/STSM) move four 8×8 core matrices per warp instruction and amortize instruction/address work across those fragments — Ben's "amortize 4:1" 1:39:18. They still consume row addresses; the point is that the row-address and instruction overhead is shared across multiple core matrices. And TK makes swizzle bit-math compile-time template arithmetic, so much of it constant-folds.
Still, yes: for kernels whose SMEM access is purely row-streaming (no tensor-core layout in play — think layernorm, elementwise, some reductions), the swizzled layout buys nothing and its address overhead is pure loss; a naive layout would be optimal 1:43:31. TK accepts being "awfully close all the time" rather than optimal every time — that tradeoff is the next section.
TK philosophy: primitives, not abstractions
Here's the distinction. An abstraction promises you can ignore what's underneath (you don't think about transistors when writing Python). A primitive is a named, composable handle on something real underneath — using it correctly requires knowing what it does 0:01:14. TK's tiles, layouts, and ops are primitives: rt_bf<16,64> is a specific arrangement of core matrices across a warp's registers; load_async is cp.async; warpgroup::mma_AB is a wgmma with a defined synchronization contract. Nothing is hidden — things are just named.
Why does this matter? Three things follow.
- The failure mode is explicit. Because nothing synchronizes for you implicitly, your memory-consistency obligations (§6–7) are exactly the hardware's. If TK were an abstraction that "handled" sync, it would either over-synchronize everything or hide races that are already difficult to observe directly.
- Extension is expected. Ben's own MoE kernel drops to raw CUDA mid-kernel for a broadcast op TK lacked 2:03:32 — a register vector is still represented as an array, so ordinary CUDA remains available. The intended usage is to read the closest existing op, copy it, and adapt it.
- The performance/readability trade is deliberate. CUTLASS can generate very efficient SASS ("when I read their code I'm always impressed" 1:39:49) but after nine months Ben still can't fluently read it 1:40:37. TK's bet: the useful margin in 2024+ is not squeezing the last 5% from already-optimized attention kernels; it is making new kernels practical to write (less standard inference ops, new architectures), where the difference can be much larger and where a 75-line readable kernel gets written but a 3,000-line CUTLASS kernel might not 1:44:26. Research that needs a custom kernel becomes less likely if the kernel takes much longer to write than the idea itself 2:06:39.
TK types: global, shared, register
TK is a header-only C++20 library (CUDA 12.x) R9R10, organized as common / types / ops, with a serious unit-test suite — ~2,000 tests over every load/store/layout/MMA combination 1:47:35. Its type system mirrors the memory hierarchy exactly: global layouts (a typed view of device memory), shared tiles/vectors (SMEM with automatic swizzling), register tiles/vectors (warp-distributed fragments). Data flows global → shared → register → shared → global, and each hop is one op.
Global layouts (gl)
A gl doesn't own memory — it's a shape-annotated pointer: "this is a batch × depth × rows × cols tensor of bf16 at this address." In the kernel you then say which tile you want ({batch, head, seq_block, 0}) and load it in one call 1:49:31. Indexing is pure-iterator: coordinates are in units of whole tiles, not elements — "row 1" of a tile-typed access means "the second tile down" 1:56:57.
Ben explicitly declines to give the reason on stream 1:50:28, so flag this as inference rather than transcript: a likely candidate is the TMA descriptor. A TMA tensor map (CUtensorMap) supports transfers of tensors up to 5 dimensions R1R6, and TK auto-generates TMA descriptors from your gl (see below); fixing the layout at 4 named dims (batch, depth, rows, cols) means every gl maps onto one descriptor shape with room to spare, and every tile load is "fix two coordinates, transfer over the last two." A fixed arity also keeps the template machinery and address math uniform.
Unrolling is just index flattening — the same trick as reshape. A 5-D tensor [B, H, X, R, C] with contiguous strides is byte-identical in memory to the 4-D tensor [B·H, X, R, C]: element (b, h, x, r, c) is element (b·H + h, x, r, c). No data moves; you fold two indices into one when you compute the coordinate, exactly as PyTorch's view(B*H, X, R, C) would. It works whenever the dims you merge are adjacent and contiguous.
-1 versus a literal in the template. But why must compile-time dims be constructed with nullptr, of all things?The template convention: gl<bf16, -1, -1, -1, 64> declares batch/depth/rows as runtime dims (−1 = "I'll tell you when I construct it") and cols as a compile-time dim (64, baked into the type). Under the hood it's C++20 concepts choosing between a runtime_dim (holds a size_t) and a compiled_dim (a static constexpr, zero bytes at runtime) 1:51:27. Why bother: compile-time dims let the compiler fold address arithmetic into immediate operands — fewer registers, fewer ALU instructions in your hot loop (§4's issue tax again). Ben's own verdict: "not a big effect, I'm not sure it was worth it, but I like template programming" 1:52:27.
The nullptr rule is defensive API design 1:51:54: at construction, a compile-time dim's argument type is std::nullptr_t, so passing an actual number there is a compile error. The failure it prevents: you pass a runtime value for a dimension that was templated as fixed, the value gets silently ignored, and your kernel reads the wrong tensor shape. Making "this slot is compile-time" unspellable-except-as-nullptr moves that bug from runtime debugging into the type checker.
TMA transfers require a tensor map descriptor — a host-side object encoding base address, dims, strides, tile box shape, and swizzle mode. Historically, that descriptor is created via a verbose driver-API call, kept in sync with your kernel's tile types by hand, and tedious to debug when wrong. In TK you parameterize the gl with the shared-tile type you intend to load (gl<bf16, -1, -1, -1, 64, st_bf<64,64>>) and the library builds and manages the descriptor, with a compile-time template dictionary tracking which descriptors exist — so requesting a TMA transfer with a tile you never registered is a compile error, not a runtime crash 1:55:34. The swizzle mode in the descriptor is derived from the tile type too, so TMA lands data in SMEM already swizzled (§8).
Shared tiles
Parameterized by type, rows, cols (multiples of 16). You address them as if row-major — st.idx(row, col) — and the template computes the swizzled address 1:58:01. Subtiles alias into a parent tile while satisfying the same template interface as full tiles. Shared vectors are intentionally simple: contiguous arrays in shared memory, without the tile-specific layout machinery.
Register tiles: the two layouts you actually touch
mma_ABt() both Q and K fragments use row layouts, but V — consumed by mma_AB() — must be declared column-layout. What forces the asymmetry?The why generalizes and is worth writing down. In TK's row/column vocabulary, the convenient primitive is effectively A·Bᵀ: B is held in the fragment arrangement the MMA expects for that operation 1:59:57. TK exposes this as two layouts, row_l and col_l, matching the PTX guide's operand layouts R8:
mma_ABt(D, A, B, C)— computes A·Bᵀ; A and B both row-layout. Attention's S = Q·Kᵀ is this natively — that's why QK comes first everywhere and why torch defaults to A·Bᵀ conventions 2:01:25.mma_AB(D, A, B, C)— computes A·B; hardware still does A·(something)ᵀ, so B must be held pre-transposed: col-layout. In attention, O += P·V, hencev_regis declaredcol_l.
Two important footnotes from the stream: (1) the "transpose" costs nothing at MMA time — it's the same MMA instruction; the difference is upstream, in which ldmatrix variant filled the fragments (LDSM.M88 vs LDSM.M88.TRANS in the SASS 2:02:42). (2) TK considered auto-transposing for you and refused — hiding it would invite you to write code with gratuitous layout churn; making you declare it keeps the cost visible. Wrong layout = compile error, never a wrong answer 2:00:54.
Register vectors are how row-wise statistics live (max/sum vectors in softmax, means in layernorm), and they need layouts because they must line up with the tile fragments they're reduced from or broadcast into:
- naive — values dealt sequentially across threads, no duplication. For plain long-vector math (layernorm-style) where no tile alignment is needed.
- align / ortho — mirror the tensor-core tile layout along one axis or the other, with ~2× value replication so that each thread already holds every vector element its tile fragments need 2:08:34. You trade a few registers to eliminate the shuffles/syncs a row-broadcast would otherwise need every softmax step.
In practice, this choice is usually inferred by the type system. You write col_vec<rt_fl<16,64>> ("a column vector matching this tile") and the templates instantiate whichever layout (ortho, as it happens) fits; mismatches are compile errors — the static_assert Ben triggers live at 2:10:08. It is a static type check for CUDA tile layouts.
Ops: warp scope & group scope
Three things sharpen the picture:
Why warp scope is the floor: §3 — the warp is the smallest unit the hardware actually executes, and TK's register tiles are warp-distributed objects (no single thread even holds a tile). So thread-scope tile ops would be incoherent. Where an operation is physically single-threaded (TMA issue), TK elects lane 0 and the rest of the warp treats it as a collective anyway 2:15:33. Uniform scope = uniform reasoning about who participates in what.
What a group op really is: group<N> is a compile-time partition — N consecutive warps act as one worker, each doing 1/N of the operation. warpgroup is an alias for group<4> 2:17:47 (4 warps = 128 threads = the unit Hopper's wgmma wants, §13). In a warpgroup::store of a 64-row Q tile, each of 4 warps stores its own 16-row slice — a collaborative store. Register ops are absent from group scope because register files don't span warps: there's no physical way for warp 1 to touch warp 0's registers, so a "group register op" would decompose into independent warp ops. (The one special case: warpgroup::mma_*, where the wgmma instruction itself spans 4 warps' register files — a hardware feature, not a TK convention.)
The sync discipline that comes with groups: a group op does not synchronize its warps unless you ask 2:18:18 — after a collaborative load you still need group<N>::sync(id) before anyone consumes the whole tile. Defaulting to the broadest barrier fails twice 2:19:20: a barrier inside an if-branch that only some warps take is a deadlock, and defensive over-syncing usually means the dependency has not been identified precisely. "You will write correct code by knowing what your code is doing."
Producer–consumer & the LCSF template
The interesting step is why it wins over the §7 style (every warp both loads and computes, "multistage" pipelining):
- Registers are the scarce resource, and the two jobs need opposite amounts. Loader warps need almost none; compute warps want the max. Hopper lets warps reallocate — producers call
setmaxnreg.dec, consumers.inc(visible in TK's matmul: "producer reduces its own number of registers… Hopper only" 2:26:30) — so specialization converts idle loader registers into bigger accumulator tiles. - Decoupling maximizes overlap. In multistage, one warp's compute phase delays its own next load-issue. With specialization, producers run a pure load loop gated only by "is a buffer free," consumers a pure compute loop gated only by "is a buffer full" — a classic bounded queue. Ben credits the rope kernel's speed to exactly this "maximal decoupling" 2:31:10.
- And yes, explicitly a programming model 2:24:34: the code contains a literal
if (warpid() >= NUM_CONSUMER_WARPS) { producer path } else { consumer path }. The compiler discovers nothing; you declare it.
producer::{setup, load, store} and consumer::{setup, compute, finish}; the template generates the pipelining, ring indices, phase-bit bookkeeping, and semaphore initialization 2:22:37. Changing pipeline depth is changing one number.The clean rule underneath is that the template family has two members:
- LCSF (load-compute-store-finish): output is produced every iteration (rope: each tile in, each tile out). Then storing is steady-state pipeline work, so it belongs to the warps whose job is global-memory traffic: producers run
store()in the loop. - LCF (no producer store): output appears once, at the end (matmul: the accumulator lives in consumer registers for the whole K-loop). Only consumers have the data, so the epilogue is theirs:
consumer::finish()writes accumulators → SMEM → TMA store to global 2:26:30.
So for matmul specifically: there is no producer store phase — finish (consumers) is the only output path. Rope has both roles active every iteration. The discriminator is simply: who holds the bytes when they're ready to leave?
This is the DeepSeek MLA prefill case: Q/K are 576-dim, of which V/O are the first 512; the same 576/512 head-dim split appears in DeepSeek's own FlashMLA kernels R7. Why it breaks the standard partitioning: the §7 kernel gives each warp its own 16×64 Q tile, resident in registers. At D=576, one 16-row Q tile is 16×576×2 B = 18 KB ≈ 144 registers' worth per thread just for Q — plus K fragments, plus a 16×512 output accumulator. The 255-register ceiling (§3) is exceeded, so the kernel cannot keep multiple full Q rows resident per warp under the previous partitioning 2:39:55.
The repartition Ben performs live: the four warps stop being four independent workers on four Q tiles and become one collective worker on a single 16-row Q tile, split along the head dimension — each warp owns a 16×144 slice (576/4), computes partial Q·Kᵀ scores on its slice, and the partials get reduced before softmax. Loads become warpgroup:: collective ops; tile types split into a 576-wide QK tile and 512-wide V/O tile. One subtlety: the 576-wide absorbed kernel dimension is the partitioning problem, not the softmax scale dimension. For DeepSeek MLA, the scale Ben folds with log₂e remains based on the original QK dimension, 192, so the constant is derived from 1/sqrt(192), not 1/sqrt(576) 2:49:06. The useful generalization is that unusual head dimensions often do not need new algorithms; they need new partitionings. A primitives library makes that repartitioning a local edit rather than a complete rewrite.
Pre- vs post-Hopper
| Pre-Hopper (Ampere/Ada: A100, 3090, 4090) | Hopper (H100/H200) | |
|---|---|---|
| Async global→shared | cp.async (SM80+), per-thread, waited with load_async_wait | cp.async.bulk.tensor — TMA: descriptor-driven, issued by one thread, completes onto an mbarrier |
| MMA | Warp-synchronous mma: operands loaded to registers manually, per-warp | wgmma: warp-group (128-thread) scope, asynchronous, can source B — and in common variants A — from shared-memory descriptors |
| Sync object | __syncthreads/__syncwarp + cp.async groups | Those, plus mbarriers (TK: "semaphores") for TMA producer/consumer readiness; WGMMA has separate fence/commit/wait groups |
| Store path | Manual STS/STG | TMA stores; store_add_async (async reduction into global) "especially useful" |
Blackwell/B200 keeps related async-copy ideas but adds TCGen05/TMEM as a new tensor-core programming model, so the Hopper column should not be read as the full Blackwell model.
TK semaphore = PTX mbarrier, renamed R8. An mbarrier is a small synchronization object living in shared memory: it has an expected arrival count, a transaction byte count, and a phase bit. Threads can arrive on it, asynchronous copy machinery can reduce its transaction count as bytes complete, and waiters proceed only when both counts for the current phase have reached zero. Ben calls it a semaphore because that's what it behaves like, and because he wants the name portable to AMD/MLX backends where the implementation will differ 2:29:17.
"One-directional" is the contrast with __syncthreads(): a syncthreads makes everyone both signal and wait — a rendezvous. An mbarrier decouples the two for a particular handoff: one side signals readiness, the other waits for it 2:29:42. In a full bounded-buffer protocol the roles alternate across paired barriers — producers may wait for "finished/free" and consumers may signal "finished/arrived" — but each barrier itself remains directed. That decoupling is exactly what a producer–consumer pipeline needs.
"TMA load happens onto an mbarrier": when you issue a TMA load you pass it an mbarrier; you also tell the barrier how many bytes to expect (tma::expect). When the copy engine finishes landing data in SMEM, it completes that transaction count against the barrier. No thread polls the transfer — consumers wait on the barrier and wake when the bytes are physically there 2:29:42. This is the mechanical meaning of "maximal decoupling" in §12, and it's the piece cp.async lacks (cp.async completion is tracked per issuing thread, which is why §7 needed the syncthreads).
What "manual loads into registers" means: pre-Hopper, the MMA instruction reads its operands from registers only. So your inner loop is: ldmatrix K-fragment from SMEM into registers → mma → repeat, one warp at a time, synchronously (the warp waits for its own MMA). The register file is a mandatory waypoint, and each warp independently owns a small MMA.
What wgmma changes — and why it's the bigger deal:
- Scope: one wgmma is executed by four warps jointly; the accumulator is striped across all 128 threads' registers. Your unit of work — "a worker" — stops being a warp and becomes a warp group. That forces the repartitioning Ben describes: tile sizes, who loads what, how you index workers — the skeleton of the kernel changes 2:35:20.
- Operand source: wgmma can read B, and in common variants A, directly from shared memory via descriptors; other variants keep A as register fragments. Compared with pre-Hopper warp MMA, much more of the operand path can bypass the explicit SMEM→register staging loop, and the SMEM→tensor-core path becomes load-bearing for the whole SM. This is what "Hopper depends a lot on the pipeline shared memory to tensor core" means.
- Asynchrony: wgmma is issued and runs in the background; you fence and
mma_async_wait()later — MMAs themselves become pipeline stages you overlap, like loads.
TMA, by contrast, mostly substitutes for cp.async at the same points in the kernel: same broad code shape, different call plus a barrier 2:34:05. That's why porting pre↔post Hopper is "copy the kernel, promote warp ops to warpgroup ops, rethink partitioning" — not a rewrite, but not one portable kernel either 2:34:49.
The DeepSeek workflow & the copy kernel
harness.impl and a PyTorch wrapper. What does each side do, and is that two-loop workflow still the right one today?harness.impl is a bare C++ test driver, conditionally included into the kernel file by a macro: a main() that memory-maps a file of reference inputs/outputs (pre-generated by a small Python/NumPy script), copies inputs to the GPU, launches the kernel, checks the result against the reference, and times it 2:37:24. It compiles with plain nvcc + a Makefile in seconds. The PyTorch wrapper is the production entry point: a pybind11/torch-extension shim that accepts torch::Tensors and launches the same kernel.
"Make sure the bindings are done correctly" = validating the glue the harness never exercised: dtype checks, contiguity/stride assumptions (is that tensor actually [B,H,S,D] row-major, or did someone .transpose() it into a strided view?), device placement, shape → grid-dim math, stream handling. Stride assumptions are a common source of wrapper bugs, so bindings need separate tests.
Two-loop workflow: inner loop (seconds) — edit kernel, make, run harness; outer loop (minutes) — build the torch extension, test the binding. Still true today? Directionally yes: a from-scratch torch.utils.cpp_extension.load_inline build with real templated CUDA is still tens of seconds to minutes, though the ecosystem has improved the outer loop (better extension caching, and newer no-compile paths for calling kernels from Python). The underlying principle is evergreen regardless of tooling: keep the iteration loop for the hard part (the kernel) decoupled from the slow machinery around it.
Ben's method, pulled into one place
Threaded through the stream is a complete working method, collected here in one place:
- Write the copy kernel first 0:55:44. Full layouts, full pipeline, all synchronization, zero compute; verify bytes in = bytes out. Most debugging effort lives in dataflow; verify it early.
- Fill in compute inside the verified memory machine. The point of the copy-kernel workflow is to separate layout and synchronization bugs from the arithmetic, so the compute pass starts from a known-good data movement scaffold.
- Know your synchronization scope. Every sync in the kernel should answer a specific cross-thread dependency you can name (§7 gives two concrete examples).
- Iterate on a harness, ship through a wrapper (this section).
- Profile in SASS, read the stalls. Long Scoreboard = memory dependency (fix with pipelining); Stall Wait + math-pipe throttle = you're near the roofline (stop) 0:36:48. One caveat: NCU pins H100 SM clocks below their real boost while HBM runs at full rate, so kernels look slightly less memory-bound under the profiler than in production 1:25:07.
- Do not rely on "impossible" cases staying impossible. "If you think you're going to get away with 'that would never happen' — you're wrong. It will happen" 0:56:56.
__syncthreads() in TK's demo 4090 attention kernel; (2) write the copy kernel for a simple op end-to-end with a harness; (3) read the PTX ISA chapter on mma/ldmatrix fragment layouts with §8 open beside it; (4) profile anything in Nsight Compute and account for every stall reason ≥ 1%; (5) then the ThunderKittens repo's kernels/ directory, simplest first — rope, then matmul, then attention.References
- NVIDIA, "NVIDIA Hopper Architecture In-Depth" — SM/quadrant structure, 144/132/114 SM configurations, 60/50 MB L2, and TMA up to 5-D tensors.
- NVIDIA, CUDA C++ Programming Guide: compute capabilities and asynchronous copies — compute-capability tables (255 registers/thread max, 32 shared-memory banks × 4 B), async copy, TMA, and thread block clusters.
- Chips and Cheese, "NVIDIA's H100: Funny L2, and Tons of Bandwidth" — independent measurement of the split-L2 far-partition latency.
- IEEE Spectrum, "The Secret to Nvidia's AI Success" (interview with Bill Dally) — instruction fetch/decode overhead ≈ 20× the ~1.5 pJ of the arithmetic; complex-instruction amortization (IMMA overhead at ~16% of the math).
- Shah, Bikshandi, Zhang, Thakkar, Ramani & Dao, "FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision", arXiv:2407.08608 — FA-2 at ~35% H100 utilization, FA-3 up to 740 TFLOPs/s (75%); see also the accompanying blog post for the MUFU throughput arithmetic (§2's 16 ops/SM/clock).
- NVIDIA, CUDA Driver API:
cuTensorMapEncodeTiled— the TMA descriptor object and its dimensionality/swizzle parameters. - DeepSeek, FlashMLA — production MLA kernels with the 576 (QK) / 512 (VO) head-dimension split.
- NVIDIA, PTX ISA —
mma/ldmatrixfragment layouts,mbarrier,wgmma,cp.asyncandcp.async.bulk.tensor. - Spector, Arora, Singhal, Fu & Ré, "ThunderKittens: Simple, Fast, and Adorable AI Kernels", arXiv:2410.20399 — the tile/template abstractions, occupancy-vs-pipelining measurements, and kernel benchmarks discussed throughout.
- Hazy Research, ThunderKittens repository — the code walked through on stream (
lcsf.cuh, kernel gallery, unit tests); see also "GPUs Go Brrr". - NVIDIA, "Inside Volta" — independent thread scheduling, per-lane program counters, and the
_syncprimitive family.