Part V — The hard parts

Dynamic shapes & control flow

Every pass in Parts III and IV bought its speedup by knowing the exact shapes. Fusion (lesson 05) grouped ops because it knew the tile counts; the memory planner (06) sized an arena to the byte; layout assignment (07) matched tensor-core fragments to concrete dimensions; the scheduler (08) and autotuner (10) searched tile sizes against one specific M, N, K and cached the winner under those numbers. That whole machine assumed the shapes are constants known at compile time. Production serving violates the assumption on the first request: batch size changes every microbatch, sequence length grows by one token every decode step, inputs arrive ragged, and a router sends a different number of tokens to each expert. If a fresh compile fires on every distinct shape you get a compile storm that buries the runtime. This lesson is how a compiler keeps the speedup while the shapes move.

1 · Why the shapes never hold still

Define the problem precisely. A static shape is a tensor dimension whose value is a compile-time constant — (batch=8, seq=512, d=4096). A dynamic shape (a.k.a. symbolic shape) is a dimension whose value is only known at run time and varies between calls. Earlier passes loved static shapes because every cost they minimized — bytes in a fused tile, slots in the arena, the autotune cache key — is a function of the exact numbers. Four forces make those numbers move in any real deployment:

Two of these are shape dynamism (the dimensions vary) and two shade into control-flow dynamism (which ops even run depends on the data — §5). Both break the static-shape contract. The compiler now has exactly three families of answer: specialize per shape and eat recompiles, reason symbolically and compile once, or normalize the shapes (bucket/pad) so only a few distinct ones ever reach the compiler.

2 · Full specialization vs symbolic shapes

Full specialization is the naive correct answer: treat every dimension as a constant, compile a fresh, fully-optimized kernel for each shape you see, and key it in the autotune cache (lesson 10). It is the fastest per call — the kernel knows its trip counts, can unroll loops, and was tuned at exactly this size. It is also correctness-trivial: nothing varies inside a compiled artifact. The catch is the compile bill. PyTorch compile and Inductor compile times run from roughly 1–60 s per graph depending on size and max-autotune; XLA HLO compiles are comparable. Pay that on the first appearance of every shape:

The decode row is fatal: you would spend seconds compiling for each step that takes microseconds to run. So full specialization is right only when the shape set is tiny and stable — which training is and serving is not.

Symbolic shapes are the principled fix. Instead of baking a dimension as the constant 512, the compiler introduces a shape variable — call it s0 — and compiles a single kernel that is correct for all values of s0. Derived dimensions are tracked as symbolic expressions: if the KV cache is length s0 and you append one token it becomes s0+1; a reshape that splits s0 into heads carries s0/h as a constraint the compiler must prove divides evenly. This is exactly TorchDynamo's symbolic shapes machinery (it uses SymInt and a small symbolic solver, sympy, to discharge the constraints) and XLA's dynamic dimensions. The trade is real but small: a symbolic kernel can't hard-unroll a loop whose bound is a variable, and the autotuner must pick one config that is good across the size range rather than perfect at one size, so you give up maybe 5–20% of peak — in exchange for compiling once instead of a thousand times.

static  (recompile each call):   matmul[M=512, N=4096, K=4096]   key=(512,4096,4096) → MISS → compile
                                 matmul[M=513, N=4096, K=4096]   key=(513,4096,4096) → MISS → compile
                                 matmul[M=514, ...]              ... a new compile every step

symbolic (compile once):         matmul[M=s0,  N=4096, K=4096]   one kernel, valid for all s0 ≥ 1
                                 guard: s0 ≥ 2  (because s0=1 was specialized away — see §3)

In PyTorch you reach this with the dynamic flag on torch.compile: dynamic=True opts every dimension to symbolic from the first call; dynamic=False forces full specialization (one compile per shape); the default is adaptive — specialize on the first shape, then if a second distinct value of a dimension appears, promote that dimension to symbolic and recompile once. The deeper mechanics of this live in gpu_kernels · 23 torch.compile.

3 · Guards and the recompile trigger

How does a compiled artifact know it is still valid? A guard is a cheap boolean predicate, attached to a compiled graph, that asserts every assumption the compiler relied on: dtypes, devices, the rank of each tensor, which dimensions are static and what their values are, and the constraints on the symbolic ones (s0 ≥ 2, s0 % 8 == 0). Before every call the runtime evaluates the guard set in a few microseconds. Guard passes → cache hit → replay the compiled kernel. Guard fails → cache miss → recompile, specializing to the new fact (or relaxing a static dim to symbolic), and add the new artifact to the cache. Guards are the bookkeeping that makes "specialize but stay safe" possible: the optimizer is allowed to assume anything it likes, as long as it writes the assumption down as a guard so a violation is caught instead of producing a silently wrong result. This is the same preserve-semantics invariant from lessons 03–04, now enforced dynamically.

The trap everyone hits: the 0/1 specialization gotcha. Dimensions equal to 0 or 1 are special — a size-1 dimension broadcasts, and a great many rewrites are only valid when a dim is exactly 1 (a "squeeze," a broadcast, dropping a unit batch). So compilers always specialize on 0 and 1 rather than treat them as symbolic. The consequence: if your very first call has batch=1 (extremely common — you smoke-test with a single example, or decode generates one sequence), the guard becomes "batch == 1" baked in. The next call with batch=8 violates it and forces a recompile, and now you carry two artifacts and a guard s0 ≥ 2 on the symbolic one. The fix is to warm up with a representative shape that is not 0 or 1 (and not the longest, either) so the compiler generalizes from a typical case. Watch for it with TORCH_LOGS="recompiles,dynamic" — every recompile prints the guard that fired and the shape that broke it.

4 · Bucketing and padding in serving

Symbolic shapes are elegant, but two realities push serving toward a blunter tool. First, some kernels really are faster fully specialized — a fixed trip count unrolls and the autotuner nails one size. Second, CUDA-graph capture (the launch-overhead killer in reduce-overhead mode) needs stable shapes; a fully dynamic graph can't be captured. The workhorse compromise is bucketing with padding.

Bucketing: pick a small set of canonical sizes — say sequence buckets {128, 256, 512, 1024, 2048} — and compile (and CUDA-graph-capture) one artifact per bucket. Padding: round each incoming request up to the nearest bucket and pad the extra positions (masked so they don't affect the result). A request of length 300 runs on the 512-bucket kernel with 212 padded positions. You now compile a handful of times instead of thousands, every call is a cache hit, and you keep CUDA graphs — at the cost of wasted FLOPs on the padding.

The waste is exactly quantifiable. With buckets and a uniform spread of true lengths, the average padded length is roughly halfway up to the next bucket, so coarse buckets waste more. Concretely, a length-300 request in a 512 bucket does 512/300 ≈ 1.7× the necessary attention work. The design knob is bucket granularity: more, tighter buckets → less padding waste but more compiles and more memory pinned by CUDA graphs; fewer, coarser buckets → cheap compile and capture but fat padding. Production stacks tune this empirically and pad attention with masks so the math stays exact. See how real engines wire variable-length batching and KV-cache growth — vLLM's paged KV and continuous batching, SGLang's RadixAttention and prefix-cache reuse — in vllm · lessons and sglang · lessons; both are, at heart, strategies to keep the compiled kernels' shapes under control while the workload churns.

5 · Data-dependent control flow

The last form of dynamism isn't a shape at all — it's which ops run. A if x.sum() > 0:, a while not converged: loop, a beam-search step, an MoE dispatch: the branch taken or the iteration count depends on a value computed at run time. The capture frontends from lesson 02 have two ways to deal with it, and only one keeps the graph whole.

The lazy way is a graph break (lesson 02): Dynamo can't statically resolve the branch, so it ends the current graph, runs the Python if eagerly, and starts a new graph after. Correct, but it severs the graph in two — fusion can't cross the break, and CUDA-graph capture stops there — so you lose the very optimizations Parts III–IV bought. A branch in the middle of a hot loop can erase most of the speedup.

The right way is a control-flow op: a real graph node whose operands are the predicate and the sub-graphs for each branch, so the whole thing stays inside one capturable, optimizable graph. PyTorch exposes torch.cond(pred, true_fn, false_fn, operands) for branches and torch.while_loop / associative scan for loops; JAX has lax.cond, lax.while_loop, and lax.scan (and lax.scan is how you express a decode loop as one rolled graph op instead of an unrolled chain of thousands of nodes). The constraint these ops impose is the price: both branches of a cond must produce the same output shapes and dtypes (the compiler picks a layout and arena slot once, for both), and a scan body's carry must be shape-stable across iterations. That constraint is exactly what lets the downstream passes — fusion, memory planning, layout — treat the branch as a single typed value instead of giving up.

# Graph break (severs the graph — fusion/CUDA-graph stop here):
if x.sum() > 0:          # data-dependent Python branch → Dynamo graph-breaks
    y = f(x)
else:
    y = g(x)

# Control-flow op (one graph node — stays fused & capturable):
y = torch.cond(x.sum() > 0, f, g, (x,))   # both f,g must return same shapes/dtypes

# JAX: a 2,048-step decode as ONE rolled op, not 2,048 unrolled nodes
final_carry, outs = jax.lax.scan(decode_step, init_carry, xs)

6 · Drive it: the shape-specialization cache

The widget streams a batch of inference requests with varying sequence lengths past four strategies. Watch the same request stream produce wildly different bills. Specialize-each compiles on every new length — a compile storm whose total time is dominated by compilation, not inference. Bucket-to-N compiles only the buckets, hits the cache after, and pays a modest padding tax (the red caps on each bar). Pad-to-max compiles once but every request runs at the maximum length — minimal compiles, maximal waste. Symbolic compiles exactly once and pads nothing. The KPIs make the §2–4 trade literal: recompiles, cache hit-rate, wasted padding FLOPs, total wall time.

Re-roll the stream a few times. Specialize-each's bill tracks the number of distinct lengths, not requests — a stream of 28 requests with 25 distinct lengths is 25 compiles. Bucketing collapses that to the bucket count no matter how many distinct lengths appear, and the padding caps shrink as you add buckets. Symbolic stays at one compile forever. That is the recompile-versus-waste trade, made into a knob.

Failure modes & checklist

Checkpoint

Where this points next

You can now keep the Part III–IV speedups while batch size, sequence length, and control flow all move at run time: symbolic shapes to compile once, guards to specialize safely, bucketing to bound the compile count, and control-flow ops to keep branches and loops inside the graph. But notice what every one of those handled — the forward graph. Everything from capture (02) through autotuning (10) and this lesson has compiled inference: the path from inputs to outputs. Training needs the backward graph too — the gradients — and you never wrote it. Hand-writing a backward for a fused, planned, laid-out, dynamically-shaped forward is hopeless; the compiler must build the backward graph from the forward and then run the same fusion, planning, and layout passes on it. That makes differentiation itself a compiler pass — lesson 12, Autodiff as a compiler pass.

Interview prompts

• Why does autoregressive decode break a statically-specialized compiler so badly? (§1–2 — the KV-cache/sequence length grows by one every step, so every step is a new shape, a new autotune cache key, a cache miss, and a recompile; you spend seconds compiling per step that runs in microseconds — a compile storm.)
• What is a symbolic shape and what does it cost? (§2 — a dimension carried as a variable s0 with constraints (s0+1, s0%8==0) so one kernel is valid for all sizes; you compile once instead of per shape, paying ≈5–20% peak because variable loop bounds can't hard-unroll and one autotune config must serve the whole range.)
• What is a guard, and how does recompilation get triggered? (§3 — a cheap boolean predicate recording the optimizer's assumptions (dtypes, ranks, static dim values, symbolic constraints); checked in µs before each call. Pass → replay cached kernel; fail → recompile and re-specialize. It's the dynamic form of preserve-semantics.)
• Explain the 0/1 specialization gotcha. (§3 — dims of size 0 or 1 enable broadcast/squeeze rewrites only valid at exactly that size, so compilers always specialize on them; warming up at batch=1 bakes a batch==1 guard that the first real batch violates, forcing a recompile. Warm up with a representative non-0/1 shape.)
• When do you bucket-and-pad instead of using symbolic shapes, and what's the trade? (§4 — when you need per-shape kernel speed or CUDA-graph capture, which needs stable shapes; compile a few buckets and round requests up. Trade: cache hits + capture vs wasted FLOPs on padding; tighter buckets cut waste but add compiles and pinned memory.)
• How should a compiler handle if x.sum() > 0: in a hot loop? (§5 — capturing it as a Python branch graph-breaks, severing fusion and CUDA-graph capture; rewrite as a control-flow op torch.cond (both arms returning the same shapes/dtypes) so it stays one optimizable graph node.)
• Why is a decode loop expressed as lax.scan rather than an unrolled Python loop? (§5 — scan is one rolled graph op with a shape-stable carry, so the graph is a few nodes the compiler can fuse and capture; unrolling 2,048 steps produces 2,048× the nodes and a hopeless compile, and a data-dependent length can't be unrolled at all.)