GPU Kernels for ML Engineers

One linear path from "thread, block, warp" to running a model on a GPU — both directions a real engineer needs. The forward half (serving): turning tokens into a stream. The backward half (training): turning a loss into gradients that fit in memory. Plus the optimization loop that makes either fast, and an interview-conditions capstone. Thirty-three lessons, strictly linear.

How to read this track

Part I (01–09): CUDA primitives. Part II (10–17): serving — the forward pass, what vLLM/SGLang actually do. Part III (18–24): the optimization loop — PyTorch dispatch, profiling, Triton, torch.compile, performant patterns. Part IV (25–30): training kernels — the backward pass: matmul backward, fused cross-entropy, norm/activation backward, FlashAttention backward, fused optimizers. Part V (31–33): the same CUDA under interview conditions. Each lesson assumes only the lessons before it.

Mental stance

A kernel is where bytes move through HBM, SRAM, registers, and tensor cores — and that lens applies whether you're serving a model (forward: which work exists, which runs together, how it streams) or training one (backward: what to save vs recompute, where the gradient reductions go, what bounds the memory peak). Optimizing the wrong layer — or the wrong half — is the most common mistake in this domain. The general engineer owns both halves.

The dependency chain

The track is a directed line, not a wheel. CUDA primitives generate the kernel ideas; the kernel ideas generate the memory-layout ideas; the memory-layout ideas generate the scheduling ideas; the scheduling ideas generate the framework. If you skip a step the next step looks like a folk trick instead of a forced move. Part IV then runs the whole line backwards: every gradient kernel is the reverse-mode twin of a forward op, and the new forced move is that activation memory — not FLOPs — is what you run out of.

Part I · CUDA from first principles

Hardware execution model, memory pyramid, your first kernel, the coalesced-access trap, tile-and-share, warp divergence, the reduction template, occupancy as a budget, and the tensor-core instruction. By the end you can read any kernel in Part II.

GPU execution model

A GPU is a hardware lattice with a specific shape. Threads, warps, blocks, SMs — and how every CUDA concept maps onto a piece of the lattice.

Memory hierarchy

HBM, L2, shared memory, registers — the pyramid of "places," with bandwidths. Why the kernel game is about moving the right bytes to the right place at the right time.

First kernel — vector add

The smallest useful CUDA program. __global__, launch syntax, host/device split, the boundary check. Walk this and you can read any kernel.

Coalesced memory access

A warp's 32 threads can land in one memory transaction or up to 32 — depending on stride. The single biggest beginner trap.

Shared memory & tiled matmul

Tile A and B into SMEM, reuse each byte O(N) times, drop HBM traffic by ~N, become compute-bound. The reason FlashAttention works.

Warps, divergence, sync

32 threads run the same instruction. When they disagree, the warp serialises. Branch warp-aligned or use predication.

Reductions

Warp shuffles, shared memory, atomics — the parallel-tree template behind softmax, normalization, and every "produce one number from many."

Occupancy

More warps in flight → better latency hiding, but each costs registers and SMEM. Maximum occupancy isn't always fastest.

Tensor cores

One hardware instruction does a 16×16 matrix multiply. Every meaningful matmul on a modern GPU lands here.

Part II · LLM serving on top of Part I

Where the primitives become FlashAttention, paged KV, radix prefix reuse, continuous batching, CUDA graphs, and the framework that turns HTTP requests into batches kernels can execute.

GPU as a bandwidth machine

Roofline synthesis of Part I, with H100 numbers. The accountant you run before asking for a custom kernel.

Transformer forward as a kernel chain

A forward pass is a sequence of GEMMs and an attention. Compute KV cache size from first principles and see exactly where bytes go each token.

Attention asymmetry & FlashAttention

Prefill compute-bound, decode bandwidth-bound, same math. FlashAttention's online softmax in seven lines, and what it does not fix.

Paged KV: virtual memory for caches

Why contiguous allocation wastes HBM, how block tables fix it, what indirection costs the kernel.

Prefix reuse: hash caches and radix trees

When prompts share prefixes, share the KV. Hash caches, radix tries, eviction, cache-aware routing.

Scheduling and graph capture

Continuous batching, chunked prefill, CUDA graphs — making the scheduler's traffic look regular and cheap to launch.

Beyond attention: GEMM, quant, MoE, sampling

The rest of the decode chain. Quantization as packing-plus-kernel, MoE dispatch, GPU sampling.

The serving framework: synthesis

HTTP → tokenizer → queue → scheduler → worker → kernel → sampler → stream. Where time goes, how to debug, vLLM and SGLang as compositions.

Part III · From measurement to optimization

Now that you know what kernels exist (Parts I–II), this part is the loop: trace PyTorch → kernel, profile, read the profile, write Triton or compile, tune the PyTorch-level patterns that surround everything. The toolkit for making real code faster.

PyTorch op → kernel launch

The dispatch chain (Python → ATen → CUDA → SM). Where the 10–30 µs per op goes, and why eager mode is sometimes 5× slower. Sync points.

Anatomy of a fused kernel

A working Triton RMSNorm read line by line: program model, masked loads, on-chip reduction, autotune. The template that recurs in every fused kernel.

Profiling toolkit: three views

torch.profiler asks "which line." nsys asks "where are the gaps." ncu asks "why is this kernel slow." Three tools, three questions, three overhead profiles.

Reading a profile: roofline in practice

Three numbers (DRAM%, SM%, occupancy) place a kernel on the four-quadrant grid. Three worked examples; the named stall reasons and their cures.

Writing Triton kernels

The DSL in one table; a fused GELU+linear with autotune; the design loop; when Triton wins and when libraries beat it. All the production gotchas.

torch.compile: compiler-driven fusion

Dynamo → AOTAutograd → Inductor. Graph breaks, dynamic shapes, cache, modes. When compile beats hand-Triton and how to verify the speedup.

Performant PyTorch patterns & synthesis

Seven habits that decide whether your kernels run healthy: async H2D, allocator, contiguous, mixed precision, autograd, dataloader, sync discipline. Plus the closing synthesis.

Part IV · Training kernels — the backward pass

The other half of the engineer. Parts I–III were the forward pass and how to make it fast; this part is what happens when the model has to learn. Each lesson is the gradient kernel of a forward op you already know — and the through-line is that training is bottlenecked by memory (the activations the backward must consume), not by the extra FLOPs. Save-vs-recompute, the two-GEMM rule, fused loss, and the optimizer that consumes every gradient.

The backward pass as a kernel chain

Autograd as a tape replayed in reverse; the 6N FLOP rule (backward ≈ 2× forward); why the bottleneck is activation memory, not compute; and the save-vs-recompute decision every op makes.

Matmul backward — the two-GEMM rule

One forward GEMM becomes two: dX = dY·Wᵀ and dW = Xᵀ·dY. Why dW contracts over the batch dimension, needs split-K, and demands fp32 accumulation. The transpose you re-stride instead of materialize.

Fused cross-entropy — the memory kernel

The logits are the biggest tensor in the model. Never materialize softmax + its gradient: fuse loss + dlogits, chunk the lm-head matmul over tokens, online-softmax it. The single biggest training-memory win.

Norm & activation backward

dx recomputes the row statistic instead of storing it; dγ/dβ reduce across every token in the batch (atomics vs deterministic two-pass). Activation derivatives fused into the matmul epilogue.

FlashAttention backward

The flagship. Recompute the N² scores from the saved log-sum-exp; the D = rowsum(dO∘O) per-row correction; dK/dV accumulate cleanly over the K-loop while dQ needs atomics. Recompute-for-memory at its most valuable.

Fused optimizers & the training step

Adam is pure bandwidth and a launch-overhead trap → multi-tensor fusion. Mixed-precision loss scaling and fp32 master weights, the 16-bytes-per-param floor, gradient checkpointing, and the full step assembled.

Part V · CUDA under interview conditions

The Part I primitives, now drilled the way a kernel interview asks for them: the declarations and host/device call flow you write from memory, the optimized snippets worth memorizing, and the algorithm/data-structure thinking behind reductions, scans, histograms, sparse traversal, top-k, and work queues. Applied practice, not a new dependency in the chain.

CUDA common operations & call flow

The declarations and host/device call flow you should be able to write cold: __global__/__device__, cudaMalloc/cudaMemcpy, streams and events, grid-stride loops, the launch boundary as the only grid-wide barrier.

Highly optimized CUDA snippets

The snippets worth knowing by heart: coalesced loads, shared-memory tiling, warp-shuffle reductions, atomics, and the small patterns that separate a textbook kernel from a fast one.

Algorithms, data structures & the line of thinking

How to reason out a kernel on the spot: reductions and scans, histograms, CSR/SpMV sparse traversal, top-k, and BFS over a graph frontier — with the atomics that make each correct.

What each lesson adds

Lesson	New first-principles idea	New trade-off you can quantify
01	Threads run as warps inside blocks inside SMs.	SM-residency vs warps in flight.
02	Memory closer to compute is smaller and faster.	Bandwidth at each tier.
03	Launch syntax + host/device split.	Launch overhead vs work.
04	Coalescing — one transaction vs many.	Stride vs transaction count.
05	Tile and reuse — turns memory-bound into compute-bound.	Tile size vs SMEM cap.
06	Warp execution and divergence.	Branch alignment vs predication cost.
07	Parallel reduction tree.	Warp shuffle vs SMEM vs atomic.
08	Occupancy is a per-SM budget.	Registers/thread vs warps/SM.
09	Tensor cores do the matmul; everything feeds them.	Tile shape vs mma fragment shape.
10	Time = max(bytes/BW, FLOPs/peak) + launch.	Fusion saves a round trip; how many?
11	KV cache size = 2·L·H·d·b per token.	Per-step bytes vs FLOPs at a given context length.
12	Prefill and decode have different arithmetic intensities.	FlashAttention HBM savings as f(T, d).
13	KV memory is virtual memory with page tables.	Block size: tail vs table vs indirection.
14	Prefix sharing is workload structure, not a kernel feature.	Prefill compute saved at a given shared fraction.
15	Scheduler makes irregular traffic look regular.	Token budget; CUDA graph win at small batch.
16	Quantization is a layout decision.	Weight bandwidth saved vs dequant overhead.
17	End-to-end = API + queue + prefill + decode + stream.	Where 1 ms of latency improvement pays back.
18	Each PyTorch op pays a 10–30 µs CPU-side dispatch tax.	Step-level CPU vs GPU time as f(ops, fusion).
19	Every fused kernel is prologue → load → math → epilogue.	Fused HBM bytes vs unfused, with launch overhead.
20	Three tools answer three different questions.	Tool choice vs symptom; overhead vs coverage.
21	DRAM% × SM% × occupancy place a kernel on a 2×2 grid.	Bottleneck class → first fix to try.
22	Triton is tile-level; CUDA is thread-level.	Triton vs library vs no-op decision.
23	Dynamo + AOTAutograd + Inductor fuse for you.	Predicted speedup as f(launch share, fusion potential, graph breaks).
24	Seven habits keep the GPU fed.	Where the next hour of work goes given a current profile.
25	Backward = the forward chain reversed; training is activation-memory-bound.	FLOPs (6N rule) vs activation peak; recompute trade.
26	One forward GEMM → two backward GEMMs.	dW contracts the batch dim → split-K + fp32.
27	The logits are the biggest tensor; fuse the loss.	3×\|logits\| naive vs one chunk; ~24× memory cut.
28	Recompute the row stat; weight-grads reduce across tokens.	Atomics vs two-pass workspace; determinism.
29	Recompute the N² scores from the saved log-sum-exp.	O(N²) memory avoided vs ~2.5× recompute FLOPs.
30	Adam is bandwidth + launch-bound; fuse it.	Per-tensor vs foreach launches; the 16P floor.
31	The call-flow API is a small, memorizable surface.	Which transfers/streams overlap vs serialize.
32	Fast snippets reuse the same five primitives.	Coalescing/SMEM/shuffle wins per snippet.
33	Atomics turn a parallel idea into a correct kernel.	Contention vs correctness in scans, top-k, BFS.

Primary references

Topic	Reference
vLLM paged kernel	vLLM paged attention kernel note
vLLM backend routing	vLLM attention backend feature support
vLLM serving architecture	vLLM architecture overview
SGLang attention backends	SGLang attention backend docs
SGLang prefix reuse	SGLang RadixAttention concept doc
SGLang gateway/router	SGLang Model Gateway docs
Kernel library layer	NVIDIA FlashInfer overview
FlashAttention original	Dao et al., FlashAttention
FlashAttention-2 (backward)	Dao, FlashAttention-2: better parallelism & work partitioning
Fused cross-entropy (training memory)	Liger Kernel: efficient Triton kernels for LLM training
Mixed-precision training	Micikevicius et al., Mixed Precision Training (loss scaling, fp32 master)