Steps for Writing a Speed-of-Light GEMM kernel: Part 1

TL;DR

• Goal: build a sequence of increasingly optimized BF16 matrix multiplication kernels in CuTeDSL for B200, ultimately reaching 99% of cuBLAS performance on several matrix dimensions.
• Introduce the GPU programming model and how it could be utilized for matrix multiplications.
• Kernel 1 (naive matmul): assign one thread to each output element and have it walk across the entire K dimension. This baseline achieves 4 TFLOP/s.
• A CuTeDSL mental model: how CuTeDSL describes where data is stored, divides that data into tiles, and maps those tiles onto GPU operations.
• Kernel 2 (tensor cores, TMA, and TMEM): give each CTA a tile of C, use TMA and barriers to move matching tiles of A and B into shared memory, issue tcgen05 MMA instructions that accumulate into TMEM, and use a per-thread epilogue view to move the result through registers into global memory. These changes raise throughput to 155 TFLOP/s.
• Kernel 3(swizzling): identify shared memory bank conflicts as the next bottleneck, then change the shared memory layout so simultaneous reads are distributed across banks. This layout-only change raises throughput to 285 TFLOP/s.

Goal

Over the past 6 months or so, I dove deeply into the world of GPUs, performance engineering, and kernels. I undertook this project to learn the core concepts of these areas and to apply my knowledge to build something useful. Thus, I wrote a series of 10 matmul kernels on B200, each progressively more performant than the last, eventually reaching 99% SOTA on several different dimensions.

In this blog and the following parts, I would like to document the steps needed to write a speed-of-light GEMM kernel, introducing the hardware-specific optimizations and CuTeDSL libraries that were used. Starting from a simple correct kernel, in this series, I will progressively introduce Blackwell tensor core tiling, TMA data movement, shared memory layouts, multistage pipelines, warp specialization, persistent scheduling, and Gilbert paths. For each iteration, I show the bottleneck, the code change, and its measured effect.

Intro to Matrix Multiplication

To say that matrix multiplication operations are common in modern-day AI models would be an understatement. Linear layers, attention, mixture of experts, and many more algorithms common in these models are at their core matrix multiplication operations. So, writing more efficient matmul implementations would greatly improve the performance of current models.

Given matrices \(A \in \mathbb{R}^{M \times K}\) and \(B \in \mathbb{R}^{K \times N}\), matrix multiplication produces \(C = AB \in \mathbb{R}^{M \times N}\), with each output element given by

Naively, we would implement a matrix multiplication like this:

# Naive matrix multiplication.
function matmul(A, B):
    M, K = shape(A)
    _, N = shape(B)
    C = zeros(M, N)

    for i in range(M):
        for j in range(N):
            for k in range(K):
                C[i, j] += A[i, k] * B[k, j]

    return C

These three nested loops describe the required arithmetic, and it's easy to note that this implementation is \(O(MNK)\) or \(O(N^3)\) if we take \(N = M = K\). This means that the naive time required to complete a matrix multiplcation grows extremely quickly as the dimension size increases. As a result, there has been extensive effort to speeding up matmul operations.

The GPU Programming Model

As detailed in the previous section, general matrix multiply operations (GEMMs) are extremely compute-intensive. However, looking at the algorithm again, we can notice that each entry into C (detailed in Figure 1 below) does not depend on each of the other entries. This therefore motivates the idea of parallel programming, where we can calculate each entry of the final matrix concurrently. So, hardware built for massive parallelization, such as GPUs, are generally used for large matrix multiplications.

There are many blogs out there on GPU fundamentals/architecture [example: link], so I'll only go over the fundamentals here. Like all modern computers, Blackwell GPUs have both memory units (caches, global memory) and compute units (ALUs, tensor cores). A general diagram is given below:

The GPU is split into many different streaming multiprocessors (SMs), each of which could be thought of as its own mini processor. Every SM contains its own tensor cores, shared memory, warp scheduler, and register file. Then, the L2 cache and off-device global memory are shared across all SMs. Note that the shared memory in each SM is the smallest but has the fastest load speed, while the L2 cache is larger but slower, and global memory is the largest but is also the slowest. As a result, we would like our desired data to be in shared memory as often as possible while minimizing loads from global memory. This will inform many design choices down the line.

The biggest difference between most everyday computers and modern GPUs is the existence of tensor cores. Tensor cores are specialized hardware built for matrix multiplications. They execute MMA operations of the form \(D = AB + C\) on small, fixed-shape tiles, performing many fused multiply-adds in parallel. Blackwell's tcgen05.mma supports shapes of up to 256*256*16.

In addition, because of the massive parallelism made possible by the GPU architecture, we must think differently when doing GPU programming compared to traditional programming. In traditional programming, we often think about processing data via loops. However, in GPU programming, we must change our mental model to think about threads instead. We could think of one thread as one individual doing a series of tasks sequentially (for example a multiplication). In a GPU, there are hundreds of thousands of individual threads that could do work simultaneously. So, when processing data, our mental model should instead be to assign data to be processed to individual threads.

Finally, we will go over the organization of threads in a GPU. The number of threads that could be scheduled in a single instruction is called a warp, which consists of 32 threads. Groups of 4 warps form a cooperative thread array (CTA) or thread block, which is the granularity at which we assign threads to a SM and allocate resources to them. Note that it is possible for multiple CTAs to be scheduled onto the same SM, but each individual CTA must be scheduled onto only 1 CTA.

Kernel 1 + Specifications

The code for this kernel can be found here.

We begin with the specifications for these kernels. First, the entries of the input and output matrices are all stored as bf16, which is a 16-bit format, following the Modular blog (which this one is modeled after). However, when accumulating over many entries in a matrix multiplication, rounding errors in bf16 could become large. As a result, we accumulate the intermediate products in fp32.

Furthermore, every reported runtime uses the same measurement protocol. We compile the CuTeDSL kernel before benchmarking, run 10 warmup iterations, and then place CUDA events around 100 consecutive kernel launches, taking the average latency. The warmup launches initialize the CUDA execution path and allow transient effects such as clock rampup to settle, while timing many launches together reduces the influence of host-launch overhead and run-to-run noise. Excluding JIT (just-in-time) compilation is also important because compilation is a one-time setup cost rather than part of the kernel's steady-state execution time. This protocol therefore measures warmed, steady-state throughput rather than first-launch latency.

We begin with the most direct CuTeDSL translation of the definition of matrix multiplication. A two-dimensional grid covers the output matrix, and each GPU thread would be responsible for an element, \(C_{ij}\). That thread walks serially across the entire \(K\) dimension, loads one value from row \(i\) of \(A\) and one value from column \(j\) of \(B\), and accumulates their product. Note that the GPU does not have enough threads to cover the entire output matrix at once, so we assign threads to output values in waves.

This naive implementation achieves 4 TFLOPS (4 trillion floating-point operations per second), which is well below the SOTA implementation.

A CuTeDSL Mental Model

Before introducing tensor cores, it is useful to understand the programming model used by the rest of the kernels. CuTeDSL is a Python-embedded GPU DSL from NVIDIA's CUTLASS project. Its central idea is to separate what data means, where it is stored, and which hardware operation will consume it. Instead of manually calculating every address used by every thread, we describe tensors, layouts, tiles, and operations. CuTe then composes those descriptions into the thread- and instruction-level mappings required by the kernel.

A tensor is storage plus a layout

A CuTe Tensor is not simply a multidimensional allocation. Conceptually, it combines an engine, which is a reference pointer into global, shared, or register memory, with a Layout. The layout maps a logical coordinate to an offset compared to that reference pointer. For example, a row-major M × K matrix has shape (M, K) and strides (K, 1), so coordinate (m, k) maps to offset m * K + k.

# Conceptual CuTe tensor: logical coordinates are separate from storage.
A_layout = make_layout(shape=(M, K), stride=(K, 1))
A = make_tensor(global_memory_pointer, A_layout)

A[m, k]  # The layout converts (m, k) into a memory offset.

This separation is important because the same logical tile can have different physical layouts in global memory, shared memory, tensor memory, or registers. CuTe layouts can also be hierarchical: a mode may itself contain several nested modes. It's also important to note that a layout only describes indexing; creating or transforming one does not by itself move any data.

Tiling creates views

Most GPU kernels repeatedly operate on smaller regions of much larger tensors, as we will see in kernel 2 and later kernels. CuTe represents this by applying a tiler to a tensor. Operations such as local_tile logically divide a tensor and select the tile owned by the current CTA. The result is another tensor view whose coordinates begin at the selected tile, but whose engine still refers to the original storage.

CuTe uses None in a coordinate to keep the global view instead of constructing a new view. This is how code can select one CTA coordinate while retaining the complete M, N, or K tile that the CTA will process. More advanced operations such as flat_divide and group_modes reshape these logical modes, but they follow the same rule: they change how a tensor is viewed and indexed, not where its values physically live.

Atoms describe hardware operations

A layout describes data, but it does not say how the GPU should operate on that data. CuTe represents a small hardware operation with an atom. An MMA atom, for example, describes one matrix multiply accumulate instruction and the mapping between its participating threads and values. A copy atom does the same for a data movement instruction. The atom therefore acts as a contract whose operands must be presented with the shapes and layouts expected by that hardware instruction.

For example, Kernel 2 begins with a Blackwell BF16-to-FP32 MMA operation of shape 64 × 256 × 16, which we will go over in more detail later. Calling make_tiled_mma turns that operation into a TiledMma, which is CuTe's reusable description of the MMA across all participating threads and values. This tiled operation allows us to replicate a smaller atom across a larger tile.

mma_op = tcgen05.MmaF16BF16Op(
    BF16,
    FP32,
    (64, 256, 16),
    tcgen05.CtaGroup.ONE,
    tcgen05.OperandSource.SMEM,
    tcgen05.OperandMajorMode.K,
    tcgen05.OperandMajorMode.K,
)
tiled_mma = cute.make_tiled_mma(mma_op)

Once the operation is known, CuTe can derive the views that it expects based on the operation.

Generic operations are specialized by their atoms and layouts

CuTe's top-level operations are intentionally generic. For example, cute.gemm receives a TiledMma and already-partitioned fragments (tensors with boundary rules), then emits the MMA operation represented by that tiled object. The call is short because the instruction shape, operand types, operand layouts, and thread-value mapping were established earlier:

cute.gemm(
    tiled_mma,
    C_fragment,
    A_fragment[k_block],
    B_fragment[k_block],
    C_fragment,
)

CuTeDSL has a compile-time stage and a runtime stage

CuTeDSL code runs in two stages, but it is not compiled twice. First, calling cute.compile runs the meta-stage on the CPU. CuTeDSL traces the host function, labeled @cute.jit, and the device function, labeled @cute.kernel, evaluates everything already known, and emits code for everything that will only be known later. Since we know some values, e.g. matrix dimensions, before any values are passed in, the compiler could use certain optimizations that would not be possible if these static values were not static.

A later invocation reuses the specialized compilation in the first stage. Compilation happens again only when a required static value changes or no cached specialization exists. The idea is to compile once per set of static values, then execute many times with different dynamic values.

To summarize, CuTeDSL's role is to make the detailed mappings composable: once we choose an operation and compatible layouts, it derives the tensor views needed to feed that operation correctly so the programmer does not need to worry about deriving addresses.

Kernel 2: Tensor Memory, Tensor Cores, and Tiling

The code for this kernel can be found here.

The naive kernel 1 is extremely slow. This is primarily because each accumulation step is requiring 2 loads (one for matrix A and one for matrix B) from global memory, as well as one write back to global memory. Since loads/stores from global memory are extremely slow, we must leverage the GPU's memory hierarchy to have as many memory accesses as possible from faster kinds of memory and as few as possible from slow kinds of memory.

The first optimization we will make is called loop tiling. In loop tiling, we load a small chunk of both matrix A and matrix B from global memory to a smaller, faster memory such as shared memory, using it as a cache. Then, the threads will do the necessary processing for that block. Finally, once the threads finish, the current block is written back to global memory, and the next block is loaded into shared memory.

Kernel 2 pseudocode

The pseudocode below follows the control flow of the CuTeDSL implementation. Each CTA owns a 64 × 256 output tile and reduces over K in 64-element tiles.

kernel_2(A, B, C):
    # Each CTA computes one tile of C.
    shared_A, shared_B = allocate_shared_memory()
    tmem_C = allocate_tensor_memory()

    # Initialize the barrier to 0, which means that it is not ready.
    tma_barrier = 0
    mma_barrier = 0

    # Move across K, accumulating one pair of A and B tiles at a time.
    for each k_tile:
        # TMA asynchronously loads the next tiles into shared memory.
        TMA.load_async(A_tile, shared_A, signal=tma_barrier)
        TMA.load_async(B_tile, shared_B, signal=tma_barrier)
        wait(tma_barrier)

        # Tensor cores add all block products from these tiles into C in TMEM.
        for each matching block in shared_A and shared_B:
            MMA.async_accumulate(shared_A_block, shared_B_block, tmem_C)
        MMA.commit(signal=mma_barrier)
        wait(mma_barrier)

    # Move the completed tile out of TMEM and write it to global memory.
    registers = TMEM.load(tmem_C)
    C_tile = cast_BF16(registers)
    global_memory.store(C, C_tile)

    free_tensor_memory(tmem_C)

Tensor Memory Accelerator (TMA)

To load tiles from global memory into shared memory, we use the Tensor Memory Accelerator (TMA). TMA is a dedicated hardware engine introduced in the Hopper generation. It is located inside every SM and does the global memory to shared memory transfer asynchronously. In the TMA operation, we pass in the coordinates of the source tile, a TMA descriptor that tells the TMA the tile's base pointer, data type, dimensions, etc., and the destination. Essentially, we are telling the TMA the layout of the tile we want it to load and where to load it to. Then, a single thread can issue a single TMA instruction for an entire multidimensional tile, and the hardware would calculate the individual addresses and perform the bulk transfer.

# Issue the TMA loads for one A tile and one B tile.
cute.copy(
    tma_atom_a,
    tAgA[(None, k_tile_idx)],  # Global memory A tile
    tAsA[(None, 0)],           # Shared memory destination
    tma_bar_ptr=ab_mbar_ptr,
)
cute.copy(
    tma_atom_b,
    tBgB[(None, k_tile_idx)],  # Global memory B tile
    tBsB[(None, 0)],           # Shared memory destination
    tma_bar_ptr=ab_mbar_ptr,
)

However, since TMA transfers are asynchronous, issuing a load does not mean that its data is immediately safe to use. So, we use a memory barrier to ensure that all of the threads will wait on the barrier until the tile is completely copied into shared memory. On a lower-level, we have a parity phase (ab_mbar_ptr) for the TMA operations and a parity phase (mma_mbar_ptr) for the MMA operation. Each thread also has its own local copy of both parity variables-- tma_phase and mma_phase. For the TMA phase, we have a counter that keeps track of the number of bytes that arrives, and when it reaches our expected amount, completes and releases the waiting threads. Then, when the threads are released, they check if their local phase matches that of the barrier, and if so, it consumes the data from the TMA or MMA. Note that we have a 2-level barrier to make sure that the threads and the barrier are in sync, and that the threads will not wrongfully consume data produced by the next cycle.

Tensor Memory and tcgen05

One major improvement that was introduced in Blackwell GPUs is tensor memory (TMEM). TMEM is an on-chip memory designed to store tcgen05 operands and results, where tcgen05 is the MMA instruction for Blackwell. In TMEM, there are 128 lanes, each with 512 32-bit columns, so the total size is 256 KiB. The smallest allocation is 32 columns across all 128 lanes, or 16 KiB.

Before Blackwell, tensor-core MMA accumulators were distributed across registers owned by the threads participating in the MMA operation themselves. These accumulators competed for register space with other threads performing address calculations, control flow, and other general-purpose work, increasing register pressure. Accumulator ownership was also tied to a thread itself, which made it harder for other threads to access the contents of the a register. So, Blackwell introduces TMEM as a separate, CTA-scoped storage space designed specifically for these accumulators. The tcgen05.mma instruction can write results directly into TMEM.

# local_tile creates this CTA's global-memory views; it does not copy data.
gA = cute.local_tile(A, cta_tile, block_coordinate_for_A)
gB = cute.local_tile(B, cta_tile, block_coordinate_for_B)
gC = cute.local_tile(C, cta_tile, block_coordinate_for_C)

# TiledMma derives fragments with the layouts expected by the MMA atom.
A_fragments = tiled_mma.make_fragment_A(shared_A)
B_fragments = tiled_mma.make_fragment_B(shared_B)
C_fragment = tiled_mma.make_fragment_C(FP32_TMEM_tile)

# The first MMA initializes C; every later MMA accumulates into it.
tiled_mma.set(ACCUMULATE, False)

# Runtime loop: reduce over every 64-element K tile.
for k_tile_idx in cutlass.range(number_of_K_tiles):
    wait_until_A_and_B_are_in_shared_memory()

    # BK=64 BF16 values; each MMA consumes 32 bytes = 16 values, so 64 / 16 = 4.
    num_k_blocks = cute.size(A_fragments, mode=[2])  # Statically 4.
    for k_block_idx in cutlass.range_constexpr(num_k_blocks):
        k_block = (None, None, k_block_idx, 0)
        cute.gemm(
            tiled_mma,
            C_fragment,
            A_fragments[k_block],
            B_fragments[k_block],
            C_fragment,
        )
        tiled_mma.set(ACCUMULATE, True)

    commit_MMA_batch()
    wait_until_MMA_finishes()

# A tiled copy maps the completed TMEM fragment into registers.
cute.copy(tmem_to_register_copy, C_fragment, registers)
gC.store(registers.to(BF16))

At a high level, the three local_tile calls select the portions of A, B, and C owned by this CTA, while TiledMma describes how the operands should be presented to the tensor cores. The shared memory tile has BK = 64, meaning that the K-dimension of the tile is 64. However, the BF16 tcgen05 MMA consumes 32 bytes along K, or 16 BF16 values, so four MMA instructions are required to consume the complete tile. Each cute.gemm adds one of those four block products into the same FP32 C_fragment in TMEM. The first MMA initializes it, and every subsequent MMA accumulates into it across the full K dimension. The partial result therefore remains in TMEM throughout the reduction instead of moving through registers after every MMA.

Using the Per-Thread View to Write Back to Global Memory

Our final objective is to store the output of the tile back into global memory.

Once the reduction is complete, the CTA owns a 64 × 256 FP32 output tile in TMEM. The first step is to load pieces of that tile into each thread's registers. At that point, every thread has a collection of FP32 output values, but it still needs to write each value to the correct address in global C. The view constructed in this section is designed to solve that registers to global memory mapping, pairing every value in a thread's register fragment with its final destination in C.

CuTe builds this destination view from the layout of the TMEM load. For each 64 × 64 piece, Ld16x256bOp(x8) determines which accumulator values land in each thread's registers. partition_S describes those values on the TMEM side, while partition_D applies the same ordering to gC. The first register value therefore matches the first location in the thread's global, or gC, view, the second value matches the second location, and so on. The thread can then store its complete fragment without calculating a separate global address for every register. In the code, we use the flat_divide function to build this view for each 64 × 64 tile.

# View the 64 × 256 result as four matching 64 × 64 pieces.
epilogue_tile = sm100_utils.compute_epilogue_tile_shape(...)  # Returns 64 × 64.
global_C_subtiles = cute.flat_divide(
    CTA_global_C_view,
    epilogue_tile,
)
tmem_C_subtiles = cute.flat_divide(
    CTA_TMEM_C_view,
    epilogue_tile,
)

# Describe one warp's TMEM → register load.
# Each repetition loads 16 rows × 8 FP32 columns; x8 covers 16 × 64.
tmem_load_atom = cute.make_copy_atom(
    tcgen05.Ld16x256bOp(tcgen05.Repetition.x8),
    cutlass.Float32,
)
tiled_tmem_load = tcgen05.make_tmem_copy(
    tmem_load_atom,
    one_TMEM_subtile,
)

# Give this thread matching views on both sides of the transfer.
thread_load = tiled_tmem_load.get_slice(thread_id)

# The accumulator values that this thread loads from TMEM.
thread_TMEM_view = thread_load.partition_S(tmem_C_subtiles)

# The matching locations that this thread stores into global C.
thread_global_C_view = thread_load.partition_D(global_C_subtiles)

On a warp- and thread-level view, the mapping can be understood as a sequence of nested splits. First, the four warps divide the current 64 × 64 subtile into four horizontal 16 × 64 bands, so warp w handles rows 16w through 16w + 15. Within one warp's band, tcgen05.ld.16x256b processes a 16 × 8 slice at a time, which we can view as two 8 × 8 blocks stacked vertically. In each block, lane_id // 4 selects one of eight rows and lane_id % 4 selects one of four adjacent column pairs. Each lane receives its selected pair from both the upper block and the corresponding row of the lower block, giving it four FP32 values from the slice. Together, the 32 lanes cover the entire 16 × 8 slice. To cover the entire subtile, the x8 repetition shifts this same mapping across eight adjacent 16 × 8 slices, allowing each warp to cover one 16 × 64 band. As a result, the four warp bands then span all 64 × 64 elements.

More precisely, let t be the lane ID within a warp. Inside each 8-column group, lane t owns rows t // 4 and t // 4 + 8 of its warp band, and the adjacent column pair beginning at 2 * (t % 4). For example, lane 0 owns rows 0 and 8 with columns 0–1, while lane 6 owns rows 1 and 9 with columns 4–5. The same positions repeat in each of the eight column groups.

# Coordinates owned by one thread.
cta_row = 16 * warp_id + lane_id // 4 + 8 * row_half
cta_col = 64 * subtile_idx + 8 * column_group \
        + 2 * (lane_id % 4) + pair_element

global_row = block_m * 64 + cta_row
global_col = block_n * 256 + cta_col

# row_half and pair_element are 0 or 1; column_group is 0 through 7.

Once this view has been constructed, the registers to global memory writeback follows directly. For each subtile, the warp collectively loads its subtile values from TMEM into its register fragment, converts them to BF16, and stores them through the corresponding partition_D view. Although one thread's locations are spread across the subtile, the views of all 32 threads interleave to reconstruct the warp's complete 16 × 64 output band.

# Repeat for each of the four 64 × 64 output subtiles.
for each subtile:
    my_TMEM_values = thread_TMEM_view[subtile]
    my_global_C_locations = thread_global_C_view[subtile]

    # Step 1: load this thread's accumulator values into FP32 registers.
    cute.copy(tiled_TMEM_load, my_TMEM_values, register_fragment)

    # Step 2: convert those register values to the output type.
    bf16_values = register_fragment.load().to(BF16)

    # Step 3: write them to their matching addresses in global C.
    my_global_C_locations.store(bf16_values)

To summarize, we utilize matrix tiling to take advantage of the memory hierarchy, TMA to perform the global memory -> shared memory load asynchronously, detail how tensor memory is used, and go through how we store the output back into global memory. Timing this kernel, we get a value of 155 TFLOPS, which is a massive improvement over our naive kernel. However, this is still only 8.8% cuBLAS, and we have many more optimizations to go.

Kernel 3: Swizzling

The code for this kernel can be found here.

Note that in kernel 2, the consumer threads performing the MMA will make multiple shared memory reads simultaneously. Intuitively, this could cause performance problems. Shared memory is divided up into 32 banks, and we could think of each bank as an independent SRAM structure. However, each bank only has one read port with a size of 4 bytes (note that adding more read ports would be expensive from a hardware perspective). As a result, we could only service one read per clock cycle. When multiple read or write requests (e.g. 2 threads from the same warp) come into the same shared memory bank, a bank conflict occurs, and we will take multiple cycles to service the reads. Obviously, to improve performance, we want to minimize bank conflicts as much as possible.

Note that these bank conflicts happen very often in kernel 2 due to the tile shape. Each shared memory row holds 64 BF16 values, which is exactly 32 four-byte words, or one complete trip through the 32 banks. The next row begins exactly 32 words later, so each index in the first row will coincide with the same index in the second row. Since an MMA needs the same K positions from multiple rows of A or B at once, the unswizzled layout naturally will send some concurrent fetches to the same bank.

To avoid bank conflicts, we use swizzling. Swizzling removes these conflicts by changing the physical shared memory address associated with each logical tile coordinate. Without swizzling, the bank pattern repeats in every row of 32-bit words. Swizzling makes the bank depend on both the row and word column, shifting this pattern between rows. Hardware performs that shuffle by XORing selected bits from the row with the address bits that normally choose the bank. In the simplified example below, four words from word column 0 therefore remain in the same logical column but are physically stored in Banks 0 through 3. Each illustrated word contains two adjacent BF16 values.

Kernel 3 applies this optimization almost entirely through a layout change rather than a new data-movement step. make_smem_layout_a and make_smem_layout_b select the 128-byte K_SW128 layout. When sA and sB are created, CuTe's built-in smem.allocate_tensor method accepts the swizzle through its swizzle= argument, so the kernel does not manually calculate XOR-adjusted addresses. The same layout is passed into the TMA construction, allowing TMA to write the tile in the correct physical arrangement and MMA to read it with a matching descriptor. Consequently, the main-loop cute.copy and cute.gemm calls remain unchanged, and no separate rearrangement or “unswizzle” instruction is required.

With swizzling, our new kernel achieves 285 TFLOPS, almost doubling the performance from kernel 2.

Part 2 of this blog will be coming soon!

Acknowledgements

I would like to thank Han Guo for his help in the making of this blog. I would also like to thank to thank the team at Modular. The high-level structure of this blog was inspired by their blog.