Writing high-performance matrix multiplication kernels for Blackwell

In this guide, we’ll progressively iterate on a matrix multiplication kernel. The first implementation will be very simple, but also quite slow. However, in just a few simple steps it can be modified into a state-of-the-art kernel, matching or exceeding highly optimized implementations such as cuBLAS and CUTLASS.

Implementation

TensorCore utilization

% of cuBLAS utilization

The cuBLAS baseline is obtained by measuring the performace of jax.dot. The CUTLASS performance is measured by taking the best result from the following cutlass_profiler invocation (excluding sparse matmuls):

At each step, we will showcase either the full implementation of the kernel, or the difference between the code listings shown in the previous and current steps. Full implementations can be found in our test file. You can also find the a full standalone optimized kernel implementation in the Pallas ops package.

0. Basic kernel#

We begin with a simple single-CTA (block) single-warpgroup example. For convenience, we split the tuning parameters of the kernel into a separate class:

tile_m, tile_n and tile_k specify the size of the matmul performed at every step of the pipeline. In general, tile_k should ideally be equal to 128 divided by the bytewidth of the input element type. max_concurrent_steps specifies the depth of memory prefetch in the compute/memory pipeline, which is frequently called the number of stages in other implementations.

The kernel implementation begins with a bit of setup:

We unpack the config variables for easier access, set the tiling and swizzling transforms to get the SMEM data format to match what’s expected by MMA instructions.

The kernel implementation itself is relatively short. The first part sets up a compute/memory pipeline using plgpu.emit_pipeline. At each step, the compute function (do_mma) consumes a (tile_m, tile_k) slice of LHS and (tile_k, tile_n) slice of RHS. As mentioned before, we specify transforms, as well delay_release=1. This last parameter ensures that the input windows (a_smem, b_smem) passed into do_mma will not be overwritten at least until the next invocation of do_mma completes. This is necessary because we only await the completion of the MMA from the one step in the following step, which is why arrive_barrier_slot and wait_barrier_slot flip between 0 and 1 at each invocation.

The kernel itself ends with an epilogue. We await the completion of the last MMA issued by the pipeline before doing anything. Then, we load the final accumulator from TMEM, write it to SMEM (remembering to call plgpu.commit_smem), and copy it back to GMEM using TMA.

What remains is to actually turn the kernel into a function that can be called with JAX arrays. We use {py:func}`plgpu.kernel <jax.experimental.pallas.mosaic_gpu.kernel> for that. The grid is for now simply 2D and iterates over the output tiles. We allocate intermediates used by the kernel:

1. The TMEM buffer used as an accumulator

2. The SMEM buffer used to stage the accumulator before its copy to GMEM

3. The barrier used to await the completion of MMA operations.

Omitting the setup code, that’s just 50 lines! Unfortunately, it’s not very fast just yet, but it does achieve half the utilization of cuBLAS already!

1. Warp specialization#

The kernel above uses a single warpgroup to do everything: from fetching the data, through issuing MMA operations, to storing the results into GMEM. While one would think that the asynchronicity in TensorCore execution should allow us to overlap the overheads of async copies (TMA) and control-flow, it does not seem to be the case.

A common solution to this problem in the Hopper generation of GPUs was to utilize warpgroup specialization. In Pallas terms, plgpu.kernel can be called with num_threads=2, meaning that each program in the grid would result in two calls to the body. The thread index is then often queried using lax.axis_index and used to select one of multiple different roles, such as only issuing async copies or only running the MMA operations.

This solution also works in the Blackwell generation, but it is in fact even simpler. Since both the async copy (TMA) as well as the tcgen05 MMA instruction only require a single CUDA lane to issue them, we don’t even need to use multiple warpgroups. We can simply break up a single warpgroup into four warps and specialize those!

In Pallas, this can be achieved using pl.core_map with a plgpu.WarpMesh. For each Pallas thread that calls such a core_map, the body will be invoked exactly four times. The core_map synchronizes all warps both at entry at exit. Note that only scalar operations are allowed in the body.

This will be the biggest rewrite to this kernel we’ll perform in this whole sequence, which is why we’ll list the entire kernel source once again.

The kernel has exactly the same structure as before: we first perform the compute, which is followed by the epilogue. The epilogue remains the same (we only use a different barrier to await the completion), so we will not discuss it further.

The plgpu.emit_pipeline call and the do_mma function has been replaced by a single pl.core_map invocation. You can see that immediately after entering its body, each Pallas thread (now representing a warp!) finds out which of the four threads it is. We then use thread with index 0 to only issue async copies that fetch the MMA operands in a loop, while thread with index 1 enters another loop in which it repeatedly calls plgpu.tcgen05_mma.

One interesting aspect here is the synchronization. We keep an array of load_barriers, each tracking progress of an outstanding GMEM->SMEM copy. The compute thread must await their completion before feeding the respective operands to the MMA operation. Going in the other direction, the thread responsible for async copies must await the completion of MMAs that consume operands before it can overwrite the memory by issuing another async copy. This is tracked through consumed_barriers. Finally, when the compute thread is done issuing all MMA operations, it calls plgpu.tcgen05_commit_arrive(mma_done_barrier), requesting the TensorCore to complete the mma_done_barrier once all the MMA operations complete.

We can now turn our attention to the plgpu.kernel definition. The only difference to the prior version is that we explicitly allocate two additional SMEM buffers that hold the MMA operands (previously they were implicitly allocated by plgpu.emit_pipeline), as well as the additional barriers. Note that the load_barriers have num_arrivals=2, because we issue two async copies on the same barrier. orders_tensor_core is necessary to specify on barriers that are meant to indicate the completion of TensorCore operations.

This relatively simple modification already gives us a meaningful bump in performance, getting us up to almost 70% of cuBLAS performance.

2. Tiled epilogue#

This time, we turn our attention away from the compute portion of the kernel and instead focus on its epilogue. We can improve its efficiency by pipelining the copy from TMEM to SMEM together with a transfer from SMEM to GMEM. To do this, we change our scratch_shapes to allocate two smaller buffers instead of an SMEM window that can hold the entire output (which also decreases our SMEM usage):

Then, in the kernel, we loop over the output columns in chunks of epilogue_tile_n, and progressively send out the output to GMEM:

3. Collective (2CTA) MMA#

If you benchmark our latest kernel, you’ll quickly find out that it can’t utilize its compute units well, because they are constantly waiting on the memory to deliver the MMA operands. This means that our kernel is memory bound, because it has too low arithmetic intensity: the number of flops we perform for each byte we load is too small.

One very effective trick of the Blackwell architecture that allows us to double our arithmetic intensity are collective MMAs. The core idea is quite simple: we use a cluster of two blocks (on two SMs) to compute a single matmul. Each block only loads half of each operand, but the MMA operation exchanges the data from SMEM of each block as its running.

We’ll start with the kernel configuration changes again:

We add the cluster parameter to plgpu.kernel to indicate that we intend to have pairs of programs collaborate (as CUDA block clusters). We also append collective=True to our TMEM allocation, to ensure that it will be allowed to be used by collective MMAs and double its number of columns (to cluster_tile_n).

Another notable change is that our pair of blocks will ultimately compute a 4x larger output tile, which is why we shrink the grid correspondingly.

We first update the entry of the kernel:

The only changes here are that we use cluster_tile_m and cluster_tile_n to compute the slice of the output the two blocks will collectively compute, and we also check if the current invocation corresponds to the first (leader) block in the cluster. This is important, because only the leader block is supposed to issue MMA instructions:

You can see a few modifications here. First of all, both blocks must issue the async copies. In both blocks we request a copy of the full window for the whole cluster, but the addition of collective_axes="cluster" indicates that the load is performed jointly by both blocks. partitioned_axis= specifies which axis of the operand should be split across the cluster. We split the LHS rows and RHS columns.

Secondly, as mentioned before, we additionally predicate the _compute body so that only the leader block runs MMA instructions. All tcgen05 calls additionally get a collective_axis= argument, to indicate that the completion of MMAs should complete the barriers in both blocks in the cluster.

Finally, we apply a small modification to the epilogue. Even though the two blocks in the cluster collectively compute a result of shape (cluster_tile_m, cluster_tile_n), each individual block only holds a result of shape (tile_m, cluster_tile_n). We change the output slicing code to need to slice out the right out_gmem_window:

4. Persistent kernel#

Our next step is to make the kernel persistent. This means that we’ll only launch however many clusters we can actually run concurrently on the GPU (SM count divided by 2), and we’ll have each cluster loop over a fixed number of output tiles. This technique allows us to better amortize block (de)initialization costs (since they are only performed once on each SM) and achieve a small degree of overlap between the SMEM to GMEM copy in the epilogue with the compute on the next output tile.

The change is relatively simple. We utilize the plgpu.nd_loop helper to specify that our iteration space is (m_iters, n_iters), but we also request that it should be split accross the cluster grid using the collective_axes= argument.

The only meaningful modification in the compute portion of the kernel body is to ensure that the first few waits on consumed_barriers in the memory warp are only skipped when processing the first output tile (as indicated by loop_info.local_index == 0). When processing the second (or later) tile, the SMEM buffers were used to compute the previous output tile, so we need to ensure that those computations have completed before we overwrite them:

Finally, we modify the kernel epilogue by appending a single line:

As the comment indicates, since TMEM loads are asynchronous, we must await their completion before we move on to the next output tile and overwrite our TMEM allocation by issuing another MMA.

5. Dedicated epilogue warpgroup#

While persistence was useful by itself, it also unlocks another optimization. When the single Pallas thread in our kernel finishes the compute portion of the kernel, it performs the entire epilogue. However, this means that it can’t issue any more work for the TensorCore until it’s done!

This leads us to a simple solution: just use 2 Pallas threads (warpgroups)! The first one will only focus on fetching the MMA operands and issuing the MMA operations, while the second one will only perform the epilogue! Of course, to enable them to run concurrently, we need to double-buffer the TMEM used for the accumulator, and use additional barriers to synchronize:

The kernel begins similarly to what we had before. We renamed acc_tmem to acc_tmem_slots and switch between its halves as we step through the loop over the output tiles:

The compute portion is additionally predicated on wg_idx == 0. There are also two important changes to how we use the barriers. First of all, if we want to reuse our TMEM allocation for MMA (which happens only for loop_info.local_index >= 2), we need to wait on the store_done_barrier for the TMEM half we want to reuse (as indicated by acc_slot). Secondly, once we want to request the TensorCore to arrive on the mma_done_barrier upon completion, we again need to select one of the two barriers that corresponds to the currently used half of TMEM.

Finally, we modify the epilogue, by only having the second warpgroup execute it, and by making the warpgroup signal the completion of the store by arriving on the store_done_barrier associated with the half of TMEM it used.

6. Grid tiling#

Our final change to this kernel is to change the order in which we produce the output blocks to better utilize L2. As mentioned before, the compute units are extremely fast compared to the memory system and so we could use all the help we can get to try to keep them busy.

Our strategy for this is inspired by CUTLASS and works as follows. First, you select which of the two dimensions in your iteration space is the faster changing (we call it grid_minor_dim). Then, you select the tile size along that dimension (grid_tile_width). Instead of traversing the whole minor dimension of the grid before incrementing the more major index, we do it every time we traverse grid_tile_width elements. Once we run out of elements, we move on to the next tile. But there’s a twist! Instead of jumping to the beginning of the second tile, we start from the end and work our way back. This ensures that as we switch the tiles, we can reuse some of the recent blocks of one of the operands.

Since this strategy is so common, we provide a helper for it: plgpu.planar_snake. When using the helper, the changes to the kernel are quite trivial:

This simple trick is incredibly effectful and is crucial in achieving state of the art performance.

Final kernel#

You’ve reached the end of this tutorial, congratulations! In the previous sections, we focused only on the differences between the different kernels and rarely listed the complete source. This is useful to hide the irrelevant details when extending the implementation, but it can also be helpful to see the full source. So here it is! The whole implementation is less than 150 lines and reaches SOTA performance (at least on the shape used in our benchmarks).