How a GPU Runs Deep Learning: Memory, the Roofline, and a First Kernel

PyTorch from the Inside treated the GPU as a device that work is sent to and waited on. This post opens it. The reason to care is practical: nearly every performance question in the rest of the series, why decoding is slow, why FlashAttention exists, why quantization speeds up serving, why batch size matters, has the same answer, and it’s about memory, not arithmetic. The plan:

• what’s inside a GPU, using the H100 as the example, with numbers from NVIDIA’s whitepaper and datasheet;
• how work is organized into threads, warps and blocks;
• the roofline, which predicts whether an operation is limited by math or by memory from one number, its FLOPs per byte;
• real kernels, written and timed.

Table of contents

Many slow cores, one wide pipe

A CPU is built for latency: a few powerful cores, large caches and elaborate prediction, so that one sequence of instructions finishes as fast as possible. A GPU is built for throughput: thousands of simple cores running the same instructions on different data, so that a huge amount of independent work finishes as fast as possible, even though each piece is slow. Deep learning is almost entirely the second kind of work: the same multiply-add on millions of numbers.

The H100 SXM, from NVIDIA’s Hopper whitepaper and datasheet:

(The 2022 whitepaper gave preliminary figures of 1,000 and 60 TFLOP/s and 3 TB/s; these are the datasheet’s final values, and the datasheet’s 1,979 TFLOP/s for BF16 assumes sparsity, twice the dense 989.)

Read it top to bottom and the memory gets bigger, farther away and slower. Registers are the fastest, private to one thread. Shared memory is a small on-chip scratchpad that the threads of one block share, and the programmer manages it by hand. The L2 cache is shared by the whole chip. HBM is the main memory, 80 GB, and the only level big enough to hold a model. Every number a kernel uses has to come up this hierarchy from HBM, and the way to go fast is to bring each number up once and use it many times while it’s close.

The last two rows are the other key fact. The tensor cores are specialized units that do small matrix multiplications, and they provide almost all of the H100’s arithmetic, 15 times what its ordinary cores can do. Only work shaped as a matrix multiplication can use them, which includes convolutions and attention once they’re lowered to matrix multiplies. Everything else, the elementwise operations and reductions of Dropout and Normalization, runs on the ordinary cores.

Threads, warps, blocks, grids

A GPU program, a kernel, is written from the point of view of one thread, and launched on many at once. The threads are organized in three levels:

• Threads each run the kernel with their own index, which they use to decide which data to work on.
• Warps are groups of 32 threads that execute the same instruction at the same time, as long as they don’t take different sides of a branch. An SM schedules warps, not threads, and keeps up to 64 of them resident so it can switch to another warp whenever one is waiting for memory. That switching is how a GPU hides memory latency: not with big caches, but with enough other work. How many warps actually fit, the occupancy, depends on each thread’s registers and shared memory: an SM’s 65,536 registers over 2,048 threads is 32 per thread, so a kernel that needs 128 registers per thread gets only 16 warps. (Apple calls a warp a SIMD group, also 32 threads.)
• Blocks are groups of threads, up to 1,024, that run on one SM and can share its shared memory and wait for each other at a barrier. (Metal: threadgroup, and threadgroup memory.)
• The grid is all the blocks of one launch. Blocks are independent, which lets the hardware spread them over however many SMs it has.

This series is written on an Apple M4 laptop, which has no NVIDIA GPU, so the kernels here are written in Apple’s Metal Shading Language and run through PyTorch’s torch.mps.compile_shader. Metal’s model is CUDA’s with different names, and the concepts, the code shapes and the conclusions carry over directly; each section gives the CUDA name. The laptop is also a useful contrast: same principles, about 300 times less compute and 30 times less memory bandwidth.

The kernels live in one file, kernels.metal. It starts with the simplest possible kernel, adding two vectors, one thread per element:

#include <metal_stdlib>
using namespace metal;

// Metal launches exactly `threads` threads, so these kernels skip bounds
// checks; CUDA rounds the grid up to whole blocks and needs `if (i < n)`.

// One thread per element: the GPU's "hello world".
kernel void vector_add(device float* out, device const float* a, device const float* b,
                       uint i [[thread_position_in_grid]]) {
    out[i] = a[i] + b[i];
}

// Every thread copies one element, but reads it from position (i * stride) mod n.
// Stride 1: neighbouring threads read neighbouring addresses.
kernel void strided_copy(device float* out, device const float* a, constant uint& n,
                         constant uint& stride, uint i [[thread_position_in_grid]]) {
    out[i] = a[(i * stride) % n];
}

// Matrix multiply, one thread per output element, reading straight from memory.
kernel void matmul_naive(device const float* A, device const float* B, device float* C,
                         constant uint& n, uint2 p [[thread_position_in_grid]]) {
    float acc = 0.0f;
    for (uint k = 0; k < n; k++)
        acc += A[p.y * n + k] * B[k * n + p.x];
    C[p.y * n + p.x] = acc;
}

// The same, but each 16 x 16 group of threads first copies a tile of A and a
// tile of B into fast on-chip threadgroup memory, then reuses each loaded
// number 16 times. Assumes n is a multiple of 16.
kernel void matmul_tiled(device const float* A, device const float* B, device float* C,
                         constant uint& n, uint2 p [[thread_position_in_grid]],
                         uint2 t [[thread_position_in_threadgroup]]) {
    threadgroup float As[16][16];
    threadgroup float Bs[16][16];
    float acc = 0.0f;
    for (uint k0 = 0; k0 < n; k0 += 16) {
        As[t.y][t.x] = A[p.y * n + k0 + t.x];
        Bs[t.y][t.x] = B[(k0 + t.y) * n + p.x];
        threadgroup_barrier(mem_flags::mem_threadgroup);     // wait until the tiles are loaded
        for (uint k = 0; k < 16; k++)
            acc += As[t.y][k] * Bs[k][t.x];
        threadgroup_barrier(mem_flags::mem_threadgroup);     // wait before overwriting them
    }
    C[p.y * n + p.x] = acc;
}

// A bigger tile, 32 x 32: twice the reuse of each number read from memory.
kernel void matmul_tiled32(device const float* A, device const float* B, device float* C,
                         constant uint& n, uint2 p [[thread_position_in_grid]],
                         uint2 t [[thread_position_in_threadgroup]]) {
    threadgroup float As[32][32];
    threadgroup float Bs[32][32];
    float acc = 0.0f;
    for (uint k0 = 0; k0 < n; k0 += 32) {
        As[t.y][t.x] = A[p.y * n + k0 + t.x];
        Bs[t.y][t.x] = B[(k0 + t.y) * n + p.x];
        threadgroup_barrier(mem_flags::mem_threadgroup);     // wait until the tiles are loaded
        for (uint k = 0; k < 32; k++)
            acc += As[t.y][k] * Bs[k][t.x];
        threadgroup_barrier(mem_flags::mem_threadgroup);     // wait before overwriting them
    }
    C[p.y * n + p.x] = acc;
}

// One level further: each thread keeps a 4 x 4 block of outputs in registers,
// so a 16 x 16 group computes a 64 x 64 tile of C. Per step of k, a thread
// loads 4 numbers of A and 4 of B from threadgroup memory and does 16
// multiply-adds with them, instead of 2 loads per multiply-add.
// Assumes n is a multiple of 64.
kernel void matmul_registers(device const float* A, device const float* B, device float* C,
                             constant uint& n, uint2 g [[threadgroup_position_in_grid]],
                             uint2 t [[thread_position_in_threadgroup]]) {
    threadgroup float As[64][16];
    threadgroup float Bs[16][64];
    float acc[4][4] = {};
    uint row0 = g.y * 64, col0 = g.x * 64, id = t.y * 16 + t.x;
    for (uint k0 = 0; k0 < n; k0 += 16) {
        for (uint e = id; e < 1024; e += 256) {              // 256 threads load each slice's 1,024 numbers, 4 apiece
            As[e / 16][e % 16] = A[(row0 + e / 16) * n + k0 + e % 16];
            Bs[e / 64][e % 64] = B[(k0 + e / 64) * n + col0 + e % 64];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint k = 0; k < 16; k++) {
            float a[4], b[4];
            for (uint i = 0; i < 4; i++) a[i] = As[t.y * 4 + i][k];
            for (uint j = 0; j < 4; j++) b[j] = Bs[k][t.x * 4 + j];
            for (uint i = 0; i < 4; i++)
                for (uint j = 0; j < 4; j++)
                    acc[i][j] += a[i] * b[j];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
    for (uint i = 0; i < 4; i++)
        for (uint j = 0; j < 4; j++)
            C[(row0 + t.y * 4 + i) * n + col0 + t.x * 4 + j] = acc[i][j];
}

Full script: kernels.metal

That file holds every kernel in this post. vector_add is the whole of the first one: each thread reads its index i and handles element i. In CUDA it would read blockIdx.x * blockDim.x + threadIdx.x; Metal hands it over directly as thread_position_in_grid. On the Python side, gpu_util.py compiles the file once and times each kernel in the two ways PyTorch from the Inside separated: one call at a time, waiting for each, for big kernels where the wait is negligible, and many calls queued with one wait at the end, which is what a training loop sees, for everything else:

import time
import torch

lib = torch.mps.compile_shader(open("kernels.metal").read())

def timed(fn, trials=50):
    """Median seconds per call, waiting for the GPU before each clock reading."""
    fn(); torch.mps.synchronize()
    times = []
    for _ in range(trials):
        t0 = time.perf_counter(); fn(); torch.mps.synchronize()
        times.append(time.perf_counter() - t0)
    return sorted(times)[trials // 2]

def timed_queued(fn, reps=200):
    """Seconds per call when `reps` calls are queued back to back and the GPU
    is waited on once: throughput, as a training loop sees it."""
    fn(); torch.mps.synchronize()
    t0 = time.perf_counter()
    for _ in range(reps):
        fn()
    torch.mps.synchronize()
    return (time.perf_counter() - t0) / reps

Full script: gpu_util.py

Memory decides: bandwidth and coalescing

Adding two vectors does one addition per 12 bytes moved (two 4-byte reads, one 4-byte write), so its speed is set entirely by how fast memory moves:

import torch
from gpu_util import lib, timed

n = 1 << 26                                             # 64M floats, 268 MB per array
a, b = torch.randn(n, device="mps"), torch.randn(n, device="mps")
out = torch.empty_like(a)
t = timed(lambda: lib.vector_add(out, a, b))
print(f"vector_add: read 2 arrays, write 1, in {t * 1e3:.1f} ms = {3 * n * 4 / t / 1e9:.0f} GB/s")
t = timed(lambda: out.copy_(a))
print(f"PyTorch copy: {2 * n * 4 / t / 1e9:.0f} GB/s")
# Odd strides are coprime with n, so every element is still read exactly once.
for stride in (1, 3, 9, 33, 1025):
    t = timed(lambda: lib.strided_copy(out, a, n, stride))
    print(f"strided_copy, stride {stride:4d}: {2 * n * 4 / t / 1e9:5.0f} GB/s")

Full script: bandwidth.py

vector_add: read 2 arrays, write 1, in 8.8 ms = 91 GB/s
PyTorch copy: 92 GB/s
strided_copy, stride    1:    93 GB/s
strided_copy, stride    3:    48 GB/s
strided_copy, stride    9:    17 GB/s
strided_copy, stride   33:    10 GB/s
strided_copy, stride 1025:     9 GB/s

• About 93 GB/s is this GPU’s memory bandwidth. The hand-written vector_add, PyTorch’s own copy, and a plain stride-1 copy all land at 91–93 GB/s, against Apple’s 120 GB/s specification for the chip’s memory, which it shares with the CPU. An H100’s 3,350 GB/s is about 30 times as much, spec for spec.
• The access pattern matters as much as the amount. The strided copy moves exactly the same bytes, every element read once, but with neighboring threads reading addresses far apart. At stride 3 bandwidth halves; at stride 33 and beyond it’s a tenth.

That’s coalescing. Memory is read in chunks of consecutive bytes, not one number at a time; on NVIDIA GPUs the chunk is a 32-byte sector. When the 32 threads of a warp read 32 consecutive 4-byte numbers, the hardware fetches four sectors, 128 bytes, and every byte is used. When they read 32 scattered numbers, it fetches up to 32 sectors, 1,024 bytes, and uses 4 of every 32: eight times the traffic for the same numbers. At small strides the loss is smaller, because a sector fetched for one warp is still in the cache when a neighboring warp needs it. This is why PyTorch from the Inside’s non-contiguous tensors matter, and why kernels are designed so that neighboring threads touch neighboring addresses.

The roofline

Every operation has an arithmetic intensity: the FLOPs it performs divided by the bytes it has to move between memory and the chip. The roofline model (Williams, Waterman and Patterson, 2009) says the speed it can reach is the lower of two ceilings:

Below a certain intensity, the ridge point $= \text{peak} / \text{bandwidth}$, memory can’t deliver numbers fast enough to keep the cores busy, and the operation is memory-bound. Above it, the cores are the limit: compute-bound. For the H100, the ridge is $989 / 3.35 = 295$ FLOPs per byte. An operation has to do nearly 300 FLOPs for every byte it reads just to use the tensor cores fully.

Measured on the laptop GPU:

import torch
from gpu_util import timed_queued as timed

# Operations from very little math per byte to a lot, measured on this GPU in
# float32, each queued 200 times and waited on once (throughput). Intensity = FLOPs / bytes that must move between memory and the chip.
rows = []
x = torch.randn(1 << 24, device="mps"); y = torch.randn_like(x)
t = timed(lambda: x + y)
rows.append(("vector add", len(x), 3 * 4 * len(x), t))
for n in (4096, 8192):
    M, v = torch.randn(n, n, device="mps"), torch.randn(n, device="mps")
    t = timed(lambda: M @ v)
    rows.append((f"matrix-vector, n={n}", 2 * n * n, 4 * (n * n + 2 * n), t))
for n in (64, 128, 256, 512, 1024, 2048, 4096):
    A, B = torch.randn(n, n, device="mps"), torch.randn(n, n, device="mps")
    t = timed(lambda: A @ B)
    rows.append((f"matmul, n={n}", 2 * n ** 3, 3 * 4 * n * n, t))
# The smallest matmul again, but 1,024 of them batched into one call.
A, B = torch.randn(1024, 64, 64, device="mps"), torch.randn(1024, 64, 64, device="mps")
t = timed(lambda: torch.bmm(A, B))
rows.append(("1,024 matmuls n=64, bmm", 1024 * 2 * 64 ** 3, 1024 * 3 * 4 * 64 * 64, t))
for name, flops, bytes_, t in rows:
    print(f"{name:22s} intensity {flops / bytes_:8.2f} FLOP/byte   {flops / t / 1e9:7.1f} GFLOP/s   "
          f"{bytes_ / t / 1e9:6.1f} GB/s")

Full script: roofline.py

vector add             intensity     0.08 FLOP/byte       6.1 GFLOP/s     73.2 GB/s
matrix-vector, n=4096  intensity     0.50 FLOP/byte      39.6 GFLOP/s     79.2 GB/s
matrix-vector, n=8192  intensity     0.50 FLOP/byte      39.7 GFLOP/s     79.5 GB/s
matmul, n=64           intensity    10.67 FLOP/byte      25.3 GFLOP/s      2.4 GB/s
matmul, n=128          intensity    21.33 FLOP/byte     273.8 GFLOP/s     12.8 GB/s
matmul, n=256          intensity    42.67 FLOP/byte    1325.1 GFLOP/s     31.1 GB/s
matmul, n=512          intensity    85.33 FLOP/byte    2778.1 GFLOP/s     32.6 GB/s
matmul, n=1024         intensity   170.67 FLOP/byte    2749.2 GFLOP/s     16.1 GB/s
matmul, n=2048         intensity   341.33 FLOP/byte    2877.5 GFLOP/s      8.4 GB/s
matmul, n=4096         intensity   682.67 FLOP/byte    2662.8 GFLOP/s      3.9 GB/s
1,024 matmuls n=64, bmm intensity    10.67 FLOP/byte     728.5 GFLOP/s     68.3 GB/s

• Low intensity rides the memory roof. The vector add (0.08 FLOP/byte) and the matrix-vector product (0.5) run near the memory bandwidth, 73–80 GB/s in this run, and their FLOP rates are tiny as a result.
• High intensity reaches the compute roof. A matrix multiply of size $n$ does $2n^3$ FLOPs on $3n^2$ float32 numbers, so its intensity is $n/6$ (in BF16, $n/3$) and grows with size. From $n = 512$ up it runs at 2.7–2.9 TFLOP/s, about the best this GPU reached in this run (the kernel ladder below measured PyTorch’s matmul at up to 3.1 TFLOP/s).
• Small matmuls hit neither roof. At $n = 64$ the intensity is 10.7, which the roofline says should give about 1 TFLOP/s, and even queued the measurement is 25 GFLOP/s. Each call has a fixed dispatch cost, and 4,096 outputs are too little work to occupy the GPU. Batching fixes it: the same 64 × 64 multiply done 1,024 times in one torch.bmm call runs at 729 GFLOP/s. (Timed one call at a time and waiting for each, as PyTorch from the Inside warned against, the small matmuls look slower still, because every measurement then includes a round trip to the GPU.)

Drag the intensity to compare the two machines:

The two ridges are the point. The laptop’s is about 31 FLOPs per byte; the H100’s is 295 for BF16 matrix math on the tensor cores. (For FP32 work on its ordinary cores, like Dropout and Normalization’s norms, the H100’s ridge is 67 / 3.35 = 20.) The laptop’s roof is drawn from its best measured matmul and best measured copy in this run, so its peak point sits on the roof by construction, and a faster run can land a little above it; the H100’s is from the datasheet. The faster the arithmetic gets relative to memory, the more operations fall on the memory side of the ridge. Decoding one token at a time is a matrix-vector product for every weight matrix, intensity about 1 in BF16, far left of either ridge, which is why The KV Cache and the cost sketch in AI Engineering, from the Inside found decode bound by memory.

A matrix multiply, from naive to tiled

A matrix multiply is where the intensity can be high, but only if the kernel reuses what it reads. Four versions of the same kernel, all in kernels.metal above:

• Naive. One thread per output element $C_{ij}$, which reads row $i$ of $A$ and column $j$ of $B$ straight from memory: $2n$ reads for $2n$ FLOPs. Nothing is reused on purpose.
• Tiled. Each block of $16 \times 16$ threads computes a $16 \times 16$ tile of $C$. It walks along the shared dimension 16 at a time: the block loads a $16 \times 16$ tile of $A$ and one of $B$ into shared (threadgroup) memory, every thread waits at a barrier until both are in, and then each thread uses 16 numbers from each tile. Every number read from main memory is used 16 times.
• Tiled 32. The same with $32 \times 32$ tiles, 1,024 threads per block: twice the reuse.
• Registers. One level further. Each thread computes a $4 \times 4$ block of outputs, kept in its registers, so a block of $16 \times 16$ threads covers a $64 \times 64$ tile of $C$. For each step along the shared dimension, a thread loads 4 numbers of $A$ and 4 of $B$ from threadgroup memory and does 16 multiply-adds with them.

In pictures, for one tile of $C$:

import torch
from gpu_util import lib, timed_queued as timed

for n in (512, 1024, 2048):
    A, B = torch.randn(n, n, device="mps"), torch.randn(n, n, device="mps")
    exact = A.cpu().double() @ B.cpu().double()        # the reference, in float64 on the CPU
    row, worst = [f"n = {n:4d}:"], 0.0
    for name, run in [("naive", lambda C: lib.matmul_naive(A, B, C, n, threads=(n, n), group_size=(16, 16))),
                      ("tiled", lambda C: lib.matmul_tiled(A, B, C, n, threads=(n, n), group_size=(16, 16))),
                      ("tiled 32", lambda C: lib.matmul_tiled32(A, B, C, n, threads=(n, n), group_size=(32, 32))),
                      ("registers", lambda C: lib.matmul_registers(A, B, C, n, threads=(n // 4, n // 4),
                                                                   group_size=(16, 16))),
                      ("PyTorch", lambda C: torch.mm(A, B, out=C))]:
        C = torch.empty(n, n, device="mps")
        t = timed(lambda: run(C))
        worst = max(worst, (C.cpu().double() - exact).abs().max().item() / exact.abs().max().item())
        row.append(f"{name} {2 * n ** 3 / t / 1e9:5.0f} GFLOP/s")
    print("   ".join(row) + f"   (all within {worst:.0e} of float64)")

Full script: matmul_ladder.py

n =  512:   naive   237 GFLOP/s   tiled   518 GFLOP/s   tiled 32   561 GFLOP/s   registers  1325 GFLOP/s   PyTorch  2869 GFLOP/s   (all within 1e-06 of float64)
n = 1024:   naive   291 GFLOP/s   tiled   514 GFLOP/s   tiled 32   581 GFLOP/s   registers  1308 GFLOP/s   PyTorch  3081 GFLOP/s   (all within 1e-06 of float64)
n = 2048:   naive   269 GFLOP/s   tiled   491 GFLOP/s   tiled 32   517 GFLOP/s   registers  1201 GFLOP/s   PyTorch  2867 GFLOP/s   (all within 2e-06 of float64)

(All agree with a float64 reference to about $10^{-6}$, the rounding of float32.) Tiling nearly doubles the naive kernel, from about 240–290 to about 500 GFLOP/s. A bigger tile barely helps. Registers help a lot, to 1.2–1.3 TFLOP/s. PyTorch’s own matmul is about twice that again. The simple counting of reuse explains less of this than it might seem to:

# How tiling changes a matmul's arithmetic intensity, counted as float32
# reads from main memory with no help from caches: an upper bound on traffic.
# Each output element takes n multiply-adds (2n FLOPs) over a row of A and a
# column of B: 2n numbers read, if nothing is reused.
ridge_here, ridge_h100 = 2878 / 93, 989e3 / 3350      # measured laptop GPU; H100 datasheet
print(f"ridge points: this GPU {ridge_here:.0f} FLOP/byte, H100 {ridge_h100:.0f} FLOP/byte\n")
print(f"{'C tile':>9s} {'reads per output':>17s} {'intensity':>13s} {'threads, one output each':>26s}")
for T in (1, 16, 32, 64, 128):
    # A T x T tile of outputs reads a T-row strip of A and a T-column strip of
    # B: 2nT numbers for T*T outputs, so each number read is used T times.
    reads_per_output = 2 / T                          # in units of n numbers
    intensity = 2 / (reads_per_output * 4)            # FLOPs per byte
    label = "naive" if T == 1 else f"{T} x {T}"
    print(f"{label:>9s} {reads_per_output:>12.3f} x n {intensity:>8.2f} FLOP/B {T * T:>18,d}")

Full script: tile_intensity.py

ridge points: this GPU 31 FLOP/byte, H100 295 FLOP/byte

   C tile  reads per output     intensity   threads, one output each
    naive        2.000 x n     0.25 FLOP/B                  1
  16 x 16        0.125 x n     4.00 FLOP/B                256
  32 x 32        0.062 x n     8.00 FLOP/B              1,024
  64 x 64        0.031 x n    16.00 FLOP/B              4,096
128 x 128        0.016 x n    32.00 FLOP/B             16,384

A $T \times T$ tile of outputs uses each number it reads $T$ times, so its intensity against main memory is at most $T/4$ FLOPs per byte, if the caches help not at all. (That’s a different number from the roofline’s $n/6$: $n/6$ is the operation’s intensity if every number were read exactly once, and $T/4$ is what one kernel design actually achieves.) By that bound the naive kernel, at 0.25 FLOP/byte and 93 GB/s, couldn’t exceed 23 GFLOP/s, and it runs at 270. The caches are doing most of the reuse: in a warp, 16 threads want the same element of $A$, and neighboring warps want the same rows. By the same bound, 16 × 16 tiles cap at 372 GFLOP/s, and the tiled kernel beats that too, so main memory isn’t what’s holding it at 500.

What is, is the inner loop. In the tiled kernel every multiply-add needs two loads from threadgroup memory, whatever the tile size, which is why the 32 × 32 tile changes little. The register version breaks that: 8 loads feed 16 multiply-adds, and it runs 2.5 times faster. That’s the general pattern. Each level of tiling raises the intensity against the slower level it reads from: shared-memory tiles against main memory, register tiles against shared memory, and on NVIDIA GPUs the tensor cores’ own matrix fragments against registers. Fast libraries also overlap loading with computing: while the cores work on one slice, the next is already being copied in (asynchronously on Hopper, with its tensor memory accelerator), so the cores never wait at a barrier. FlashAttention uses exactly these tricks.

The table’s last column is the other constraint. With one output per thread, a tile can’t be bigger than the block: at most 1,024 threads, so 32 × 32. A $128 \times 128$ tile, the size the H100’s libraries use, only exists with register tiling, each thread owning many outputs. And shared memory doesn’t hold two square tiles; it holds thin slices along the shared dimension, several in flight at once. How big can they be here?

import torch

# How much threadgroup (shared) memory will this GPU give one block?
for kb in (16, 32, 64):
    source = f"""
    #include <metal_stdlib>
    using namespace metal;
    kernel void probe(device float* out, uint i [[thread_position_in_grid]],
                      uint t [[thread_position_in_threadgroup]]) {{
        threadgroup float buf[{kb * 256}];                   // kb kilobytes of float32
        buf[t] = i;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        out[i] = buf[(t + 1) % 32] + buf[{kb * 256 - 1}];
    }}"""
    try:
        out = torch.zeros(64, device="mps")
        torch.mps.compile_shader(source).probe(out, threads=64, group_size=32)
        torch.mps.synchronize()
        print(f"{kb} KB of threadgroup memory: ok")
    except Exception as e:
        print(f"{kb} KB of threadgroup memory: refused ({str(e).split('\"')[1][:60]})")

Full script: smem_probe.py

16 KB of threadgroup memory: ok
32 KB of threadgroup memory: ok
64 KB of threadgroup memory: refused (Threadgroup memory size (65536) exceeds the maximum threadgr)

This GPU gives a threadgroup at most 32 KB. The H100’s 228 KB of shared memory per SM is sized for pipelines of such slices: for a $128 \times 128$ BF16 tile, a slice of $A$ and one of $B$, 64 deep, take 32 KB, and three to five of them in flight fit comfortably. BF16 halves the bytes of every slice, one more reason lower precision is faster.

Why matrix shapes matter

A matrix-multiply kernel splits the output into tiles and hands them out to the SMs. Take a kernel with fixed 128 × 128 tiles, one per SM at a time, on an H100’s 132 SMs: the tiles run in waves of 132, and the last wave is paid for in full even if it’s nearly empty:

import math
# Arithmetic for an H100: a matmul kernel splits the output into 128 x 128
# tiles and hands them to 132 SMs in waves. A last, partly empty wave is paid
# for in full.
for n in (1024, 1408, 1536, 4096, 4097, 4224):
    per_side = math.ceil(n / 128)
    tiles = per_side ** 2
    waves = math.ceil(tiles / 132)
    useful = (n / (per_side * 128)) ** 2              # edge tiles are partly padding
    print(f"n = {n}: {tiles:4d} tiles -> {waves:2d} waves; SMs busy {tiles / (waves * 132):4.0%}, "
          f"useful work in the tiles {useful:4.0%}, overall {useful * tiles / (waves * 132):4.0%}")

Full script: waves.py

n = 1024:   64 tiles ->  1 waves; SMs busy  48%, useful work in the tiles 100%, overall  48%
n = 1408:  121 tiles ->  1 waves; SMs busy  92%, useful work in the tiles 100%, overall  92%
n = 1536:  144 tiles ->  2 waves; SMs busy  55%, useful work in the tiles 100%, overall  55%
n = 4096: 1024 tiles ->  8 waves; SMs busy  97%, useful work in the tiles 100%, overall  97%
n = 4097: 1089 tiles ->  9 waves; SMs busy  92%, useful work in the tiles  94%, overall  86%
n = 4224: 1089 tiles ->  9 waves; SMs busy  92%, useful work in the tiles 100%, overall  92%

• Wave quantization. $n = 1408$ makes 121 tiles, one wave, 92% busy. $n = 1536$ makes 144, which is 12 more than one wave holds: two waves, and the second is 91% empty. For this kernel, a matrix 9% larger could take nearly twice as long.
• Partial tiles. $n = 4097$ needs one more row and column of tiles than $n = 4096$, almost all padding. The same work costs 9 waves instead of 8.

Real libraries soften both. They choose among tile shapes for each problem, split the shared dimension across SMs, and schedule tiles persistently, so the cliffs are smaller than this table, but they’re still visible in benchmarks (NVIDIA’s matrix-multiplication performance guide shows them). Multiples of the tile size, 128 here, avoid the partial tiles; they don’t prevent a nearly empty last wave, which depends on the tile count against the SM count. The same logic is why a language model’s vocabulary is often padded: nanoGPT pads GPT-2’s 50,257 tokens to 50,304, a multiple of 64. (On this laptop, run-to-run timing varies by more than these effects, so the example is arithmetic, not a measurement.)

Where this leads

• The KV Cache is this post’s roofline applied to decoding: every generated token is a pass over the weights at intensity about 1.
• FlashAttention is a tiled kernel for attention: it keeps tiles of queries, keys and values in shared memory so the full attention matrix never goes to HBM.
• Quantization from Scratch is intensity by other means: fewer bytes per weight, so a memory-bound operation does more FLOPs per byte moved.
• vLLM Architecture batches requests in serving, which turns many matrix-vector products into one matrix-matrix product, moving decode to the right on the roofline.

Takeaways

• A GPU is thousands of simple cores behind a memory hierarchy: registers, shared memory, L2, HBM, each bigger and slower. Going fast means bringing each number up once and reusing it while it’s close.
• Kernels run as threads in warps of 32, grouped into blocks that share an SM’s shared memory, spread over a grid. Warps waiting on memory are swapped for ones that aren’t.
• Coalescing: neighboring threads should read neighboring addresses. The same bytes read scattered ran at a tenth of the speed; a scattered warp moves up to eight times the bytes it uses.
• The roofline: attainable speed is the lower of peak compute and intensity × bandwidth. The H100’s ridge is 295 FLOPs per byte, so most operations other than large matrix multiplications are memory-bound. Tiny operations are bound by per-call overhead instead; batching them into one call (1,024 small matmuls in one bmm) recovers most of the speed.
• Tiling works at every level of the hierarchy. Shared-memory tiles doubled a naive kernel, a bigger tile added little because the inner loop was bound by shared-memory loads, and a register level added 2.5×. Fast libraries stack all the levels, overlap loads with compute and use tensor cores.
• Tile counts and SM counts make shapes matter: multiples of the tile size avoid partial tiles, but a nearly empty last wave depends on the tile count against the SM count.

References

Start with the CUDA programming guide’s chapters on the programming model and memory hierarchy, then the roofline paper. CS336 lectures 5–6 apply both to language models.

• NVIDIA, NVIDIA H100 Tensor Core GPU Architecture (whitepaper, 2022): SMs, tensor cores, registers, shared memory, L2.
• NVIDIA, H100 Tensor Core GPU (datasheet): the final figures, 989 TFLOP/s dense BF16, 67 TFLOP/s FP32, 3.35 TB/s.
• NVIDIA, CUDA C++ Programming Guide (NVIDIA docs): threads, blocks, warps, shared memory and coalescing.
• NVIDIA, Matrix Multiplication Background User’s Guide (NVIDIA docs): tile and wave quantization, and dimension alignment.
• Williams, Waterman and Patterson, Roofline: An Insightful Visual Performance Model for Multicore Architectures (Communications of the ACM, 2009): the roofline model.
• Apple, Metal Shading Language Specification (Apple docs): threadgroups, SIMD groups and threadgroup memory, the Metal names for CUDA’s terms.
• Stanford CS336, Language Modeling from Scratch, lectures 5–6 (2025): GPUs, kernels and the roofline for language models.