The Stack Below the Stack (Part 2): Below Python

The Stack Below the Stack, a 3-part series on how modern LLM inference actually works, told through a single DeepSeek V4 dtype bug.

• Part 1 · Physics of a request : why the first token is a different problem from every token after it, and why batching exists.
• Part 2 (this post) · Below Python: what actually runs under vllm serve, and why the escape hatches failed.
• Part 3 · Serving at scale : recipes, fleets, traffic shapes, and the bug resolved layer by layer.

Catch-up: Part 1 established the physics: every request is two workloads on one GPU. Prefill is compute-bound and sets TTFT; decode is bandwidth-bound and sets TPOT; we raise the batch $B$ toward the crossover $B^{\star} \approx 300 \cdot N_{\text{total}}/N_{\text{active}}$ to pay for the weight read, and continuous batching keeps $B$ full.

This part goes under that: the compute cores, the kernels, and the compilers where the opening bug actually lived, a int32/int64 dtype mismatch in a fused MoE router that --enforce-eager and disabling torch.compile could not touch. By the end, we will see exactly why those knobs missed.

Part 1 covered what a request costs; this part is where that cost is spent, one layer at a time. A quick orientation of the pieces: the SMs (streaming multiprocessors) are the GPU’s compute engines, the blocks that actually do the math. HBM (high-bandwidth memory) is its main memory pool, where the weights and KV cache live, and decode is slow because feeding the SMs from HBM is the bottleneck.

The unit of GPU work is a kernel, a single function launched across thousands of threads. The fastest optimizations (fusion, and the attention kernels Flash-Decoding and FlashAttention-4) are all about moving fewer bytes and keeping more of the SMs busy. A layer up, CUDA Graphs and torch.compile speed up how kernels are launched, but neither can see the type contracts written inside a hand-written kernel, and that blind spot is the one layer below anything the opening knobs could reach.

1. Why a GPU isn’t just a fast CPU

The natural instinct when reading GPU specifications is to interpret them as a CPU on steroids. An H200 has 132 streaming multiprocessors, 141 GB of HBM at 4.8 TB/s, and roughly 2 petaflops of FP8 compute. Surely, the reasoning goes, it’s just a very fast CPU?

It’s not. It’s a fundamentally different kind of machine, built for a fundamentally different workload.

Being an F1 fan, the analogy I find most helpful is this: think of a CPU as a Formula 1 car and a GPU as a freight train. A CPU is built to make one task as fast as possible, with aggressive branch prediction, deep out-of-order pipelines, multi-level caches, and prefetchers. A GPU, on the other hand, is built to move enormous amounts of data through arithmetic units in bulk. Its per-task latency is poor, but its throughput is staggering. If we have one person to move, take the F1 car. If we have ten thousand, take the train.

1.1 SIMT: the GPU threading model

The train analogy only helps if we know how work is scheduled onto the rail cars. NVIDIA’s model is SIMT (Single Instruction, Multiple Thread), and its scheduling unit is the warp: a group of 32 threads that share one program counter and execute the same instruction each cycle, just on different data.

The hardware never schedules individual threads, only whole warps, and it hides memory latency by keeping many warps in flight, so that when one warp stalls waiting on HBM another ready warp runs in its place. That “lockstep” execution is what makes matmul and batched decode a natural fit.

It also sets up two failure modes:

• warp divergence (lanes in a warp take different branches)
• insufficient parallelism (too few warps to hide latency, as in batch-1 attention before Flash-Decoding).

Decode occupancy and fused attention both live or die on this model.

Closest CPU analogy: SIMD (Single Instruction, Multiple Data, as in AVX-512), where one instruction applies the same operation to a whole vector of values at once, crossed with a massively oversubscribed thread pool that hides latency by always having more ready work to switch to.

The difference is scale and how divergence is punished. On a CPU, a mispredicted branch costs a single core a handful of cycles; when lanes in a warp diverge, the GPU serializes both branches and runs each with the inactive lanes masked off, so the cost is paid across all 32 lanes rather than one.

1.2 Tensor cores

An SM actually carries two different kinds of compute units, and it helps to keep them separate.

• CUDA cores are the general-purpose ALUs: each does ordinary scalar floating-point and integer math, roughly one multiply-add per clock, and they are what run the SIMT threads from the last section (elementwise ops, activations, address arithmetic, and so on).
• Tensor cores sit on the same SM but are specialized: each one consumes a small tile of two matrices (16×16 or larger) and emits the whole fused multiply-accumulate for that tile in a single instruction, which is exactly the operation a matmul is built from.

They are not separate chips, just two unit types inside every SM. The reason the distinction matters is that most of the FLOPs on the datasheet come from the tensor cores, not the CUDA cores: the H200’s roughly 989 BF16 teraflops (or 1,979 FP8 teraflops) are tensor-core numbers, far above what its CUDA cores alone could reach.

Transformer inference is almost entirely “keep those units busy with the right precision and layout.” Fusion, FlashAttention, and FlashAttention-4 are all, in different ways, tactics for that. The software analogy is AES-NI on a CPU: one instruction that replaces what would otherwise be hundreds of individual operations. Figure 1 is the contrast in one diagram:

1.3 Why fabrics matter as much as FLOPs

Peak FLOPs only matter once arithmetic intensity is high enough to reach the compute roof. The other half of the hardware story is the interconnect hierarchy: where the bytes live and how fast they can move between those places. The pattern that matters is that every step outward from the SMs costs us bandwidth, usually close to an order of magnitude at a time, and that cliff is what decides which kind of traffic we can afford to send across which boundary.

At the top, on-chip SRAM feeds the SMs at effectively terabytes per second, and HBM (the GPU’s own memory) delivers a few TB/s: 3.35 on an H100, 4.8 on an H200. Step off the chip and the drops begin. Within a single 8-GPU node, GPUs talk to each other over NVLink at roughly 900 GB/s bidirectionally per GPU through the NVSwitch crossbar, fast enough that bandwidth-heavy patterns can treat the eight GPUs almost as one accelerator. (That “almost as one” is a bandwidth statement, not full cache coherence; hardware coherence is a Grace Hopper NVLink-C2C property, not standard H100 peer access.)

Between two nodes, we fall to InfiniBand, 400 to 800 Gbps per link on current NDR/XDR generations (older HDR is 200 Gbps), and the unit matters here: that is bits per second, so dividing by 8 leaves only about 50 to 100 GB/s, roughly a 10× cliff below NVLink. The CPU-to-GPU link is PCIe at around 64 GB/s on a 16-lane Gen 5 connection, the slow host path we keep off any per-token critical path.

In familiar terms, the memory hierarchy looks a lot like a service-tier diagram (Figure 2):

Decode’s bandwidth story is exactly this hierarchy: each token, weights (and growing KV) stream from HBM into on-chip SRAM so the SMs can do a little math, then we do it again. If HBM cannot feed SRAM fast enough, the SMs sit idle. That is what “bandwidth-bound” means in one sentence.

Every parallelism choice is fundamentally a choice about which traffic crosses which fabric. Tensor parallelism needs NVLink-class bandwidth for per-layer all-reduces (every GPU summing its partial result with all the others, once per layer); pipeline parallelism can survive InfiniBand because it only ships activations at layer boundaries; expert parallelism’s all-to-all (each GPU shipping a different slice to every other) is somewhere in between and is exactly why DeepSeek’s recipe prefers DP+EP on H200. The fused MoE router only appears on that EP path. Part 3 defines these collectives properly; here the point is just that each one stresses a different fabric.

That hardware only does useful work when we launch work onto it. The unit of that work is a kernel, and almost every “mysterious” inference speedup we will hear about is either a better kernel or fewer launches.

2. What a kernel is, and why it matters

A kernel is a function, written in CUDA C++ or one of its analogs, that runs across thousands of threads on the GPU. We launch it with a grid of thread blocks (each block around 128 to 256 threads), and inside the function, threadIdx, blockIdx, and blockDim give us our coordinates so we can compute which slice of the input belongs to us. Threads in the same block can communicate through shared memory, a software-managed scratchpad on each multiprocessor, typically 100 to 228 KB. Across blocks, communication happens only through global HBM, atomics, or by ending and relaunching the kernel.

Because those thousands of threads run concurrently, a kernel is also a small concurrency problem. Whenever threads share data through shared memory, they have to synchronize: a thread that reads a slot another thread is meant to fill first must wait behind a barrier (__syncthreads() in CUDA) until that write lands. Forget the barrier and we have a race condition, where the reader sometimes sees stale data purely because of the order the hardware happened to schedule the warps, an order the kernel does not control. That is what makes races so nasty: one can pass every test on one machine and only surface when the timing shifts. Hold onto that idea; it is the hidden cause behind a class of production nondeterminism bugs that Part 3 comes back to.

Each kernel launch costs CPU time, because the driver has to validate arguments, transfer them, and schedule the work onto a queue. NVIDIA’s measurements put the floor near 9.6 microseconds for a trivial kernel; in real workloads it’s 20 to 200 microseconds. That sounds small until we count the launches.

A single transformer layer might involve an RMSNorm, three linear projections, a rotary positional embedding (RoPE), an attention call, an output projection, residual adds, another norm, a router, MoE dispatch, MoE compute, MoE combine, and an output residual. A 60-layer model is easily more than 1,000 kernel launches per token. At 20 microseconds each, that’s 20 milliseconds of pure CPU-side overhead, which on a fast GPU is more than the entire compute budget for that token. 😮

The web-app version of this problem is N+1 query syndrome: we query for a list, then loop and make N more queries to fetch related items. It’s the same disease with the same cures: do less, batch what’s left, and eliminate round-trip overhead wherever we can’t batch.

2.1 Kernel fusion

The first cure is to merge operations. This is similar to the optimization in compilers called loop fusion or deforestation: when we chain operations over a list (map f . map g . filter h), we don’t materialize the intermediate lists; instead, we stream each element through all three operations in a single pass.

On a GPU, softmax(matmul(A, B) + C) shouldn’t write the matmul result to slow memory, read it back to add C, write it again, and then read it back to apply the softmax. Instead, we fuse it into one kernel that keeps intermediates in registers and shared memory, never touching HBM between operations.

What this helps us achieve in practice is:

• One kernel launch instead of three, which kills the CPU overhead.
• Intermediate values live in registers and never touch HBM, which kills the memory-bandwidth cost.
• The whole sequence moves closer to compute-bound, which finally uses the silicon we paid for.

This is what “fused” means in names like fused_moe, fused_attention, and fused_topk_bias_router. Each is a hand-written kernel that combines what would otherwise be 3 to 10 separate operations. Figure 3 shows the HBM traffic difference:

The most famous example is FlashAttention. Standard attention materializes the full $N \times N$ attention matrix in HBM, reads it back for softmax, writes again, then reads once more for the value-weighted sum. That intermediate alone is:

$$ \text{bytes per layer} = N^{2} \cdot b $$

For $N = 32{,}768$, FP16 ($b=2$), we get:

$32{,}768^{2} \cdot 2 = 2{,}147{,}483{,}648\ \text{bytes} \approx 2\ \text{GB per layer}$.

Across 80 layers in a 70B model that is on the order of 160 GB of attention intermediates for a single forward, before we even count weights or KV.

FlashAttention (Dao et al., 2022) tiles the inputs into small blocks, computes the softmax online with a numerically stable running-max-and-sum trick, and keeps everything inside fast on-chip memory. The large intermediate is never written. So instead of that 2 GB per layer (about 160 GB across the model), the extra memory attention needs grows with $N$ rather than $N^2$: FlashAttention keeps only a running max and sum per query row, on the order of a few hundred KB per layer, and that little bit of state lives on-chip rather than in HBM.

In other words, the ~160 GB of attention intermediates effectively drops to nothing. FlashAttention reaches around 70% of GEMM (general matrix multiply) throughput on modern GPUs, which means attention runs almost as fast as a pure matrix multiply.

The idea that makes all of this work is online softmax (here “online” means an incremental, streaming algorithm, not the internet; the ordinary version needs the whole row at once). Normally, softmax has to see a whole row of scores at once: first to find the largest value (we subtract it before taking $e^x$, or the exponentials overflow), then to add up all the exponentials we divide by. That “see the whole row at once” requirement is exactly what forces the full $N \times N$ matrix into memory.

Online softmax reaches the same answer without ever holding a full row. It walks the scores one tile at a time and carries just two running numbers: the biggest score seen so far, and the running sum of exponentials. When a later tile turns up a bigger maximum, it rescales the running sum so the totals from earlier tiles still line up. By the last tile, those two numbers match what a plain softmax would have produced, but we never kept more than one tile on-chip. That is why the scores never touch HBM, and Flash-Decoding reuses the very same trick to stitch KV chunks together in the next subsection.

Note: FlashAttention has gone through several hardware-tuned revisions since, and the rest of this post uses the standard shorthand for them: FA2, FA3, and FA4 (FlashAttention-2, -3, and -4), each one re-tuned for a newer GPU generation. The original 2022 release is the one usually just called “FlashAttention”; FA4 gets its own treatment in section 2.3.

Concretely, FlashAttention streams tiles of Q, K, and V through SRAM and never lets that N×N matrix touch HBM (Figure 4):

vLLM’s day-zero DeepSeek V4 implementation ships at least four new fusions of this shape:

• Fused compressor, RMSNorm, RoPE, and KV-cache insertion: 1.4 to 3× speedup.
• Inverse-RoPE plus FP8 quantization: 2 to 3×.
• Fused Q-norm, KV-RoPE, and K-insert: 10 to 20×.
• Fused MoE router (the one the bug lives in): new for V4.

Every fused kernel is a contract written in C++, and every contract has type invariants. The DeepSeek crash was in topk_softplus_sqrt, one of these fused MoE routers: the opening bug, one layer down.

2.2 Flash-Decoding: attention that works at batch size 1

Kernel fusion helps whenever intermediates thrash HBM. Decode has a sharper problem: even a perfectly fused attention kernel leaves most of the GPU idle when the query length is 1.

To see why, start from what a GPU needs to stay busy: many independent pieces of work it can run at the same time, enough to fill all 132 SMs on an H100. FlashAttention gets that parallel work from two places: the batch (each sequence is independent) and the query positions (each token’s query is independent).

Prefill and training are swimming in both. Prefill’s query is the entire prompt, so it has thousands of positions to work on at once, and training runs large batches. Decode has neither. It generates one token at a time, so the query length is just 1, and the batch is often small as well, at low load or on long-context requests. That leaves almost nothing to spread across the SMs, so FlashAttention barely touches the GPU: at batch size 1, FA2 uses under 1% of an H100’s SMs.

This is the decode half of the prefill/decode split from Part 1 colliding with the hardware from section 1: the datasheet FLOPs are all there, but with one query and a tiny batch, there is no work to hand them.

Flash-Decoding (Tri Dao, Daniel Haziza et al., 2023, now shipping in FlashAttention ≥ 2.2) finds parallelism in the one dimension decode has to spare: the KV sequence length itself, which can run to hundreds of thousands of tokens. The algorithm has three steps (Figure 5):

1. Split the KV cache into C chunks along the sequence dimension.
2. Process each chunk independently, computing its partial attention scores plus a local log-sum-exp: a running max-and-sum normalizer that records just enough about the chunk’s softmax to merge it with the others later without error.
3. Combine all C chunks in a second kernel, using those local log-sum-exp values to stitch the partial softmax numerators back into the correct global result.

The trick that makes this exact rather than approximate is the same online softmax identity FlashAttention uses for block tiling. Once we know the log-sum-exp of each chunk’s partial scores, we can reassemble the true global softmax from those partial results without ever materializing the full $N×N$ matrix. The chunks supply the partial log-sum-exps; the reduce kernel applies the identity.

The payoff is C-way parallelism across the sequence length even at batch size 1, enough to keep all 132 SMs busy. On CodeLlama-34B at 64K context, that lifts attention throughput to roughly 50× faster than FA2 at the same batch size, and overall decode throughput up to 8× higher. Past 128K context, Flash-Decoding stops being an optimization and becomes load-bearing.

2.3 Blackwell and FlashAttention-4: when the bottleneck shifts

Note: If we are shipping on Hopper (H100/H200) today, we can skim this subsection: FA2/FA3 remain the defaults there; FA4’s Blackwell-specific paths do not apply until we move to B200/GB200.

Flash-Decoding assumed Hopper’s bottleneck shape, where the tensor cores are fast enough that the whole job is keeping them fed. Blackwell (B200/GB200) changes which bottleneck we’re fighting, and that shift is why FA4 is a genuine redesign rather than “FlashAttention with a higher version number.”

Going from H100 to B200, the tensor cores (the matmul hardware) got about 2.25× faster (1 to 2.25 BF16 PFLOPs). Still, two things they lean on did not budge: the exponential unit (MUFU.EX2, the special-function hardware every softmax uses to compute $e^x$) and shared-memory bandwidth. When the matmul engine speeds up, but the parts feeding it stay put, the bottleneck simply moves. So on Blackwell the matmuls are no longer the limit; attention now waits on the softmax exponential in the forward pass and on shared-memory traffic in the backward pass.

The familiar “math per byte fetched” budget from Part 1 now has a sibling: a “math per special-function op” budget. 🤓

FlashAttention-4 (Zadouri et al., 2026) is the algorithmic and kernel response to that shift. Its one job is to keep the tensor cores fed, and three things now get in the way: the exponential unit is starved, shared memory is overloaded, and some of the softmax bookkeeping sits needlessly on the critical path. The five techniques below map onto exactly those three problems: 1 and 2 feed the starved exponential unit, 3 and 4 relieve the shared-memory crunch, and 5 trims redundant work off the critical path.

Note: A CTA (cooperative thread array) is NVIDIA’s name for a thread block, the group of threads that share on-chip memory.

1. Ping-pong pipelining. Attention alternates between two kinds of work that live on different hardware: matmuls on the tensor cores and softmax on the special-function (MUFU) units. FA4 keeps two query tiles moving inside each CTA and staggers them, so while one tile is busy with softmax, the other runs its matmul. Because the two tiles never reach for the same unit at the same instant, neither the tensor cores nor the MUFU units sit idle waiting on the other.
2. Software-emulated exponential. Every softmax has to raise a number to a power, and on Blackwell that step runs on a single piece of hardware, the MUFU.EX2 exponential unit, which is now the bottleneck. FA4 stops leaning on it alone and computes the exponential a second way in parallel, on the ordinary floating-point (FMA) units, using a standard numerical recipe (a Cody-Waite range reduction plus a short degree-3 polynomial) to keep the answer accurate. With two paths producing exponentials instead of one, throughput roughly doubles, and FMA capacity that would otherwise sit idle gets put to work.
3. Tensor Memory (TMEM). The running totals attention carries as it works normally sit in shared memory, which is exactly the resource under strain. Blackwell adds a separate 256 KB per-SM scratchpad, TMEM, dedicated to those accumulators, so parking them there instead takes load off shared memory (SMEM) and unclogs the backward pass.
4. 2-CTA MMA mode. Normally each thread block loads its own copy of an input matrix into shared memory. Here, two neighboring CTAs in the same cluster pair up, share one copy, and work together on a single large 256×256 tensor-core multiply (an MMA, the tensor-core instruction). Loading that operand once instead of twice halves the shared-memory traffic for it, and it also halves the global atomic updates needed to accumulate dQ, the query gradient computed in the backward pass.
5. Conditional online softmax rescaling. As FlashAttention streams through the scores, every time it meets a new maximum it rescales its running output, and each rescale is extra vector work sitting on the critical path. In practice, most new maxima are only slightly larger and barely change the result. FA4 rescales only when the maximum jumps by more than a set threshold $\tau$ and skips the rest. The answer stays exact because the final normalization still uses the true running statistics.

The result is 1,613 TFLOPs/s on B200 at 71% utilization, roughly 1.3× faster than cuDNN 9.13 and 2.7× faster than Triton. For comparison, FlashAttention-2 on Hopper reached about 35% utilization of the H100.

One interesting tidbit: FA4 is written entirely in CuTe-DSL, CUTLASS’s Python-embedded kernel DSL. The kernels lower straight from Python to PTX (the portable GPU assembly section 3 covers in full), skipping C++ templates altogether, and compile times drop by 20-30× versus the C++ CUTLASS path.

Personally, I find this interesting, and I wonder if anyone else is thinking along the same lines: the three-ways-to-write-a-kernel hierarchy from the next section is shifting: the Python-DSL tier is now turning out kernels that match or beat hand-written C++ on shipping hardware.

Blackwell’s tensor cores got so fast that the non-tensor-core work (exp, SMEM traffic) is now the limiter: same queries, new cost model.

What this means for our own deployment:

• If we’re on H200 (Hopper), nothing changes: FA2/FA3 stay the defaults, and FA4’s Blackwell-specific paths don’t apply to us.
• If we’re on B200 (Blackwell), FA4 is what we want, and cuDNN ≥ 9.13 already folds in many of its tricks. vLLM and SGLang pick the backend for us, so we shouldn’t have to choose a kernel by hand.
• Longer term, expect more of the same. Each new generation widens the gap between matmul throughput and everything feeding it, so this hardware-software co-design isn’t a one-off; it’s the new normal.

3. The CUDA compilation pipeline

A kernel does not run on the SMs as we wrote it; in fact, it reaches the hardware through a two-stage compile:

• First, CUDA C++ (or Triton) is compiled ahead of time into PTX, a portable assembly-like IR that is not tied to any one GPU generation.
• Two, at load time, the host driver JIT-compiles that PTX into SASS, the actual machine code for the specific GPU in the box.

That split explains two failures we actually hit. First, a single cu130 container can still die with “PTX JIT compilation failed” on an old driver: the driver is too old to translate the newer PTX. Second, it is why production images ship fatbins, which carry prebuilt SASS for the GPUs we expect and keep PTX as a fallback for the rest. Figure 6 is the pipeline behind those failure modes:

Let us walk the two stages in a little more detail:

• nvcc compiles our CUDA C++ down to PTX, which stays portable across GPU generations the same way Java bytecode stays portable across JVM versions: it is a stable, abstract assembly language that names operations without committing to any one chip’s instruction encoding.
• Then, at load time, a JIT inside the GPU driver lowers that PTX to SASS, the real machine code for one specific architecture (sm_90 for Hopper, sm_100 for Blackwell). SASS never carries across generations, which is why PTX sits between them as the forward-compatibility layer: ship the PTX, and a future driver can still generate SASS for a GPU that didn’t exist when we compiled.

That is why a production binary is usually a fatbin. The name is literal: it is a fat binary in the same sense as a macOS universal binary, one file that carries several architecture-specific builds side by side. It packs precompiled SASS for the architectures we already expect, plus PTX as a fallback for the ones we don’t. At load time, the driver reaches for the exact-match SASS if it is present and JITs from the PTX otherwise, so the only run that pays the compile cost is the first one on an unfamiliar GPU.

This is what makes the container tag cu130 (CUDA 13.0) more than cosmetic: a newer toolkit knows about newer architectures and emits PTX that uses their features. To actually run Blackwell-specific instructions we need both ends recent enough, a CUDA 13 toolchain to emit the PTX and a host driver new enough to JIT it. When only one side is current, the load-time JIT is where it breaks, the same class of failure as shipping a binary that needs a newer runtime than the host provides. It wasn’t the cause of this bug, but it is one of the most common ways a container that runs on one box dies on another.

3.1 Three ways to write a kernel

For years, picking a kernel language was a clean trade-off: the closer we wrote to the metal, the higher our performance ceiling, but the slower we shipped. As the FA4 results above show, that rule is breaking down at the top.

So in some sense, we should read the three tiers below as a historical ranking more than a fixed one: raw CUDA on top, Triton in the middle, higher-level DSLs climbing fast. The deciding question is shifting from “C++ or nothing?” to “which abstraction reaches the newest hardware features first?”

• Raw CUDA C++ gives us the full instruction set: __global__ functions, manual shared-memory management, inline PTX, and CUTLASS (NVIDIA’s C++ template library for matrix-multiply primitives). The performance ceiling is still the highest for many hand-tuned paths, but the engineering cost is steep. A lot of production tensor-core code still lives here, and it is exactly the tier Graphs and Inductor won’t rewrite. The DeepSeek router bug lived here too.

• Triton is the dominant middle ground: a Python-embedded DSL where we decorate a function with @triton.jit, write with block-pointer abstractions, and let the compiler handle scheduling, vectorization, memory coalescing, and register allocation within a block. The rough rule of thumb is 80% of CUDA’s performance for 20% of the effort on typical transformer ops, and that’s fair. Most of vLLM’s fused kernels are Triton, and torch.compile emits Triton.

• Above Triton sit higher-level tools such as CUTLASS templates, CuTe-DSL, TileLang, ThunderKittens, and NVIDIA’s cuTile, which compose kernels from building blocks. CuTe-DSL is no longer aspirational: it shipped FA4. ThunderKittens (Hazy Research, Stanford) is the interesting bet for writing new attention variants fast without hand-writing CUDA.

4. CUDA Graphs: record once, replay cheaply

How kernel fusion helps solve the memory-bound problem, CUDA Graphs help solve the overhead-bound problem. The idea is simple. Instead of the CPU dispatching more than 1,000 separate kernel launches per inference step, we record that whole sequence once, every kernel address, argument, and shape, and then replay the entire recording with a single API call.

NVIDIA’s own measurement makes the payoff concrete. Launched one at a time, each kernel costs ~9.6 microseconds of CPU dispatch, so a 1,000-kernel forward pass burns roughly 20ms of pure launch overhead before the GPU does any useful work. Replayed as one graph, that same pass costs about 10 microseconds of submission total, and stays nearly flat no matter how many kernels it contains.

Crucially, this collapses the host-side submission cost only, not GPU execution: the kernels still run and take their normal time on the SMs. What disappears is the launch overhead that was starving a fast GPU in the gaps between kernels, which hurts most when that overhead rivals the token’s actual compute (small batches, short decode steps).

Another way to think about it: CUDA Graphs are prepared statements for GPU launch sequences.

Of course there is no free lunch, and there is always something else. 😄

The catch here is that a captured graph is static. Pointer addresses, argument values, and launch configurations are all baked into the recording, so if our batch size changes, the graph is invalid. Frameworks work around this by capturing one graph per batch shape and dispatching to the nearest match.

vLLM captures CUDA graphs at startup for a discrete set of batch sizes, typically powers of two and other common values, and pads actual batches up to the nearest captured size. Its architecture has a graph-capture dispatcher with several modes (NONE, PIECEWISE, FULL, FULL_DECODE_ONLY, FULL_AND_PIECEWISE) that trade off how much of the model gets captured against how much GPU memory the captures consume.

What --enforce-eager does, in vLLM’s own words, is “disable CUDA graph and always execute the model in eager mode.” No recordings, and every kernel launches individually. This was the first escape hatch I tried for the bug. It made no difference.

There is a related cost that shows up at startup, not at steady state. Capturing those graphs and compiling what Inductor emits is a large share of why a cold vLLM process can take on the order of a minute before it can serve a request.

Serverless GPU platforms attack that with process snapshots (CRIU for CPU state plus GPU-memory restore) so a warm replica comes back in seconds instead of replaying the whole capture. Modal, as one example took a ~1 GiB-model vLLM cold start from about 96s down to about 14s on their stack. Short version: CUDA Graphs are free at replay time and expensive at capture time, which is why --enforce-eager exists as a debug lever and why production deployments usually pay the warm-up once and keep the process alive.

5. torch. compile and its blind spot

CUDA Graphs record a fixed launch sequence that we already wrote. torch.compile sits one layer above that and tries to generate the sequence: PyTorch 2.x’s graph-capture-and-codegen pipeline. If Graphs are prepared statements, torch.compile is closer to a JIT that rewrites our SQL into a better plan and then prepares it.

It is structurally a four-stage compiler in its own right (Dynamo, AOTAutograd, PrimTorch, Inductor). It was also the second thing I reached for to sidestep the opening bug, after --enforce-eager, and it made no difference either, for a reason that only makes sense once we see where custom ops sit in its pipeline. Figure 7 is that pipeline:

Let us walk the four stages in order.

TorchDynamo is the front end: it hooks CPython’s frame-evaluation API. It catches each Python function just before its bytecode runs, then interprets that bytecode symbolically to build an FX graph of the tensor operations. When it hits something it can’t trace (data-dependent control flow, an unsupported library call, anything that needs Python’s full dynamism), it takes a graph break: it compiles the part it has, drops back to ordinary Python for the tricky region, and picks tracing back up afterward.

AOTAutograd runs PyTorch’s autograd engine ahead of time instead of during execution: it traces the backward pass symbolically, builds one combined forward-and-backward graph, then splits it back into separate forward and backward pieces each backend can optimize on its own. PrimTorch then boils PyTorch’s roughly 2,000 operators down to about 250 primitives, so a backend has a far smaller set of operations to support.

TorchInductor is the third stage does the real work. It takes the FX graph, lowers it through a Python-like loop-level IR, fuses operators, and emits Triton for the GPU and C++/OpenMP for the CPU. Count them up, and that is three compilers stacked at runtime (FX → Triton → the driver’s PTX→SASS step), the same idea as a JavaScript engine compiling a hot path, applied to tensor graphs.

Here’s the detail that matters for our bug: torch.compile treats custom C++ ops registered through torch.ops as opaque, a black box it cannot see inside. It sees the op’s declared schema and the input dtypes, but not the C++ implementation, so any type invariants the kernel relies on are entirely its own problem. That is exactly why disabling compile did nothing for the MoE router bug.

There’s one notable exception to the opaque-custom-op limitation, and it’s specific to attention. FlexAttention (PyTorch 2.5, late 2024) gives us a flex_attention API where we write the score modification (causal mask, sliding window, document masking, ALiBi, and so on) as a plain Python function, and the compiler JIT-compiles it into an optimized attention kernel, now targeting the FA4 backend on Blackwell. We write:

from torch.nn.attention.flex_attention import flex_attention, create_block_mask

def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

block_mask = create_block_mask(causal_mask, B, H, Q_LEN, KV_LEN)
out = flex_attention(query, key, value, block_mask=block_mask)

The compiler sees inside our attention variant and fuses it into a single kernel with the score function inlined, with no custom C++ op boundary and no opaque dispatch. The FlexAttention paper reports it landing within 0.7-1.4× of hand-written FlashAttention-2. On the gpt-fast inference path the win is model- and context-dependent: about 2.04× versus SDPA for LLaMA-3.1-8B at 16k context, and 0.99-1.66× for LLaMA-3.1-70B as context grows ( PyTorch FlexAttention for Inference ; Dong et al., arXiv:2412.05496 ). On the FA4 backend, it stays competitive with hand-written kernels while remaining fully composable with torch.compile. It doesn’t fix the general problem, since arbitrary custom ops stay opaque, but for the single most important custom op in the inference stack, attention, it dissolves the boundary entirely.

Part 2 recap: the unit of GPU work is the kernel, and a kernel is a contract written in C++ with its own dtype and layout invariants. Everything above it optimizes around that contract without seeing inside: fusion and Flash-Decoding cut the bytes and launches a kernel costs, CUDA Graphs replay a fixed launch sequence cheaply, and torch.compile regenerates that sequence, but none of them rewrite what happens within a custom op. That blind spot is exactly why --enforce-eager, disabling compile, and blanket casts all failed on the DeepSeek router: the mismatch lived one layer below anything those knobs could touch. FlexAttention is the one place the boundary dissolves, and only for attention.

Keep reading

Part 2 of 3. Next, Part 3 · Serving at scale : how a recipe like DP+EP + FP8 KV + chunked prefill is built from these pieces, fleets, disaggregation, speculative decoding, reasoning-effort traffic, and the opening bug resolved layer by layer.

• ← Part 1 · Physics of a request
• Part 2 (this post) · Below Python
• Part 3 · Serving at scale →

References & Further Reading