Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 6: Kernels, Triton, XLA

HBM scales while on-chip memory remains constant

Modern NVIDIA GPUs (B200s) maintain roughly 100-200 SMs with ~256KB shared memory per SM, but High Bandwidth Memory (HBM) capacity grows significantly, widening the speed gap between fast local memory and slow global memory.

Registers and shared memory minimize HBM access

Fast on-chip memory (registers and L1/shared memory) local to each SM enables threads to communicate and compute without expensive round-trips to high-latency HBM, which is the primary bottleneck in kernel optimization.

Thread blocks map to SMs for shared memory access

While the programming model presents a clean hierarchy of threads, thread blocks, and grids, thread blocks specifically schedule onto Streaming Multiprocessors to enable fast communication via shared memory rather than slow HBM.

Warps execute in lockstep causing divergence penalties

Threads are grouped into 32-thread warps that must execute identical instructions simultaneously, causing serialized execution when threads take different branches (control divergence), which severely reduces throughput.

Warp scheduling hides latency via zero-cost switching

SMs maintain multiple resident warps and switch between them instantly when one stalls waiting for HBM, allowing the hardware to hide memory latency by keeping compute units busy with other warps.

Register usage limits thread occupancy

Each thread can use at most 255 registers, meaning threads using many registers reduce the total number of concurrent threads (occupancy), though fewer threads doing more work via thread coarsening can sometimes improve efficiency.

Shared memory bank conflicts serialize access

Shared memory is divided into 32 banks, and when multiple threads in a warp access the same bank simultaneously, the hardware must serialize these accesses, creating delays that require swizzling techniques to mitigate.

Uncoalesced global memory wastes bandwidth

When threads in a warp access scattered HBM locations rather than contiguous cache lines (128 bytes), the hardware executes multiple memory transactions instead of one, drastically reducing effective bandwidth utilization.

Block quantization causes SM underutilization

Since thread blocks cannot split across SMs, launching a block count that doesn't evenly divide the total SM count (e.g., 160 blocks on 148 SMs) leaves some SMs idle during the final execution wave.