From 1% to 94% of cuBLAS, one kernel at a time

A visual walkthrough of Simon Boehm's CUDA matrix multiplication worklog. Each kernel fixes one specific bottleneck. This page explains what that bottleneck is, what the fix looks like on the hardware, and why the number moves. It starts from zero: threads, warps, blocks, grids, and the memory hierarchy are all explained before the first kernel.

Source article: siboehm.com/articles/22/CUDA-MMM (December 2022). Code: github.com/siboehm/SGEMM_CUDA. All throughput numbers are the article's measurements on an NVIDIA RTX A6000 multiplying two 4092 x 4092 fp32 matrices. Diagrams and text on this page are original; code sketches are simplified rewrites, not the repository's code.

The whole story in one chart. GFLOP/s achieved by each kernel, as a fraction of cuBLAS. Click a bar to jump to that kernel. Kernels 7 and 8 are missing because the article dropped them (they fixed shared memory bank conflicts but ran slower overall).

The shape of the climb is the lesson. The first three kernels are about not wasting memory bandwidth: reading global memory in the pattern the hardware wants, then caching in shared memory. Kernels 4 to 6 are about arithmetic intensity: doing more multiply-adds per byte loaded, by giving each thread a tile of outputs and keeping operands in registers. Kernels 9 and 10 are about matching the hardware's real structure: tuning tile sizes per GPU and organising work around warps, the unit the scheduler actually sees.

If you already know what a warp is and how shared memory differs from global memory, skip to the problem. Otherwise the primer below defines every term used later. The final section maps every level of this hierarchy onto what Triton does for you, and shows where the flash-linear-attention (FLA) kernels sit on top of it.

Primer: the vocabulary

CUDA describes a computation in a hierarchy that exists for correctness (grid, block, thread) and a hardware hierarchy that exists for performance (SM, warp scheduler, warp, lane). The two overlap but are not the same thing. Most of the optimisations later come from understanding where they differ.

Kernel, thread, block, grid

Kernel

__global__ void f(...)

A function that runs on the GPU (the device). The CPU (the host) launches it with f<<<gridDim, blockDim>>>(args). The launch returns immediately; the GPU executes it asynchronously. One launch creates one grid.

Thread

threadIdx.x/y/z

The smallest unit. Every thread runs the same kernel code but with its own values of threadIdx and blockIdx, its own registers and its own program counter. Kernel code is written from the point of view of a single thread.

Block (thread block)

blockIdx, blockDim, up to 1024 threads

A group of threads that are guaranteed to run on the same SM at the same time. Threads in one block can share data through shared memory and can synchronise with __syncthreads(). Threads in different blocks cannot cheaply talk to each other.

Grid

gridDim

All the blocks of one launch. Blocks are independent and the hardware may run them in any order, on any SM. Both gridDim and blockDim are 3D vectors (x, y, z), though most kernels use one or two dimensions.

Left: one block, each square is a thread, colour is the warp it belongs to (consecutive linear thread ids in groups of 32). Right: the grid of blocks needed to cover an M x N matrix with 32 x 32 tiles. Blocks that hang over the edge still launch all their threads; the ones outside the matrix do nothing. That waste is called tile quantisation.

Inside a kernel, a thread computes its global position from the built-ins. For a 2D layout the usual formula is x = blockIdx.x * blockDim.x + threadIdx.x and likewise for y. The value of threadIdx.x ranges from 0 to blockDim.x - 1; blockIdx.x ranges from 0 to gridDim.x - 1.

Warp: the unit the hardware actually executes

A warp is a group of 32 threads that the hardware executes together, in lockstep, one instruction at a time (this is the SIMT model, single instruction multiple threads). Warps do not appear anywhere in CUDA source code. They are formed by taking the threads of a block in order of their linear thread id and cutting every 32:

linearId = threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z)
warpId   = linearId / 32      // which warp in the block
laneId   = linearId % 32      // position inside the warp, 0..31

Because threadIdx.x is the fastest-varying dimension, threads with consecutive threadIdx.x land in the same warp. This one fact decides which memory accesses get combined (kernel 2) and which shared memory accesses collide (kernels 7 and 8).

Three consequences of the warp being the execution unit:

  • A memory instruction issued by a warp produces 32 addresses at once. The memory system services them as a set, and how many transactions it needs depends on how those 32 addresses are laid out. This is coalescing.
  • If threads in a warp take different branches, the warp executes both paths with some lanes masked off. This is divergence, and it wastes issue slots.
  • Scheduling, register allocation and latency hiding all happen per warp, not per thread.

Streaming multiprocessor (SM) and warp schedulers

The GPU is a collection of SMs. The RTX A6000 has 84 of them. Each SM holds a register file, an on-chip data store split between L1 cache and shared memory, and four warp schedulers. Every cycle each scheduler looks at the warps assigned to it, picks one that is ready (its operands have arrived), and issues that warp's next instruction to the execution units. When a warp is waiting on memory, the scheduler simply picks another warp. This is how GPUs hide latency: not with big caches and out-of-order execution like a CPU, but by keeping many warps in flight and switching between them for free.

One SM, simplified. Blocks are assigned to an SM as a whole; their warps are distributed across the four schedulers. Resident warps share the register file and the shared memory. If any of those three resources is exhausted, no more blocks fit and the scheduler has fewer warps to choose from.

Relevant limits on the A6000 (compute capability 8.6), used in the calculations later:

SMs
84
threads per block, max
1024
threads per SM, max
1536 (48 warps)
32-bit registers per SM
65536, allocated in units of 256 per warp
shared memory per SM
up to 100 KB (of a 128 KB L1/shared pool), 48 KB per block by default
global memory
48 GB GDDR6, about 768 GB/s
fp32 peak
the article budgets about 30 TFLOP/s (the datasheet number is 38.7)

Occupancy is the number of warps resident on an SM divided by the maximum (48). It is limited by whichever of registers, shared memory or thread count runs out first. Higher occupancy gives the scheduler more warps to switch between, which helps hide latency, but it is not the only way to hide latency: a single warp with lots of independent instructions (instruction-level parallelism, ILP) also keeps the pipeline busy. The later kernels deliberately trade occupancy for ILP and register reuse.

Memory hierarchy

Bandwidth and capacity move in opposite directions. Every kernel after the first is an attempt to serve as many multiply-adds as possible from the fast, small levels and to touch global memory as rarely as possible.

Registers

per thread, fastest

Scalar variables in kernel code live here. Fixed-size arrays with compile-time indices can too (float acc[8] becomes 8 registers). A thread can use up to 255. More registers per thread means fewer warps fit on the SM.

Shared memory (SMEM)

per block, on the SM, __shared__

Programmer-managed scratchpad. Roughly 16x the bandwidth of global memory and far lower latency. Physically the same silicon as L1, split by configuration. Divided into 32 banks of 4-byte words; a warp instruction that hits the same bank at different addresses is serialised (a bank conflict).

L1 and L2 cache

hardware managed

L1 sits in each SM next to shared memory. L2 is shared by the whole chip (6 MB on the A6000). Both cache global memory transparently. The naive kernel already gets most of its data from L1, which is why the gain from shared memory alone (kernel 3) is smaller than you might expect.

Global memory (GMEM)

device DRAM, cudaMalloc

The 48 GB on the card. Highest latency (hundreds of cycles) and lowest bandwidth. Accessed in 32-byte sectors; a 128-byte line is four sectors. All inputs start here and the output must end here.

Row-major layout and strides

Matrices are stored as one flat array. Row-major means element (row, col) of an M x N matrix is at index row * N + col: consecutive columns in a row are neighbours in memory, consecutive rows are N floats apart. That gap (the stride) is why reading down a column is expensive and reading along a row is cheap.

Hover a cell to see where it lives in the flat array. A 4 x 6 matrix: moving one column right moves 1 float (4 bytes); moving one row down moves 6 floats (24 bytes). For the article's 4092-wide matrices, one row down is 16368 bytes.

FLOPs, arithmetic intensity and the roofline

FLOP and FMA

One floating point operation. A fused multiply-add d = a*b + c is one instruction (FFMA in SASS) but counts as two FLOPs. A matmul of size M x N x K does 2MNK FLOPs.

Arithmetic intensity (AI)

FLOPs performed per byte moved across some memory boundary. Raising it is the theme of kernels 3 to 6: do more math per byte loaded from global memory, then per byte loaded from shared memory.

Roofline

Attainable FLOP/s is at most the smaller of peak compute and (AI x memory bandwidth). At low AI the kernel is memory-bound and sits on the sloped part; at high AI it is compute-bound and sits under the flat roof.

PTX and SASS

PTX is NVIDIA's virtual assembly; SASS is the real machine code for one GPU generation. Instruction names used later: LDG (load global), STG (store global), LDS (load shared), FFMA (fp32 fused multiply-add). A .128 suffix means one instruction moves 128 bits (four floats).

Roofline for the A6000. Horizontal marks show each kernel's achieved throughput. Only kernels 1, 2 and cuBLAS are placed on the x axis, using their measured global memory traffic; the rest are shown against the compute roof only. Kernels 1 and 2 sit far below the memory slope because their accesses waste most of each transaction; from kernel 5 onward the gap to the roof is instruction overhead, latency and scheduling rather than bandwidth.

The problem: SGEMM

SGEMM computes C = alpha * A @ B + beta * C in single precision, with A of size M x K, B of size K x N and C of size M x N. Each output element is a dot product of one row of A with one column of B, of length K. The article benchmarks M = N = K = 4092.

Napkin math for 4092 x 4092

FLOPs
2 x 4092^3 + 4092^2 = about 137 GFLOP
minimum bytes read
3 matrices x 4092^2 x 4 B = 201 MB
minimum bytes written
4092^2 x 4 B = 67 MB
time at 30 TFLOP/s
4.5 ms
time to move 268 MB at 768 GB/s
0.35 ms

So if a kernel moved only the minimum data it would be compute-bound by a factor of 13. Put differently: a kernel can afford to read global memory up to about 13 times more than the minimum and still be compute-bound. That is the budget the whole worklog is spent inside. cuBLAS reads about 500 MB, less than twice the minimum. The naive kernel, with zero caching, would read 548 GB.

One output of C needs a full row of A and a full column of B. There are M x N outputs and each pair of inputs is reused many times: row i of A contributes to all N outputs in row i of C. Every optimisation below is a way of exploiting that reuse at some level of the memory hierarchy.
1

Naive kernel

One thread per output element, straight from global memory.

309 GFLOP/s1.3% of cuBLAS

The most direct mapping of the problem onto the thread hierarchy: launch a grid of 32 x 32 blocks, one block per 32 x 32 tile of C, one thread per element. Each thread runs the K-long dot product on its own and writes one value. Nothing is shared, so no synchronisation is needed.

// host side
dim3 grid(ceil_div(M, 32), ceil_div(N, 32));
dim3 block(32, 32);           // 1024 threads
k1_naive<<<grid, block>>>(M, N, K, alpha, A, B, beta, C);

// device side, from one thread's point of view
__global__ void k1_naive(int M, int N, int K, float alpha,
                         const float *A, const float *B, float beta, float *C) {
  int row = blockIdx.x * blockDim.x + threadIdx.x;   // note: x picks the row
  int col = blockIdx.y * blockDim.y + threadIdx.y;
  if (row < M && col < N) {
    float acc = 0.f;
    for (int k = 0; k < K; ++k)
      acc += A[row * K + k] * B[k * N + col];
    C[row * N + col] = alpha * acc + beta * C[row * N + col];
  }
}
Hover an element of C. The thread that owns it walks the highlighted row of A and column of B, one pair per iteration of k. Shown with M = N = K = 8 and 4 x 4 blocks so that the grid structure is visible.

Why it is slow

Look at which threads form a warp. With blockDim = (32, 32), a warp is 32 threads with the same threadIdx.y and threadIdx.x = 0..31. In this kernel threadIdx.x selects the row. So in a single iteration k, the 32 lanes of a warp read A[row*K + k] for 32 different rows: 32 floats that are each 4092 x 4 = 16368 bytes apart. The memory system has to fetch 32 separate 32-byte sectors to deliver 128 bytes of useful data. Meanwhile all 32 lanes read the same B[k*N + col], which is at least cheap (one broadcast).

The profiler shows the consequence: about 15 GB/s of global memory throughput on a card capable of 768 GB/s. The next kernel changes nothing except which thread owns which element.

2

Global memory coalescing

Make the 32 lanes of a warp read 32 neighbouring floats.

1986 GFLOP/s8.5% of cuBLAS

When a warp issues a load, the hardware collects the 32 addresses and merges those that fall in the same 32-byte sector (or 128-byte line) into one transaction. If the 32 lanes read 32 consecutive, aligned floats, that is 128 bytes served by a single 128-byte request. If they read 32 floats scattered across 32 different lines, it is 32 requests for the same amount of useful data. This merging is called coalescing and it is done at run time by the hardware, not by the compiler: the two kernels compile to identical SASS.

The fix is to flip which index the fast-varying thread id controls. Make the block one-dimensional (1024 threads) and derive row and column arithmetically so that consecutive threads get consecutive columns:

dim3 block(32 * 32);          // same thread count, now 1D

int row = blockIdx.x * 32 + threadIdx.x / 32;   // changes every 32 threads
int col = blockIdx.y * 32 + threadIdx.x % 32;   // changes every thread
// dot product loop is unchanged

Now a warp is 32 threads with the same row and 32 consecutive columns. Per iteration k: every lane reads the same A[row*K + k] (one broadcast), and the lanes read B[k*N + col .. col+31], 128 contiguous bytes, one transaction.

The 32 lanes of one warp, at one iteration k. Top: which elements of A and B each lane reads. Bottom: the flat memory of A and B in 32-byte sectors, with touched sectors filled. Kernel 1 touches 32 sectors of A to get 32 floats; kernel 2 touches 1 sector of A and 4 sectors of B for the same work.

Global memory throughput rises from 15 GB/s to 110 GB/s and the kernel gets 6.4x faster. Note that alignment also matters: the 32 addresses must start on a sector boundary for the minimum transaction count. Also note that lanes do not have to access addresses in lane order; any permutation of the same 32 consecutive addresses coalesces just as well.

3

Shared memory cache blocking

Load a chunk once per block, reuse it 32 times from on-chip memory.

2980 GFLOP/s12.8% of cuBLAS

A block computing a 32 x 32 tile of C needs a 32 x K slab of A and a K x 32 slab of B. Every element of that slab of A is used by 32 threads in the block (one per column of the tile). So instead of each thread reading from global memory independently, the block cooperates: it walks along K in chunks of BK = 32, and for each chunk all 1024 threads together copy a 32 x 32 piece of A and a 32 x 32 piece of B into shared memory, one element per thread. After a barrier, each thread does 32 multiply-adds against the shared copies, then the block moves to the next chunk.

__shared__ float As[32 * 32];
__shared__ float Bs[32 * 32];
int tRow = threadIdx.x / 32, tCol = threadIdx.x % 32;   // tCol is the fast index

// move base pointers to this block's tile
A += blockIdx.x * 32 * K;  B += blockIdx.y * 32;  C += blockIdx.x * 32 * N + blockIdx.y * 32;

float acc = 0.f;
for (int k0 = 0; k0 < K; k0 += 32) {
  As[tRow * 32 + tCol] = A[tRow * K + tCol];   // coalesced: tCol varies fastest
  Bs[tRow * 32 + tCol] = B[tRow * N + tCol];
  __syncthreads();                             // wait until the whole chunk is in SMEM
  A += 32;  B += 32 * N;                       // advance to next chunk
  for (int k = 0; k < 32; ++k)
    acc += As[tRow * 32 + k] * Bs[k * 32 + tCol];
  __syncthreads();                             // nobody overwrites SMEM while others still read
}
C[tRow * N + tCol] = alpha * acc + beta * C[tRow * N + tCol];
One block, one 32 x 32 tile of C (drawn at 8 x 8 for legibility, so BK = 8 here). Each step copies the next chunk of A and B into shared memory, then every thread accumulates its partial dot product from the copies. After K / BK steps the tile is complete.

Two barriers per chunk are needed. The first guarantees the chunk is fully written before anyone reads it. The second guarantees everyone has finished reading before a fast thread starts overwriting the buffer with the next chunk.

Occupancy check for this kernel

This kernel uses 8 KB of shared memory per block (two 32 x 32 float arrays), 37 registers per thread and 1024 threads per block. Feeding that into the SM limits:

  • Shared memory: (8192 + 1024 bytes runtime overhead) per block into 102400 bytes per SM allows 11 blocks.
  • Threads: 1024 per block into 1536 per SM allows only 1 block.
  • Registers: 37 x 32 = 1184 per warp, rounded up to the 256 allocation unit gives 1280; x 32 warps = 40960 per block; 65536 / 40960 allows 1 block.

One block per SM, 32 warps out of 48 possible: 67% occupancy. That is decent, so occupancy is not the bottleneck. You can check other configurations in the calculator below.

Occupancy calculator using the A6000 limits from the primer. Each bar is the number of blocks that would fit per SM if only that resource mattered; the smallest wins. Try 128 threads with 8 KB and 37 registers to see the later kernels' regime.

The real bottleneck: shared memory instruction pressure

The inner loop compiles to two shared loads and one FMA:

ld.shared.f32  %f1, [As + ...];
ld.shared.f32  %f2, [Bs + ...];
fma.rn.f32     %f3, %f2, %f1, %f3;

Two memory instructions per math instruction. The profiler's warp-state sampling shows warps stalled on MIO throttle: the queue that feeds shared memory instructions is full. The kernel is compute-bound on paper but in practice it is bound by how fast it can issue LDS instructions. The fix is not more caching; it is fewer loads per FMA, which means each thread must produce more than one output so that loaded values can be reused from registers.

4

1D blocktiling

Each thread computes a column of 8 outputs and reuses one B value 8 times.

8475 GFLOP/s36.5% of cuBLAS

New tile shape: a block now owns a BM x BN = 64 x 64 tile of C and walks K in chunks of BK = 8. Each thread computes TM = 8 vertically adjacent outputs, so the block needs 64 x 64 / 8 = 512 threads. Shared memory holds a 64 x 8 chunk of A and an 8 x 64 chunk of B: 1024 floats, 4 KB.

The important change is the inner loop. For each k in the chunk, the thread loads one value from Bs (its column) into a register and then reuses it for all 8 of its rows, loading one As value per row:

float acc[TM] = {0};                          // 8 accumulators, live in registers
for (int k0 = 0; k0 < K; k0 += BK) {
  // cooperative GMEM -> SMEM copy of the 64x8 and 8x64 chunks (one float each)
  As[aRow * BK + aCol] = A[aRow * K + aCol];
  Bs[bRow * BN + bCol] = B[bRow * N + bCol];
  __syncthreads();
  A += BK;  B += BK * N;
  for (int k = 0; k < BK; ++k) {
    float b = Bs[k * BN + tCol];               // loaded once ...
    for (int i = 0; i < TM; ++i)
      acc[i] += As[(tRow * TM + i) * BK + k] * b;   // ... used 8 times
  }
  __syncthreads();
}
for (int i = 0; i < TM; ++i)
  C[(tRow * TM + i) * N + tCol] = alpha * acc[i] + beta * C[(tRow * TM + i) * N + tCol];
Hover a thread (a column strip in the C tile). Its 8 outputs need 8 rows of the As chunk and 1 column of the Bs chunk. Drawn at BM = BN = 16, BK = 4, TM = 4 to keep it readable; the idea is identical at 64 / 8 / 8.

Counting loads per output

Per k in the chunk a thread issues 1 + TM = 9 shared loads and does TM = 8 FMAs. That is 1.125 loads per FMA instead of 2. Per output element over the full K: K/32 global loads and 9K/8 shared loads, versus K/16 and 2K in kernel 3. The warp-stall profile confirms far fewer cycles lost to the memory pipeline, and throughput almost triples.

Sidenote: the compiler would have done this anyway

If the two inner loops are written in the natural order (outputs outer, k inner) with no explicit b register, the generated code is just as fast. Both loop counts are compile-time constants, so nvcc fully unrolls them, notices that the same Bs element is loaded 8 times, and keeps it in a register. With this loop order it also emits 128-bit LDS.128 for the As loads, because for one output row the eight k values sit side by side in memory. The Bs loads stay 32-bit, since consecutive k are BN floats apart. In kernel 5 the k loop moves outermost and the situation flips: the Bs slice a thread needs is contiguous and the As slice is strided, which is what kernel 6 fixes by transposing As.

5

2D blocktiling

Each thread computes an 8 x 8 square as an outer product held entirely in registers.

15972 GFLOP/s68.7% of cuBLAS

Tile constants become BM = BN = 128, BK = 8, TM = TN = 8. A block of 128 x 128 / 64 = 256 threads owns a 128 x 128 tile of C; each thread owns an 8 x 8 square of it, 64 accumulators in registers.

Loading the chunk into shared memory now takes several trips per thread: As is 128 x 8 = 1024 floats and Bs is 8 x 128 = 1024 floats, so each of the 256 threads copies 4 floats of each. The copy loop strides through the chunk so that, at every step, the 32 lanes of a warp still read 32 contiguous floats from global memory (coalesced).

The inner loop is the interesting part. For each k in the chunk, a thread copies the 8 As values it needs (a column slice) and the 8 Bs values it needs (a row slice) into registers, then does a rank-1 update of its 8 x 8 accumulator block: 64 FMAs from 16 loads.

float acc[TM * TN] = {0};   // 64 accumulators
float regA[TM], regB[TN];    // 16 operand registers

for (int k = 0; k < BK; ++k) {
  for (int i = 0; i < TM; ++i) regA[i] = As[(tRow * TM + i) * BK + k];
  for (int j = 0; j < TN; ++j) regB[j] = Bs[k * BN + tCol * TN + j];
  for (int i = 0; i < TM; ++i)
    for (int j = 0; j < TN; ++j)
      acc[i * TN + j] += regA[i] * regB[j];     // outer product, all in registers
}
Hover a thread (an outlined square in the C tile). For the chosen k, the highlighted As column slice and Bs row slice are copied into regA and regB; their outer product updates the thread's whole square. Drawn at BM = BN = 16, BK = 4, TM = TN = 4.

Why a square beats a column

A thread computing an r x c block of outputs needs r + c loads per k for r x c FMAs. For a column (c = 1) the ratio is (r+1)/r, which can never get below 1. For a square it is 2/r, which halves every time r doubles. That is the whole reason 2D tiling exists.

Loads per FMA for different per-thread output shapes at one k. Filled cells are outputs, outlined strips are the operands that must be loaded to update them.

Counting again

Per output element over the full K: K/64 global loads and K/4 shared loads. Compared with kernel 4 that is another 2x fewer global loads and 4.5x fewer shared loads. Throughput doubles to 16 TFLOP/s. The block-level arithmetic intensity (FLOPs per byte moved from global memory into shared memory) is BM x BN / (2 (BM + BN)) = 32 FLOP per byte for a 128 x 128 tile, versus 8 for the 32 x 32 tile of kernel 3.

6

Vectorised memory access

Transpose As so both operand slices are contiguous, and move global data as float4.

18237 GFLOP/s78.4% of cuBLAS

Part 1: transpose As in shared memory

In kernel 5, the 8 values a thread needs from As for one k are As[(tRow*8 + i) * BK + k] for i = 0..7: eight floats spaced BK apart. A single load instruction cannot fetch them. If As is stored transposed, as As[k * BM + row], those eight values become As[k * BM + tRow*8 .. tRow*8+7]: contiguous, so two LDS.128 instructions replace eight LDS.32. The transpose is free: it is done by writing each element to the swapped position during the global-to-shared copy, which happens anyway.

Left: row-major As, the values one thread needs for k = 1 are scattered with stride BK. Right: the same chunk stored transposed, the values are adjacent and load as 128-bit vectors. The Bs slice was already contiguous in kernel 5.

Part 2: 128-bit global loads and stores

Each thread's share of the copy is four consecutive floats. Read them as one float4:

float4 a4 = reinterpret_cast<const float4*>(&A[aRow * K + aCol * 4])[0];   // LDG.E.128
As[(aCol * 4 + 0) * BM + aRow] = a4.x;   // scatter into the transposed layout
As[(aCol * 4 + 1) * BM + aRow] = a4.y;
As[(aCol * 4 + 2) * BM + aRow] = a4.z;
As[(aCol * 4 + 3) * BM + aRow] = a4.w;

reinterpret_cast<float4*>(&Bs[bRow * BN + bCol * 4])[0] =
    reinterpret_cast<const float4*>(&B[bRow * N + bCol * 4])[0];        // LDG.E.128 then STS.128

Why does the cast matter, when the compiler could have merged four adjacent 32-bit loads itself? Because a 128-bit load requires the address to be 16-byte aligned, and the compiler cannot prove that about a float* passed in as an argument. The reinterpret_cast to float4* is the programmer's promise that the pointer is aligned. Shared memory needs no such promise; the compiler owns its layout and vectorises those loads on its own.

Together the two changes bring another 2.3 TFLOP/s. The profiler now lists three things: shared memory bank conflicts, occupancy higher than needed, and no overlap between loading the next chunk and computing on the current one.

7, 8

Bank conflicts (the two kernels that were dropped)

They removed the conflicts and still ran slower, so the article skips them. The concept still matters.

not reportedslower than kernel 6

Shared memory is built from 32 banks, each 4 bytes wide, interleaved: byte address a lives in bank (a / 4) % 32. In one cycle each bank can serve one 32-bit word. When a warp issues a shared load, the 32 lanes' addresses are grouped by bank. If two lanes want different words from the same bank, the hardware replays the instruction for the second one. This is an n-way bank conflict, and it multiplies the cost of that instruction by n. Lanes reading the same word are fine (broadcast).

The outer-product kernels read As and Bs in patterns that depend on the tile sizes and on TM, TN, and those patterns produce conflicts. Kernels 7 and 8 rearranged the shared memory layout so that a warp's accesses land in 32 distinct banks. They succeeded at that but the extra index arithmetic and the changed access pattern cost more than the conflicts did, so the article moved on. cuBLAS does avoid conflicts, which is part of the remaining gap.

32 lanes each read the float at index lane x stride. Lines connect lanes to banks. Stride 1 is conflict-free. Stride 2 puts two lanes on every even bank (2-way). Stride 32 puts all 32 lanes on bank 0 (32-way, the pattern of reading a column of a 32-wide row-major array). Stride 33 is conflict-free again, which is why padding a row by one float is the classic fix.
9

Autotuning

Same kernel, search the five tile parameters instead of guessing them.

19721 GFLOP/s84.8% of cuBLAS

By now the kernel has five template parameters. BM, BN, BK set how much of A and B is cached in shared memory per step. TM, TN set how much of that a thread pulls into registers. They interact with every hardware limit at once: shared memory capacity, register file size, occupancy, coalescing, the float4 copy loop, bank conflict patterns. Nobody can reason their way to the optimum, so the article does what every production library does: enumerate the sensible configurations and time them.

What "sensible" means

  • The thread count is fixed by the tile: (BM / TM) x (BN / TN) threads, which must be a multiple of 32 and at most 1024.
  • Each thread copies one float4 per pass of the load loop, so BM x BK and BN x BK must be divisible by 4 x threads.
  • Shared memory (BM + BN) x BK x 4 bytes must fit, and registers per thread (TM x TN accumulators plus TM + TN operands plus addressing) must leave room for enough warps.
  • TM and TN should be multiples of 4 so the register slices load as 128-bit vectors.

About 400 configurations survived those filters and were benchmarked with a script. On the A6000 the winner was BM = BN = 128, BK = 16, TM = TN = 8: only BK changed from kernel 6, for a gain of about 8% in the benchmark table. On an A100 the winner was BM = BN = 64, BK = 16, TM = TN = 4, and running the A6000's best configuration there would have left 6% on the table. The optimum is a property of the GPU, not of the algorithm, which is why compilers like Triton ship an autotuner and why cuBLAS ships hundreds of pre-tuned kernels.

Tile parameter explorer. Presets: kernel 3 is 32 / 32 / 32 / 1 / 1 (BK equal to BM), kernel 4 is 64 / 64 / 8 / 8 / 1, kernel 5 is 128 / 128 / 8 / 8 / 8, kernel 9 is 128 / 128 / 16 / 8 / 8. The derived quantities are exact; the "fits" checks use the A6000 limits, and the float4 divisibility check only applies from kernel 6 onward (earlier kernels copied one float per thread).
10

Warptiling

Insert the warp as an explicit tiling level between block and thread.

21779 GFLOP/s93.7% of cuBLAS

So far the loop structure has two tiling levels: the block owns BM x BN and each thread owns TM x TN, with threads laid out in a plain row-major grid over the block tile. That layout ignores the fact that the 32 threads of a warp are scheduled together, share a register cache, and conflict with each other (and only each other) in shared memory. Warptiling adds a level: the block tile is divided among warps, each warp tile is divided among its 32 lanes, and each lane still computes TM x TN squares.

The new parameters

WM, WN
size of the tile owned by one warp; BM / WM x BN / WN warps per block
WMITER, WNITER
how many times the warp iterates over its tile; each iteration covers a sub-tile of WSUBM x WSUBN = (WM / WMITER) x (WN / WNITER)
TM, TN
the per-thread square inside a sub-tile; 32 lanes x TM x TN = WSUBM x WSUBN
warpId, laneId
threadIdx.x / 32 and threadIdx.x % 32; the lane's row and column in the sub-tile are laneId / (WSUBN / TN) and laneId % (WSUBN / TN)

Each thread therefore accumulates WMITER x WNITER squares of TM x TN. For each k in the chunk it loads WMITER x TM values from As and WNITER x TN values from Bs into registers, then updates all its squares with outer products. The k loop is kept outermost inside the chunk, so everything inside it is independent work the scheduler can overlap.

for (int k = 0; k < BK; ++k) {
  for (int wm = 0; wm < WMITER; ++wm)                 // A slices for each sub-tile row
    for (int i = 0; i < TM; ++i)
      regA[wm * TM + i] = As[k * BM + warpRow * WM + wm * WSUBM + laneRow * TM + i];
  for (int wn = 0; wn < WNITER; ++wn)                 // B slices for each sub-tile column
    for (int j = 0; j < TN; ++j)
      regB[wn * TN + j] = Bs[k * BN + warpCol * WN + wn * WSUBN + laneCol * TN + j];
  for (int wm = 0; wm < WMITER; ++wm)                 // warp-level matmul
    for (int wn = 0; wn < WNITER; ++wn)
      for (int i = 0; i < TM; ++i)
        for (int j = 0; j < TN; ++j)
          acc[(wm * TM + i) * (WNITER * TN) + wn * TN + j] += regA[wm * TM + i] * regB[wn * TN + j];
}
Three levels of tiling in one 128 x 128 block tile, using BM = BN = 128, four warps of WM = WN = 64, WMITER = 1, WNITER = 4 (so WSUBM x WSUBN = 64 x 16), TM = 8, TN = 4. Hover a warp tile to see its four sub-tiles, or a lane's square to see all four squares that lane accumulates. Every lane in a warp reads the same rows of As (bounded by the warp tile), which is the locality the warp-level register cache rewards.

Why it is faster

  • Explicit parallelism at each hardware level. Blocks run in parallel on different SMs; warps run in parallel on the four schedulers of an SM (and interleave on one scheduler); the independent FMAs inside a thread overlap in the pipeline (ILP). Warptiling makes the middle level a deliberate choice instead of an accident of thread numbering.
  • Register cache locality. Recent NVIDIA GPUs have a small operand cache in front of the register file. Tighter per-warp tiles mean the same As and Bs registers are reused in consecutive instructions.
  • Bank conflicts are a per-warp phenomenon, so controlling which addresses a warp touches together is the lever for reducing them.
  • It is the shape tensor cores want. A warp-level matrix multiply on a WSUBM x WSUBN sub-tile maps directly onto the warp-wide mma instructions that tensor cores execute. This kernel is one step from that.

After re-tuning the enlarged parameter space, throughput reaches 21.8 TFLOP/s, within 6% of cuBLAS. Something that did not help: thread block swizzling (remapping blockIdx to C tiles so that concurrently running blocks share rows of A in L2). L2 hit rate was already about 80%, and the swizzle produced no measurable gain, so it was removed.

0

The cuBLAS reference

Not one kernel but a library of hundreds, chosen at run time by shape.

23250 GFLOP/s100%

Comparing kernel 10 with cuBLAS across matrix sizes shows two regimes. At 2048 and 4096 the hand-written kernel is within a few percent. At small sizes it loses badly. The reason is that cuBLAS is a dispatcher: it contains many SGEMM implementations (the article counted 16 distinct ones for square sizes up to 4096, out of a 500 MB binary) and picks one per call based on M, N, K, data type and GPU.

The trace at size 256 is instructive: cuBLAS launched a matmul kernel and a reduction kernel. That is split-K. A 256 x 256 output with 128 x 128 tiles is only 4 blocks, which cannot occupy 84 SMs. Splitting the K dimension across several blocks gives each SM something to do; each block produces a partial sum for its slice of K, and a second kernel adds the partials.

Split-K. Without it, a small output matrix means few blocks and idle SMs. With it, several blocks compute partial products for the same C tile over different ranges of k, write them to scratch memory, and a reduce kernel sums them. It costs an extra pass over the partials but multiplies the available parallelism.
11

What is left (work in progress in the article)

The remaining 6%, and what a tensor core kernel would add.

no numberwork in progress
  • Double buffering. Right now each chunk is a strict sequence: copy from global memory, barrier, compute, barrier. With two shared memory buffers, the copy of chunk i + 1 can be issued before the compute on chunk i starts, so memory latency overlaps math. CUTLASS does this at both levels: global to shared, and shared to registers.
  • Conflict-free shared memory layouts. Swizzled layouts that keep vectorised loads while spreading a warp's accesses over all 32 banks. Kernels 7 and 8 were the first attempt.
  • Hardware copy paths. On Ampere, cp.async moves data from global memory straight into shared memory without passing through registers. On Hopper, TMA and warp specialisation (some warps only load, others only compute, with different register budgets) go further.
  • Tensor cores. Everything above is fp32 on CUDA cores. With TF32 or BF16 inputs, warp-level mma instructions raise peak throughput by roughly 3x on this card and turn the problem back into a memory-bound one, where the loading machinery above becomes the entire game.
Timeline of one block. Top: the current kernels serialise load and compute for each chunk. Bottom: with two buffers, the load for the next chunk runs while the current one is being computed, and the barrier only has to wait for the data that is needed next.

Where Triton and the FLA kernels sit

The kernels above are written one level at a time: you decide which thread owns which element, you declare shared memory, you choose the register tile, you transpose by hand, you pipeline by hand. Triton flips the contract. You write a program at the block tile level (roughly kernel 3's viewpoint: one program owns a BM x BN tile and loops over K), and the compiler fills in kernels 2, 4 to 8, 10 and 11 from a handful of knobs. The flash-linear-attention library (fla-org/flash-linear-attention, the kernels behind the fla-hub models) is written entirely in Triton, so understanding what Triton takes off your hands is also understanding what FLA can and cannot control.

Vocabulary translation

program
a Triton kernel instance, launched over a grid exactly like a CUDA block; tl.program_id(axis) is blockIdx
num_warps
threads per program, in warps (default 4, so 128 threads); there is no threadIdx anywhere in Triton source
tl.constexpr
compile-time constants such as BT, BK, BV: the template parameters of kernels 4 to 10
tl.load / tl.store
a whole tile of pointers at once; the compiler decides which lane loads which element (coalescing) and whether to vectorise
tl.dot(a, b, acc)
a block-tile matmul: operands are staged into shared memory, accumulated in registers, executed on tensor cores where the dtype allows
num_stages
depth of the software pipeline for the K loop; 2 is double buffering, 3 or 4 is deeper
@triton.autotune
kernel 9 as a decorator: a list of configs (block sizes, num_warps, num_stages) benchmarked once per value of the key arguments and cached
@triton.heuristics
constexpr flags derived from the arguments at launch (is this a variable-length batch, is there an initial state), so each variant is compiled separately with dead branches removed
Every kernel in the worklog, and what does the same job in Triton. Hover a row. Green rows are handled by the compiler from a knob or automatically; amber rows are still your decision in Triton source. The split is the whole point: Triton keeps the decisions that change the algorithm (tile shapes, what to fuse, what to materialise) and automates the ones that only change the machine code.

What you still own in Triton

  • The grid decomposition: which program computes which tile, and in which order. This is kernel 1's decision and it still determines L2 locality and load balance.
  • Block tile sizes and the number of warps. Too large a tile per warp and the compiler spills accumulators to local memory, which shows up as LDL and STL instructions and a cliff in throughput. That is kernel 9's register budget, now enforced by the compiler instead of the occupancy calculator.
  • Dtypes at each boundary: FLA loads bf16, accumulates in fp32, and casts back before the next tl.dot. With fp32 inputs Triton uses TF32 on tensor cores unless told otherwise.
  • Fusion and materialisation: which intermediate tiles go back to global memory and which stay in registers. The kernels above never had this choice because a GEMM has no intermediates; a linear attention chunk kernel has several.
  • Masks for ragged edges (the tile quantisation problem from the primer) and variable-length batches.

What FLA builds on top of this

A causal linear attention layer with a decay (GLA, gated delta rule, KDA and their relatives) can be written as a recurrence over tokens with a state matrix S of size K x V per head. Computing it token by token is sequential and slow; computing it as one T x T attention matrix is quadratic. FLA's chunk kernels take the middle path: split the sequence into chunks of BT tokens (typically 64), handle everything inside a chunk as a small dense attention (a BT x BT matrix, causally masked), and pass everything between chunks through S.

Anatomy of an FLA chunk forward pass, simplified (gates and decays omitted). Two Triton kernels: the state pass walks the chunks in order and materialises S after each chunk; the output pass is then embarrassingly parallel over chunks. Every box marked with a dot is a block-tile matmul, the object this whole page is about, but with K equal to the chunk length or the head dimension rather than 4092.

Read the grid axes on that figure against the primer. program_id(2) runs over batch x heads: the outermost independent dimension, so it maps to different SMs. program_id(1) runs over chunks in the output kernel, or over tiles of the state in the state kernel. program_id(0) runs over tiles of the head dimension (BK or BV), which is what keeps the register tile of one program small enough. Each program is still a block of num_warps x 32 threads; the warp and thread tiling inside a tl.dot is chosen by the compiler.

Why FLA sits somewhere else on the roofline

  • The matmuls are small. With BT = 64 and head dimension 128, each tl.dot is a 64 x 64 x 128 or 64 x 128 x 64 product: about 1 MFLOP, versus 137 GFLOP for the article's GEMM. Block-level arithmetic intensity is bounded by the tile, so these kernels sit far to the left of the ridge point and are bound by memory traffic and latency, not FMA throughput.
  • The state pass is a scan. One program advances S through all T / BT chunks in order, so parallelism comes from batch x heads x state tiles, not from the sequence. For a small batch at long context this is the part that leaves SMs idle, the same shape of problem as cuBLAS's split-K case.
  • S gets materialised. Writing the per-chunk state to global memory (T / BT x K x V floats per head) so the output kernel can read it back is exactly the kind of traffic kernels 3 to 6 spent all their effort avoiding. Fusing the two passes trades that traffic for a longer serial chain; FLA offers both fused and chunked paths for a reason.
  • Delta-rule variants add a triangular solve inside each chunk (the solve_tril kernels), which is why those kernels carry a sub-chunk size BC and autotune over it separately.

So the levers that matter for FLA are the ones Triton leaves in your hands: chunk size, which tensors are materialised, how the grid is cut so that batch x heads x tiles fills the GPU, bf16 operands with fp32 accumulation, and the autotune space over num_warps and num_stages. The thread-level work the article spends kernels 2 to 8 on is done by the compiler, and done about as well as a careful hand-written kernel for these tile shapes. Serving stacks such as vLLM vendor FLA's Triton ops for their gated delta rule models, so the same trade-offs carry over to inference.

One-line summary. The article climbs the hierarchy from thread to warp to block. Triton starts you at the block and hides the two levels below it. FLA stacks a fourth level, the chunk, on top of the block, and its performance problems live there.

Cheat sheet

KernelChangeTile (BM, BN, BK, TM, TN)GMEM loads / outputSMEM loads / outputGFLOP/sof cuBLAS
1one thread per output, strided warp reads of A32, 32, K, 1, 12K03091.3%
2consecutive lanes read consecutive columns32, 32, K, 1, 12K (coalesced)019878.5%
3block caches chunks of A and B in SMEM32, 32, 32, 1, 1K/162K298012.8%
4thread owns a column of 8 outputs64, 64, 8, 8, 1K/329K/8847536.5%
5thread owns an 8 x 8 square, outer product in registers128, 128, 8, 8, 8K/64K/41597268.7%
6As transposed for LDS.128; float4 GMEM loads128, 128, 8, 8, 8K/64 (128-bit)K/4 (128-bit)1823778.4%
9parameters searched by benchmark128, 128, 16, 8, 8K/64K/41972184.8%
10warp tile level between block and thread+ WM, WN, WMITER, WNITERK/64below K/42177993.7%
0cuBLAS, shape-dependent dispatchvaries23250100%

Loads per output are counted for the full K loop, per thread, divided by the number of outputs the thread produces. For kernel 10 the shared loads per output depend on WMITER and WNITER: per k a thread loads (WMITER x TM + WNITER x TN) values for WMITER x WNITER x TM x TN FMAs.

The five rules the worklog teaches

  • Make the 32 lanes of a warp touch 32 consecutive, aligned words. Everything else about global memory is secondary.
  • Move data up the hierarchy once, then reuse it as many times as possible before moving more. Block tiles reuse through shared memory; thread tiles reuse through registers.
  • Count instructions, not just bytes. Two shared loads per FMA is a bottleneck even when the bytes are on-chip.
  • Give the compiler the information it cannot infer: compile-time loop bounds so it can unroll, and aligned vector types so it can emit 128-bit accesses.
  • Organise around warps. They are the unit of scheduling, of coalescing, of bank conflicts and of tensor core instructions.