A GEMM Kernel, Step by Step

The idea in one sentence

Every rung of a fast matrix multiply is the same move at a different scale: give one worker a squarer, bigger tile of the answer, because a tile of shape $m \times n$ costs $m + n$ numbers to feed and returns $m\,n$ multiplications.

The ladder we are climbing

This chapter walks a published set of kernels. Simon Boehm wrote them, measured them one at a time on an RTX A6000, and published the code and the numbers in his CUDA matmul worklog with the repository at siboehm/SGEMM_CUDA. The rung numbering and every GFLOP/s figure below are his. The drawings, the arithmetic and the one idea that ties the rungs together are ours.

The problem is one call: $C \leftarrow \alpha AB + \beta C$ with all three matrices $4092 \times 4092$ in single precision. The reference to beat is cuBLAS, NVIDIA’s own library, which reaches 23249.6 GFLOP/s on the same card.

What the problem is allowed to cost

Before writing a kernel, work out what a perfect one would cost. There are two costs and the larger one wins.

The arithmetic is fixed. Each of the $4092^2$ entries of $C$ is a dot product of length 4092, so each takes 4092 multiplications and 4092 additions. That is

floating-point operations, counting the one extra multiply-add per entry that $\alpha$ and $\beta$ cost, or 137 GFLOP.

The memory is not fixed, and that is the whole chapter. The floor is easy: you must read $A$, read $B$, read the old $C$, and write the new $C$. At four bytes per float that is $3 \times 4092^2 \times 4 = 201$ MB read and $4092^2 \times 4 = 67$ MB written, so 268 MB moved in total. Divide the work by the traffic and you get the ratio that decides everything:

Five hundred and eleven floating-point operations for every byte. That number is called arithmetic intensity, written $I$ from here on, and it is the single axis a kernel lives on. Boehm’s card offers 768 GB/s of memory bandwidth and, on the peak his roofline uses, 30 TFLOP/s of arithmetic. Sending 268 MB at 768 GB/s takes 0.35 ms. Doing 137 GFLOP at 30 TFLOP/s takes 4.6 ms. The arithmetic is thirteen times the memory, so a good matrix multiply is compute bound, and the entire difficulty is getting anywhere near that floor.

The picture that holds both facts is the roofline. Plot attainable performance against $I$. Divide the 30 TFLOP/s by the 768 GB/s and you get 39.1 FLOP per byte. Below that, memory is the ceiling and performance rises along a slope. Above it, arithmetic is the ceiling and the line goes flat. That crossing is called the ridge point, and a kernel’s whole life is spent trying to get to the right of it. NVIDIA’s data sheet quotes 38.7 TFLOP/s at the boost clock rather than Boehm’s 30, which would lift the flat part and move the ridge to 50 FLOP per byte, and would change none of the conclusions below.

Here is the whole ceiling table, and the code that prints it, so no number in this chapter has to be taken on trust.

PEAK, BW = 30e12, 768e9          # RTX A6000: Boehm's roofline ceilings

def intensity(m, n):             # FLOP per byte of a float32 m x n tile
    return m * n / (2 * (m + n))

print("  tile      FLOP/byte      ceiling GFLOP/s")
for m, n in [(1, 1), (32, 32), (64, 64), (128, 128), (128, 256)]:
    i = intensity(m, n)
    print("%4d x %-4d %10.2f %15.0f" % (m, n, i, min(PEAK, i * BW) / 1e9))
print("ridge point %.2f FLOP/byte" % (PEAK / BW))

#   tile      FLOP/byte      ceiling GFLOP/s
#    1 x 1          0.25             192
#   32 x 32         8.00            6144
#   64 x 64        16.00           12288
#  128 x 128       32.00           24576
#  128 x 256       42.67           30000
# ridge point 39.06 FLOP/byte

The intensity function in that listing is the one idea of the chapter, and the next section derives it.

Rung one: one thread, one entry

The obvious kernel gives every entry of $C$ its own thread. The thread works out which entry it owns from the launch indices of Lecture 25, runs a loop over $k$, and writes once.

__global__ void sgemm_naive(int M, int N, int K, const float *A, const float *B, float *C) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;   // the fast index picks the row
    int col = blockIdx.y * blockDim.y + threadIdx.y;
    if (row < M && col < N) {
        float acc = 0.0f;
        for (int k = 0; k < K; ++k)
            acc += A[row * K + k] * B[k * N + col];
        C[row * N + col] = acc;
    }
}

Count what one thread costs. It reads a whole row of $A$ and a whole column of $B$, which is $2 \times 4092 = 8184$ floats, or 32736 bytes, and it does 8184 operations on them. So its arithmetic intensity is

FLOP per byte, four bytes moved for every single operation. On the roofline that puts the ceiling at $0.25 \times 768 = 192$ GFLOP/s, which is 0.6% of the flat roof. The kernel is not slow because of anything it computes. It is slow because of what it asks for.

Boehm measures this kernel at 309.0 GFLOP/s, which is 1.3% of cuBLAS. Notice that 309 is above the 192 the roofline allows at an intensity of 0.25. Nothing is wrong: 0.25 counts what the kernel asks for, and the caches answer much of it without going out to DRAM, because all 32 lanes of a warp here share one column of $B$ and read it at the same moment. His profiler puts the real global memory throughput at 15 GB/s, far under the 768 the card can do, so bandwidth is not what holds this kernel back. What holds it back is the shape of the requests, and that is the next rung.

Rung two: turn the warp sideways

The first fix costs two lines and changes nothing else. In the kernel above threadIdx.x walks row, so the 32 lanes of a warp sit in 32 different rows. Every entry they touch in $A$, and every entry they write in $C$, is $N$ floats from its neighbour, which is 32 separate 32-byte segments per instruction. Bind the fast thread coordinate to the column instead and the same warp reads 32 consecutive entries of a row of $B$ and writes 32 consecutive entries of $C$, which is the coalesced case of Lecture 25 exactly.

// the fast thread coordinate must run along the fast matrix coordinate
int row = blockIdx.y * BLOCKSIZE + threadIdx.x / BLOCKSIZE;
int col = blockIdx.x * BLOCKSIZE + threadIdx.x % BLOCKSIZE;

Both versions do identical arithmetic in an identical order. One asks the memory system for four transactions per warp and the other asks for 32. Boehm’s profiler shows global memory throughput going from 15 GB/s to 110 GB/s, and the kernel from 309.0 to 1986.5 GFLOP/s, a factor of 6.4, for the price of choosing which index gets threadIdx.x. It is the largest single jump on the ladder, and it is the drawing from Lecture 25 cashed in.

Rung three: a desk the block shares

Coalescing fixed how the traffic is shaped. Now reduce how much of it there is.

The block is the level that has a desk. Its threads share one multiprocessor’s shared memory, and they can be made to wait for each other. So stop letting every thread fetch its own operands. Cut $A$ and $B$ into chunks, have the block cooperatively load one chunk of each into shared memory, have everyone compute out of the desk, then move to the next chunk.

__shared__ float As[BM * BK];
__shared__ float Bs[BK * BN];

float acc = 0.0f;
for (int chunk = 0; chunk < K; chunk += BK) {
    As[ty * BK + tx] = A[ty * K + tx];         // each thread brings one number
    Bs[ty * BN + tx] = B[ty * N + tx];
    __syncthreads();                            // the desk is now complete

    for (int k = 0; k < BK; ++k)
        acc += As[ty * BK + k] * Bs[k * BN + tx];
    __syncthreads();                            // nobody overwrites before everyone is done

    A += BK;                                    // walk both chunks along k
    B += BK * N;
}

The two __syncthreads() calls are the rhythm of every GPU matrix multiply, and each one prevents a different accident. The first stops a thread computing from a desk somebody has not finished loading. The second stops a fast thread loading the next chunk over numbers a slow thread is still reading.

Now count the traffic for a whole block tile of $BM \times BN$ entries. Over the full $k$ the block reads $BM \times K$ floats of $A$ and $K \times BN$ floats of $B$, and it produces $BM \times BN$ entries at $2K$ operations each. So

The $K$ cancels, which is the first pleasant surprise: the depth of the multiply does not matter, only the shape of the tile. At $BM = BN = 32$ that is 8 FLOP per byte, up from 0.25, a thirty-two-fold cut in requested traffic. Boehm measures 2980.3 GFLOP/s, 12.8% of cuBLAS.

Rungs four and five: a column, then a square

Rung three still gives each thread one entry of $C$. The desk made the global traffic cheap and did nothing about the traffic between the desk and the registers, which is now the bottleneck. Per multiply-add, a thread reads one number from As and one from Bs. Two shared-memory loads buy one operation.

Apply the same idea one level down. Let a thread own several entries of $C$ at once, holding them in its registers. If the thread owns a column of $TM$ entries, then for each $k$ it loads $TM$ values of $A$ and one value of $B$, and does $TM$ multiply-adds. That is $TM + 1$ loads for $TM$ operations, already better than 2 for 1. If it owns a square of $TM \times TN$, it loads $TM + TN$ values and does $TM \times TN$ multiply-adds.

float acc[TM][TN] = {0.0f};                      // 64 registers when TM = TN = 8
for (int k = 0; k < BK; ++k) {
    for (int i = 0; i < TM; ++i) a_frag[i] = As[(ty * TM + i) * BK + k];
    for (int j = 0; j < TN; ++j) b_frag[j] = Bs[k * BN + tx * TN + j];

    for (int i = 0; i < TM; ++i)                 // 16 loads above, 64 products here
        for (int j = 0; j < TN; ++j)
            acc[i][j] += a_frag[i] * b_frag[j];
}

The measured numbers follow the shape of the argument without matching it. The column version, $TM = 8$ with $BM = BN = 64$, reaches 8474.7 GFLOP/s, 36.5% of cuBLAS. The square version, $TM = TN = 8$ with $BM = BN = 128$, reaches 15971.7 GFLOP/s, 68.7%. Between those two the block tile doubled its intensity, from 16 FLOP per byte to 32, and the thread tile went from 0.9 products per shared-memory load to 4. The kernel got 1.9 times faster. Relieving one bottleneck moves the pressure somewhere else, every time.

The cost of a square is registers. Sixty-four accumulators per thread, at 256 threads to a block, is 16384 registers before anything else is counted, a quarter of the multiprocessor’s whole file. That caps the chip at four such blocks, so 32 warps out of 48, an occupancy of 67% at best and lower in practice. This is the trade Lecture 25 promised: fewer warps in flight, far less data in motion, and a kernel twice as fast.

Rung six: four floats in one instruction

A small rung with a large name. A load of four consecutive floats can be issued as a single 128-bit instruction instead of four 32-bit ones, by casting the pointer to float4.

// one 128-bit load instead of four 32-bit loads
float4 tmp = reinterpret_cast<const float4 *>(&A[row * K + col])[0];

Two things have to be true. The address must be 16-byte aligned, which is Lecture 23’s padding argument reappearing, and the four floats must be consecutive in the direction you want to read them. Boehm transposes $A$ as it lands on the desk, so that reading down a column of $A$ becomes reading along a row, and then both operands can be taken four at a time. That gets 18237.3 GFLOP/s, 78.4% of cuBLAS.

The last rung: the whole hierarchy at once

The last rung puts the warp back in the picture, and it is the drawing this part of the course exists for.

So far the block tile was carved directly into thread tiles. But threads do not execute independently. Thirty-two of them move as one warp, and a warp is the unit that issues shared memory loads and arithmetic. If the 32 threads of a warp own thread tiles scattered across the block tile, then each of their shared-memory reads touches a different region and their register operands cannot be reused between them. Give the warp its own rectangle instead, and the operands it loads serve all 32 of its threads.

That is three nested tilings of the same matrix $C$, one per level of Lecture 25’s picture.

Every number in that figure comes from Amanzhol Salykov’s RTX 3090 worklog, whose kernel uses a $128 \times 128 \times 8$ block tile with 256 threads, a $32 \times 64$ warp tile and an $8 \times 8$ thread tile. Check that they agree: $128/32 = 4$ warp tiles down and $128/64 = 2$ across, so 8 warp tiles, which is exactly the 8 warps in 256 threads. Inside a warp tile, $32/8 = 4$ thread tiles down and $64/8 = 8$ across, so 32 thread tiles, which is exactly the 32 lanes of a warp. The hierarchy tiles the matrix with nothing left over, twice.

The shared memory follows from the same numbers. One chunk pair is $128 \times 8$ floats of $A$ and $8 \times 128$ of $B$, so $2 \times 256 \times 8 \times 4 = 16384$ bytes when double buffered, which is the 16 KB Salykov’s kernel reserves.

Boehm’s warp-tiled kernel reaches 21779.3 GFLOP/s, 93.7% of cuBLAS, and it is the last rung.

The fringe and the last wave

Two costs are left that have nothing to do with the kernel and everything to do with the shape of the launch. Neither is a rung. Both are geometry, and both stay invisible in the GFLOP/s number until the day the matrix changes size.

The first one wastes space. A tile size is chosen before the matrix is known, so the grid of tiles has to cover the matrix, and the last row and column of tiles hang over the edge. At $4092$ with $128 \times 128$ tiles the overhang is nothing to worry about: 32 tiles a side cover 4096, the last band is 124 wide out of 128, and the launch computes 0.2% more positions than the answer has. Shrink the matrix and the same arithmetic stops being harmless. A $300 \times 300$ multiply needs 2.34 tiles a side, so it launches 3, the grid covers $384 \times 384$, and 61% of what it computes is real.

The overhanging tiles are not cheap. Their threads march the same loop over $k$ and compute the same $128 \times 128$ block of results, and the results outside the matrix are dropped at the store. NVIDIA’s Matrix Multiplication Background User’s Guide calls this tile quantization, and gives a case where the launch does 1.5 times the arithmetic for 0.39% more answer.

The second one wastes time. A block runs on one multiprocessor from start to finish, so the number of blocks the machine can hold at once is fixed: 84 multiprocessors times however many blocks the registers and the shared memory allow. Rung five put that at four for a $128 \times 128$ tile with 256 threads, so the card holds 336 of these blocks together. That many at once is called a wave. The 4092 problem has $32 \times 32 = 1024$ tiles to get through, which is 3.05 waves. The machine cannot run 0.05 of a wave, so it runs four, and the fourth one carries 16 blocks.

Over the whole launch that is 76% of the machine, and the loss arrives in jumps. A launch of 336 blocks takes one wave. A launch of 337 takes two, for 0.3% more work. Aleksa Gordić’s anatomy of high performance matmul kernels shows the same cliff on an H100 PCIe, which has 114 multiprocessors: 114 blocks finish in one wave, and 115 blocks take nearly twice as long with a single block running through the second wave while the rest of the chip waits.

Blocks are not really released in step. A multiprocessor takes a new one as soon as an old one retires, so the boundary between waves is soft and 76% is the pessimistic end of the range. It is close when every block does the same work, which is the case here. Notice also which way the loss runs. If the registers allowed only two of these blocks per multiprocessor, a wave would be 168, and the same 1024 tiles would waste less, because a coarse wave rounds up by more than a fine one.

Both of these are the same fringe. One of them rounds the answer up to whole tiles and the other rounds the launch up to whole waves, and each is a fixed shape meeting a size that was never chosen to fit it. That is one more reason the tile size cannot be reasoned out from first principles.

Six parameters, found by search

Six parameters have appeared: $BM$, $BN$, $BK$, $TM$, $TN$, and the threads per block. They are constrained by each other and by the hardware, and the best combination is not derivable. It depends on the register file, the shared memory, the number of multiprocessors and the ratio of bandwidth to arithmetic, and every one of those changes between GPU models.

So the honest method is search. Boehm compiles every legal combination and times it, which finds $BM = BN = 128$, $BK = 16$, $TM = TN = 8$ on his card, and moves his kernel from 18237.3 to 19721.0 GFLOP/s, 84.8% of cuBLAS, with no new idea at all. On his ladder the search is kernel 9 and the warp tiling of the previous section is kernel 10, which is the order the figure below uses. This is also why cuBLAS is a large library rather than a function: it ships many kernels and picks between them at run time by shape and by device.

What is left

Ninety-four percent is where this ladder stops and where the published tuned kernels start. Four things stand between the two, and each is one paragraph here and a week of work in practice.

Bank conflicts. Shared memory is 32 banks wide. If the 32 lanes of a warp read 32 addresses that land on the same bank, the hardware serialises them and the read costs 32 times what it should. Salykov pads the leading dimension of the shared tile from 128 floats to 132, which shifts each row by four banks and breaks the collision, and he splits the thread tile into four pieces for the same reason.

Double buffering. In the loop of rung three, the block loads a chunk, waits, computes, waits, and loads the next. The arithmetic units are idle during both waits. Keep two chunks in shared memory instead, compute from one while the next is arriving, and the waiting disappears into the computing. It costs twice the shared memory, which is why Salykov’s kernel reserves 16 KB for what would otherwise be 8 KB.

Asynchronous copies. On Ampere and later, cp.async moves data from global memory straight into shared memory without landing in registers on the way, so the copy does not occupy the thread that issued it. Salykov uses it in his larger kernel and reports better speed and lower power than his own synchronous version.

Split-K. When $M$ and $N$ are small and $K$ is large, there are not enough tiles of $C$ to fill the machine. Splitting the $k$ range across several blocks and summing their partial results gives the chip something to do, at the price of an extra pass over $C$.

Salykov benchmarks with the GPU clocks locked and the L2 cache flushed between replays, so that the numbers mean something, and reports that his kernels beat cuBLAS on an RTX 3090 under those conditions. He is equally careful about the cost: his 128x128x8 kernel averages 3% to 4% faster than cuBLAS while drawing 12% more power, which pushes the card into its power limit at the largest sizes and takes some of the speed back. The larger kernel, the one using cp.async, wins on both counts. His sgemm.cu repository is the one to read after Boehm’s.

Those four are the end of this ladder, and the ladder holds two things fixed the whole way up. The multiplying is done by the plain arithmetic units, and a thread owns the numbers it multiplies. The fastest kernels on current hardware give up both. One instruction hands a whole tile to a tensor core, and a copy engine fetches the next tile on one thread’s say-so. The warps that are left over spend their time running the pipeline between those two, which is why a modern kernel reads more like a factory floor than a loop. Lecture 26.5 builds that kernel.

Where this is going

Look back at what actually happened. The block tile, the warp tile and the thread tile are the same construction three times, at three scales, each one choosing a rectangle of $C$ and pulling its operands one level up the memory hierarchy. The formula $m\,n / (m + n)$ governed all three. The code did not say that once. It said it three times, in three different notations, with six tuning parameters threaded through by hand.

That repetition is a sign of missing algebra. There is an object here that nobody named: the function from a coordinate to an address, together with the rules for cutting it into tiles and composing the pieces. Naming it is what CUTLASS and CuTe do, and it turns the six parameters above into arithmetic on layouts. That is Lecture 27, and it closes the part by making the whole of Lecture 23 into an algebra.