2805 Bowers Ave, Santa Clara, CA 95051 | 408-730-2275research@colfax-intl.com

Tutorial: Matrix Transpose in CUTLASS

The goal of this tutorial is to elicit the concepts and techniques involving memory copy when programming on NVIDIA® GPUs using CUTLASS and its core backend library CuTe. Specifically, we will study the task of matrix transpose as an illustrative example for these concepts. We choose this task because it involves no operation other than copying data from one set of addresses to another, allowing us to study in isolation those aspects of memory copy optimization, such as coalesced accesses, that can be separated from workloads also involving computation.

Our treatment takes inspiration from Mark Harris’s Efficient Matrix Transpose tutorial, which we recommend for an in-depth discussion of the matrix transpose problem that does not directly involve the abstractions from CuTe that we use here. Conversely, our tutorial can also serve as an introduction to these abstractions for the reader already familiar with Harris’ tutorial. In any case, we will review the key ideas from that tutorial here before explaining how to implement the corresponding optimized solution using CuTe.

Review of Coalesced accesses:

In many compute workloads, particularly those found in ML/AI applications, we work with multidimensional arrays known as tensors. Since computer memory is inherently one-dimensional, these tensors must be linearized, or organized into this one-dimensional space. As a consequence, adjacent elements in some dimensions of a tensor might not be adjacent in memory. We say that a dimension is contiguous when elements adjacent in that dimension are also adjacent in memory. A block of consecutive elements in a contiguous dimension is also called contiguous.

Accesses — i.e., reads or writes — to a contiguous block of memory are called  coalesced, whereas accesses to non-contiguous blocks of memory are called  strided. Coalesced accesses typically offer faster performance than strided ones because they align more effectively with the GPU’s memory architecture, allowing for more efficient data caching and retrieval. For this reason, optimizing for coalesced memory access is highly desirable when programming for GPUs.

Certain workloads, however, necessitate strided accesses and cannot be implemented otherwise. The matrix transpose — or more generally, tensor permute operations — are prime examples where strided accesses are unavoidable. In such cases, it is crucial to minimize the performance impact of these less efficient access patterns. One standard technique is to only perform the strided accesses at the lower and faster levels of the GPU memory hierarchy, which we now recall.

For the purposes of our discussion, there are three programmable levels to the GPU memory hierarchy. From the highest to lowest level, we have global, shared, and register memory.

The global memory (GMEM), i.e. High-Bandwidth Memory (HBM), is the largest of the three, and also the slowest one to read from or write into. For instance, an NVIDIA H100 Tensor Core GPU has 80 GB of GMEM. Strided access here will have the worst effect on performance.

Next is the shared memory (SMEM), which is much smaller but considerably faster than GMEM. An NVIDIA H100 Tensor Core GPU, for example, has up to 228KB of SMEM per streaming multiprocessor (SM). For readers more familiar with memory architectures, we note that SMEM is physically carved out of the L1 cache. SMEM is shared amongst all threads within the same cooperative thread array (CTA), and each CTA operates within its own segment of SMEM. Strided access here is still suboptimal, but is greatly preferred over strided access in GMEM.

Finally, we have the register memory (RMEM), which is dedicated to individual threads.

In this tutorial, memory accesses only comprise copying numbers (e.g., 32-bit floats) around, either from one level to another, or between different locations in the same level.

The naive transpose method discussed in Harris’ tutorial starts out with strided accesses in  GMEM -transpose-> GMEM. He then improves upon this by first copying data GMEM to SMEM so that we have  GMEM -> SMEM -transpose-> SMEM -> GMEM. This way, the strided load happens in SMEM whereas both GMEM accesses are coalesced.

CuTe way:

We now discuss how to implement these two approaches using the CuTe library. We start with the naive method, mainly to demonstrate what not to do.

Data in the CuTe framework is abstracted as cute::Tensor objects. A CuTe tensor consists of a pointer (in the C sense) to the first element of the tensor, as well as a cute::Layout object, which describes the offset of each element in the tensor with respect to the first element through defining shape and stride integer tuples. For example, for a row-major matrix of dimensions M by N, we would define the layout to have shape (M, N) and stride (N, 1).

For a cute::Layout, we note that one of the available options when defining the layout of a new tensor is to specify whether it is row-major or column-major in terms of the strides (GenRowMajor or GenColMajor). In a column-major matrix, adjacent elements within a column are contiguous while adjacent elements across columns are strided in memory. By default, CuTe uses column-major layouts. More generally, we could specify strides for each dimension of the shape of the layout.

One easy way to implement transpose is to simply define the input as column major, output as row-major and let CuTe figure out the copy.

using namespace cute;
int M = 2048, N = 2048;
float *d_S, *d_D;
// Allocate and initialize d_S and d_D on device (omitted).

// Create the row major layouts.
auto tensor_shape = make_shape(M, N);
auto tensor_shape_trans = make_shape(N, M);
auto gmemLayoutS = make_layout(tensor_shape, GenRowMajor{});
auto gmemLayoutD = make_layout(tensor_shape_trans, GenRowMajor{});

// Create the row major tensors.
Tensor tensor_S = make_tensor(make_gmem_ptr(d_S), gmemLayoutS);
Tensor tensor_D = make_tensor(make_gmem_ptr(d_D), gmemLayoutD);

// Create a column major layout. Note that we use (M,N) for shape.
auto gmemLayoutDT = make_layout(tensor_shape, GenColMajor{});

// Create a column major view of the dst tensor.
Tensor tensor_DT = make_tensor(make_gmem_ptr(d_D), gmemLayoutDT);

An important note here that although we have three tensors, we only have two actual copies of the data. This is because tensor_D and tensor_DT both use the data in d_D — they are two different views on the same data. We will be using the column-major view for our transpose kernel, but using the row-major view when we are verifying the transpose result.

Next, we need to determine how to divide the input Tensor into smaller chunks that we can distribute over the CTAs. We can do this using the cute::tiled_divide method.

using namespace cute;
using b = Int<32>;
auto block_shape = make_shape(b{}, b{});       // (b, b)
Tensor tiled_tensor_S  = tiled_divide(tensor_S, block_shape); // ([b,b], m/b, n/b)
Tensor tiled_tensor_DT = tiled_divide(tensor_DT, block_shape); // ([b,b], m/b, n/b)

Here, we specify the tile size to be 32 by 32. The values for tile size is an important tuning parameter, and should be tuned for each specific workload. In fact 32 by 32 is not the optimal value for the transpose kernel, and we will tune it before benchmarking.

tiled_divide creates a tensor with the same data but a different layout, i.e., a different view of the data. In our case, for tensor_S we start out with a 2D matrix of size (M, N). cute::tiled_divide with tile size of b generates a view of a 3D matrix of size ([b,b], M/b, N/b); b by b matrices in a M/b by N/b grid.

This view makes it much easier to access the correct tile once inside the kernel.

Tensor tile_S = tiled_tensor_S(make_coord(_, _), blockIdx.x, blockIdx.y);
Tensor tile_DT = tiled_tensor_DT(make_coord(_, _), blockIdx.x, blockIdx.y);

Here, placing make_coord(_, _) as the first argument takes the entire first dimension, while specifying the integer values of the second and third dimensions as the block indices takes the corresponding slice of the tensor. (For those familiar with numpy: the underscore (_) in CuTe is equivalent to the colon (:) notation there.) In other words, tile_S represents the entire b by b matrix located at grid point (blockIdx.x, blockIdx.y). Note that we don’t swap blockIdx.x and blockIdx.y when slicing into tiled_tensor_DT since we already took the column-major view with shape (M, N) (by contrast, if we instead took a tiled divide of tensor_D, we would need to swap the block indices and then use different thread layouts for source and destination in local_partition below). We can then get the part assigned to a specific thread with:

auto thr_layout =
      make_layout(make_shape(Int<8>{}, Int<32>{}), GenRowMajor{});
Tensor thr_tile_S = local_partition(tile_S, thr_layout, threadIdx.x);
Tensor thr_tile_DT = local_partition(tile_DT, thr_layout, threadIdx.x);

Here, we launched the kernel with 256 threads per CTA and chose a thread layout so that loads from gmem are coalesced whereas stores to gmem are uncoalesced (as we emphasized above, no matter which thread layout is chosen there will be uncoalesced accesses). Finally we can use cute::copy to copy the data from thr_tile_S to thr_tile_DT.

Tensor rmem = make_tensor_like(thr_tile_S);
copy(thr_tile_S, rmem);
copy(rmem, thr_tile_DT);

Now we can benchmark this against a pure copy kernel. The code for the copy kernel is based on the CUTLASS’s tiled_copy example, so we will leave unpacking it as an exercise for the reader. Additionally, we have found empirically that the tile size 32 by 1024 gave the best performance for our workload.

Just as we saw in Harris’ post, the speed of this naive method is not very good. This is because this copy is a strided copy from GMEM -> GMEM. In order to confirm this, let’s profile this transpose using NVIDIA Nsight™ Compute. This profiling tool can detect problems in code that lead to lowered performance. Profiling the naive transpose, the summary page of the GUI shows us:

Nsight Compute has a wide range of tools to help with optimization, but full exploration of Nsight is beyond the scope of this article. For this article we will simply be looking at the summary page. In the above summary page, we see that indeed the problem of uncoalesced accesses constitutes the primary reported issue.

Next, we study the improved algorithm: copy the data from GMEM into SMEM first, then do the transpose, then copy back from SMEM to GMEM.

To move the strided access to SMEM, we need a Tensor that uses SMEM. We will use CuTe to allocate an array_aligned object in a CTA’s SMEM.

using namespace cute;
using CuteArray = array_aligned<Element, cosize_v<SmemLayout>>;

extern __shared__ char shared_memory[];
CuteArray &smem = *reinterpret_cast<CuteArray*>(shared_memory);

Here, smemLayout is the Layout for the SMEM used in a single tile. We can now create a tensor whose data pointer is shared_memory:

Tensor sS = make_tensor(make_smem_ptr(smem.data()), smemLayout);

One important note here is that we must ensure the SMEM tensor is small enough to fit on a single SM. In other words, the size of smemLayout multiplied by the number of bytes per Element must be less than the total SMEM capacity on a single SM. Beyond that, we have occupancy considerations to contend with depending on the SMEM used per CTA.

Now we can repeat the column-major view trick we did with the data in GMEM, except that this time we apply it to SMEM. We create two different views of SMEM — one row-major and the other column-major.

using namespace cute;
using b = Int<32>;
auto block_shape = make_shape(b{}, b{});       // (b, b)

// Create two Layouts, one col-major and one row-major
auto smemLayout = make_layout(block_shape, GenRowMajor{});
auto smemLayoutT = make_layout(block_shape, GenColMajor{});

// Create two views of smem
Tensor sS  = make_tensor(make_smem_ptr(smem.data()), smemLayout);
Tensor sD = make_tensor(make_smem_ptr(smem.data()), smemLayoutT);

Finally, we can use the cute::copy to copy from GMEM to SMEM and then back from SMEM to GMEM. Note here that S and D are tiled_divide‘s of tensor_S and tensor_D and tS and tD are thread layouts chosen to ensure coalesced accesses to GMEM (in fact, they are both equal to thr_layout from above!).

// Slice to get the CTA's view of GMEM.
Tensor gS = S(make_coord(_, _), blockIdx.x, blockIdx.y); // (bM, bN)
Tensor gD = D(make_coord(_, _), blockIdx.y, blockIdx.x); // (bN, bM)

// Create the thread partitions for each Tensor.
Tensor tSgS = local_partition(gS, tS, threadIdx.x);
Tensor tSsS = local_partition(sS, tS, threadIdx.x);
Tensor tDgD = local_partition(gD, tD, threadIdx.x);
Tensor tDsD = local_partition(sD, tD, threadIdx.x);

// Copy GMEM to SMEM.
cute::copy(tSgS, tSsS);

// Synchronization step. On SM80 and above, cute::copy
// does LDGSTS which necessitates async fence and wait.
cp_async_fence();
cp_async_wait<0>();
__syncthreads();

// Copy transposed SMEM to GMEM.
cute::copy(tDsD, tDgD);

Now when we benchmark, we get a much better result.

Nonetheless, we are still a way off from the copy result. Profiling the code again, we can spot the next issue at hand — memory bank conflicts.

Memory Bank conflicts:

The strided SMEM version gets much better performance than the naive version, but it still does not match the copy performance. A large part of this discrepancy is due to memory bank conflicts. On most NVIDIA GPUs, the shared memory is organized into 32 memory banks. Only one thread in a warp is able to access a memory bank at a time; this is true for both read and write accesses. Hence, if multiple threads try to access the same memory bank, the accesses are serialized. This is called a bank conflict. For a more in-depth discussion about bank conflicts, we recommend Lei Mao’s excellent blog post.

In more detail, elements are assigned in 32 bits to memory banks in a round-robin format. The first 32 bits are assigned to 0, the next to 1, and so on, until the 33rd set of 32 bits is assigned to bank 0 again. So in a 32 by 32 (row-major) tile of float, each column maps to the same memory bank. This is the worst case scenario; with 32 threads in a warp, this causes a 32-way bank conflict.

Mark Harris’ tutorial solves this issue by padding the rows by 1 number. This offsets the elements, causing every element in a column to fall in a different bank. We could replicate this workaround in CuTe by using non-default strides. The CuTe Layout contains information about stride, which defines the offset between elements in each dimension. We can add padding by setting the stride for the columns to be 33 instead of 32. In code, this can be done simply by:

auto block_shape = make_shape(Int<32>, Int<33>); // (b, b+1)

// Create two Layouts, one col-major and one row-major
auto smemLayout = make_layout(block_shape, GenRowMajor{});
auto smemLayoutT = make_layout(block_shape, GenColMajor{});

However, this wastes memory for the extra 32 numbers in SMEM. In this article we will implement an alternate solution — swizzle.

Swizzle and Layout Composition:

In order to discuss swizzle, we first need to expand on a CuTe Layout. A Layout is not simply a container that stores information about a Tensor’s structure, but is a function that map one coordinate to another. For example, take a column-major tensor A with M rows and N columns. Given coordinates (4,5) — 4th column, 5th row — this Layout for A will map the tuple (4,5) to the integer 5M+4. This is the index for the element at coordinates (4,5) in the 1D pointer to the data. This abstracts away the often confusing coordinate math that comes with working with higher dimensional Tensors.

Normally, the coordinate calculation is done using just the stride of the Tensor, which defines the offset in 1D memory space between adjacent elements in a dimension. For example using the same Tensor A, the stride is (1,M). Elements in a column are adjacent to each other, i.e., with offset of 1, while elements in a row are offset by M.

CuTe provides tools for more complex coordinate mapping functions. One such tool is Swizzle. The details of swizzling are beyond the scope of this tutorial, and we refer curious readers to NVIDIA’s PTX documentation.

By defining an appropriate swizzling function, CuTe programmers can access data the same way they would in the non-swizzling case, without worrying about bank conflicts. CuTe abstracts away the swizzling details by baking in swizzle as a property of a tensor’s layout using the composition operation.

Composition, as its name suggests, creates a functional composition of the layout arguments. Specifically, when a programmer accesses data in a swizzled tensor in SMEM — say by calling tensor(i) in CuTe, where the logical index i is what they think the accessing location is — they actually access the data at swizzle_function(tensor(i)).

Returning to transpose, the swizzle function we need is Swizzle<5,0,5>. The number 5 here refers to the number of bits in the mask. Per the CuTe documentation, this function modifies the lower 5 bits by taking the xor with the upper 5 bits (the mask). Then with this pattern applied to a 32 by 32 set of addresses, no two elements in a column map to the same memory bank, thereby avoiding all bank conflicts. We add this swizzle pattern to our Layouts.

auto tileLayoutS = make_layout(block_shape, GenRowMajor{});
auto smemLayoutS_swizzle = composition(Swizzle<5, 0, 5>{}, tileLayoutS);

We also note that other data storing patterns in SMEM will require different swizzling functions. We encourage the reader to experiment with the generic swizzle functions exposed via CuTe and pick what works best for them.

Transposing via Layout Composition:

Above, we have discussed how to transpose a tile in SMEM by defining a column-major Layout of the tile. Here, we show an alternate method using Layout composition. Specifically, we make a Layout composed out of the swizzled LayoutS and LayoutD.

auto tileLayoutD = make_layout(block_shape_trans, GenRowMajor{});
auto smemLayoutD_swizzle = composition(smemLayoutS_swizzle, tileLayoutD);

The trick here is that both of these layouts are defined to be row-major, but CuTe uses column-major by default including for layout algebra. We now claim that  composition(tileLayoutS,tileLayoutD) equals

auto tileLayoutDT = make_layout(block_shape_trans, GenColMajor{});

To explain, let the block dimensions be bM and bN, so tileLayoutS and tileLayoutD have Shape:Stride given by (bM,bN):(bN,1) and (bN,bM):(bM,1), respectively. Then we have:

tileLayoutS(tileLayoutD(x,y)) = tileLayoutS(bM*x+y).

Now to compute what the integer bM*x+y maps to under tileLayoutS, it is convenient to represent it as a coordinate pair in the domain shape (bM,bN). But since CuTe algebra for mapping 1D indices to coordinates in a shape is done column-major (or left-to-right), we see that bM*x+y corresponds to the coordinate (y,x). Thus, we get:

tileLayoutS(bM*x+y) = tileLayoutS((y,x)) = bN*y+x.

This shows that the composed Layout function equals that for the Layout (bN,bM):(1,bN), which validates the claim. Finally, we note that in the presence of post-composition with a swizzle function, pre-composition leaves the same swizzle in-place, thus avoiding some code duplication.

Our swizzled solution gets us close to the performance of the copy kernel, just as Mark Harris’ article did.

With performance approaching the bandwidth limit, we are nearing the hardware limitation. When profiling the swizzle version, the summary page shows:

We see that we have resolved the memory bank conflict issue. The last reported issue on long scoreboard stalls may be disregarded as we are profiling a completely memory-bound kernel.

TMA:

Note that data transfer between GMEM and SMEM constitutes by far the majority of time spent in our transpose kernel. The Tensor Memory Accelerator (TMA) is a feature introduced in the NVIDIA Hopper™ architecture that can be used in place of regular load and store instructions between GMEM and SMEM, thereby potentially improving the performance of our transpose kernel. We studied the usage of TMA for this tutorial, and found a mixed set of results, which we describe in this section.

To review, TMA is a dedicated asynchronous memory copy unit for copying multi-dimensional data from GMEM to SMEM and vice-versa. In the TMA model for async copy, instead of having the threads/warps in a CTA cooperatively copy a portion of the source tensor to the target tensor, one elects a single thread in the CTA to issue the load or store TMA instruction. While the instruction executes in the async proxy, threads are free to do other independent work. Barrier objects and synchronization primitives (fence, arrive and wait) are used to synchronize the data movement with computation that relies on the data. When used in conjunction with a software pipelining scheme, TMA allows overlap memory copy  instructions to overlap with computations, which helps to hide latency. However, since the transpose kernel only does memory copy, we don’t have an opportunity to demonstrate this benefit of TMA in this tutorial.

To clarify the performance of TMA for memory copy in isolation, we first studied the performance of TMA load and store copy kernels versus other alternatives, such as CuTe’s TiledCopy tutorial, which does 128-bit vectorized loads and stores passing through RMEM only. We found that in this situation, TMA performed on par with this simpler alternative (after tile size tuning for both), with both near the memory bandwidth spec of the device. This outcome was in line with our expectations — indeed, we have no reason to expect TMA to outperform in the situation of pure memory copy according to a regular pattern.

In contrast, a naive attempt at using TMA for both load and store in a transpose kernel, with the same tile sizes chosen as above, performed worse than our best-performing version. This was due to the presence of bank conflicts! The immediate problem was that TMA only supports a restricted set of swizzle functions (intended for use in conjunction with WGMMA); for instance, see this section of the CuTe codebase. In particular, it doesn’t support the Swizzle<5,0,5> function that we used above, which makes it less straightforward to completely eliminate bank conflicts. Note however that we have no reason to believe this is an essential problem, but we chose not to pursue this line of investigation further in light of our benchmarking with the copy kernel. Moreover, when trying out a version with TMA store only and with 128-bit vectorized load into registers and then writing to SMEM, we found that it performed only slightly under par even though the profiler still reported shared store bank conflicts (but avoiding bank conflicts for the TMA store from SMEM to GMEM).

Because of these mixed results, we do not describe in detail the mechanics of how to use TMA, deferring this to a future blog post in which we will aim to study TMA in a context more suited to its strengths.

Conclusion:

In this tutorial, we introduced readers to a number of fundamental GPU memory concepts and how to program for them using the CuTe library, by way of implementing an efficient matrix transpose kernel.

Starting with coalesced reads and writes, we touched upon the concepts of CuTe Layouts and Tensors, bank conflicts, swizzle functions, and TMA. Except for TMA, we have seen how a good understanding of these concepts is necessary to implement an efficient transpose kernel. In a subsequent post, we plan to study TMA in a setting where it is important for optimization.

To conclude this tutorial, we present the runtimes of the various kernels that we discussed. We included the JustCopy kernel as the headroom of what can be achieved, as well as both a naive PyTorch implementation (given by invoking contiguous() on torch.transpose) and one using torch.compile, in order to demonstrate the magnitude of the efficiency gains available through writing these low-level kernels.

The source code for all of these kernels as well as the benchmarking script is available in a Colfax Research GitHub repository.

Edit (05/07/24): Added JIT compiled version of PyTorch transpose for reference (h/t @CHHillee).

Share this:

Discover more from Colfax Research

Subscribe to get the latest posts sent to your email.

Type your email…

Subscribe

Posted

in

Comments

7 responses to “Tutorial: Matrix Transpose in CUTLASS”

The shape comment is incorrect in this line.
“`
Tensor tiled_tensor_DT = tiled_divide(tensor_DT, block_shape); // ([b,b], n/b, m/b)
“`
It should be ([b,b], m/b, n/b).

Fixed, thanks!

Thanks for sharing!

Regarding the code snippet below:

“`
Tensor rmem = make_tensor_like(thr_tile_S);
copy(thr_tile_S, rmem);
copy(rmem, thr_tile_DT);
“`

Why do we need to explicitly define a register fragment instead of simply writing:
“`
copy(thr_tile_S, thr_tile_DT);
“`

Both code snippets will achieve the same copy functionality. However, for the copy through an explicitly defined register fragment the compiler nvcc will generate SASS with all LDG instructions before STG while the second will have LDG and STG interleaved, which can increase occurrence of long scoreboard stalls (from gmem latency) and thereby degrade performance.

I pushed an update to the cfx-article-src repo that includes both copy kernels for comparison.

Thanks for the reply!

If I get it correctly: It’s quite similar to the pipelining (or double buffer) optimization when writing GEMM kernel (but with the help of compiler instead of manually coding), which reduces the instruction latency to achieve maximum memory throughput. Please correct me if I’m wrong.

Yes, that’s a good analogy, though I would say the coding is ‘manual’ since writing it in the two different ways we described generates different SASS. In the analogy, this would be akin to choosing the number of buffers.

Thanks, this really helps me understand it!

Leave a Reply Cancel reply

Your email address will not be published. Required fields are marked *

Comment *

Name

Email

Website

Save my name, email, and website in this browser for the next time I comment.

Δ

Copyright © 2023-2024 Colfax International. All rights reserved.


2805 Bowers Ave, Santa Clara, CA 95051 | 408-730-2275research@colfax-intl.com

Delivering 1 PFLOP/s of Performance with FP8 FlashAttention-2

We recently released an update to our FlashAttention-2 forward pass implementation on NVIDIA Hopper™ architecture that incorporates a number of new optimizations and improvements, including a software pipelining scheme and FP8 support. In this article, we will explain a challenge with achieving layout conformance of register fragments for WGMMA instructions that we encountered in the process of fusing back-to-back mixed-precision GEMMs with FP8 operands and FP32 accumulator. After explaining why this problem surfaces for FP8 but not FP16 precision operands, we’ll describe a performant method to achieve layout conformance through using a concise combination of two NVIDIA® CUDA® intrinsics – namely, the byte permute and shuffle sync instructions. Our solution takes some inspiration from a blog post of Manish Gupta that studied a similar problem in the context of mixed-input GEMMs. Furthermore, we’ll discuss another complication that arises with the majorness of the V matrix for FP8 WGMMA. Finally, we will present some FLOPs/s benchmarks collected from runs with synthetic data on the NVIDIA® H100 Tensor Core GPU with SXM5 board form factor. In particular, for head dimension 256 and large sequence length, we obtain over 1 petaflop/s of performance. This is consistent with other claims on FP8 fused attention performance in the same regime, e.g. by HippoAttention.

Recollections on FlashAttention-2

FlashAttention-2 is a memory-aware algorithm for multi-head attention (MHA) that was introduced by Tri Dao last year, building on earlier work of Dao and his collaborators. It serves as blueprint for implementing MHA as a fused CUDA kernel (FMHA), in which intermediate steps of the attention computation

aren’t read back to HBM (i.e., global memory), but rather kept in shared memory or in registers. The key idea of FlashAttention-2, as well as the original FlashAttention, is to leverage tiling of the input tensors Q, K, V in conjunction with the online-softmax algorithm, which allows one to circumvent keeping the entire S matrix live for the softmax calculation.

Apart from the algorithm itself, there are a number of interesting challenges in terms of optimization to get a performant implementation on modern GPUs, in particular those built on NVIDIA Hopper architecture like the H100 GPU. For example, to make optimal use of the Tensor Cores on an H100 GPU for carrying out the matrix multiplications in FMHA, it’s important to target the new Hopper-specific WGMMA instructions as well as the Tensor Memory Accelerator (TMA) for loads from global to shared memory. We described how to accomplish this using tools from NVIDIA’s open-source CUTLASS library in our paper from last December. There, we chose to work with half-precision FP16 input tensors Q, K, V for the head-to-head performance comparison with Dao’s base implementation written for NVIDIA Ampere architecture. However, in practice we want to transition to even lower precision types, in keeping with the philosophy of model quantization. Therefore, as a first step we need to understand how to accommodate FP8 input tensors in the design of the FMHA kernel.

In what follows, when we discuss invoking WGMMA (that is, wgmma.mma_async in PTX) to multiply matrices such as Q·KT or P·V, we will implicitly be referring to tiles of these matrices. For example, we could choose 64×64 tiles.

The challenge with FP8: layout conformance of WGMMA register fragments

WGMMA is warpgroup-level, so involves 128 threads (4 warps) to carry out the MMA. When invoking a WGMMA instruction, accumulation occurs in those threads’ registers. It is then paramount to understand the thread ownership pattern, i.e. register fragment layout, of the entries of the result matrix, as prescribed by WGMMA. For example, we have the following layout for a slice of the FP32 accumulator tile, which is extracted from Figure 119 in the PTX documentation:

In Figure 1, we took the subset of the 64×N output tile consisting of the two rows 0 and 8 and columns 0-15. Figure 1 then indicates that these entries are evenly divided among threads 0-3. Moreover, the entries are arrayed contiguously in each thread’s registers according to the d index. This is the register fragment layout that would appear as part of the wgmma.mma_async calls for computing either Q·KT or P·V in FMHA, for instance, and regardless of the precision type of the operands given FP32 accumulator.

When invoking WGMMA to do a matrix multiplication A·B, you have the option of arranging for the operand A to either be in shared memory or in registers (by contrast, operand B must always be in shared memory). For the second GEMM P·V in FMHA, we want to use the option of keeping the first operand in registers for performance reasons, avoiding unnecessary writes and reads with shared memory. In order to compute a correct result, we then need to match the FP32 accumulator register fragment layout to either the FP16 or FP8 operand A register fragment layout. To be more precise, we should first downcast the entries held in registers in place to the desired precision format, and then potentially do some register movement, both within and across threads.

This sounds complicated, and it turns out to involve some non-trivial maneuvers in the case of FP8, but let us first point out that one doesn’t need to do any register movement at all in the case of FP16. This is because the layouts in question are then identical: just compare Figure 119 and Figure 118 in the PTX documentation. In fact, this makes it easy to take our FP16 FMHA kernel and allow for the input tensors Q and K to be FP8 while still fixing the input tensor V to be FP16. We call this the FP8 hybrid FMHA kernel, which is akin to the FP8 version of FMHA in the Triton tutorial.

Let’s now consider the case at hand, with all three tensors Q, K, V in FP8 precision. Consider the following extract from Figure 122 in the PTX documentation:

This is the thread ownership pattern, or layout conformance requirement, that must be satisfied prior to invoking an FP8 WGMMA. For example, according to Figure 2, we need to arrange that thread 0 has these 8 entries in its registers in order:

T0d0, T0d1, T1d0, T1d1, T0d2, T0d3, T1d2, T1d3.

In particular, observe that we will need to exchange data among threads to achieve layout conformance.

The Solution: Byte Permute and Shuffle Sync

Our solution in code is exposed on our repo via ReorgCFp8toAFp8, which we invoke in the consumer path of the pipelined implementation immediately before the second gemm call. Invoking that method on the CUTLASS Tensor tSrSPrec (the FP8 downcasted register fragment Tensor) accomplishes the following data movement, going from top to bottom:

In Figure 3, each box represents two entries from Figure 1 of the same thread and color combined. Keeping in mind that the entries have been downcasted to FP8, each box is then 16 bits wide. We then accomplish the register data movement in code by invoking a combination of the following two CUDA intrinsics:

An added wrinkle: transposing the V matrix offline

We’ve been thinking about Q, K, V as matrices for the attention formula, but for the attention layer the Q, K, V are actually 4-dimensional tensors, with dimensions given by the batch size B, sequence length S, numbers of heads H, and head dimension D. For our FP16 FMHA kernel, we conformed to a certain convention regarding how these tensors are arrayed in memory. Namely, we supposed that they are packed in the (B, S, 3, H, D) format. Then when loading tiles from global to shared memory via TMA, this convention forces the 2-dimensional Q, K, V tiles to be contiguous in the head dimension and strided in the sequence length dimension.

Given a gemm call to multiply A·BT for an (M×K)-matrix A and an (N×K)-matrix B, we say that the A resp. B operand is mn-major if it is contiguous in the M resp. N dimension (or outer dimension), and k-major if is instead contiguous in the K-dimension (or inner dimension). Then given the (B, S, 3, H, D) convention, for the first GEMM Q·KT the 2nd operand is k-major, whereas for the second GEMM P·V the 2nd operand is mn-major. Fortuitously, for FP16 precision this distinction turns out to be immaterial since WGMMA accepts both k-major or mn-major for its 2nd operand. However, this is no longer the case for FP8 precision. Therefore, we need to either transpose the V tensor as a preprocessing step before calling our FP8 FMHA kernel (but not the FP8 hybrid version!), or otherwise handle the transpose elsewhere, such as fusing it to the epilogue of the projection that creates the V tensor.

FLOPS benchmarks

Figure 4 shows the FLOPs/s improvements we see from moving to lower precision types. Runs were conducted with synthetic data drawn from a normal distribution with mean 0 and variance 1. To interpret the figure correctly, recall from above that FP8 Hybrid refers to the FMHA kernel in which the Q and K tensors are FP8 and V is FP16, and note that the TFLOPs/s reported for FP8 don’t include the cost of transposing V since that is a pre-processing step that is expected to be accounted for elsewhere.

In our experiments, we have tuned the number of pipeline stages as well as the tile sizes used in WGMMA independently for each kernel: we have configurable parameters QBLKSIZE and KBLKSIZE for tiling along the sequence length of Q and K, V respectively (note that we never divide along the head dimension). We see that the largest improvement going from FP8 Hybrid to FP8 occurs in the case of head dimension 256. This happens because keeping V as FP8 reduces pressure on the shared memory, enabling us to set both QBLKSIZE and KBLKSIZE to 128.

Apart from the set of parameters chosen for Figure 4, we can also find parameters that showcase over 1 petaflop/s of performance through increasing the sequence length. For example, for sequence length 8448 = 64*132 we have:

dgxuser@PM-DGX-H100:~/kernels$ ./fmha_fp8_pipe_128x128xCTA256 \
--batch-size=4 --seq-length=8448 --head-size=256 --iterations=1000
Using device 0: NVIDIA H100 80GB HBM3  (SM90, 132 SMs)
M = 8448
N = 8448
K = 256
QBLK = 128
KBLK = 128
L = 32 : 8 * 4
CUTE_FMHA:     [1028661.0]Gflop/s  (2.2735)ms

Or if we let sequence length be 16896 = 2*8448, then we have:

dgxuser@PM-DGX-H100:~/kernels$ ./fmha_fp8_pipe_128x128xCTA256 \
--batch-size=2 --seq-length=16896 --head-size=256 --iterations=1000
Using device 0: NVIDIA H100 80GB HBM3  (SM90, 132 SMs)
M = 16896
N = 16896
K = 256
QBLK = 128
KBLK = 128
L = 16 : 8 * 2
CUTE_FMHA:     [1057474.8]Gflop/s  (4.4230)ms

For these examples, multiples of 132 were chosen to account for wave quantization effects.

Finally, for ease of reproducibility we give our autotuned parameters for the different kernels:

In Table 1, NOPIPE refers to replacing the software pipelining by our original copy/gemm overlapping scheme described in section §6 of our paper. Also, with NOPIPE we don’t use the QINRMEM option, while with the default pipelined version (STAGES=3) we do use QINRMEM.

Addendum: accuracy loss with FP8

As always with moving to lower precision types, the concomitant effects on accuracy loss have to be understood and managed correctly. We leave this important topic to future work. For now, we can report the root-mean-square error taken between the output matrix and a reference calculation in Table 2 (with the same synthetic data and other parameters fixed as in Figure 4):

In Table 2, the reference calculation is done on the Q, K, V tensors generated in the given precision type and then upcasted to FP32.

Acknowledgments

We would like to thank Tri Dao for helpful comments on an earlier draft and Ying Zhang from the PyTorch team for pointing out the wave quantization optimization to us.

Edit (03/06/2024): Included more thorough autotuning for auxiliary parameters and updated FLOPS numbers based on this.Edit (03/10/2024): Added reference to similar work by HippoAttention on FP8 fused attention.

FP8-FMHA-blog.pdf

Share this:

Discover more from Colfax Research

Subscribe to get the latest posts sent to your email.

Type your email…

Subscribe

Posted

in

Comments

Leave a Reply Cancel reply

Your email address will not be published. Required fields are marked *

Comment *

Name

Email

Website

Save my name, email, and website in this browser for the next time I comment.

Δ

Copyright © 2023-2024 Colfax International. All rights reserved.


Developing CUDA Kernels for Accelerated Matrix Multiplication on
NVIDIA Hopper Architecture using the CUTLASS Library
GANESH BIKSHANDI† and JAY SHAH†
We explain how to develop NVIDIA CUDA kernels for optimized general matrix multiplication (GEMM) on NVIDIA Hopper
architecture using the template collection CUTLASS and its core library CuTe. Our main contribution is to provide an
implementation of a GEMM kernel that uses the Tensor Memory Accelerator (TMA) and Warp Group Matrix-Multiply-
Accumulate (WGMMA) operations introduced with NVIDIA Hopper architecture.
1
INTRODUCTION
The problem of devising computationally efficient algorithms for matrix multiplication is of fundamental
importance due to the basic role of this algebraic operation in scientific computing and machine learning. Indeed,
it is the ability to massively parallelize matrix multiplication on the GPU that has in large part enabled the rise of
AI in the modern world (in the form of ever-more sophisticated neural networks).
As a closed-source and broadly scoped solution, NVIDIA already delivers high-performance implementations
of general matrix multiplication (GEMM) in the cuBLAS library and convolutions (CONV) in the cuDNN library.
However, in view of the ongoing rapid evolution of GPU architectures and powerful new AI models, it is becoming
increasingly important for end-users to be able to implement GEMM and CONV algorithms for their specific
applications by writing their own CUDA code. Custom code is mostly needed in the epilogue of a GEMM kernel,
where another operation has to be fused before writing out the result matrix.
For example, Flash Multi-Head Attention (FMHA) is a recent optimized implementation of attention layers in
transformer deep learning models such as Large Language Models (LLMs) [1]. As part of that proposed attention
algorithm, one wants to fuse the matrix multiplications with the softmax nonlinear operations in order to exploit
data locality and thereby achieve better performance, as opposed to a straightforward implementation without
fusion. Implementing FMHA using cuBLAS is not possible as cuBLAS does not expose the thread block or
thread-level operations, where the fusion would need to happen.1
To address these needs, NVIDIA has provided a suite of CUDA C++ template abstractions for implementing
GEMM and related computations through their open-source library CUTLASS, integrated with the CuTe core
library and backend since version 3.0. CUTLASS eases the development of optimized kernels by bringing to
bear modern software engineering practices otherwise lacking in plain CUDA, such as OOPs, C++ templates
(generics), abstract data types (e.g., Tensors, Layouts, and Shapes), and re-usable design patterns.
Nonetheless, it can be challenging for the CUDA programmer to fully leverage the capabilities afforded to
them by CUTLASS. In this publication and the accompanying code, we explain and implement a CuTe-based
GEMM kernel for Hopper architecture that in particular exploits the new Tensor Memory Accelerator (TMA) and
Warp Group MMA (WGMMA) operations to improve performance. To make this publication more self-contained,
we also include a larger analysis of the GEMM problem; other insightful sources on this include [3] and [5]. In
future work, we plan to integrate our GEMM kernel into FMHA and other LLM kernels.
1NVIDIA has announced the cuBLASDx API as a version of cuBLAS that will provide custom fusion support [2], but this isn’t widely available
yet.
†Colfax Research. A copy of this paper is available at https://research.colfax-intl.com/nvidia-hopper-gemm-cutlass/.
Date: December 14, 2023. Email: research@colfax-intl.com.
2
•
Ganesh Bikshandi and Jay Shah
Complete working code for the GEMM kernel explained in this paper is available at: https://github.com/
ColfaxResearch/cutlass-kernels.
2
INITIAL RUNDOWN ON GEMM PERFORMANCE
The H100 PCIe GPU provides a maximum compute performance of 51 Teraflops for FP32 execution and 378
Teraflops for TF32 execution [6].2 The goal for a performant GEMM implementation is to achieve as high a
percentage of this theoretical maximum as possible. With reference to the functionality exposed by CUTLASS,
performance is sensitive to many parameters such as:
 0
 50000
 100000
 150000
 200000
 250000
 300000
Performance (GFLOPS)
K=32768
K=4096
K=512
K=256
K=128
K=64
Fig. 1. CUTLASS SGEMM (TF32) Performance on Hopper (H100 PCIe) GPU for 𝑀=𝑁=4096 and 𝐾∈{64,128,256,512,4096,32768}
given different tiling shapes, warp counts and stage counts. Sorted in decreasing order of performance.
• Shape of matrices (tall, skinny, or square);
• Layout of matrices (row or column);
• Number of warps per thread block;
• Thread block shape;
• Thread cluster shape;
• Number of pipeline stages in the software pipelining optimization;
• Chosen precision (TF32 vs FP32 vs FP16 vs FP8);
• Usage of special MMA instructions like WGMMA or TMA.
Figure 1 shows the performance variation with different values chosen for TF32 precision.3 The results were
obtained using cutlass_profiler, a tool provided by CUTLASS that generates different GEMM kernels off of a
base kernel by varying parameters and then measures the resulting running time and GFLOPS (among other
metrics).
From our empirical studies, we discovered that the best GEMM kernels try to optimize for maximum Tensor
Core utilization and maximum possible GPU occupancy using software pipelining. In Figure 2, we display the
2The NVIDIA datasheet lists 756 TFlops for TF32 with sparsity.
3Performance is evaluated with respect to the matrix-multiply operation𝐶= 𝐴·𝐵for 𝐴and 𝐵generic 𝑀×𝐾and 𝐾×𝑁matrices, respectively.
In general, the variables 𝑀, 𝑁, and 𝐾always refer to these dimensions.
Developing CUDA Kernels for GEMM on Hopper using CUTLASS
•
3
performance of four such versions of GEMM for single-precision TF32 on H100 PCIe GPU. To the best of our
knowledge, these kernels have been optimally tuned given the available specifications.
 0
 50
 100
 150
 200
 250
 300
cuBLAS
CuTe
CUTLASS,
TMA+WGMMA
CUTLASS,
TMA+WGMMA+WS
 0
 20
 40
 60
 80
 100
Performance (TF32 TFLOPS)
Efficiency (%)
Performance
215.6
170.0
228.5
249.5
Efficiency
54
42
57
62
Fig. 2. SGEMM (TF32) Performance for 𝑀=𝑁=𝐾=4096.
The four versions listed here are:
(1) cuBLAS: A version from the cuBLAS library in TF32 execution mode (set using NVIDIA_TF32_OVERRIDE=1).
(2) CuTe: A hand-implemented version (presented in this paper) using CuTe, a collection of C++ CUDA
template abstractions for defining and operating on hierarchically multidimensional layouts of threads and
data. This implementation uses TMA for loads and WGMMA for matrix operations. Software pipelining (cf.
§5) is not used in this version.
(3) CUTLASS, TMA+WGMMA: A version shipped with the CUTLASS library that uses TMA loads and
WGMMA instructions along with software pipelining optimization. This version served as our guide in
implementing GEMM kernels using CuTe.
(4) CUTLASS, TMA+WGMMA+WS: An improved version of (3) that also uses Warp Specialization (WS) and
Thread Clusters. This version, while still not the best possible, utilizes other key differentiators of Hopper
architecture.4 We have chosen to include this to showcase deeper optimizations available for Hopper.
Discussion on using cuBLAS versus CUTLASS has sometimes been framed as trading off the superior general
performance of cuBLAS for the customizability of CUTLASS.5 However, Figure 2 shows that CUTLASS is now
more than competitive with cuBLAS; even our custom version, which implements only a small subset of all
possible optimizations, comes close in performance. We also note that the best CUTLASS kernel we studied
achieves close to 280 TeraFlops in performance. This observed discrepancy in performance between cuBLAS
and CUTLASS is consistent with NVIDIA’s own benchmarking [4].
At any rate, the upshot is that good performance for a custom lightweight GEMM kernel is achievable thanks
to CUTLASS given some effort put in by the CUDA programmer, which bodes well for the performance of fused
4The idea of warp specialization itself is of course not new to Hopper, but it is only implemented in CUTLASS starting with Hopper [11].
5See for example https://github.com/NVIDIA/cutlass/issues/109 for public comment.
4
•
Ganesh Bikshandi and Jay Shah
kernels down the line. Furthermore, in §6 we will see that our CuTe program outperforms both cuBLAS and
CUTLASS kernels for Batched-GEMM with 𝐾= 64 and 𝐿= 96 for batch_count.
3
MATRIX MULTIPLICATION ON THE GPU
To orient the reader, we give a rapid review of some of the basic ideas involved in writing a GEMM kernel in
CUDA.
3.1
The GPU Memory Hierarchy and CUDA Thread Hierarchy
To begin with, understanding the GPU memory hierarchy is crucial for optimizing GEMM kernels. The GPU
has three distinct levels in its memory hierarchy, proceeding from larger and slower to smaller and faster memory:
(1) HBM (High Bandwidth Memory) or Global Memory (GMEM).
(2) Shared memory (on-chip) (SMEM).
(3) Register memory (RMEM).
On the other hand, the CUDA programming model has the thread hierarchy. This is comprised, from coarser
to finer groupings, of grids, thread blocks (i.e., cooperative thread arrays or CTAs), and threads.6 GMEM is
available across the entire grid; SMEM is per streaming multiprocessor (SM), to which a thread block is assigned;
and individual threads have their own RMEM. In this way, CUDA exposes the GPU memory hierarchy to the
programmer [8]. Also note that for the H100 GPU, 32 threads are grouped into one warp for parallel execution on
a SM; the size of a warp is fixed by the particular GPU architecture and thus not under the programmer’s control.
The matrix multiplication algorithms of interest to us are written to be aware of this hierarchical structure.
More precisely, they decompose the top-level matrix multiplication into multiple sub-matrix multiplications (or
tiled matrix multiplications). Each decomposition step made in the algorithm corresponds to moving across one
of the levels in the CUDA thread hierarchy and GPU memory hierarchy. A typical example of how one might
assign tiling shapes to different levels of the hierarchy is given in Figure 3 (taken from [3]).7
Fig. 3. A typical tiling hierarchy used by a GEMM kernel optimized for NVIDIA GPU.
Figure 3 should be read as follows:
6On top of this, Hopper introduces thread block clusters as intermediate between grids and thread blocks. The thread blocks within a cluster,
spread over different SMs, can access each other’s shared memory (distributed shared memory).
7In practice, there are many possible choices for how one might tile the matrices (cf. Figure 1).
Developing CUDA Kernels for GEMM on Hopper using CUTLASS
•
5
• At the HBM or global memory level, the 𝐴matrix is divided across the rows (𝑀dimension) and the 𝐵
matrix is divided across the columns (𝑁dimension). The result matrix 𝐶is divided along both the 𝑀and
𝑁dimensions. A row panel of 𝐴and a column panel of 𝐵are assigned to a thread block, which computes a
tile of output 𝐶.
• Each thread block further tiles the row panel of 𝐴and column panel of 𝐵along the 𝐾dimension; this ensures
that the tiles fit in shared memory. The corresponding tiles of 𝐴and 𝐵are then multiplied, accumulating
the result in the tile of the 𝐶matrix assigned to the thread block.
• Further optimization then proceeds recursively by sub-tiling each of the tiles for warp and thread-level
computation. Sub-tiles are moved into register memory and the actual element-wise multiplication is
always done at the last level of tiling.
• Note that the sub-tile chosen in the warp tile to be assigned to a thread does not have a contiguous shape;
rather, we have 32 threads assigned to the sub-tiles in each quadrant of the warp tile, so that each thread
sub-tile is spread over all four quadrants. This is done to exploit memory coalescing for parallel thread
execution.
In the remainder of this section, we go over some (pseudo)code that describes both the naive Matmul and tiled
Matmul algorithms.
3.2
Naive Matrix Multiplication
Recall that a naive matrix multiplication can be written using a triply nested loop in C++ as follows:
for (int i = 0; i < M; ++i)
for (int j = 0; j < N; ++j)
for (int k = 0; k < K; ++k)
C[i][j] += A[i][k] * B[k][j];
However, writing in CUDA means we can parallelize one or two or all the loops. CUDA naturally provides
2D parallelism in the form of thread blocks and grids when a CUDA kernel is launched, which makes it easy to
parallelize the 𝑖and 𝑗loops. The following code snippet illustrates this principle.
__global__ void NaiveGemmKernel(int M, int N, int K, float const *A, int lda, float
const *B, int ldb, float *C, int ldc) {
int i = threadIdx.x + blockIdx.x * blockDim.x;
int j = threadIdx.y + blockIdx.y * blockDim.y;
if (i < M && j < N) {
for (int k = 0; k < K; ++k)
C[i + j * ldc] += A[i + k * lda] * B[k + j * ldb];
}
}
dim3 block(16, 16);
dim3 grid( (M + block.x - 1) / block.x, (N + block.y - 1) / block.y);
NaiveGemmKernel<<< grid, block >>>(M, N, K, A, lda, B, ldb, C, ldc);
6
•
Ganesh Bikshandi and Jay Shah
The 𝑘loop is now the main loop of the matrix multiplication, while 𝑖and 𝑗are not expressed “explicitly” in
the CUDA code as they are inferred from the kernel launch parameters given for the thread block and grid
dimensions. The 𝑘loop is not parallelized, but CUTLASS provides an option to parallelize along that dimension
also, via the 𝑘-split algorithm. For our discussion, we will suppose that the 𝑘loop is not parallelized.
3.3
Outer Product Summation and Tiling
The naive matrix multiplication in §3.2 requires the matrices 𝐴and 𝐵to be repeatedly brought inside the cache
or shared memory. To obviate this issue, the 𝑘loop can be permuted to be the outside loop, and this leads to the
“outer product” version of matrix multiplication in which the result matrix 𝐶is computed as a sum of 𝑘many
outer products (of columns of 𝐴and rows of 𝐵).
for (int k = 0; k < K; ++k) // K dimension now outer-most loop
for (int i = 0; i < M; ++i)
for (int j = 0; j < N; ++j)
C[i][j] += A[i][k] * B[k][j];
One problem with computing the result matrix 𝐶via outer product summation is that this requires 𝐶to be
live all the time, which for large enough dimensions can quickly exceed all available on-chip memory. A way to
circumvent this problem is “tiling” or “blocking”. The computation can first be tiled into smaller units that can fit
into on-chip memory, after which another set of loops iterates over the tiles:
for (int m = 0; m < M; m += MT) // iterate over M dimension
for (int n = 0; n < N; n += NT) // iterate over N dimension
for (int k = 0; k < K; ++k)
for (int i = 0; i < MT; ++i) // compute for one tile
for (int j = 0; j < NT; ++j) {
int row = m + i;
int col = n + j;
C[row][col] += A[row][k] * B[k][col];
}
Note that we have not yet tiled along the 𝐾dimension; this comes next.
3.4
Hierarchical or Recursive Matrix Multiplication
We now refer back to Figure 3 and its multiple levels of tiling. Matrix multiplication optimized for the GPU
is implemented in an hierarchical or recursive fashion. At every level of the recursion, the program copies a
tile from one memory to another (e.g., global to shared memory). The actual multiply-accumulate-add happens
only at the leaf level (the last level), which in Figure 3 is thread-level. At the leaf level, we have a choice to
either use the standard multiply-accumulate-add of the SIMT core or specialized Tensor Core instructions. An
implementation optimized for Hopper architecture will choose to recurse only till warpgroup-level (instead of
thread-level) and use the special WGMMA instructions that use the Tensor Cores more optimally [10, §9.7.14].
Pseudocode illustrating this is displayed below:
Developing CUDA Kernels for GEMM on Hopper using CUTLASS
•
7
// MT, NT, KT = dimensions at threadblock level
// MW, NW, KW = dimensions at warp level
// Loop1A: threadblock-level concurrency
Loop1A: for each m, n in M, N with step MT, NT
Loop1B: for each k in K with step KT
Move a chunk of A from GMEM to SMEM (As)
Move a chunk of B from GMEM to SMEM (Bs)
// Loop2A: warp-level concurrency
Loop2A: for each mm, nn in MT, NT with step MW, NW
Loop2B: for each kk in KT with step KW
Move a chunk of As from SMEM to RMEM (Ar)
Move a chunk of Bs from SMEM to RMEM (Br)
// run mma and accumulate in registers
// further recursion is hidden by mma call
mma(Ar, Br, accum)
Note that the equivalent of the parallelization step handled by launching the CUDA kernel corresponds to the
outermost loop in this pseudocode. Since parallelization occurs over chunks of the 𝐴and 𝐵matrices that now are
no longer single rows and columns, converting this pseudocode into working CUDA code is not straightforward.
We will see that the CUTLASS API provides a method for doing this via local_tile.
3.5
Fusing Operations with Matrix Multiplication
For use in real-life workloads, a GEMM kernel will often involve additional operations to be fused with the
above Matmul, such as a linear scaling step. The place where that is done (at the end of the outermost loop) is
referred to as the epilogue. Typically, the accumulator will be in registers in the epilogue. To understand the basic
reason for kernel fusion, consider linear scaling: even though this could be done separate to the initial Matmul,
fusing in the epilogue helps reduce the total bandwidth requirement, since otherwise the linear scaling would
necessitate a re-read of the original 𝐶and the resulting 𝐷matrix from global memory.
Although the idea is simple, as a software engineering task there are many issues to be aware of when writing
a fused kernel. For one, there are multiple elementwise operations that need to be fused with Matmul (e.g., ReLU)
in a typical AI-intensive workload. Secondly, the difficulty of fusing various operations can vary greatly with the
operation; for example, softmax is harder to handle than linear scaling.
4
DEVELOPING A GEMM KERNEL USING CUTLASS AND CUTE
Developing and maintaining a performant Matmul kernel is difficult mainly because the shape of the tiles at
each level must be tuned for different GPU architectures, different precision types of the underlying data, and
different matrix sizes. When new levels in the CUDA thread hierarchy are introduced (e.g., the thread cluster
level introduced with Hopper architecture), the kernel code must also be changed with additional loops and other
code to handle that. The unroll parameters need to be tuned by the compiler too. Apart from all this, there are
specialized instructions like WGMMA in Hopper that can boost the performance of the leaf Matmul step. Finally,
fusing operations like linear scaling or softmax with Matmul introduces more complexity into the equation.
8
•
Ganesh Bikshandi and Jay Shah
To handle all of this, one has the CUTLASS library. CUTLASS is a C++ templates-based library that provides a
very high-level interface for defining the shape of the tiles at each level of the memory hierarchy. It also provides
an interface called Epilogue that allows the programmer to fuse operations like linear scaling after the completion
of a leaf Matmul. Moreover, CUTLASS defines default leaf level kernels itself and most often selects highly tuned
kernels for the GPU architecture based on the tiling parameters supplied during compile time, the specified
architecture and the precision.
4.1
CUTLASS API Basics
A basic listing of CUTLASS-based Matmul is described in Listing 1. Apart from CuTe, CUTLASS has the
following 3 important APIs for GEMM, each corresponding to a distinct level of the GPU memory hierarchy [12]:
(1) Device API;
(2) Kernel API;
(3) Collective API.
The Collective API embodies a thread block or a cluster of thread blocks (from Hopper architecture onwards).
Collective APIs can be used to construct a GEMM as well as the epilogue to be fused with GEMM. The default
epilogue simply writes out the accumulator of GEMM from register memory to global memory. CUTLASS defines
several other typical operations such as linear scaling and clamping; other device-side function call operators
may also be used to perform custom operations.
The Kernel API embodies the entire grid. It thus schedules the collectives and is responsible for tiling the input
matrices into row and column panels, loading the references to them and invoking the GEMM and the epilogues.
Fusion of epilogues with GEMM happens at the Kernel API level.
The Device API is the highest-level API. It is invoked from the Host (i.e., CPU) and does not have any detail
about the specifics of the Matmul implementation. This API is used by host-side .cu code to invoke CUTLASS’s
GEMM kernels, much like cuBLAS API.
The CUTLASS API relies on C++ templates and meta-programming. Many compile-time values can be set to
default (or auto) values, while the optimal values are chosen to be best possible by the CUTLASS implementation.
Most of the compile-time parameters described in the listing are self-explanatory. The constructed GEMM kernel
can be executed with input and output arguments just like any C++ functor objects.
using ElementA = float; // Element type for A matrix operand
using ElementB = float; // Element type for B matrix operand
using ElementC = float; // Element type for C and D matrix operands
using ArchTag = cutlass::arch::Sm90; // Tag indicating the SM
using OperatorClass = cutlass::arch::OpClassTensorOp;
// Operator class tag
using TileShape = Shape<_128,_128,_32>; // Threadblock-level tile size
using ClusterShape = Shape<_1,_2,_1>;
// Shape of the threadblocks in a cluster
Collective API
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
ElementA, RowMajor, 4,
Developing CUDA Kernels for GEMM on Hopper using CUTLASS
•
9
ElementB, ColumnMajor, 4,
ElementAccumulator,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAuto,
cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
using CollectiveEpilogue = typename cutlass::epilogue::collective::
CollectiveBuilder<
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementC, ElementC,
ElementC, ColumnMajor, 4,
ElementC, ColumnMajor, 4,
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
Kernel API
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int>, // Indicates ProblemShape
CollectiveMainloop,
CollectiveEpilogue
>;
Device API
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
Listing 1. Constructing a GEMM Kernel for SM90 (Hopper) Architecture
4.2
Tiled Matrix Multiplication Using CuTe
Even though CUTLASS provides APIs for optimized GEMM and fusing operations with GEMM, developing
fused kernels like FMHA requires lower level APIs of CUTLASS, as the fusion is not straightforward. Such a
custom kernel has been shipped as part of CUTLASS [9]. However, that kernel does not use the new Hopper
features and is mostly customized for SM80 architecture. It is somewhat cumbersome to redevelop this for SM90.
CuTe is another API within the CUTLASS API that provides even more flexibility to develop GEMM kernels. It
specifically introduces the concept of Shapes and Layouts, using which programmers can define the different
levels of tiling explicitly. Additionally, it provide APIs to:
(a) Convert matrices in to tensors and partition them;
(b) Access the tiles of a tensor that belong to a thread block (local_tiles);
(c) Make a local partition of a tensor that belongs to a thread within a thread block (local_partition);
(d) Copy between GEMM, SMEM and RMEM (copy);
(e) Multiply tensors with special Matmul instructions like WGMMA (gemm);
10
•
Ganesh Bikshandi and Jay Shah
(f) Synchronize between thread clusters;
(g) Make special swizzle layouts for shared memory.
Listing 2 shows a basic CuTe based Matmul CUDA kernel, obtained from CUTLASS repository [14].8
template <class MShape, class NShape, class KShape,
class TA, class AStride, class ABlockLayout, class AThreadLayout,
class TB, class BStride, class BBlockLayout, class BThreadLayout,
class TC, class CStride, class CBlockLayout, class CThreadLayout,
class Alpha, class Beta>
__global__ static
void
gemm_device(MShape M, NShape N, KShape K,
TA const* A, AStride dA, ABlockLayout blockA, AThreadLayout tA,
TB const* B, BStride dB, BBlockLayout blockB, BThreadLayout tB,
TC
* C, CStride dC, CBlockLayout blockC, CThreadLayout tC,
Alpha alpha, Beta beta)
{
using namespace cute;
using X = Underscore;
// Shared memory buffers.
__shared__ TA smemA[cosize_v<ABlockLayout>];
__shared__ TB smemB[cosize_v<BBlockLayout>];
auto sA = make_tensor(make_smem_ptr(smemA), blockA);
auto sB = make_tensor(make_smem_ptr(smemB), blockB);
// Represent the full tensors.
auto mA = make_tensor(make_gmem_ptr(A), make_shape(M,K), dA);
auto mB = make_tensor(make_gmem_ptr(B), make_shape(N,K), dB);
auto mC = make_tensor(make_gmem_ptr(C), make_shape(M,N), dC);
// Get the appropriate blocks for this thread block.
auto MT = size<0>(sA);
auto NT = size<0>(sB);
auto KT = size<1>(sB);
auto gA = local_tile(mA, make_shape(MT, KT), make_coord(blockIdx.x, _));
auto gB = local_tile(mB, make_shape(NT, KT), make_coord(blockIdx.y, _));
auto gC = local_tile(mC, make_shape(MT, NT), make_coord(blockIdx.x, blockIdx.y);
// Define partitioned views of GMEM and SMEM for COPY
8The code has been adjusted in places for clarity. We also display the kernel only, not the main program that launches the kernel.
Developing CUDA Kernels for GEMM on Hopper using CUTLASS
•
11
auto tAgA = local_partition(gA, tA, threadIdx.x);
auto tAsA = local_partition(sA, tA, threadIdx.x);
auto tBgB = local_partition(gB, tB, threadIdx.x);
auto tBsB = local_partition(sB, tB, threadIdx.x);
// Define partitioned views of SMEM for GEMM.
// Partition sA (M,K) by the rows of tC.
auto tCsA = local_partition(sA, tC, threadIdx.x, Step<_1, X>{});
// Partition sB (N,K) by the cols of tC.
auto tCsB = local_partition(sB, tC, threadIdx.x, Step< X,_1>{});
// Partition gC (M,N) by the tile of tC.
auto tCgC = local_partition(gC, tC, threadIdx.x, Step<_1,_1>{});
// Allocate the accumulators (RMEM).
auto tCrC = make_fragment_like(tCgC);
// Clear the accumulators
clear(tCrC);
Matmul Begin
// Data is copied from GMEM to SMEM using the COPY views.
// gemm(.) operates on the GEMM views.
auto k_max = size<2>(tAgA);
for (int k = 0; k < k_max; ++k) {
// Copy GMEM to SMEM.
copy(tAgA(_,_,k), tAsA);
copy(tBgB(_,_,k), tBsB);
cp_async_fence();
cp_async_wait<0>();
__syncthreads();
// Compute GEMM on SMEM.
// Accumulate to registers.
gemm(tCsA, tCsB, tCrC);
__syncthreads();
}
Matmul End
// Epilogue fusion goes here.
for (int i = 0; i < size(tCgC); ++i)
{
tCgC(i) = tCrC(i);
}
12
•
Ganesh Bikshandi and Jay Shah
}
Listing 2. Basic GEMM using CuTe
The key part of the Matmul computation is listed in the loop. The kernel computes the matrix multiplication of
𝐴and 𝐵𝑇resulting in 𝐶. The matrices 𝐴and 𝐵are tiled as shown in Figure 4. The tiling of 𝐶was already discussed
earlier in the naive Matmul implementation. The key differences between naive Matmul and this Matmul are:
(1) Matrix elements are brought from GMEM to SMEM first using an async copy operation;
(2) The result matrix 𝐶is stored in RMEM and finally written back to GMEM in the epilogue;
(3) The computation is tiled along the 𝐾dimension also. This enables step (1), as 𝐴and 𝐵tiles are small enough
to fit in shared memory compared to the full row or column panel.
The two key APIs of CuTe to be emphasized in Listing 2 are:
(1) local_tile: extracts the tiles local to a thread block into a tensor.
(2) local_partition: extracts the elements local to a thread in a thread block into a tensor.
Once local tiles are obtained using the CuTe API, the corresponding tiles of 𝐴and 𝐵are multiplied using the
GEMM API and the result is accumulated in the 𝐶matrix (in registers). After the last tile along the 𝐾dimension
is processed, the result 𝐶is then written to GMEM.
An important feature in CuTe is the view of a tensor (in the sense of the C++ concept). During the copy
operation, where the data is read from global to shared memory, a view based on AThreadLayoutA (tA) and
BThreadLayout (tB) is used for input tensors. Such a view is created to improve coalescing for global memory
loads, for example. However, during the gemm operation, a view based on CThreadLayout (tC) is used. Such
a thread-to-data mapping improves the performance of matrix-multiply computation but may not result in
coalesced stores to global memory. The original shared memory can be read and written using the different views.
Thus, the thread layouts of the copy and gemm operations are decoupled so that the best choice for each operation
can be chosen by the user.
Fig. 4. Graphical representation of Tiled CuTe GEMM kernel with SIMT Core. The optimized implementation in §4.3 will
instead use TMA for loading tiles of 𝐴and 𝐵from GMEM to SMEM, and WGMMA for computing the tiles of 𝐶.
Note that we have used no-transpose for 𝐴and transpose for 𝐵(commonly referred to as NT layout), as
that is easy to implement using SIMT mul-add. In contrast, the versions shown in Figure 2 all use TN layout.
Nevertheless, this version will serve as a basis for the version of Matmul that we present next.
Developing CUDA Kernels for GEMM on Hopper using CUTLASS
•
13
4.3
Incorporating TMA and WGMMA instructions from NVIDIA Hopper Architecture
The listing 2 in the preceding section uses a plain SIMT mul-add operation as the atom to compute the product
of two tiles. On the other hand, modern NVIDIA GPUs such as the H100 provide Tensor Cores that accelerate
mixed-precision computation several-fold. Additionally, Hopper architecture also introduces the Tensor Memory
Accelerator (TMA), which can transfer large blocks of data efficiently between global memory and shared memory.
To utilize TMA and Tensor Cores, two important changes are required:
• The copy API call should be changed to include the TMA copy atom;
• The gemm API call should be changed to include the MMA atom – for Hopper, we choose WGMMA.
The WGMMA atom schedules data of size 64×8 (other sizes are possible too) split across 128 threads atomically
as one operation on a tensor core. The mapping of rows and columns of the resultant 𝑀𝑇×𝑁𝑇tile of 𝐶to threads
is more complex than those listed in the preceeding sections. CuTe provides utilities that define the mapping,
so users need not understand the details of the complex thread layout. At the same time, while implementing
complex fused kernels in the epilogue, programmers can use CuTe APIs to unravel the thread layout so as to
obtain the correct row and column of individual elements of the tile.
Typically, while using WGMMAs the loads of the 𝐴and 𝐵tiles are executed using TMA. TMA loads offer better
performance than cp.async. WGMMA operations require that the shared memory allocated for the tiles of the 𝐴
and 𝐵matrices be in a certain “swizzled” format [10], which is supported by TMA. TMA loads are asynchronous
and are invoked from one thread (typically thread 0), while the other threads wait on a cuda::barrier for the
operation to be completed. Thus, TMA loads require producer-consumer style synchronization between the
threads of a warp.
Listing 3 shows the relevant part of the new GEMM kernel that uses TMA and WGMMA. The copy and gemm
views are not shown in the new listing for brevity.9
....
for (int k = 0; k < size<1>(tAgA); ++k)
{
.....
//copy A and B from GMEM to SMEM using COPY views.
if (threadIdx.x == 0)
{
/// Initialize shared memory barrier
....
copy(tma_copy_a, tAgA(_,k), tAsA);
copy(tma_copy_b, tBgB(_,k), tBsB);
}
__syncthreads();
/// Wait on the shared memory barrier.
....
__syncthreads();
9CuTe provides APIs to get the correct views for TMA copy and WGMMA gemm operations.
14
•
Ganesh Bikshandi and Jay Shah
warpgroup_fence_operand(tCrC);
cute::gemm(wmma_atom, tCrA, tCrB, tCrC);
warpgroup_commit_batch();
warpgroup_wait<1>();
__syncthreads();
}
.....
Listing 3. GEMM using TMA+WGMMA
As is shown in Figure 5, the TMA+WGMMA version provides almost a sevenfold improvement compared to
the CuTe basic version, since it uses the Tensor Cores. However, there is clear room for improvement in terms of
Tensor Core utilization.
 0
 100
 200
 300
 400
 500
TF32
FP16
 0
 20
 40
 60
 80
 100
 0
 2
 4
 6
 8
 10
Performance (TFLOPS)
Tenscore Utilization (%)
cuBLAS
CuTe
CuTe-Basic
Tensor Core Utilization
Fig. 5. cuBLAS vs CuTe vs CuTe-Basic (TF32 and FP16 precision).
5
ADDITIONAL OPTIMIZATIONS
Recall from §2 that the best CUTLASS kernel for SGEMM delivers around 280 TFLOPS, while cuBLAS delivers
around 215 TFLOPS. CUTLASS implements many more optimizations to achieve this superior level of performance.
To list a few from the documentation [11], we have:
(1) Software Pipelining – Software pipelining is a technique to hide memory latency where memory accesses
and math instructions are executed concurrently, while always accounting for the dependencies between
these steps. The CUTLASS implementation uses multiple buffers at both the thread block and warp level.
(2) Warp Specialization – With optimizations like software pipelining, different threads or groups of threads
naturally have distinct roles. Some are producers that load data, while others are consumers that run the
MMA instructions. The idea of warp specialization is to spatially partition the warps in a thread block into
two groupings of producers and consumers.
Developing CUDA Kernels for GEMM on Hopper using CUTLASS
•
15
(3) Persistent Kernels – Persistent kernels is a CUDA design pattern that aims to avoid kernel launch and
configuration overhead by keeping the kernel persistent on the GPU across multiple calls. In CUTLASS,
this involves having persistent thread blocks compute multiple output tiles over their lifetime.
(4) Two co-operative consumer warp groups – WGMMA allows the operand 𝐴tile to be in register memory
too instead of shared memory. However, that restricts the tile size of 𝐴due to limited register space. Splitting
the tile size across the 𝑀dimension into two and assigning to two different consumer warp groups allows
for larger tile sizes and eases register pressure.
(5) Warp-Specialized Persistent Ping-Pong kernel – The two consumer warp groups from (4) are each
assigned to a different output tile. This allows for the epilogue of one consumer warp group to be overlapped
with the math operations of the other consumer warp group, thus maximizing tensor core utilization. There
is also synchronization on the side of the producer warp groups.
From our empirical studies, point (5) in particular is largely responsible for the gap between the fourth column
in Figure 2 and the indicated 280 TFLOPS number for the best measured CUTLASS kernel.
6
BATCHED-GEMM
The AI workflow that we are targeting does not involve multiplying large square matrices. Instead, it involves
large square matrices decomposed as products of matrices with small 𝐾(e.g., 64 or 128), and with batch count
𝐿> 1 (e.g., 64 or 96); cf. [1, §2.2]. Such a scheme is popularly known as Batched-GEMM. Our CuTe program can
be extended to handle Batched-GEMM by simply setting the third dimension of the grid to be 𝐿. We then use
blockIdx.z when using the local_tile operation inside the CUDA kernel, as shown in listing 4.
....
auto gA = local_tile(mA, make_shape(MT, KT), make_coord(blockIdx.x, _, blockIdx.
z));
auto gB = local_tile(mB, make_shape(NT, KT), make_coord(blockIdx.y, _, blockIdx.
z));
auto gC = local_tile(mC, make_shape(MT, NT), make_coord(blockIdx.x, blockIdx.y,
blockIdx.z);
....
Listing 4. Batched-GEMM kernel using CuTe
Performance of such a Batched-GEMM using CuTe is shown in Figure 6. Surprisingly, the CuTe program
outperforms both cuBLAS and CUTLASS, even though it does not use any of the additional optimizations that
CUTLASS uses as listed in §5.
At the same time, all of the programs deliver sub-optimal performance for Batched-GEMM. From Figure 1,
we can infer that, for small 𝐾(64 or 128), the performance of GEMM is sub-optimal in general. In particular, it
appears that the optimizations listed in §5 need to be adapted to the case of smaller matrices and Batched-GEMM,
or else different methods should be brought to bear.
7
CONCLUSION AND FUTURE WORK
To summarize, we discussed how to develop a CuTe-based GEMM kernel for NVIDIA Hopper architecture that
uses the Tensor Memory Accelerator (TMA) and Warp Group MMA (WGMMA) operations. Our CuTe program
16
•
Ganesh Bikshandi and Jay Shah
 0
 20
 40
 60
 80
 100
cuBLAS
CuTe
CUTLASS
Performance (TF32 TFLOPS)
47.9
53.9
51.1
Fig. 6. Batched-SGEMM with 𝑀=𝑁=4096, 𝐾=64 and 𝐿=96.
achieved close to 80% of the performance of the standard cuBLAS GEMM kernel with this one single optimization.
For Batched-GEMM, our CuTe program outperformed both cuBLAS and CUTLASS.
Currently, we are working to integrate our GEMM kernel with Flash Multi-Head Attention (FMHA) [1], a
popular attention layer used in Large Language Models (LLMs). Some important challenges to be addressed are:
• FMHA necessitates changes to the basic GEMM kernel described in this paper. The most significant change
is that all panels of the 𝐵matrix are assigned to the same thread block, nullifying any parallelism along
that dimension in the grid. On the other hand, LLMs use batched GEMM, introducing parallelism along the
third dimension 𝐿of the grid.
• The online-softmax [13] computation has to be fused with the result of the GEMM (tCrC) in registers,
which involves atomic operations across threads for computing the maximum and sum of each row.
• FMHA uses a smaller 𝐾dimension (64 or 128) in comparison to 𝑀and 𝑁. From Figure 1, we know that
GEMM performance for small 𝐾values tends to be very sub-optimal. We plan to adapt the optimizations
listed in §5 to the setting of matrices with small 𝐾and Batched-GEMM.
REFERENCES
[1] FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. Tri Dao. July 17, 2023. https://arxiv.org/abs/2307.08691.
[2] New cuBLAS 12.0 Features and Matrix Multiplication Performance on NVIDIA Hopper GPUs. Roman Dubtsov, Evarist Fomenko and
Babak Hejazi. February 1, 2023. https://developer.nvidia.com/blog/new-cublas-12-0-features-and-matrix-multiplication-performance-on-
nvidia-hopper-gpus/
[3] CUTLASS: Fast Linear Algebra in CUDA C++. Andrew Kerr, Duane Merrill, Julien Demouth and John Tran. December 5, 2017.
https://developer.nvidia.com/blog/cutlass-linear-algebra-cuda/.
[4] CUTLASS 3.2 – Performance. https://github.com/NVIDIA/cutlass#performance.
[5] CuTe dense matrix-matrix multiply tutorial. https://github.com/NVIDIA/cutlass/blob/main/media/docs/cute/0x_gemm_tutorial.md.
[6] NVIDIA H100 Tensor Core GPU Datasheet. https://resources.nvidia.com/en-us-tensor-core/nvidia-tensor-core-gpu-datasheet.
[7] NVIDIA Hopper Architecture In-Depth. Michael Andersch, Greg Palmer, Ronny Krashinsky, Nick Stam, Vishal Mehta, Gonzalo Brito and
Sridhar Ramaswamy. March 22, 2022. https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/
Developing CUDA Kernels for GEMM on Hopper using CUTLASS
•
17
[8] CUDA Refresher: The CUDA Programming Model. Pradeep Gupta. https://developer.nvidia.com/blog/cuda-refresher-cuda-programming-
model/
[9] xFormers: A modular and hackable Transformer modelling library. Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan
Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza. 2022. https:
//github.com/facebookresearch/xformers.
[10] Parallel Thread Execution ISA Version 8.2. https://docs.nvidia.com/cuda/parallel-thread-execution/index.html.
[11] Efficient GEMM in CUDA. https://github.com/NVIDIA/cutlass/blob/main/media/docs/efficient_gemm.md.
[12] CUTLASS 3.0 GEMM API. https://github.com/NVIDIA/cutlass/blob/main/media/docs/gemm_api_3x.md
[13] Online normalizer calculation for softmax. Maxim Milakov and Natalia Gimelshein. 2018. https://doi.org/10.48550/arXiv.1805.02867.
[14] https://github.com/NVIDIA/cutlass/blob/main/examples/cute/tutorial/sgemm_nt_1.cu.