Outperforming cuBLAS on B200

This article will outline iteratively implementing Blackwell features to optimize a matmul kernel in Brain Floating Point 16 (bf16) on varying matrix shapes on a B200 NVIDIA GPU. Our final kernel reaches 106% of cuBLAS's performance on M=N=K=8192. These optimizations arise from the increasing asymmetric hardware scaling of tensor core performance relative to memory bandwidth and other general purpose computation units on NVIDIA GPUs. We will begin with a brief summary of features new to Blackwell aswell as B200 specifications then begin iteratively optimizing our matmul kernel.

Blackwell Architecture

The main new hardware features in Blackwell data-center GPUs are the introduction of Tensor Memory, Cluster Launch Control, and the 5th generation MMA instruction that enables MMA's spanning across peer CTAs.

1. Tensor Memory

   Tensor memory arises from issues with the Hopper architecture. Hopper MMA instructions (WMMA and WGMMA) computed operands and stored accumulator results in the register file, namely in the registers of each thread respectively. The issue with this architectural choice is that many other operands and instructions also require the usage of the register file, with index calculations being the most notable, which leads to immense contention for registers between MMA operands, results, and instructions. Overall this harms performance as operands are spilled into more expensive memory regions like local memory which increases the latency of fetching these operands for operations. With Blackwell datacenter GPUs like the B200, MMA results are now accumulated in a new region of memory called tensor memory, instead of registers.

   Lane (2KB each)

   Addressing direction

   The tensor memory region is per SM, with tensor memory being organized as a 2-dimensional matrix that is 128 lanes (rows) and 512 columns in size. Each cell within this matrix is 32-bits or 4-bytes in size. Making the total memory capacity:

   $$128 \text{ lanes} \times 512 \text{ cols} \times 4 \text{ bytes} = 256 \text{ KB per SM}$$

   Another key detail is the access pattern of tensor memory.

   Lane (2KB each)

   Warp boundary

   As opposed to other memory regions, tensor memory requires an entire warpgroup (4 warps) for full access of all lanes and columns, also the warpgroup must start at a warpgroup aligned index (multiple of 4). With each warp in the warpgroup being responsible for accessing 32 lanes. Therefore an entire warpgroup is required for the epilogue stage of matrix multiplication for writing computed values from tensor memory back to HBM. Another notable detail is that tmem must be deallocated by the kernel explicitly, as it is managed by the programmer, as opposed to smem that will be automatically deallocated.

2. Cluster Launch Control

   Cluster Launch Control (CLC) is a Blackwell hardware mechanism that gives CTAs the ability to cancel queued CTAs that haven't begun executing yet and steal their work by taking the cancelled CTA's index. Unlike non-CLC persistent kernels that launch as many CTAs as there are SMs, CLC kernels launch as many CTAs as there are output tiles. This work stealing approach allows for CTAs to dynamically receive the next tile to work on as opposed to static assignment. Overall this improves load balancing as CTA run-times can exhibit variability so CLC can dynamically adjust for such variability via scheduling.

3. 2-CTA MMA

   Blackwell tensor cores can be used with the new tcgen05.mma PTX instruction. One new feature of tcgen05.mma is that alongside the new cta_group operand, it enables larger MMA shapes that span across 2-CTAs within a thread block cluster.

B200 Specifications

Developed on the TSMC 4NP process node, the B200 has the following specifications:

• Power Target: 1000W
• Maximum Clock Speed: 1.965 GHz
• Base Clock Speed: 700 MHz
• Streaming Multiprocessor Count: 148
• L2 Cache: 126 MB
• Die Arrangement: 2 Dies
• VRAM: 288 GB HBM3E @ 8 TB/s
• Shared Memory per SM: 227 KB

Initial Kernel

We will now walk through iteratively optimizing a matrix multiplcation kernel starting with the initial kernel that leverages the 2 main workhorses of NVIDIA GPUs, the Tensor Memory Accelerator (TMA) and Tensor Cores. The basic design of a CUDA matmul is to bring data from global memory (HBM) into shared memory, with its lower latency accesses, and then perform a matrix multiply accumulate on that data with Tensor Cores.

We will iterate using tiles of dimensions BMxBK for A and BKxBN for B along the K dimension. Where at each point we will perform a load using the TMA of a BMxBK tile of A and a BKxBN tile of B into smem. Then we will issue an MMA to compute the output tile from the A and B tiles. Both TMA and MMA are only issued by one thread, so we will choose an arbitrary warp, warp0, and elect one thread from it for both the TMA and MMA issues. Once our entire tile over K loop is finished, we can then write out output BMxBN tile back to C in HBM. Both TMA loads and MMA instructions are asynchronous, meaning that they work in the background and return control back to the issueing thread before the operation has completed. Given the asynchronous nature of TMA and MMA, we will use mbarriers to track and wait for their completion. This is needed on the TMA side as we need to ensure that smem is fully populated with A and B tiles in order for our MMA to have the operands it needs, and it is needed on the MMA side as we need to ensure that our MMA is fully completed with using the A and B tiles in smem, before we can load new A and B tiles into smem. In the TMA loads, we need an mbarrier expect/wait for (BM + BN) * BK * 2 bytes to arrive in order to signify that our TMA load has successfuly moved both A and B tiles into smem. For synchronizing the MMA, after completing the loop over the BK dimension, we will perform a tcgen05_commit in order to link the execution of this MMA to our mbarrier, then we can wait for the MMA to finish.

At this starting point we are reaching about 40% of cuBLAS's performance.

2-CTA MMA

Our first optimization is to leverage 2-CTA MMA. 2-CTA MMA provides the ability for a pair of CTAs in the same thread block cluster to work together to produce 2 output tiles of C in one MMA instruction.

The benefit is that we can create a larger MMA tile, which in turn means that we can increase the arithmetic intensity of our operations. Arithmetic intensity is number of FLOPs / number of byte accesses. As tensor core performance grows faster than memory bandwidth, it is critical to get the most operations out of any operands fetched. Instead of doing 2 128x256 MMAs to get a 256x256 output tile, we can directly do 1 256x256 MMA to get the same result. For our purposes we will have each CTA contain a unique row of A and half of the shared column of B, with half of the 256x256 output tile being accumulated in the tmem of each CTA.

By having each CTA load half of B, we are saving a redundant load of B that occurs in 1-CTA MMA where the exact same column of B is loaded into the SMEM of both CTAs. The 1.5x arithmetic intensity benefit can be seen below:

In terms of actually implementing this optimization, we must first understand thread block clusters.

First introduced in Hopper, thread block clusters are a new level of hierarchy above CTAs. Clusters ensure that CTAs within a cluster are scheduled onto the same GPU Processing Cluster which gurantees that they will be on neighbouring SMs. The benefit of clusters is in being able to leverage Distributed Shared Memory, where CTAs are able to read from and write to the shared memory of other CTAs within the same cluster. So with thread block clusters we can store half of our B column across 2 CTAs while B is still accessible, which prevents the need for any data overlap.

We can enable thread block clusters by adding the __cluster_dims__(x, y, z) compile-time kernel attribute to our kernel. Cluster dimensions are not modifiable after launching the kernel. In our case we only need 2 CTAs per cluster so we will launch with the cluster dimensions of (2, 1, 1).

Then we find the rank of the CTA through the %cluster_ctaid.x special register which contains the CTA identifier or rank within a cluster, with only 2 CTAs in our cluster it will return either 0 or 1, and we will choose CTA0 as the leader. In terms of changes to our kernel, the core logic remains the same but we need to adjust synchronization and which mbarriers to use.

Now that we are operating at the cluster level, __syncthreads() is no longer sufficient as we need to synchronize prologue and tmem dealloc across CTAs, so we will use the cluster equivalent which is: barrier.cluster.arrive.release.aligned and barrier.cluster.wait.acquire.aligned. The only TMA change is that both CTAs will now arrive on CTA0's tma_mbar and only CTA0 will have to wait for smem to finish loading before proceeding. We accomplish this by using the mapa instruction to map the local shared memory address of tma_mbar_addr on CTA1 to the tma_mbar_addr on CTA0. On the MMA side we now issue our MMA and tcgen05 instructions with the argument cta_group = 2 which encodes that we are using 2-CTA MMA and we will have to tcgen05_commit with multicast in order both mma_mbar in CTA0 and CTA1 to track its completion. The multicast instruction takes a cta_mask operand which specifies which CTAs participate, 0b11 selects both so we will pass that as our cta_mask.

This is a visual representation of how our new kernel executes.

We get a slight performance boost even with the arithmetic intensity increase as we have additional cluster-level synchronization. However the main benefit is that by decreasing the size of A and B tiles in smem, we can hold more A and B tiles in smem. Which leads us into our next optimization.

Pipelining

In our previous approach execution was serialized, the kernel would either be waiting for a TMA load or waiting on an MMA, yet the TMA and tensor cores have the ability to run simultaneously. Both units are stalling waiting on each other because they are working on the same tile, but if they could work on different tiles, then we could eliminate these stalls almost entirely. We can first leverage this by pipelining or pre-loading tiles of A and B into smem with the TMA then run our normal loop of TMA load then MMA. The difference being that we would be loading tile i + 3 while doing MMA on tile i that was already preloaded into smem.

In terms of implementation we change our monolothic K / BK loop to first load our pipeline_depth tiles of A and B into smem. We then have our normal loop where we iterate over K, MMA remains the same but our TMA loads are now offset. As we preloaded tiles with TMA loads, we are left with tiles of A and B in smem that have yet to be consumed so we have our final loop to consume all tiles left in smem.

Now we are reaching about 70% of cuBLAS's performance.

Warp Specialization

As only 1 thread of 1 warp is needed to issue a TMA load or MMA instruction we can create dedicated warps that either only issue TMA loads (producers) or only issue MMAs (consumers). Since TMA loads and tensor core MMAs execute asynchronously, they will only be constrained by the state of the current tile.

In order to maximally minimize stall times, we need to maximize the amount of tiles that we can store in smem. Each SM has 227 KB of smem, so total tiles stored in smem is simply 227 KB / (A_tile_size + B_tile_size). However, the default maximum smem size for a CTA is 48 KB, so we will use cudaFuncSetAttribute on the host side to set the MaxDynamicSharedMemorySize to our required smem.

On the device side, in order to manage the state of the smem buffers we will have a queue_size number of full_mbar and empty_mbar mbarriers. Each mbar will track the state of a tile (empty or full). We also need one mbarrier to denote when the last consumer has finished so that the epilogue can begin.

We can assign any arbitrary warp to be a producer/consumer and we can reuse these warps for our epilogue warpgroup, so we only need 4 warps. On the kernel side we will change which mbar our producer and consumer will wait on and arrive/commit on.

1. Producer: Waits on empty_mbar of the current tile, arrives on full_mbar of the current tile.
2. Consumer: Waits on full_mbar of the current tile, arrives on empty_mbar of the current tile.

On the producer side we need to adjust our TMA loads for the current tile that we are on which is based on the current buffer or stage. We will also have to wait for our current tile to be empty, which is signaled by the empty_mbar_addr mbarrier. For the first queue_size iterations we don't have to wait as the tiles are default empty when we launch the kernel. On completion of the TMA load the producer will arrive on the full_mbar which signals that the current tile is full and ready for any consumer.

The consumer has the same adjustment for the current tile/stage but will wait on mbar_full and commit to mbar_empty once it has consumed the tile.

Now we are reaching about 82.8% of cuBLAS's performance on average, a 14.3-point gain over v3.

Persistent Kernel

The execution of our warp-specialized kernel looks like this. (note that each block represents time not number of TMA loads/MMA instructions issued and that the actual issue/wait is not synced this is just shown for brevity).

Throughout the execution of the kernel we can see that the producer and consumer are issueing TMA loads and MMA instructions, except for during the epilogue phase. During the epilogue, both the producer and consumer can only wait as there is no data to load or compute on, which means that our TMA and tensor cores are completely idle.

The solution to this problem is a persistent kernel. Instead of launching 1 kernel for each BMxBN output tile in C, we can have each kernel produce total_output_tiles / total_kernels_launched tiles. Each kernel will have a new outer-loop that assigns the next output tile to work on. So when the kernel reaches the epilogue stage, the epilogue warpgroup can begin the epilogue for output tile i, while the producer can begin loading tiles for output tile i + 1. With this approach our new execution will look like this:

The 2 main benefits are that:

1. The prologue costs (smem/tmem allocation, mbarrier initialization) and kernel launch overhead are paid once across the 148 kernels launched, rather than once per output tile.
2. The TMA can begin loading A and B tiles into smem for the next output tile during the epilogue of the current output tile keeping the TMA busy.

However, we currently cannot issue any MMA instructions during the epilogue as the tmem buffer is still being used for the epilogue. Which brings us to our next optimization that is needed to fully overlap our main loop with the epilogue.

Double-Buffering Tensor Memory

As we have previously noted, tmem consists of 128 lanes and 512 columns. Given that our BM value is 128 and BN is 256, we can store 2 output tiles/buffers in tmem. With this approach we can have 2 128x256 tiles in flight, with each tile either in the epilogue state or in the mainloop/MMA accumulation state.

Now when the kernel reaches the epilogue stage, tensor cores can work on the next tile by accumulating in the other empty tmem buffer. Our new execution looks like this:

In terms of implementing this into our kernel, we first have to change the number of kernels that we launch. Instead of launching (M / BM) * (N / BN) kernels we will instead launch as many kernels as SMs, which on the B200 is 148.

Inside the kernel we have to add an additional loop, with the producer and consumer operating as the inner loop over K and the new outer loop iterating over output tiles. Each CTA will grab a tile, then run the full main loop with the output tile accumulated into tmem, signals the epilogue, then immediately starts a new main loop with the next tile.

Next in order to coordinate the main loop and outer loop handoff, we need a second set of mbarriers which are tmem_full and tmem_empty. We will have one of each per elilogue stage, which in our case is 2. The consumer signals tmem_full after committing the last MMA of a tile, the epilogue waits on it before reading, and then signals tmem_empty once tmem reads are complete so that the consumer can reuse that buffer for accumulation of the next tile.

The last piece is the scheduler itself, which is just a counter. Each cluster advances a local wave iterator, multiplies by the number of clusters in flight, and adds its cluster ID to get a globally unique tile index. Converting that flat index back into (block_row, block_col) is a single div-mod, and we terminate once the index runs past the total tile count. Every cluster independently computes which tiles belong to it, so there are no atomics or cross-cluster coordination needed.

Now we are reaching about 88.9% of cuBLAS's performance on average, a 0.7-point gain over v5.

Epilogue Optimizations

Currently our epilogue writes from registers directly to HBM. The issue, along with these writes being uncoalesced, is that these loads utilize the Load/Store Unit (LSUs) which has finite throughput. Since throughput is limited, the LSU ends up stalling waiting to issue the next store. We can bypass this issue by using the TMA, which is a seperate hardware unit that does not utilize the LSU. By routing our stores through TMA loads, we can sidestep this LSU bottleneck. As the TMA cannot read from registers, we must change our epilogue data path from tmem -> registers -> HBM to tmem -> registers -> smem -> HBM. However, naively staging the entire BMxBN tile in smem would cost multiple pipeline stages given the smem footprint. The solution is to pipeline the approach by breaking the output tile into chunks and executing the 3 stage process iteratively.

To avoid stalling between chunks, we can double buffer the smem staging area. While the TMA drains buffer n to HBM, the epilogue warps can then write the n + 1 chunk into buffer n + 1, which will keep the TMA hot.

1. Tensory Memory to Registers

   Instead of issueing and waiting on each load individually, we cna batch all inital tmem loads then wait with a single tcgen05.wait. Each load is independent so this batching avoids unnecessary serialization. Also once all of the values are loaded from tmem to registers, this tmem buffer is now empty so we can signal to any waiting consumer warp that it can freely write to this buffer.

   if (warp_id == 0) tma_store_wait<NUM_TMA_STORE_STAGES - 1>();
   
   float tmp[LOADS_PER_CHUNK][8];
   #pragma unroll
   for (int n = 0; n < LOADS_PER_CHUNK; ++n)
   {
       tcgen05_ld(tmp[n], tmem_row + chunk * STORE_N + n * 8);
   }
   tcgen05_wait_ld();
   
   if (chunk == num_chunks - 1)
   {
       tcgen05_before_thread_sync();
       const int tmem_empty_cta0 = (tmem_empty_addr + wave_stage * 8) & 0xFEFFFFFF;
       if (elect_sync())
       {
           mbarrier_arrive_cluster(tmem_empty_cta0);
       }
   }
   
   barrier_sync(EPILOGUE_BAR, EPILOGUE_THREADS);

2. Registers to Shared Memory

   The packed array is scoped inside the loop in order to reuse those registers across each iteration rather than keeping all iterations packing registers live at once.

   #pragma unroll
   for (int n = 0; n < LOADS_PER_CHUNK; ++n)
   {
       nv_bfloat162 packed_bf[4];
       #pragma unroll
       for (int i = 0; i < 4; ++i)
       {
           packed_bf[i] = __float22bfloat162_rn({tmp[n][i * 2], tmp[n][i * 2 + 1]});
       }
       const int swizzled_n = n ^ (row & 7);
       nv_bfloat16 *write_ptr = store_base + store_stage * BM * STORE_N + row * STORE_N + swizzled_n * 8;
       *reinterpret_cast<int4*>(write_ptr) = *reinterpret_cast<int4*>(packed_bf);
   }
   
   __syncwarp();
   tma_store_fence();
   barrier_sync(EPILOGUE_BAR, EPILOGUE_THREADS);

3. Shared Memory to HBM

   Finally, we can have 1 thread per warpgroup issue the TMA load of the output to HBM.

   if (warp_id == 0 && elect_sync())
   {
       const int src = store_smem + store_stage * store_buf_size;
       tma_2d_smem2gmem(src, &C_tmap, block_col * BN + chunk * STORE_N, block_row * BM);
       tma_store_commit();
   }
   
   store_stage ^= 1;

Hilbert Curves

Our persistent scheduler assigns tiles by striding each cluster's index by 74 (the total number of clusters) . This means each cluster jumps across the output grid with a stride that's large enough to evict the A and B tiles it just loaded before it can reuse them. As A and B tiles used by output tile i are likely not being used for output tile i + 74. We can see the impact by profiling our kernel with M=N=K=8192 using Nsight Compute:

With a 49.23% L2 hit rate about half of our requests are going straight to HBM instead of from cache. We can solve this by adjusting our tile traversal to make consecutive tiles in each cluster's schedule, also be spatially close in the output grid. This way A and B blocks can stay resident in L2 cache. Hilbert curves solve exactly this problem. A Hilbert curve maps 1D indexs to 2D coordinates while maintaining spacial locality. Specifically it gurantees that consecutive blocks will always have a Manhattan distance of exactly 1. Meaning that the next tile will always be in the same row, same column, or diagonally adjacent to the current tile. Here is a Hilbert curve traversing a grid:

The Hilbert curve algorithm is standard so we won't go into it, but we will take a look at the scheduler. To integrate this into our kernel we first have to precompute the full tile schedule on host and then upload it to a buffer in HBM so that each cluster can retrieve its indexes without any computation.

Since Hilbert curves require a power of 2 grid, we have to round up to the smallest power of 2 that covers both dimensions. We then walk every index along the curve, convert it into 2D coordinates, and skip any that are outside of the grid. As we walk we will pack valid tiles into a single int value with the row value being in the upper 16 bits and the column value being in the lower 16 bits.

On the device side, we update the scheduler to become a single gmem read that only reads packed coordinates instead of performing any computation. Each cluster indexes into its contigous slice of the precomputed buffer and unpacks the coordinates of the tiles. Once there are no tiles left, the scheduler will return false and the kernel will stop executing.

With this change on host and to our scheduler, we are able to get a 15% increase in our L2 hit rate as illustrated by NCU:

Now we are reaching about 91.9% of cuBLAS's performance on average, a 3.0-point gain over v6.

Cluster Launch Control

At M=N=K=16384 our current kernel's gap vs cuBLAS widens compared to smaller matrix dims. 2 issues are causing this. The first is the tail effect. As our output tiles are dispatched to SMs in waves, when our output tile count is not an exact multiple of SM count, our final wave will only partially saturate our SMs. This causes each SM, along with its TMA and tensor cores, to sit idle resulting in wasted computing capability. The second is runtime variance. Even SMs within the same full wave will finish executing at different times due to variables like L2 cache state and inter-die latency. With a static scheduler, the tile to SM assignment is fixed at launch so fast SMs have no way to process more tiles, which would effectively hide the latency caused by slower SMs. Cluster launch control fixes both by naturally load balancing across SMs. When an SM finishes a tile, it can claim the next output tile via CLC by getting its CTA id and mapping it to block_row and block_col. Tiles will flow to whichever SM is free first, so the partial tail wave can be grabbed by idle SMs rather than blocking to wait for slower SMs.