Skip to content

Prefill and Decode Have Opposite Bottlenecks

#llm-inference #quantization #speculative-decoding #fp8 #model-compression #gpu-kernels

The two halves of a decode step ​

Prefill and decode share a weight matrix and almost nothing else.

Prefill pushes thousands of tokens through the same weights, so every byte you load from HBM gets reused hundreds of times. The kernel is dominated by matrix multiplies and the movement of activation data between them. Decode produces one token at a time, reads the entire weight matrix and the entire KV cache to do it, and reuses each of those bytes exactly once. Prefill is compute-bound. Decode is bandwidth-bound. Treating them with a single kernel configuration is how you end up with a serving job that looks fine on paper and slow in production.

FuTransformer quantifies the gap on a single H100. Prefill hits 1.29 PFLOP/s, which is roughly two thirds of the H100's dense FP8 peak. A full Llama-3.1-8B decode step at batch one takes 2.8 ms, against 8.1 ms for the compiled transformers implementation. That ratio is the whole argument for splitting the kernel path in two.

What 2.8 ms per step means in practice: about 357 tokens per second for a single stream before any batching. Nobody perceives typing lag at that rate. The 8.1 ms baseline is 123 tokens per second, still fine for chat, but you're burning three GPUs' worth of time to serve one user's stream.

What FuTransformer changes at the kernel level ​

The implementation borrows CODA's GEMM-plus-epilogue structure and retargets it at inference. Several pieces of the Llama block that normally cost a separate kernel launch get folded into the epilogue: the RMSNorm gain becomes a fixed scale baked into the weights at load time, FP8 output gets written directly with delayed scaling instead of a cast pass, and prefill runs without padding, so no MMA throughput is wasted on pad tokens. KV-cache results land directly in cache layout rather than a scratch tile that gets copied afterward.

Decode is a different design entirely. Projection weights are INT4, and the KV cache is kept in an INT4 copy. Weights and keys are stored in Tensor Core fragment layouts so dequantization happens as the MMA operands load. There's no full intermediate tile in BF16 sitting in memory waiting to be consumed.

Then there's the block-diagonal Hadamard transform applied to the cache. Quantization error in a KV cache is usually dominated by a handful of channels with large magnitudes, which drag the quantizer's range out and crush everything else into a few levels. A fixed orthogonal rotation mixes those channels so no single one dominates. It's cheap because the transform is fixed, and it can be folded into the read path.

One number in the paper deserves more attention than it usually gets. At the 8B shape for the SwiGLU gate/up projection, FuTransformer's BF16 path is already 1.13x faster than CODA's. The FP8 path is 2.02 to 2.18x faster than that same BF16 path. So the format switch contributes most of the win, and the kernel engineering contributes a smaller but still real slice. If you're deciding where to spend next quarter, the format conversion is the bigger lever.

What the speed costs in quality ​

Quantization always costs something. The question is whether the cost lands where your workload notices it.

On WikiText-2, BF16 gives a perplexity of 6.24. INT4 weights with plain round-to-nearest give 7.23. The same INT4 format with GPTQ gives 6.99. The rotated INT4 cache adds another 0.20 on top.

That 6.24 to 7.23 jump is about 16% worse perplexity. For a chat assistant the damage is mostly invisible. For long-form summarization or code generation over a full repository, small perplexity differences compound across thousands of tokens and show up as dropped constraints. GPTQ recovers roughly a third of the loss and costs nothing at runtime, which makes skipping it hard to justify. The 0.20 from the cache rotation is the price of halving KV-cache bytes, and at long context that trade is usually worth taking because cache bandwidth is the wall you're actually hitting.

The detail I find most reassuring: the kernels match a quantized reference computation within 0.001 perplexity. The CUDA implementation isn't leaking error on top of what the quantization format already costs.

Speculative decoding meets agentic context ​

Speculative decoding rests on an assumption that agentic workloads quietly break: that a small draft model can guess the target model's next several tokens. The guess gets verified in parallel, so the only thing that matters is acceptance length, the number of draft tokens that survive verification.

Under long, continuously evolving contexts and frequent multi-turn interaction, existing draft models lose up to 80% of their acceptance length. Keeping a historical KV cache for the draft model to fix this introduces its own overhead, which eats the speedup you were chasing. You end up paying memory and latency to maintain a shadow copy of conversation history that drifts out of sync with the target model anyway.

DLite drops the history entirely. It reuses contextual hidden states that the target model already computed, then does speculative prediction locally, without maintaining any explicit historical context. A lightweight sequential head restores causal dependencies within each block. The architecture is heterogeneous enough that it needs a two-stage training recipe, but the payoff is 2 to 6 accepted tokens per verification pass, up to 70% better than DSpark.

An acceptance length of 2 to 6 means each expensive target-model forward pass yields 2 to 6 tokens. That's your decode multiplier. Track acceptance length, not single-turn tokens per second, because it's the number that collapses when a real conversation gets long.

Quick Take: The compression techniques that matter in production are the ones matched to a specific bottleneck, not the ones with the best headline compression ratio.

Four ideas, four different failure points ​

These papers aren't variations on a theme. Each one removes a different kind of cost, and each one has a distinct way of biting you.

TechniqueWhat it reducesWhere it pays offReported gainMain risk
FP8 prefill, INT4 decodeActivation bytes in prefill, weight and cache reads in decodeDense single-node serving on Hopper1.48x prefill, 2.1-3.3x decodePerplexity 6.24 to 7.23 without GPTQ
History-free speculative decodingNumber of target-model forward passesMulti-turn agents with long contexts2-6 accepted tokens, up to 70% over DSparkDraft head must be retrained with the target
Budget-aware optimizerHidden width that survives structured pruningAny model you plan to prune after trainingBeats Muon under MLP channel pruningRetraining cost if you already trained with Muon
Haar-wavelet sparsificationBytes per input vector in repeated inner productsEdge and distributed inference with fixed weightsFewer coefficient pairs than dimensionsPermutation metadata storage and transfer

Read the last column carefully. All four methods work. None of them are free.

Muon trains well and prunes badly ​

Muon is a strong optimizer, and it has a clean interpretation as steepest descent under a spectral-norm geometry. It also produces Transformers that fall apart when you prune MLP hidden channels after training. That's an awkward combination, because the teams most likely to prune are the ones training the biggest models, and they're the ones most likely to reach for a better optimizer.

BMuon's diagnosis is geometric. Muon has no mechanism that makes channel importance comparable across blocks, so hidden channels get used redundantly. For linear factorization models, Muon admits minimum-cost exact representations that carry redundant hidden width. Nothing pushes it toward a width-minimal solution. When you prune afterward, you're cutting channels the optimizer never had a reason to keep distinct.

BMuon modifies the spectral geometry so hidden channels across all blocks compete for a shared allocation budget. The resulting oracle is exactly normalized steepest descent under a family of norms interpolating between Muon's geometry and structured sparse geometries. Under a sufficiently scarce budget, every minimum-cost representation becomes width-minimal.

The practical shift is in the knob you turn. Instead of a sparsity-penalty coefficient that has to be retuned per setting, you set a global channel-allocation budget, which is an interpretable number: how many MLP channels do you want to keep in total. A minorization-maximization surrogate makes the exact allocation problem tractable during training.

Results: on ImageNet-1K with ViT-S/16, BMuon beats Muon under gradual structured MLP pruning, and its learned channel allocations prune better than weight magnitude. On a 1B-parameter LLaMA-style model trained on FineWeb, it beats Muon under post-training MLP channel pruning. If you plan to prune, the optimizer choice belongs in the pruning plan from day one.

Key numbers from the four papers, in one place. Decode latency of 2.8 ms per step versus 8.1 ms in compiled transformers for Llama-3.1-8B at batch one. Prefill throughput of 1.29 PFLOP/s and a 1.48x speedup over compiled PyTorch FP8. WikiText-2 perplexity moving from 6.24 to 7.23 with round-to-nearest INT4 weights, or 6.99 with GPTQ. Acceptance length of 2 to 6 tokens for DLite against an 80% degradation for draft models on evolving contexts. A 1B-parameter LLaMA-style model trained on FineWeb for the BMuon pruning comparison.

Sparsifying the repeated inner product ​

SETAREH targets a case the serving stack usually ignores: a fixed coefficient vector dotted against a stream of input vectors, over and over. That's the shape of embedding lookups, fixed linear heads, retrieval scoring, and a lot of what runs on edge devices where memory traffic is the entire budget.

The method applies the same permutation sequence and Haar transforms to both the fixed coefficients and each incoming input. Because both sides get the identical orthogonal treatment, the inner product is preserved exactly before any truncation. Then it keeps only a subset of the transformed coefficient-input pairs. What you get is a controllable approximation error in exchange for representing fewer pairs than the original dimension.

The sorting step is what makes this more than plain wavelet thresholding. It concentrates energy so the retained subset carries most of the product.

Evaluation covers a synthetic Gaussian baseline and a real apparent-age estimation task using a CNN on APPA-REAL. In both, the inner products can be approximated with substantially fewer transformed pairs than the original vector dimension.

The catch is stated plainly in the paper: the practical benefit depends on approximation accuracy and on what it costs you to store or transmit the permutation metadata. If the metadata travels with every query, you can erase the savings. This works when the permutation is fixed and lives on both sides of the wire. It's not a drop-in for attention.

Common pitfalls ​

You will hit these. I've hit most of them.

Quantizing and testing on short prompts. INT4 weight error is roughly flat across sequence length, but the rotated cache cost of 0.20 perplexity is not. Validate at your real context length, not at 2K because it ran faster.

Benchmarking decode at batch 32. At batch one, PyTorch's INT4 projection kernel is competitive. By batch 32 it loses to BF16, because it targets a single query row and can't amortize the weight read. If your production traffic is batched, that crossover is the only measurement that matters, and it's the one most benchmarks skip.

Pruning a Muon-trained model without budget-aware training. The pruning evaluation looks reasonable on ImageNet-scale vision and then the model degrades on your deployment distribution. The failure is in the optimizer's geometry, so no amount of pruning-criterion tuning fixes it.

Assuming speculative decoding acceptance is a fixed property. It's a function of context length and turn count. A draft model that accepts 4 tokens on a fresh prompt can drop toward 1 by turn eight, and that's exactly when your latency budget is tightest.

Mixing prefill and decode in one kernel configuration. Padding a prefill batch to a fixed shape burns MMA cycles on tokens that don't exist. Running decode with a prefill-shaped epilogue adds a full intermediate tile. These are separate code paths for a reason.

What I ran into in practice ​

I watched a serving job regress after a quantization upgrade last year, and the cause was dumber than any of the above. My team had tuned the INT4 config against a 4K-context eval set, shipped it, and then took complaints from users running 60K-token documents. Perplexity looked fine in the harness. The cache quantization error was doing the damage, and it only showed up once the cache was big enough that the quantizer's channel range got dominated by a handful of outliers.

Rotating the cache fixed most of it, and adding GPTQ on top of the weights recovered the rest. The lesson wasn't about quantization formats. It was that our eval set didn't resemble production traffic, and every compression technique we tested had a hidden dependency on sequence length, batch size, or turn count that the benchmark didn't capture.

I've seen the same shape of failure with speculative decoding. A colleague's agent setup showed a 2.4x speedup in a synthetic multi-turn benchmark and a 1.3x speedup in production. Acceptance length was collapsing as conversations grew, and nobody was logging it. Once we instrumented acceptance length per turn, the picture became obvious in an afternoon. If you take one thing from the DLite paper, take the diagnostic framing: the historical context is the common cause of both the acceptance drop and the cache overhead.

One thing to remember ​

Every technique in this cluster reduces one specific cost and leaves the others untouched. FuTransformer cuts bytes moved per decoded token but doesn't reduce the number of forward passes. DLite reduces forward passes but doesn't make each one cheaper. BMuon changes which weights survive pruning but doesn't change inference math at all. SETAREH cuts representation cost for fixed-weight inner products and does nothing for attention. Stack the ones whose bottlenecks don't overlap, and measure the stack end to end, because two 2x wins that fight for the same resource give you 2x, not 4x.

The bottom line ​

If you're serving Llama-class models on Hopper at batch one to 32, split your kernel path: FP8 for prefill, INT4 weights with an INT4 rotated KV cache for decode, and pay for GPTQ on the weights. The 2.8 ms versus 8.1 ms gap at batch one is a GPU-cost difference you can measure on your invoice.

If you're running agents with long, evolving conversations, stop tuning your draft model against fresh prompts. Acceptance length degrades up to 80% as context accumulates, so either move to a history-free draft that reuses target hidden states or budget for the KV-cache overhead of maintaining history explicitly.

If you're training a model you intend to prune, pick the optimizer for compressibility now. Muon-trained Transformers degrade under post-training MLP channel pruning in a way that no pruning criterion repairs, and BMuon's global channel budget gives you a number you can actually specify up front. Watch the metadata cost on any fixed-weight sparsification scheme before you deploy it to the edge, since a permutation table shipped with every query will eat the savings.