Achieve cuBLAS Performance on NVIDIA GB300 for HGEMM

Tensor core HGEMM (half-precision BF16 GEMM) kernels written from scratch, whose goal is to match cuBLAS performance on NVIDIA GB300 GPUs. Each kernel computes a row-major C = A @ B^T (A: M×K, B: N×K, C: M×N) with BF16 inputs/outputs and FP32 accumulation in tensor memory (tmem), built directly on the Blackwell tensor cores and the Tensor Memory Accelerator (TMA).

The kernels are numbered v0 … v7 and form a step-by-step optimization tutorial: every version builds on the previous one and adds a single major optimization. v0 … v5 were adapted from https://github.com/gau-nernst/learn-cuda/tree/main/02_matmul_sm100.

Performance

Measured on NVIDIA GB300 for M = N = K = 8192 (BF16):

Implementation	Performance (TFLOP/s)
cuBLAS	1963.582
v0	314.655
v1	818.752
v2	1213.041
v3	1490.823
v4	1678.140
v5	1764.173
v6	1905.022
v7	1980.925

Optimizations

v0 — baseline

Single-tile GEMM (128×256×256). One warp issues non-swizzled TMA loads of the A/B tiles into shared memory behind a single mbarrier, then runs a flat tcgen05 MMA loop (sixteen MMA_K=16 slices) accumulating over K into tmem. The FP32 result is read back from tmem, converted to BF16, and written to global C with strided (uncoalesced) per-thread stores. Load and compute are fully serial — no pipelining.

v1 — swizzled shared memory + coarser TMA

Same tile shape (128×256×256) and math as v0, but reworks all data movement so the shared layout matches how tcgen05.mma reads it.

Key points:

• Coarser, swizzled TMA. TMA uses 128B swizzling (CU_TENSOR_MAP_SWIZZLE_128B) with a transfer granularity of K=64 instead of 8, cutting TMA issues per operand from 32 down to 4. Shared memory is organized as (M or N, 64)-wide swizzled blocks (128-byte rows).
• Matching MMA descriptor. The tcgen05 shared descriptor drops the explicit LBO/SBO encoding and instead sets 128B-swizzle mode (2ULL << 61, LBO implied), so the TMA-written layout and the MMA-read layout are identical.
• Reorganized MMA loop. The same 16 MMA_K=16 slices are now iterated as a nested 4×4 loop: k1 selects the 64-wide swizzled tile, k2 the 16-wide slice inside it.

v2 — multi-stage pipeline

Adds a software pipeline (NUM_STAGES=2) that double-buffers the A/B shared tiles so global loads overlap the MMA.

Key points:

• Per-stage barriers. Each stage owns its own shared-memory region and its own tma_mbar, plus a single shared mma_mbar. They form a dependency chain: TMA copy → tma_mbar[stage] completes → tcgen05.mma reads that stage → commit → mma_mbar.
• Prefetch + overlap. The loop prefetches NUM_STAGES-1 tiles, then loads the next tile while computing the current one. The overlap is between the TMA and MMA async units; each iteration still waits on mma_mbar, so multiple MMA groups are not kept outstanding yet.

v3 — warp specialization

Splits the mainloop across two warps so load and compute run as independent producers and consumers.

Key points:

• Producer / consumer warps. warp 0's elected thread issues TMA loads; warp 1's elected thread issues tcgen05.mma. This enables a deeper pipeline (NUM_STAGES=4, smaller BLOCK_K=64).
• Full/empty handshake. Each stage has a pair of barriers: tma_mbar[stage] (TMA→compute: tile ready) and mma_mbar[stage] (compute→loader: stage reusable), giving proper pipeline back-pressure between the two warps.
• Mainloop→epilogue barrier. A dedicated mainloop_mbar signals that the final tmem accumulator is ready before the whole CTA runs the tcgen05.ld epilogue. It is kept separate from the per-stage mma_mbar (which only means "shared stage reusable") so the async MMA/commit completion is synchronized with clear, fixed-phase semantics.

v4 — 2-CTA cluster MMA

Introduces a cluster of 2 CTAs (CTA_GROUP=2) that jointly compute a wider tile.

Key points:

• Shared load barrier. BLOCK_N is split across the two CTAs; both ranks run load() and each arrive.expect_tx on the same cluster tma_mbar, which is initialized with count CTA_GROUP. The MMA therefore waits until both CTAs' data (arrivals + bytes) have landed.
• Cluster MMA. Only rank 0 issues the cta_group::2 tcgen05.mma (MMA_M = BLOCK_M * CTA_GROUP = 256), writing a 256×BLOCK_N group accumulator into shared tensor memory; in the epilogue each CTA reads back its own 128 rows via cta_rank.
• Multicast commit. The MMA commit uses multicast with cta_mask = 0b11 so the completion arrival reaches both CTAs' mma_mbar (and later mainloop_mbar), letting each CTA's loader reuse stages and start its epilogue.
• Deeper pipeline + swizzle. Pipeline depth grows to NUM_STAGES=6, with a tile-swizzle grid mapping (bid_m/bid_n) that improves L2 reuse.

v5 — persistent kernel + epilogue overlap

Turns the kernel persistent (grid sized to the SM count, e.g. 148) with an inner loop over output tiles, and splits the work across three warp roles so a second, tensor-memory-level pipeline can hide epilogue latency.

Key points:

• Three warp roles. warp 0 feeds (TMA global→shared), warp 1 computes (shared→tmem MMA), and warps ≥2 drain (tmem→global C).
• Two-level pipeline.
  1. Shared-memory K pipeline — NUM_STAGES shared stages with a tma_mbar/mma_mbar handshake: loader waits mma_mbar[stage] (stage free), issues TMA, arrives tma_mbar[stage]; compute waits tma_mbar[stage], runs the MMA, arrives mma_mbar[stage] to release the stage.
  2. Tensor-memory output pipeline — 2 tmem slots with a mainloop_mbar/epilogue_mbar handshake: compute waits epilogue_mbar[slot] (slot free), runs the full K-loop into that slot, arrives mainloop_mbar[slot]; the epilogue waits mainloop_mbar[slot], reads/casts/stores C, arrives epilogue_mbar[slot] to release the slot.
• Overlap. The loader prefetches later K-stages while compute consumes ready ones, and compute works on one tmem slot while the epilogue drains the other — so tile N's epilogue overlaps tile N+1's mainloop.

v6 — TMA store epilogue

Replaces the strided per-thread output writes with a TMA store path. In the epilogue each accumulator sub-block is read from tmem into registers, converted to BF16, written into a swizzled shared-memory staging buffer, and pushed to global C with cp.async.bulk.tensor.2d (via a C_tmap descriptor) for fully coalesced stores.

Key points:

• Store staging layout. Two C-store staging buffers (SC_SIZE = 2 * SC_STAGE_SIZE) are reserved at the start of dynamic shared memory, ahead of the A/B load stages, so store staging never overlaps the operand load buffers. Each buffer holds one 128×64 BF16 sub-block (16 KB), and a full 128×256 tile is emitted as BLOCK_N/64 = 4 TMA stores that ping-pong between the two buffers.
• Overlapped stores. Stores are committed to a bulk-async group; wait_group 1 only drains the older group before a buffer is reused, keeping one store in flight. This overlap spans the n-loop and even the next output tile (its tmem load / cast / shared-store), with a final wait_group 0 draining the tail.
• Tile-swizzle scheduling. A compute_bid mapping walks M-tiles first within groups of SWIZZLE_GROUP=16 before advancing N, so consecutive CTAs reuse the same B tile for better L2 locality; m_in_group clamps the last (partial) group.

v7 — larger tiles, M-waves, and L2 promotion

Grows each CTA's tile to a larger BLOCK_M=256 and adds an overlapped, TMA-based epilogue.

Key points:

• M-waves. Each CTA owns a 256×256 output tile (from its own TileScheduler coordinate) processed as two 128-row M-waves (WAVE_M=128, NUM_M_WAVES=2). The 512 tmem columns store the two waves' accumulators side by side (rather than being a two-slot double buffer as in v6).
• CTA-group MMA. MMA_M = WAVE_M * CTA_GROUP = 256, so a single tcgen05_mma<2> produces a 256×256 result across the cluster with rank 0/1 each covering 128 rows; each CTA runs two waves to fill its 256×256 tile.
• Leader-owned load barrier. Per K-stage every CTA loads its own A (256×64) and a B slice (128×64, split across CTAs in N). The leader CTA owns the tma_bar and registers the whole cluster's expected bytes in one arrive_expect_tx(TMA_BYTES * CTA_GROUP); the peer CTA only remote-arrives, with the barrier address normalized to CTA0 via & 0xFEFFFFFF.
• Pipelined TMA-store epilogue. The accumulator is drained wave-by-wave through a 2-stage shared staging buffer and written to global C with cp.async.bulk TMA stores, keeping one store in flight so global writes overlap the next tmem read/cast.