Matmul on sm_120 — Pranav Deepak

Intro

Assumed: NVIDIA GPU basics (warps, blocks, smem, tensor cores, roofline). Modal’s GPU glossary if you need a primer.

The tiled view

The grade school matrix multiply says: $A \in \mathbb{R}^{M \times K}$ , $B \in \mathbb{R}^{K \times N}$ , $C \in \mathbb{R}^{M \times N}$ , $C_{ij} = \sum_{k=0}^{K} A_{ik} B_{kj}$ . Pick $tM, tN, tK$ that divide $M, N, K$ , and tile $A, B, C$ into $A', B', C'$ where each element of $A'$ is a $tM \times tK$ tile of $A$ (same for $B', C'$ ). Then $C'_{ij} = \sum_{k=0}^{K/tK} A'_{ik} \, @ \, B'_{kj}$ computes the original matmul, with the scalar product $A_{ik} B_{kj}$ replaced by the tile matmul $A'_{ik} \, @ \, B'_{kj}$ , and the sum of two matrix tiles taken element-wise.

Assign one worker per $C'$ tile. The worker computing $C'_{00}$ does two tile matmuls and one element-wise add: $A'_{00} \, @ \, B'_{00} + A'_{01} \, @ \, B'_{10}$ . So: concurrent workers own tiles of $C$ , each streams $A, B$ tiles along $K$ . When writing square-ish matmuls for GPUs, we use this form of tiling recursively.

Roofline

How fast can we go? In a magical world, you’d still have to read each element of $A$ and $B$ at least once, and write each element of $C$ once. Even if compute were instantaneous, you’d still have to move $2K(M+N) + 4MN$ bytes ( $A, B$ are bf16, $C$ is fp32), which takes $t = \frac{2K(M+N) + 4MN}{B}$ seconds, where $B$ is peak memory bandwidth in bytes/second. The number of floating point operations is $2MNK$ , so:

$\text{FLOP/s} = \frac{2MNK}{t} = \frac{MNK \cdot B}{K(M+N) + 2MN}$

At $M = N = K = 8192$ on the RTX 5090 ( $B = 1.79$ TB/s), this ceiling is about $3667$ TFLOP/s. The 5090’s BF16 tensor core peak is $210$ TFLOP/s, so at this size we’re nowhere near memory-bound. The game is feeding the tensor cores.

In essence, matmuls give us $O(N^3)$ floating point operations to try and hide $O(N^2)$ bytes of memory traffic.

We’re not in magic land. Elements get reused: each element of $A$ participates in $2N$ ops, each element of $B$ in $2M$ ops (roughly). Pay the GMEM-to-shared transfer once per element, read it $O(N)$ (or $O(M)$ ) times at shared-memory speed.

Instructions and machinery

sm_120’s tensor core instruction is Ampere-style synchronous warp-level mma. Each warp holds fragments of $A$ , $B$ , and the accumulator in its registers. We use mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32: a $16 \times 8 \times 16$ ( $m \times n \times k$ ) matmul, $A, B$ in bf16, accumulator in fp32, $A$ row-major, $B$ column-major. Per thread: 8 elements of $A$ in 4 registers, 4 elements of $B$ in 2 registers, 4 elements of $C$ in 4 registers. The layouts:

$A$ ‘s fragment: treat the warp’s 32 lanes as an $(8, 4)$ row-major grid. Each lane holds a 2-vector of bf16 in one 32-bit register, so the warp holds an $(8, 8)$ tile. Each lane has 4 such registers, stacked in the order above, giving a $(16, 16)$ tile overall.

$B$ ’s fragment: lanes arranged as $(4, 8)$ col-major. Each lane holds a 2-vector of bf16 in one 32-bit register, so the warp holds an $(8, 8)$ tile. Each lane has 2 such registers, stacked one below the other, giving a $(16, 8)$ tile.

$C$ ’s fragment: lanes arranged as $(8, 4)$ row-major. Each lane’s reg0, reg1 together hold a 2-vector in the top $(8, 8)$ tile, and reg2, reg3 hold the 2-vector in the $(8, 8)$ tile directly below, giving a $(16, 8)$ output tile.

To load bf16 elements from shared memory into these register layouts, we use ldmatrix.m8n8.x{1,2,4}. The x1 variant: the first 8 lanes supply start addresses of 8 contiguous 8-vectors (the 8-vectors themselves must be contiguous in shared memory, but the 8 start addresses need not be). ldmatrix reads each 8-vector in lane-supplied order and distributes 2 elements into each lane’s 32-bit register.

The x2 and x4 variants extend this: x2 takes 16 lane-supplied addresses and fills 2 registers per lane, x4 takes all 32 and fills 4. Each group of 8 lanes’ addresses feeds one register.

We use the Tensor Memory Accelerator (TMA) for async global-to-shared loads, and mbarriers to synchronize. There are already great writeups on both (Hopper matmul blogs cover them well), so I won’t rehash here.

Thanks to gau-nernst for help throughout the sm_120 journey. I stole his PTX wrappers (here), plus some niceties: tensor map creation, kernel launcher, swizzling utilities.

Common kernel structure

Each kernel has a config file for compile-time constants, and every kernel shares the same basic structure:

Each block computes a (BM, K) @ (K, BN) -> (BM, BN) matmul, streaming K in BK tiles.

2. The problem reduces to a stream of (BM, BK) @ (BK, BN) shared-memory matmuls (As, Bs), with the (BM, BN) result stored register-to-gmem directly.

Each warp computes a (WM, BK) @ (BK, WN) matmul, streaming BK in mma_k = 16 tiles. It exploits outer-product structure: each warp loads an acc_per_warp_m-sized column of As (elements are (mma_m, mma_k) tiles) and an acc_per_warp_n-sized row of Bs (elements are (mma_k, mma_n) tiles), then issues acc_per_warp_m * acc_per_warp_n mma instructions into independent accumulator registers. The scoreboard sees the mmas are independent and pipelines them, giving instruction-level parallelism. So WM = mma_m * acc_per_warp_m and WN = mma_n * acc_per_warp_n.

Each block launches warps_per_block_m * warps_per_block_n warps, covering the full (BM, BN) tile.

I captured these common ideas into a struct, pasted here as-is:

#pragma once 
#include "../clean_abstractions/all.cuh"

template <typename Cfg>
struct GemmMaker
{
  const int lane_id;
  const int warp_id;
  const int thread_id;
  const CUtensorMap* tma_a;
  const CUtensorMap* tma_b;
  int c2_row_start;
  int c2_col_start;
  uint32_t a_s2r_off;
  uint32_t b_s2r_off;

  __device__ __forceinline__
  GemmMaker(const CUtensorMap& tma_a_, const CUtensorMap& tma_b_, int tid)
    : tma_a(&tma_a_), tma_b(&tma_b_),
      thread_id(tid),
      lane_id(tid % 32),
      warp_id(tid / 32)
  {
    int warp_start_m = (warp_id / Cfg::warps_per_block_n) * Cfg::WM;
    int warp_start_n = (warp_id % Cfg::warps_per_block_n) * Cfg::WN;

    a_s2r_off = ((warp_start_m + (lane_id % 16)) * Cfg::BK + (8 * (lane_id / 16))) * sizeof(nv_bfloat16);
    b_s2r_off = ((warp_start_n + ((lane_id % 8) + (8 * (lane_id / 16)))) * Cfg::BK + (8 * ((lane_id / 8) % 2))) * sizeof(nv_bfloat16);
    c2_row_start = warp_start_m + (lane_id/4);
    c2_col_start = warp_start_n + 2*(lane_id%4);
  }

  __device__ __forceinline__
  void load_A_g2s(int bk_idx, int block_start_m, uint32_t As_addr, uint32_t mbar_addr)
  {
    cp_async_bulk_tensor_2d(As_addr, tma_a, bk_idx * Cfg::BK, block_start_m, mbar_addr);
  }

  __device__ __forceinline__
  void load_B_g2s(int bk_idx, int block_start_n, uint32_t Bs_addr, uint32_t mbar_addr)
  {
    cp_async_bulk_tensor_2d(Bs_addr, tma_b, bk_idx * Cfg::BK, block_start_n, mbar_addr);
  }

  __device__ __forceinline__
  void store_C(float2* C2, float rc[Cfg::acc_per_warp_m][Cfg::acc_per_warp_n][4], int block_start_m, int block_start_n)
  {
    static constexpr int ldc2 = Cfg::N / 2;
    int C_row_start = block_start_m + c2_row_start;
    int C_col_start = block_start_n + c2_col_start;

    #pragma unroll
    for (int m = 0; m < Cfg::acc_per_warp_m; m++)
    {
      #pragma unroll
      for (int n = 0; n < Cfg::acc_per_warp_n; n++)
      {
        int C_row = C_row_start + (m * Cfg::mma_m);
        int C_col = (C_col_start + (n * Cfg::mma_n)) / 2;
        C2[C_row       * ldc2 + C_col] = {rc[m][n][0], rc[m][n][1]};
        C2[(C_row + 8) * ldc2 + C_col] = {rc[m][n][2], rc[m][n][3]};
      }
    }
  }

  __device__ __forceinline__
  void load_A_s2r(uint32_t ra[Cfg::acc_per_warp_m][4], uint32_t As_addr, int k)
  {
    #pragma unroll
    for (int m = 0; m < Cfg::acc_per_warp_m; m++)
    {
      uint32_t addr = As_addr + compact_swizzle<Cfg::swizzle_num>(
        a_s2r_off + ((m * Cfg::mma_m * Cfg::BK) + (k * Cfg::mma_k)) * sizeof(nv_bfloat16));
      ldmatrix_m8n8_x4_b16(ra[m], addr);
    }
  }

  __device__ __forceinline__
  void load_B_s2r(uint32_t rb[Cfg::acc_per_warp_n][2], uint32_t Bs_addr, int k)
  {
    #pragma unroll
    for (int n = 0; n < Cfg::acc_per_warp_n / 2; n++)
    {
      uint32_t addr = Bs_addr + compact_swizzle<Cfg::swizzle_num>(
        b_s2r_off + ((2 * n * Cfg::mma_n * Cfg::BK) + (k * Cfg::mma_k)) * sizeof(nv_bfloat16));
      ldmatrix_m8n8_x4_b16(rb[2 * n], addr);
    }
  }

  __device__ __forceinline__
  void mma(float rc[Cfg::acc_per_warp_m][Cfg::acc_per_warp_n][4],
          uint32_t ra[Cfg::acc_per_warp_m][4],
          uint32_t rb[Cfg::acc_per_warp_n][2])
  {
    #pragma unroll
    for (int m = 0; m < Cfg::acc_per_warp_m; m++)
    {
      #pragma unroll
      for (int n = 0; n < Cfg::acc_per_warp_n; n++)
      {
        mma_m16n8k16_row_col_f32_bf16(rc[m][n], ra[m], rb[n]);
      }
    }
  }
};

Since acc_per_warp_n is always at least 2, I fuse two (16, 8) loads into one ldmatrix.x4 of (16, 16) instead of issuing two ldmatrix.x2s. gau-nernst’s FA-5090 blog notes this matters. He fuses along k; I fuse along n because my later kernels pipeline the smem-to-rmem loads along k, and fusing along the pipelined axis would conflict with that.

Kernel 0: Naive

Note: torch.matmul launches a CUTLASS Ampere kernel on sm_120. Since clock rates and power envelopes vary, I report TFLOP/s alongside percentage against torch.matmul. All benchmarks use triton.testing.do_bench with default params and median return mode.

Kernel 0 comes in pretty straightforward:

template <class Cfg>
__global__ void matmul_kernel(
    __grid_constant__ const CUtensorMap a_map,
    __grid_constant__ const CUtensorMap b_map,
    float* C
)
{
  GemmMaker<Cfg> gemm(a_map, b_map, threadIdx.x);
  extern __shared__ __align__(1024) uint8_t smem_raw[];
  int b = blockIdx.x; 
  int t = threadIdx.x; 
  int l = t % 32;
  int w = t / 32; 
  int block_start_m = (b / Cfg::GN) * Cfg::BM; 
  int block_start_n = (b % Cfg::GN) * Cfg::BN; 
  float2 *C2 = reinterpret_cast<float2*>(C); 

  uint32_t As = static_cast<uint32_t>(__cvta_generic_to_shared(smem_raw));
  uint32_t Bs = As + Cfg::As_bytes;
  uint32_t m_bar = Bs + Cfg::Bs_bytes; 

  uint32_t ra[Cfg::acc_per_warp_m][4];
  uint32_t rb[Cfg::acc_per_warp_n][2];
  float rc[Cfg::acc_per_warp_m][Cfg::acc_per_warp_n][4] = {0.0};

  if (t == 0)
  {
    mbarrier_init(m_bar,32);
  }
  asm volatile("fence.mbarrier_init.release.cluster;");
  __syncthreads();

  int parity = 0; 
  for (int bk_idx = 0; bk_idx < Cfg::block_k_iters; bk_idx ++)
  {
    __syncthreads();
    if (w == 0)
    { 
      if (l == 0)
      {
        mbarrier_arrive_expect_tx(m_bar, Cfg::As_bytes + Cfg::Bs_bytes); 
        gemm.load_A_g2s(bk_idx, block_start_m, As, m_bar); 
        gemm.load_B_g2s(bk_idx, block_start_n, Bs, m_bar);
      }
      else
      {
        mbarrier_arrive(m_bar);
      }
    }
    __syncthreads(); 
    mbarrier_wait_parity(m_bar, parity); 

    for (int wk_idx = 0; wk_idx < Cfg::warp_k_iters; wk_idx++)
    {
      gemm.load_A_s2r(ra,As,wk_idx);
      gemm.load_B_s2r(rb,Bs,wk_idx); 
      gemm.mma(rc,ra,rb);
    }
    parity^=1;
  }
  __syncthreads();
  gemm.store_C(C2,rc,block_start_m,block_start_n);
}

K0: 194 TFLOP/s on 5090 (-14.5% vs torch), 250 TFLOP/s on 6000 Pro (-42.4% vs torch).

Kernel 1: Swizzle

The kernel body is unchanged. We turn on two swizzles via config.

TMA swizzle. We use CUTLASS-style <b_bits, m_base, s_shift> descriptors to mitigate bank conflicts during ldmatrix:

32B load dim → (1, 4, 3)
64B load dim → (2, 4, 3)
128B load dim → (3, 4, 3)

Lei Mao’s swizzle post and CuTe swizzle post are good primers on bank conflicts and swizzling, so I won’t rehash.

A friend recently asked me: if there are 32 banks of 32 bits each, 32 lanes, and each lane loads more than 32 bits, then by pigeonhole at least two lanes land in the same bank, so won’t that serialize and hurt perf? Relevant here since each lane’s ldmatrix loads an 8-element bf16 vector (128 bits). This post explains why it doesn’t; I originally found the reasoning in this microbenchmark blog, which explains the interleaving across load orderings:

bank conflict interleaving

Block swizzle. Kernel 0’s block-index-to-tile mapping is plain row-major:

int block_start_m = (b / Cfg::GN) * Cfg::BM;
int block_start_n = (b % Cfg::GN) * Cfg::BN;

Say two blocks co-schedule on one SM, and the 5090 has 170 SMs, so ~340 blocks run concurrently. With row-major mapping on a (say) 64×64 grid of (BM, BN) tiles, those 340 blocks sweep many full rows of C, which means collectively touching the entire B matrix, likely too big for L2. Column-major would do the same with A. Neither is great.

Triton-style block swizzle fixes this: tile the grid into chunks, iterate chunks column-major, iterate tiles within a chunk row-major. The active blocks now cover a compact (group_m * chunk_height, group_n * chunk_width) region, reusing both A and B rows/cols heavily in L2.

So the new block-to-tile map:

template<int group_m, int group_n, int blocks_per_group,
         int g_outer_m, int g_outer_n, int BM, int BN>
__device__ __forceinline__
void block_swizzle(int b, int &block_start_m, int &block_start_n) {
  int group_id = b / blocks_per_group;
  int local_id = b % blocks_per_group;

  int global_m = group_id % g_outer_m;
  int global_n = group_id / g_outer_m;

  int local_m = local_id / group_n;
  int local_n = local_id % group_n;

  int tile_m = global_m * group_m + local_m;
  int tile_n = global_n * group_n + local_n;

  block_start_m = tile_m * BM;
  block_start_n = tile_n * BN;
}

K1: 214 TFLOP/s on 5090 (-6.0% vs torch), 339 TFLOP/s on 6000 Pro (-21.8% vs torch).

Kernel 2: Warp-specialized

Standard warp-specialized layout. Launch one extra warp; the last warp is the producer (fires TMA loads, g2s) and all other warps are consumers (ldmatrix, mma). Empty/full mbarrier handshake: consumers signal empty after consuming, wait on full to consume; producer waits on empty, signals full on tx complete.

One note: the usual pattern starts the producer by signalling all buffers empty. I use gau-nernst’s trick instead: start the producer’s wait_parity at 1. Same result, no perf difference, just cleaner.

template <class Cfg>
__global__ void matmul_kernel(
    __grid_constant__ const CUtensorMap a_map,
    __grid_constant__ const CUtensorMap b_map,
    float* C
)
{
  GemmMaker<Cfg> gemm(a_map, b_map, threadIdx.x);
  extern __shared__ __align__(1024) uint8_t smem_raw[];
  int b = blockIdx.x; 
  int t = threadIdx.x; 
  int l = t % 32;
  int w = t / 32; 
  int block_start_m; 
  int block_start_n; 
  block_swizzle<Cfg::group_m,Cfg::group_n,Cfg::blocks_per_group,Cfg::G_outer_M, Cfg::G_outer_N, Cfg::BM, Cfg::BN>(b,block_start_m,block_start_n);
  float2 *C2 = reinterpret_cast<float2*>(C); 

  uint32_t As_base = static_cast<uint32_t>(__cvta_generic_to_shared(smem_raw));
  uint32_t Bs_base = As_base + (Cfg::As_bytes*Cfg::bk_stages);
  uint32_t empty_bar_base = Bs_base + (Cfg::Bs_bytes*Cfg::bk_stages);
  uint32_t full_bar_base = empty_bar_base + (8*Cfg::bk_stages); 

  auto As        = [&](int s) { return As_base        + s * Cfg::As_bytes; };
  auto Bs        = [&](int s) { return Bs_base        + s * Cfg::Bs_bytes; };
  auto empty_bar = [&](int s) { return empty_bar_base + s * 8; };
  auto full_bar  = [&](int s) { return full_bar_base  + s * 8; };


  if (t == 0)
  {
    for (int s = 0; s < Cfg::bk_stages; s++)
    {
      mbarrier_init(empty_bar(s),Cfg::warps_per_block_m*Cfg::warps_per_block_n*32);
      mbarrier_init(full_bar(s),32);
    }
  }
  asm volatile("fence.mbarrier_init.release.cluster;");
  __syncthreads();

  
  if (w == Cfg::dma_warp_id)
  {
    int producer_parity = 1; 
    int stage = 0; 
    for (int bk_idx = 0; bk_idx < Cfg::block_k_iters; bk_idx++)
    {
      mbarrier_wait_parity(empty_bar(stage),producer_parity);
      if(l == 0)
      {
        mbarrier_arrive_expect_tx(full_bar(stage),Cfg::As_bytes + Cfg::Bs_bytes);
        gemm.load_A_g2s(bk_idx, block_start_m, As(stage), full_bar(stage));
        gemm.load_B_g2s(bk_idx, block_start_n, Bs(stage), full_bar(stage)); 
      }
      else
      {
        mbarrier_arrive(full_bar(stage));
      }
      stage = (stage + 1) % Cfg::bk_stages; 
      if (stage == 0) producer_parity ^= 1;
    }
  }
  else 
  {
    int consumer_parity = 0; 
    int stage = 0; 
    uint32_t ra[Cfg::acc_per_warp_m][4];
    uint32_t rb[Cfg::acc_per_warp_n][2];
    float rc[Cfg::acc_per_warp_m][Cfg::acc_per_warp_n][4] = {0.0};
    for (int bk_idx = 0; bk_idx < Cfg::block_k_iters; bk_idx++)
    {
      mbarrier_wait_parity(full_bar(stage), consumer_parity); 
      for (int wk_idx = 0; wk_idx < Cfg::warp_k_iters; wk_idx++)
      {
        gemm.load_A_s2r(ra,As(stage),wk_idx);
        gemm.load_B_s2r(rb,Bs(stage),wk_idx); 
        gemm.mma(rc,ra,rb);
      }
      mbarrier_arrive(empty_bar(stage));
      stage = (stage + 1) % Cfg::bk_stages; 
      if (stage == 0) consumer_parity ^= 1;
    }
    sync_bar<Cfg::warps_per_block_m*Cfg::warps_per_block_n*32>(); 
    gemm.store_C(C2, rc, block_start_m,block_start_n);
  }

}

K2: 230 TFLOP/s on 5090 (+0.8% vs torch), 423 TFLOP/s on 6000 Pro (-2.6% vs torch). First kernel to beat torch on 5090.

Kernel 3: Fused nested pipeline

My favorite. No warp specialization. Instead, standard Ampere-style pipelining at both levels: global-to-shared (outer, over BK tiles) and shared-to-register (inner, over mma_k = 16 tiles). With naive nested pipelining, you pay the inner pipeline’s prologue and epilogue every outer iteration. The idea from the Triton pipelining talk: fuse them.

I derived the fused logic by hand on an odd instance: 10 outer jobs with 3 outer buffers, 7 inner jobs per outer with 4 inner buffers. I picked non-power-of-2 numbers deliberately pow2 everywhere lets algebraic errors hide in modular arithmetic.

below is unfused diagram (rough) below is fused diagram I strongly suggest the reader to draw something like this, with color coded buffers to derive this fusion, I am too lazy to explain this in great detail.

Rough shape:

Prologue. Fire the outer pipeline’s prologue loads, then the inner pipeline’s prologue loads off the first outer buffer.
Steady state. Inner loads run ahead of inner computes. As inner compute drains the current outer buffer, fire (a) the next outer buffer’s TMA and (b) the next inner loads on the upcoming outer buffer. The inner pipeline straddles the outer boundary.
Drain (no new outer). Same as steady state, but stop issuing outer TMAs.
Drain inner. Finish the inner compute on the last outer buffer.

template <class Cfg>
__global__ void matmul_kernel(
    __grid_constant__ const CUtensorMap a_map,
    __grid_constant__ const CUtensorMap b_map,
    float* C
)
{
  GemmMaker<Cfg> gemm(a_map, b_map, threadIdx.x);
  extern __shared__ __align__(1024) uint8_t smem_raw[];
  int b = blockIdx.x;
  int t = threadIdx.x;
  int l = t % 32;
  int w = t / 32;
  int block_start_m;
  int block_start_n;
  block_swizzle<Cfg::group_m, Cfg::group_n, Cfg::blocks_per_group,
                Cfg::G_outer_M, Cfg::G_outer_N, Cfg::BM, Cfg::BN>
               (b, block_start_m, block_start_n);
  float2* C2 = reinterpret_cast<float2*>(C);

  uint32_t As_base   = static_cast<uint32_t>(__cvta_generic_to_shared(smem_raw));
  uint32_t Bs_base   = As_base + Cfg::As_bytes * Cfg::bk_stages;
  uint32_t full_base = Bs_base + Cfg::Bs_bytes * Cfg::bk_stages;

  auto As       = [&](int s) { return As_base   + s * Cfg::As_bytes; };
  auto Bs       = [&](int s) { return Bs_base   + s * Cfg::Bs_bytes; };
  auto full_bar = [&](int s) { return full_base + s * 8; };

  uint32_t ra[Cfg::wk_stages][Cfg::acc_per_warp_m][4];
  uint32_t rb[Cfg::wk_stages][Cfg::acc_per_warp_n][2];
  float rc[Cfg::acc_per_warp_m][Cfg::acc_per_warp_n][4] = {0.0};

  // init barriers
  if (t == 0) {
    for (int s = 0; s < Cfg::bk_stages; s++) mbarrier_init(full_bar(s), 32);
  }
  asm volatile("fence.mbarrier_init.release.cluster;");
  __syncthreads();

  // prime outer bk pipeline: fire all bk_stages TMAs upfront
  auto issue_tma = [&](int bk_idx, int stage) {
    if (l == 0) {
      mbarrier_arrive_expect_tx(full_bar(stage), Cfg::As_bytes + Cfg::Bs_bytes);
      gemm.load_A_g2s(bk_idx, block_start_m, As(stage), full_bar(stage));
      gemm.load_B_g2s(bk_idx, block_start_n, Bs(stage), full_bar(stage));
    } else {
      mbarrier_arrive(full_bar(stage));
    }
  };

  #pragma unroll
  for (int s = 0; s < Cfg::bk_stages; s++) {
    if (w == 0) issue_tma(s, s);
  }
  mbarrier_wait_parity(full_bar(0), 0);

  // prime inner wk pipeline on bk stage 0: (wk_stages - 1) ldmatrix load
  #pragma unroll
  for (int i = 0; i < Cfg::wk_stages - 1; i++) {
    gemm.load_A_s2r(ra[i], As(0), i);
    gemm.load_B_s2r(rb[i], Bs(0), i);
  }

  static constexpr int full_bk_iters = Cfg::block_k_iters - Cfg::bk_stages;
  static constexpr int wk_iters      = Cfg::warp_k_iters  - (Cfg::wk_stages - 1);

  // steady state: TMA + wk-pipeline fused, wk crosses bk boundary 
  for (int bk_idx = 0; bk_idx < full_bk_iters; bk_idx++) {
    int bk_cons_stage      = bk_idx % Cfg::bk_stages;
    int next_bk_cons_stage = (bk_idx + 1) % Cfg::bk_stages;
    int parity             = ((bk_idx + 1) / Cfg::bk_stages) % 2;
    int next_bk_load_idx   = bk_idx + Cfg::bk_stages;
    int next_bk_load_stage = next_bk_load_idx % Cfg::bk_stages;
    int bk_base            = bk_idx * Cfg::warp_k_iters;

    // phase 1: inner wk-pipe on current bk_cons_stage
    for (int wk_idx = 0; wk_idx < wk_iters; wk_idx++) {
      int wk_load_idx     = (wk_idx + (Cfg::wk_stages - 1)) % Cfg::warp_k_iters;
      int wk_load_stage   = (bk_base + wk_load_idx) % Cfg::wk_stages;
      int wk_compute_stage= (bk_base + wk_idx)      % Cfg::wk_stages;

      gemm.load_A_s2r(ra[wk_load_stage], As(bk_cons_stage), wk_load_idx);
      gemm.load_B_s2r(rb[wk_load_stage], Bs(bk_cons_stage), wk_load_idx);
      gemm.mma(rc, ra[wk_compute_stage], rb[wk_compute_stage]);
    }

    __syncthreads();                                   // loads of curr bk done
    if (w == 0) issue_tma(next_bk_load_idx, next_bk_load_stage);
    mbarrier_wait_parity(full_bar(next_bk_cons_stage), parity);
    __syncthreads();

    // phase 2: inner wk-pipe crosses onto next_bk_cons_stage
    for (int wk_idx = wk_iters; wk_idx < Cfg::warp_k_iters; wk_idx++) {
      int wk_load_idx     = (wk_idx + (Cfg::wk_stages - 1)) % Cfg::warp_k_iters;
      int wk_load_stage   = (bk_base + wk_load_idx) % Cfg::wk_stages;
      int wk_compute_stage= (bk_base + wk_idx)      % Cfg::wk_stages;

      gemm.load_A_s2r(ra[wk_load_stage], As(next_bk_cons_stage), wk_load_idx);
      gemm.load_B_s2r(rb[wk_load_stage], Bs(next_bk_cons_stage), wk_load_idx);
      gemm.mma(rc, ra[wk_compute_stage], rb[wk_compute_stage]);
    }
  }

  //drain: same structure as steady state, no more TMAs issued 
  static constexpr int no_tma_end = full_bk_iters + (Cfg::bk_stages - 1);

  for (int bk_idx = full_bk_iters; bk_idx < no_tma_end; bk_idx++) {
    int bk_cons_stage      = bk_idx % Cfg::bk_stages;
    int next_bk_cons_stage = (bk_idx + 1) % Cfg::bk_stages;
    int parity             = ((bk_idx + 1) / Cfg::bk_stages) % 2;
    int bk_base            = bk_idx * Cfg::warp_k_iters;

    for (int wk_idx = 0; wk_idx < wk_iters; wk_idx++) {
      int wk_load_idx     = (wk_idx + (Cfg::wk_stages - 1)) % Cfg::warp_k_iters;
      int wk_load_stage   = (bk_base + wk_load_idx) % Cfg::wk_stages;
      int wk_compute_stage= (bk_base + wk_idx)      % Cfg::wk_stages;

      gemm.load_A_s2r(ra[wk_load_stage], As(bk_cons_stage), wk_load_idx);
      gemm.load_B_s2r(rb[wk_load_stage], Bs(bk_cons_stage), wk_load_idx);
      gemm.mma(rc, ra[wk_compute_stage], rb[wk_compute_stage]);
    }

    __syncthreads();
    mbarrier_wait_parity(full_bar(next_bk_cons_stage), parity);
    __syncthreads();

    for (int wk_idx = wk_iters; wk_idx < Cfg::warp_k_iters; wk_idx++) {
      int wk_load_idx     = (wk_idx + (Cfg::wk_stages - 1)) % Cfg::warp_k_iters;
      int wk_load_stage   = (bk_base + wk_load_idx) % Cfg::wk_stages;
      int wk_compute_stage= (bk_base + wk_idx)      % Cfg::wk_stages;

      gemm.load_A_s2r(ra[wk_load_stage], As(next_bk_cons_stage), wk_load_idx);
      gemm.load_B_s2r(rb[wk_load_stage], Bs(next_bk_cons_stage), wk_load_idx);
      gemm.mma(rc, ra[wk_compute_stage], rb[wk_compute_stage]);
    }
  }

  //  epilogue: last bk stage, no next bk to cross into 
  static constexpr int bk_idx         = Cfg::block_k_iters - 1;
  static constexpr int bk_cons_stage  = bk_idx % Cfg::bk_stages;
  static constexpr int bk_base        = bk_idx * Cfg::warp_k_iters;

  for (int wk_idx = 0; wk_idx < wk_iters; wk_idx++) {
    int wk_load_idx     = (wk_idx + (Cfg::wk_stages - 1)) % Cfg::warp_k_iters;
    int wk_load_stage   = (bk_base + wk_load_idx) % Cfg::wk_stages;
    int wk_compute_stage= (bk_base + wk_idx)      % Cfg::wk_stages;

    gemm.load_A_s2r(ra[wk_load_stage], As(bk_cons_stage), wk_load_idx);
    gemm.load_B_s2r(rb[wk_load_stage], Bs(bk_cons_stage), wk_load_idx);
    gemm.mma(rc, ra[wk_compute_stage], rb[wk_compute_stage]);
  }
  __syncthreads();

  for (int wk_idx = wk_iters; wk_idx < Cfg::warp_k_iters; wk_idx++) {
    int wk_compute_stage = (bk_base + wk_idx) % Cfg::wk_stages;
    gemm.mma(rc, ra[wk_compute_stage], rb[wk_compute_stage]);
  }

  __syncthreads();   // absolutely needed before store
  gemm.store_C(C2, rc, block_start_m, block_start_n);
}

Honestly: this kernel was just fun to write, and it made me understand pipelining way better. Whether the fusion itself actually gave perf gains is unclear. It should, but I didn’t verify :)

K3: 235 TFLOP/s on 5090 (+4.1% vs torch), 425 TFLOP/s on 6000 Pro (-2.4% vs torch).

Perf table

kernel	5090 TFLOP/s	5090 vs torch	6000 Pro TFLOP/s	6000 Pro vs torch
K0 naive	194.24	-14.5%	249.94	-42.4%
K1 swizzle	214.43	-6.0%	338.88	-21.8%
K2 warp-spec	229.78	+0.8%	422.56	-2.6%
K3 fused	235.01	+4.1%	425.24	-2.4%
`torch.matmul` (`cutlass_80_tensorop_bf16_*`)	228.00	—	435.59	—

All numbers: $M = N = K = 8192$ , bf16 inputs, fp32 accumulate, median of triton.testing.do_bench.

Notes on autotuning

My autotuning is not perfect. I brute-force over reasonable parameter combinations and got Claude to do back-of-the-envelope register-count estimates to prune dead configs before nvcc. I haven’t experimented with __launch_bounds__; since each block launches a small number of warps, the compiler might allocate more registers per thread, which could support larger acc_per_warp_m / acc_per_warp_n. Worth trying.

One hard constraint: TMA doesn’t allow copying more than 256 elements along the non-leading dimension, which caps BM and BN at 256.

Notes on persistence (and what’s coming next)

I did try a persistent kernel with a static Morton-based tile scheduler, but it only worked for square tiles and I didn’t autotune it. There also isn’t enough shared memory on sm_120 to justify storing C to shared and firing async stores. Overlapping C stores with the next tile’s A/B prologue has the most bang for buck when the C store is async. You can fire the next tile’s A/B TMA load and let it run in the background while doing a standard C store, but I haven’t gotten it to be fast.

Proper persistence, along with a non-square-grid scheduler, is part 2.