Chapter 3 · Kernel optimization
Tensor cores and low precision
3.3

Tensor cores and low precision

Every kernel in the previous section topped out well short of the GPU’s advertised peak because none of them touched a . On the same A6000 used above, switching cuBLAS from plain fp32 to TF32 or BF16 precision, so it can dispatch to tensor cores, raises its measured throughput by 2.5× and 3.5× respectively, entirely from using different hardware for the same multiply-accumulate.

Going narrower than bf16 means an 8-bit floating-point format. The interchange format published by NVIDIA, Arm, and Intel defines two 8-bit encodings rather than one: , with a 4-bit exponent and 3-bit mantissa, and E5M2, with a 5-bit exponent and 2-bit mantissa. The two trade off differently: E4M3 gives up representing infinity and most NaN bit-patterns to extend its range, topping out at a maximum normal value of 448, while E5M2 keeps full IEEE 754 behavior for those special values and reaches a maximum normal value of 57,344. The paper’s own recommendation is E4M3 for weight and activation tensors and E5M2 for gradients, since gradients tend to need the wider dynamic range.

Two ways to spend eight bits. Both FP8 formats spend one bit on the sign and split the remaining seven differently. E4M3 gives up infinity and most NaN bit-patterns to push its maximum normal value to 448; E5M2 keeps full IEEE 754 special values and reaches 57,344, the wider dynamic range the paper recommends for gradients. Exponent bits are shaded, mantissa bits highlighted.Constants from the source

FP8’s narrow range means a tensor has to be rescaled before it is cast down: a is applied first so a tensor’s largest values sit near 448 or 57,344 instead of overflowing or clustering near zero, then removed again once the FP8 matrix multiply has produced a higher-precision result. NVIDIA’s TransformerEngine library implements this bookkeeping so it does not have to be hand-rolled per model: its DelayedScaling recipe wraps a forward pass in FP8 at a chosen format, E4M3 in the README’s own example, and it targets Hopper, Ada, and Blackwell GPUs, the same generations whose tensor cores can execute FP8 matrix multiplies directly.

Because the boundary between compute-bound and memory-bound work depends on flops per byte, halving a tensor’s size in memory does two things at once: it doubles the arithmetic a tensor core can push through in the same span of time, and it halves the bytes that have to move to feed it. That is why FP8, and the narrower MXFP8 and NVFP4 formats TransformerEngine has since added for Blackwell, keep showing up wherever a kernel is waiting on bytes rather than flops, the same trade-off the KV cache made in the previous chapter, applied to the weights and activations themselves.

Feeding tensor cores at these rates also changed how data moves. The Hopper generation added the (TMA), a dedicated unit that copies tiles of multi-dimensional arrays between global and shared memory asynchronously. A single thread issues the copy against a descriptor that carries the tensor’s shape and strides, so address computation and out-of-bounds predication happen in one place instead of consuming registers in every thread, and the copy’s asynchrony is what makes warp-specialized kernel schedules practical. TMA also emits the swizzled shared-memory layouts Hopper’s tensor-core instructions expect, layouts too intricate to be worth loading by hand.

The tensor-core instruction itself went asynchronous in the same generation. Hopper’s wgmma.mma_async family is issued not by a thread or a warp but by a of four warps, because the accumulator no longer fits anywhere smaller: an m64n16k16 multiply accumulates a 64×16 tile of fp32 values, 1,024 registers, where a single thread can address at most 255. Spread across the warp group’s 128 threads that is 8 registers each, and the instruction variants scale the N dimension from 8 all the way to 256.

That constraint is a Hopper fact rather than a permanent one. Blackwell SM100 introduces the tcgen05 MMA family, which reads and writes the accumulator directly in tensor memory, a new on-chip memory dedicated to it, freeing the register file for other work. With the accumulator out of the register file the reason for spreading the instruction across 128 threads goes with it: only one thread issues a tcgen05 MMA, and two adjacent CTAs can jointly execute a single one to double the tile size. The hardware chapter covers what those instructions compute. The point here is that the register-budget argument expires with the generation that motivated it.

A worklog in the same spirit as the A6000 one above shows what those two units are worth on an H100. The A6000-style kernel, ported to bf16, reaches 32 TFLOP/s where cuBLAS reaches 716, because on Hopper the tensor cores are not optional. Rewriting the inner loop around TMA loads and wgmma lifts it to 317 TFLOP/s, larger 128×128 tiles to 423, and splitting each block into a producer warp group that issues TMA loads into a circular buffer of shared-memory tiles and a consumer warp group that drains it through the tensor cores, the pattern called warp specialization, to 498. The finished kernel, several optimizations later, outruns cuBLAS by 7% at N=4096.

That machinery is no longer exotic: DeepSeek’s open-source DeepGEMM is a tensor-core kernel library for exactly these SM90 and SM100 GPUs, covering FP8, FP4, and BF16 GEMMs plus fused mixture-of-experts kernels, compiled at runtime by a lightweight JIT module so installation needs no CUDA compilation. Its README describes borrowing concepts from CUTLASS and CuTe while avoiding heavy reliance on their templates, keeping a small set of core kernel functions readable as a learning resource, and reaching up to 1550 TFLOP/s on an H800 while matching or exceeding expert-tuned libraries across matrix shapes.

How a TMA copy says it is done#

The previous section’s pipeline needed a producer to signal a consumer that a stage was full. TMA supplies that signal in hardware, and the protocol is worth spelling out because it is the same one every warp-specialized Hopper kernel runs. Colfax’s tutorial fits the whole of it into one short kernel: two lines set the barrier up, one issues the copy, and two more make the rest of the block wait on it.

TMA load, from CUTLASS Tutorial: Mastering TMA (declarations elided)
if (threadIdx.x == 0) {
    initialize_barrier(tma_load_mbar, /* arrival count */ 1);

    set_barrier_transaction_bytes(tma_load_mbar, tma_transaction_bytes);

    auto tma_load_per_cta = tma_load.get_slice(0);
    copy(tma_load.with(tma_load_mbar),
         tma_load_per_cta.partition_S(gmem_tensor_coord_cta),
         tma_load_per_cta.partition_D(smem_tensor));
}

__syncthreads();
wait_barrier(tma_load_mbar, /* phase */ 0);

// after this line, the TMA load is finished

The barrier is an , a shared-memory object initialized with an arrival count. The count is 1 here, because a single thread issues the TMA operation, and the barrier’s starting phase is always 0. The second line does two things in one instruction, arriving on the barrier and declaring how many bytes the transfer will deliver, and that byte count is exactly the size of the tile being copied. The copy then names the same barrier as its completion mechanism, so the hardware credits bytes against it as they land.

The wait is where the shape of the code stops matching the shape of the work. Only thread 0 issued the copy, but every thread in the block waits on the barrier, which is why the __syncthreads() before the wait is not optional: it resolves the divergence the if block created. The wait blocks until the barrier’s phase bit flips, and the phase is supplied by the caller, 0 for the first use after initialization. A second load through the same mbarrier has to flip the phase to reuse it, which is precisely the bookkeeping CUTLASS’s Pipeline classes exist to take over once a kernel is running a circular buffer of stages rather than a single load. Once the wait returns, the memory model guarantees that the TMA’s writes to shared memory are visible to every thread that waited.

One detail sits outside the kernel body. The descriptor argument has to be annotated __grid_constant__ const, and a kernel copying two tensors needs its own such instance for each, which is the requirement for passing a tensor map from host to device. Stores work differently: a TMA store uses no mbarrier at all, enforcing memory consistency with a memory fence instead.