The Asynchronous Kernel: Tensor Cores & TMA

The idea in one sentence

The fastest matrix multiply on a modern GPU is a factory: one piece of hardware fetches the tiles, another multiplies them, and the threads that are left do nothing but say when.

The kernel that stopped being the best one

Lecture 26 climbed to 93.7% of NVIDIA’s own library, and every rung of that ladder asked one thread to do three jobs at once. Work out an address. Fetch the numbers sitting at it. Multiply them. The last rung got 64 results into a single thread’s registers and kept the arithmetic units fed by hand, with a rhythm of loads and barriers the programmer wrote out line by line.

Look at what that kernel does not use. It never touches a tensor core, even though the card it runs on has them. Every multiply-add is issued by a thread, for a thread, on numbers that thread loaded. That is the shape of a fast kernel from before the hardware changed underneath it.

The hardware then took two of the three jobs away. Fetching a tile became the work of a dedicated copy engine that one thread can order and then forget about. Multiplying a tile became a single instruction that 128 threads issue together and none of them owns the result of. What is left for the threads is the third job, saying when, and this chapter is about why that turned out to be the hard one.

The worklog behind the numbers here is Pranjal Shankhdhar’s H100 worklog, which walks eleven kernels from 32 to 764 TFLOP/s on one H100. The explanation that pulled the machinery into one place is Aleksa Gordic’s Inside NVIDIA GPUs. Both are worth reading after this chapter, and both are credited where their numbers appear.

The ceiling moved

Start where Lecture 26 started, with what a perfect kernel would cost, because the answer is no longer the same answer.

Peak arithmetic on a GPU is one multiplication of three numbers. How fast the clock ticks, how many tensor cores are listening to it, and how many operations one tensor core finishes in one tick. Nothing else enters, which is why the result is called a speed of light: it is what you would get if no thread ever waited for anything.

Check the multiplication on a chip whose answer is published. An A100 has 108 multiprocessors with 4 tensor cores each and a 1.41 GHz boost clock, and NVIDIA rates it at 312 TFLOP/s of dense half-precision arithmetic. Divide and you learn what one tensor core does in one tick:

Five hundred and twelve floating-point operations per core per clock, which is a round number because the hardware was built to a round number.

Now do the H100. It has 132 multiprocessors, still 4 tensor cores each, so 528 in all. NVIDIA’s Hopper architecture post says the fourth-generation tensor cores deliver “2x the MMA computational rates of the A100 SM on equivalent data types”, which turns 512 into 1024. Horace He’s post on matmul timing gives the chip’s top clock as 1.830 GHz. Multiply:

989 TFLOP/s. Lecture 26’s whole ladder lived under a 30 TFLOP/s roof.

The other half of the roofline moved much less. An H100 SXM reads its own memory at 3.35 TB/s. Divide one by the other and you get the arithmetic intensity a kernel must reach before arithmetic, rather than memory, is what limits it. With the ordinary FP32 units, which NVIDIA rates at 67 TFLOP/s on this card, that ridge point sits at 20 floating-point operations per byte. With the tensor cores it sits at 295. Same chip, same memory, and a bar almost fifteen times higher.

Now price a tile against that bar. A block tile of $m$ rows by $n$ columns marching along $k$ reads $m + n$ numbers for every step of $k$ and does $2mn$ operations on them. In bf16 a number is two bytes, so over the full depth $K$ the intensity is

which is Lecture 26’s formula with the byte count halved. The largest tile in this chapter is $128 \times 256$, and it gives $I = 32768/384 = 85.3$. That is nowhere near 295. On paper the tile alone caps you at $85.3 \times 3.35 = 286$ TFLOP/s, which is 29% of the machine.

That gap is why this chapter has a second half. Getting near the speed of light is not a question of picking a bigger tile. It is a question of arranging for most of the bytes a tile asks for to be found on the chip already, and that turns out to be a question about the order the tiles are visited in.

The tile is the instruction

Here is the change that undoes Lecture 26’s mental model.

On Hopper the matrix multiply instruction is called wgmma, for warp-group matrix multiply accumulate. A warp-group is four warps, so 128 threads, and they issue it together. For bf16 operands its shapes are written m64nNk16 with $N$ running from 8 to 256 in steps of 8, which means the instruction consumes a $64 \times 16$ block of $A$ and a $16 \times N$ block of $B$ and adds their product into a $64 \times N$ accumulator.

Take $N = 64$ and count. That one instruction performs $64 \times 64 \times 16 = 65{,}536$ multiply-adds, which is 131,072 floating-point operations. Lecture 26’s inner loop did two of them per lane per issue.

Two operands and a result, and none of them is where you would expect. Operand $B$ must be in shared memory. Operand $A$ may be in shared memory or in registers. And the accumulator, the thing you actually wanted, is spread across the whole warp-group.

That last one is the sentence that ends the thread as a unit of thought. A $64 \times 64$ accumulator holds 4096 floats. A warp-group has 128 threads. So each thread holds $4096/128 = 32$ of them, and the declaration in a real kernel is a plain array, float d[N/16][8], which at $N = 64$ is d[4][8], which is 32. No thread can read the answer. No thread can print it. Asking one thread what the product is has stopped being a question with an answer.

The widths matter for a reason a programmer meets fast. At $N = 256$ the accumulator is $64 \times 256 = 16{,}384$ floats, so 128 registers per thread, and a CUDA thread may hold at most 255 registers in all. One warp-group therefore cannot own a $128 \times 256$ output tile: two instructions of that width would want 256 accumulator registers per thread, more than the whole budget, with nothing left over for addresses or operands. That single constraint is why the fast kernels in this chapter run two consumer warp-groups rather than one, each owning half the tile.

The instruction is also asynchronous, which is where this chapter gets its title. It returns before the answer exists. Around it sit three more instructions whose only job is to say when the accumulator may be touched.

// one warp-group, one accumulator, three fences around the multiply
float d[WGMMA_N / 16][8];                     // 32 registers when N = 64

asm volatile("wgmma.fence.sync.aligned;");    // freeze d: the hardware owns it
wgmma_async(d, desc_a, desc_b);               // wgmma.mma_async...m64n64k16
asm volatile("wgmma.commit_group.sync.aligned;");   // close the batch
asm volatile("wgmma.wait_group.sync.aligned 0;");   // only here does d hold it

The two descriptors are not pointers. They are packed words that tell the tensor core where the operand sits in shared memory, how far apart its rows are, and how it was scrambled on the way in, which is the subject two sections down.

One thread orders the groceries

Now the other job that left. Getting a tile of $A$ from global memory into shared memory used to be a loop: every thread worked out an address, issued a load, and waited for it. Hopper has a unit that does the whole tile, called the Tensor Memory Accelerator, or TMA.

The setup happens once, on the host, before the kernel launches. A call to cuTensorMapEncodeTiled builds a descriptor holding the element type, the base pointer, the shape of the matrix, the distance between its rows, the shape of the tile you want, and how that tile should be scrambled on arrival. Read that list again and notice what it is. It is a layout, in the sense Lecture 27 is about to make precise, handed to hardware instead of to a loop.

Inside the kernel the copy is one call, made by one thread.

__shared__ alignas(128) bf16 slot[BM * BK];   // 128 x 64 of bf16: 16,384 bytes
__shared__ barrier bar;

if (threadIdx.x == 0) init(&bar, blockDim.x); // all 128 threads will arrive
__syncthreads();

barrier::arrival_token token;
if (threadIdx.x == 0) {
    cp_async_bulk_tensor_2d_global_to_shared(slot, &tensor_map, col, row, bar);
    token = barrier_arrive_tx(bar, 1, sizeof(slot));   // 1 arrival, 16,384 bytes
} else {
    token = bar.arrive();                     // the others promise no bytes
}
bar.wait(std::move(token));                   // opens when both counts are full

The interesting object there is not the copy. It is bar. A __syncthreads() counts threads and nothing else: it opens when everybody has arrived. This barrier counts two things, and the second one is bytes. Thread 0 arms it with the size of the transfer, the copy engine reports its progress against that number, and the wait opens only when every thread has arrived and every promised byte has landed.

Counting bytes rather than threads is what makes the copy free. The 127 threads that arrived and promised nothing are not blocked on the transfer at all. They can be given a different job, and two sections from now they are.

The XOR that unties the banks

There is a reason the descriptor carries a scrambling mode, and it is the oldest problem in shared memory.

Lecture 25 walked through the thirty-two doors: shared memory is 32 banks of 4 bytes, and when the lanes of a warp ask for different addresses in the same bank the hardware serialises them. A tensor core sharpens the problem, because it does not read a tile the way a loop does. It wants rows on one step and columns on another, out of the same block of shared memory.

Work in 16-byte chunks, which is the unit a wide load moves, and think of a tile as an 8 by 8 grid of them. Stored the obvious way, chunk $(r, c)$ lands in bank group $c$ for every row $r$. Reading a row is then perfect: eight chunks, eight different bank groups, one turn. Reading a column is the worst case the hardware has, because all eight chunks want bank group $c$ and get served one at a time, at eight times the cost.

The fix is one exclusive-or. Store chunk $(r, c)$ at position $c \oplus r$ instead of $c$. Nothing changes along a row, since exclusive-or with a fixed $r$ is a permutation of $0$ through $7$. And now a column is $c \oplus 0, c \oplus 1, \ldots, c \oplus 7$, which is a permutation for the same reason. Both directions are conflict free, out of one operation that costs nothing.

The real thing works on address bits rather than on a toy grid. CuTe writes it Swizzle<B, M, S>, which takes the $B$ address bits sitting $S$ places up and exclusive-ors them down into the $B$ bits at position $M$. Swizzle<3, 4, 3> is the common choice for tensor core operands: three bits, so eight positions, over chunks of 16 bytes, giving a pattern that repeats every 128 bytes. The pleasant part is that you rarely write it. It is a field in the TMA descriptor, and the copy engine does the scrambling as the bytes land.

The factory

Every piece is now in place, and they assemble into a shape that has nothing to do with a loop.

The copy engine can run without a thread. The tensor cores can run without a thread. What remains is to keep both busy at once, and the way to do that is to stop asking one warp-group to alternate between them. Split the block’s warp-groups by role. One of them does nothing but issue copies. The others do nothing but multiply. This is called warp specialization, and PyTorch’s write-up of the CUTLASS ping-pong kernel states the split plainly: “each warp group takes on a specialized role of either Data producer or Data consumer”.

Between them sits a queue. Shared memory is carved into a small fixed number of slots, each big enough for one pair of tiles, and the slots are used in a ring. Each slot carries its own pair of barriers, and that pair is the entire synchronisation protocol:

• the producer waits on empty[i], fills slot $i$, and signals full[i];
• a consumer waits on full[i], multiplies out of slot $i$, and signals empty[i].

Nothing else is needed. No barrier across the block, no __syncthreads() around the loop, no thread waiting on a thread. Each slot is a small contract between two groups that never look at each other.

if (warpgroup_id == 0) {                   // the producer
    setmaxnreg_dec();                      // it barely computes: give registers up
    for (int k = 0; k < k_tiles; ++k) {
        int slot = k % QSIZE;
        empty[slot].wait(phase(k));        // the consumer is done with this slot
        if (elected_one_thread) {
            tma_load(sA[slot], sB[slot], k);
            barrier_arrive_tx(full[slot], 1, BYTES_PER_SLOT);
        }
    }
} else {                                   // the consumers
    setmaxnreg_inc();                      // and they take what it gave up
    for (int k = 0; k < k_tiles; ++k) {
        int slot = k % QSIZE;
        full[slot].wait(phase(k));         // the bytes are in
        wgmma(d, sA[slot], sB[slot]);
        empty[slot].arrive();              // the slot may be refilled
    }
}

Two details in that skeleton pay for themselves. The first is setmaxnreg. A producer warp-group computes almost nothing, so its threads need almost no registers, while each consumer thread is holding 128 accumulator registers and would spill without more. The instruction moves the budget from one group to the other at run time. PyTorch’s post gives the sizes CUTLASS uses: the producer’s count goes down by 40 and each consumer’s goes up by 232.

The second is phase. A barrier that is reused every QSIZE iterations cannot be waited on by name, because a fast group would sail straight through a barrier the slow group opened one lap ago. Each barrier carries a phase bit that flips every time it opens, and a waiter names the phase it expects. Getting that wrong gives you a kernel that is right at one queue depth and silently wrong at another.

There are two ways to run the consumers and both are shipped. In the cooperative arrangement the two consumer warp-groups work on the same output tile. In the ping-pong arrangement they work on different ones, which lets one run its epilogue, the scaling and the store, while the other is still multiplying. PyTorch gives the reason: “one can be using the tensor cores for MMA while the other performs the epilogue, and then vice-versa”.

Keeping the line busy

The factory keeps one multiprocessor fed. Two things stop the whole chip from staying fed, and both are about launch geometry rather than about the inner loop.

The first is waves. Lecture 26 measured the cliff: blocks are dealt to the multiprocessors in rounds, and a launch that overshoots a round by one block pays a whole extra wave for it. The shape this chapter has been building is no exception. A $4096 \times 4096$ output cut into $128 \times 256$ tiles is $32 \times 16 = 512$ tiles, which is 3.88 waves on the 132 multiprocessors of an H100 SXM, so 4 waves, the last of which fills 116 of the 132 units. That is 97% of the machine used, and the missing 3% is worth having.

A persistent kernel takes it. Instead of launching one block per tile and letting the scheduler deal them out, launch exactly as many blocks as there are multiprocessors, keep them alive for the whole call, and have each one pull tiles from a shared counter until the work is gone. The wave boundary disappears because there are no waves. Something better comes with it: a block about to start its next tile can issue that tile’s loads while the stores of the previous one are still in flight, so the tail of one tile hides inside the head of the next. Pranjal’s kernel goes from 631 to 660 TFLOP/s in the step that makes the blocks persistent, hides the stores and groups the tiles for the cache.

Scheduling is geometry

Which raises the question that decides the last stretch, and it is a question about drawing rather than about code: in what order should the output tiles be visited?

Any order computes the same answer. The orders differ in how much the tiles in flight have in common. A tile of $C$ at row $i$, column $j$ needs panel $i$ of $A$ and panel $j$ of $B$. If eight tiles are in flight and they all sit in one row of the tile grid, they share one panel of $A$ and need eight different panels of $B$, so nine in all. If they sit in a $4 \times 2$ block instead, they need four and two. Same eight tiles, two-thirds of the operand traffic, and the L2 cache is what feels the difference.

The middle order is the one most libraries ship: walk a few rows of the tile grid at a time so that consecutive tiles stay inside a compact block. It costs a few lines of index arithmetic and takes the average from 9.86 panels to 7.02. A Hilbert curve is a space-filling curve that visits every cell of a grid while keeping every prefix compact, and it reaches 6.67. On Pranjal’s ladder the grouped order is part of the step that gets to 660 TFLOP/s and lifts his L2 hit rate to 83%, against 70% for cuBLAS on the same problem. The Hilbert curve is his last kernel and adds 6 TFLOP/s on top of 758.

Hopper adds one more level to that picture. Blocks can be grouped into a cluster, a set of blocks guaranteed to be resident on nearby multiprocessors at the same time, and the members of a cluster can read each other’s shared memory directly. NVIDIA calls that distributed shared memory. It means a panel of $B$ that two neighbouring blocks both want can be fetched once and multicast, instead of fetched twice. A cluster behaves like a wider multiprocessor with a wider desk, and on Pranjal’s ladder it is worth 704 to 734 TFLOP/s, the step that first passes cuBLAS.

The scoreboard

Here is the whole climb, on one H100 SXM, at $M = N = K = 4096$ in bf16. Every rate is Pranjal Shankhdhar’s, measured with CUDA 12.6 and averaged over eight runs. The percentages below and the comparison against the speed of light are ours.

Three things in that picture are worth saying out loud.

The first bar is 32 TFLOP/s, which is 3.2% of the machine. That is Lecture 26’s design, the one that reached 93.7% of cuBLAS on an RTX A6000, running on hardware whose fast path it does not know exists. A kernel is only as good as the machine it was written for.

The second bar is 317, and it is one step: use the tensor cores and use the copy engine. Ten times the speed, from adopting the two instructions this chapter is named after. Everything after that step is about keeping those two busy, and everything after it adds up to another factor of 2.4.

The last bar is 764 against cuBLAS at 716. Divide and the kernel is 6.7% faster than NVIDIA’s library on this shape, which is what Pranjal reports, and he is careful to add that the margin moves with size: about 2% at $N = 512$, 7 to 8% at 2048 and 4096, and 1.5% at 8192. It is also 77% of the 989 computed in Figure 1, and the missing 23% is where the honest answer lives. Part of it is the memory system, part of it is the power limit from the warning above, and the rest is the last percent that only exists below the language.

Where this is going

Count what the thread lost in this chapter. It does not work out the address, because a descriptor built on the host did that. It does not fetch, because a copy engine does. It does not multiply, because a warp-group does. It does not even own the answer. What it owns is a role and a place in a queue.

So what is a kernel written in, once it is not written in threads? Every piece of machinery above turned out to be a promise about where numbers sit. The TMA descriptor is a shape, a stride, a tile size and a scramble. The swizzle is a function from a coordinate to an address. The accumulator fragment is a map from a thread and a register number to a place in a tile. A queue slot is a tile shape with an offset.

Lecture 27 names that object. A layout is a function from coordinates to addresses, written as a shape and a stride, and the operations a kernel needs become algebra on those functions. It is also where four generations of this hardware get drawn as one timeline, each one taking another job away from the thread, which is the story this chapter is a single frame of.

If you want to write these kernels rather than read about them, the Library has the practice grounds: LeetGPU for CUDA judged on real hardware, Triton Puzzles for the block-level way of thinking, and the GPU MODE reference kernels for problems people are still competing on. Pranjal’s fast.cu repository is the code behind every number in Figure 9.