CUDA
on the Metal

What a GPU is made of, how its programming model maps onto that silicon, and a line-by-line trace of a real LLM kernel — RMSNorm.

Navigate with ← → or Space · F for fullscreen

Inside the L4: The Whole Chip

The host talks to the GPU over PCIe. On-chip, data climbs a hierarchy: the further from the cores, the larger and slower.

Inside a Streaming Multiprocessor

Every SM is a self-contained processor: its own register file, its own shared memory, and the cores that do the math. A block lives entirely inside one SM.

What You Write vs What Runs It

You program in logical terms. The hardware maps each one onto physical silicon. Understanding that mapping is the whole game.

Warps & SIMT

Zoom inside one SM. You write code for a single thread, but the hardware executes a warp — 32 threads — together. One instruction, 32 threads: Single Instruction, Multiple Threads.

A Warp at Work: Adding a Vector of 64

64 elements need 64 threads — exactly 2 warps. Step through the six instructions the SM issues to compute C = A + B.

Each warp adds its 32 pairs in a single FADD — the whole vector is done in six warp-instructions.

128 Cores, ~1,536 Threads: Latency Hiding

An SM has only 128 physical cores, yet keeps up to ~1,536 threads resident at once. Why so many? A memory load costs ~200 cycles — so instead of stalling, the SM switches the cores to another ready thread.

The Block, and Why It Caps at 1024

CUDA's way to launch threads is the block — a team that cooperates through shared memory and __syncthreads(). The whole team must fit on one SM at once, and that constraint is exactly why blocks cap at 1024.

More Data Than Threads? Loop.

2048 elements but only 1024 threads. You don't ask for more threads — each thread processes several elements in sequence: a grid-stride loop.

Kernel Scheduling: The Hardware Queue

Launch a grid of 2,048 blocks on 58 SMs. The GigaThread Engine fills all 58, and the remaining blocks wait in a hardware queue.

Walkthrough: an RMSNorm Kernel

The normalization at the heart of modern LLMs. Compute the root-mean-square of each row, scale by it, then multiply by a learned weight.

The Naive Version: Eager PyTorch

This is correct, readable, and slow. Every operation on the input is a separate CUDA kernel launch that reads from and writes back to VRAM.

class RMSNorm(nn.Module):
    def forward(self, x):
        # fp32 for numerical stability
        v = x.float().pow(2) \
              .mean(-1, keepdim=True)
        xn = x * torch.rsqrt(v + self.eps)
        return self.weight * xn

Source: chapter_1/llama_inference.py · next, we fuse all six ops into one kernel — the four phases on the silicon.

Four Phases on the Silicon

One block of 1024 threads turns one row of 2048 numbers into its normalized output. It happens in four phases — the next four slides walk each one line by line.

Phase 1 — Local Sums

Every thread squares its own elements and accumulates them in a private register. With d = 2048 and 1024 threads, each thread handles two elements via the grid-stride loop.

// per thread · tid = threadIdx.x
float local = 0.0f;        // private register
for (int i = tid; i < d; i += blockDim.x) {
    float v = x[row*d + i];   // LDG from VRAM
    local += v * v;           // FFMA, stays in reg
}
// i = tid, then tid+1024  → 2 passes

reads VRAM writes registers

Phase 2 — The Reduction

1024 partial sums must become one. Registers are private, so we reduce within each warp by shuffle, hand off through shared memory, then reduce once more.

// 1 · collapse 32 lanes → lane 0  (5 steps)
for (int off = 16; off > 0; off >>= 1)
    local += __shfl_down_sync(FULL, local, off);

// 2 · each warp's lane 0 → shared memory
if (lane == 0) smem[warpId] = local;
__syncthreads();

// 3 · warp 0 reduces the 32 partial sums
float total = (tid < 32) ? smem[tid] : 0;
if (warpId == 0)
    for (int off = 16; off > 0; off >>= 1)
        total += __shfl_down_sync(FULL, total, off);

registers → shared → registers   shuffle is ~5 cyc/step and never touches shared memory; shared memory is only the cross-warp handoff.

Phase 3 — Final Math & Constant Memory

Thread 0 now holds the grand sum of squares. It computes the single RMS scalar and publishes it so all 1024 threads can use it.

if (tid == 0) {
    float mean = total / d;        // d ← constant mem
    smem_rms = rsqrtf(mean + eps);  // eps ← constant mem
}
__syncthreads();   // RMS now visible to all

reads constant mem writes shared mem

Phase 4 — Applying the Weights

All 1024 threads wake up, read the shared RMS value, and write the normalized, weighted output back to VRAM — again two elements each.

float rms = smem_rms;          // from shared mem
for (int i = tid; i < d; i += blockDim.x) {
    y[row*d + i] =
        x[row*d + i] * rms * w[i]; // STG to VRAM
}
// w[i]: left in VRAM, hot in L1 cache

reads VRAM + shared + L1 writes VRAM

Where Data Lives Is the Optimization

Two choices in the trace look lazy but are the real performance wins.

You write threads.
The hardware runs warps.

Every kernel is a negotiation between these three forces. That's CUDA on the metal.

From Theory to 7.2× Faster

Same math, same GPU. A single fused CUDA kernel replaces PyTorch's chain of eager ops — measured across the three RMSNorm calls per transformer step.