subha

Optimized B200 Matrix Multiplication

A deep dive into building a performant matmul kernel using NVIDIA's B200 GPU with persistent kernels, pipelining, TMA, and TMEM.

Hi everyone! Today we’ll be looking into how to build a performant matrix multiplication kernel using the new NVIDIA B200 architecture :)

I’ll be walking through the code step-by-step, explaining the specific hardware features of Blackwell that allow us to achieve high utilization.

First, let’s define the problem. We want to solve C=ABC = AB where:

  • AA is of size M×KM \times K
  • BB is of size K×NK \times N
  • CC is of size M×NM \times N

1. The Strategy: Persistent Kernels

For this implementation, I used a Persistent Kernel approach.

In a standard CUDA kernel, you might launch 10,000 blocks to cover a large matrix. The GPU scheduler assigns them to SMs, retires them when done, and fetches new ones. This incurs scheduling overhead and forces repeated data fetching from global memory.

The Persistent Approach: Instead of launching a block for every tile, we launch exactly enough blocks to fill the hardware. The B200 has 148 Streaming Multiprocessors (SMs), so we launch 148 blocks.

These blocks never “finish.” Instead, they run a while loop, grabbing a claim on a workload tile (via an atomic counter or simple grid striding), processing it, and immediately moving to the next. This allows us to keep data in the L2 cache and maintain our TMEM allocations without setup overhead.

Thread Block Hierarchy

Within each of our 148 blocks, we use 256 threads divided into Warpgroups:

  • Producer Warpgroup (Warp 0): Handles TMA (loading data from Global →\to Shared Memory).
  • Consumer Warpgroup (Warp 1): Handles MMA (Matrix Multiply Accumulate using Tensor Cores).

2. Blackwell Hardware Features

To make this efficient, we leverage three specific hardware features: TMA, TMEM, and Swizzling.

TMA (Tensor Memory Accelerator)

TMA is a dedicated hardware engine that copies data from Global Memory to Shared Memory asynchronously. It frees up our threads to do math while data is moving.

TMEM (Tensor Memory)

This is a new memory type in the Blackwell architecture.

  • Old Way: Tensor cores accumulated results into Registers (RF).
  • Blackwell Way: TMEM acts as a dedicated accumulator memory shared across the SM.
  • Benefit: It isolates the accumulation traffic, preventing register file pressure.

We allocate TMEM at the start of the kernel using tcgen05.alloc:

// Allocate TMEM columns (performed by Warp 1)
if (threadIdx.x >= 32 && threadIdx.x < 64) { 
    const int addr = static_cast<int>(__cvta_generic_to_shared(tmem_base));
    // Allocate 256 columns of TMEM
    asm volatile(
        "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
        :: "r"(addr), "r"(num_cols) : "memory"
    );
}
__syncthreads();

Swizzling (The Hidden Optimizer)

One detail often overlooked is memory layout. In my host code, when creating the TMA descriptors, I specifically set CU_TENSOR_MAP_SWIZZLE_128B.

// From host code
cuTensorMapEncodeTiled(
    ..., 
    CU_TENSOR_MAP_SWIZZLE_128B, // <--- Key for performance
    ...
);

Swizzling rearranges the layout of the bytes in shared memory to avoid Bank Conflicts. Without this, multiple threads might try to access the same memory bank simultaneously, serializing access and killing bandwidth.

3. The Pipeline (Producer vs. Consumer)

We use a Multi-Stage Pipeline with a ring buffer of depth PIPE_DEPTH = 4.

  • The Producer loads into slot (X+3)%4(X+3) \% 4.
  • The Consumer computes on slot XX.

We need strict synchronization to ensure the Producer doesn’t overwrite data the Consumer is still using, and the Consumer doesn’t read data that hasn’t arrived yet. We use mbarriers for this.

The Producer (TMA Warp)

The producer’s job is to keep the pipeline full. It issues cp.async.bulk commands.

// Producer Logic (Warp 0)
for (int iter_k = 0; iter_k < num_k_tiles; iter_k++) {
    int stage_id = iter_k % PIPE_DEPTH;
    
    // 1. Wait for Consumer: Don't overwrite data if MMA is still using this stage!
    // We skip this wait for the first fill of the pipeline.
    if (iter_k >= PIPE_DEPTH) {
        mbarrier_wait(mma_mbar_base + stage_id * 8, phase ^ 1);
    }

    // 2. Flip phase if we wrapped around the ring buffer
    if (stage_id == PIPE_DEPTH - 1) phase ^= 1;

    // 3. Issue TMA Load (A and B tiles)
    // ... cp.async.bulk instructions ...
}

The Consumer (MMA Warp)

The consumer waits for data, computes, and then tells the producer “I’m done with this buffer.”

// Consumer Logic (Warp 1)
for (int iter_k = 0; iter_k < num_k_tiles; iter_k++) {
    int stage_id = iter_k % PIPE_DEPTH;

    // 1. Wait for Producer: Wait for TMA to finish loading this stage
    mbarrier_wait(tma_mbar_base + stage_id * 8, phase);

    // 2. Issue MMA (Matrix Multiply)
    // broken down into smaller chunks for the hardware instruction
    
    // 3. Commit: Tell the producer we are done with this stage
    asm volatile(
        "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
        :: "r"(mma_mbar_base + stage_id * 8) : "memory"
    );
}

4. Optimization: From 800 to 1200 TFLOPs

Initially, my kernel was stuck at 800 TFLOPs.

The Bug: My original implementation was too synchronous. In every k_tile iteration, I was waiting for the TMA to arrive, doing the math, and then waiting for the math to finish before moving to the next tile.

The Fix (Asynchronous Pipeline): I decoupled the wait signals.

  1. Consumer never waits for MMA to finish. It issues the command, commits the completion signal to the barrier, and immediately loops to the next tile.
  2. Producer only checks that barrier when it wraps around the ring buffer and needs to reuse that specific slot.

This separation allows the GPU to have TMA copies and Tensor Core math overlapping almost perfectly. This pushed performance to ~1200 TFLOPs.

5. The Epilogue: Escaping TMEM

Once the main loop finishes, our result CC sits in TMEM. We can’t write TMEM directly to Global Memory. We have to:

  1. Load from TMEM to Registers.
  2. Convert formats (Accumulation is FP32, Output is BF16).
  3. Store to Global Memory.

We use int4 vectorization to maximize write bandwidth:

// Epilogue: Direct TMEM -> Global
for (int n = 0; n < TILE_N / 8; n++) {
    // 1. Load from TMEM
    float tmp[8];
    asm volatile(
        "tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0...%7}, [%8];"
        : "=f"(tmp[0])... : "r"(addr)
    );

    // 2. Convert FP32 -> BF16
    __nv_bfloat162 out[4]; // vectorized type
    out[0] = __floats2bfloat162_rn(tmp[0], tmp[1]);
    // ...

    // 3. Vectorized Store (16 bytes = int4)
    reinterpret_cast<int4*>(out_ptr)[0] = reinterpret_cast<int4*>(out)[0];
}

Conclusion

By combining Persistent Kernels, Blackwell’s new TMEM/TMA engines, and a carefully managed asynchronous pipeline, we can extract massive performance from the B200.

If you have questions about the mma_mbar logic or the specific PTX instructions, feel free to reach out!

a sleeping pixel cat
fin