Chapter 17 Going Beyond the Transformer for Inference

17.1 State-Space and Linear-Attention Models

17.1.1 Mamba and Associative Recall Tasks

The most direct path to faster generation is changing what the model stores. Every trick applied to the KV cache-GQA, MLA, cross-layer sharing-still leaves a memory footprint that grows linearly with context length. A senior interview question in serving-system design often goes: “What architectural change would eliminate that growth entirely?” State-space models, and specifically Mamba, are the answer interviewers are probing for. This topic shows up at staff and principal level, where the expectation is that you can explain the RNN duality, name the specific failure mode of early SSMs on language tasks, and sketch a hybrid deployment strategy with numbers.

State-space models derive from classical linear dynamical systems. A continuous-time SSM maps an input signal $x ( t )$ to an output $y ( t )$ through a hidden state $h ( t )$:

Discretized with step size $\Delta$, this becomes a recurrence:

where $A$ and $B$ are the discretized matrices derived from $A$, $B$, and $\Delta$. At generation time, producing token $t + 1$ requires loading only $h_{t}$-a fixed-size vector of dimension $N$-not an ever-expanding KV cache. The state $h_{t}$ is the model’s entire memory of the past.

The same recurrence can be re-expressed as a convolution during training:

This is parallelizable via a parallel scan in $O ( T \log T )$ time, giving training efficiency comparable to a transformer.

The problem with early SSMs like S4 is that $A$ and $B$ are fixed for all inputs-every token blends into the state identically regardless of content. This fails on associative recall: given a sequence of key-value pairs followed by a query key, retrieve the matching value. Fixed linear blending cannot isolate one binding from others; the signal is diluted by all intervening content.

Mamba resolves this by making the discretization input-dependent:

When the model sees a query that matches a prior key, the input-dependent gate can sharpen the state to amplify that specific binding and suppress others. This selective SSM passes MQAR (Multi-Query Associative Recall) benchmarks that defeat all fixed-parameter SSMs.

Mamba’s selective state-space mechanism gives an RNN the content-based memory it needs to match transformer quality on language tasks while keeping generation cost proportional to hidden state size rather than context length.

The argument is bandwidth. At generation step $t$ a standard transformer reads the entire per-sequence KV cache: $O ( S \cdot d \cdot L )$ bytes, growing with every generated token. The attention arithmetic intensity during generation is always approximately 1-the KV cache is sequence-specific, so batching $B$ requests together does not amortize the memory read; the batch dimension cancels in the intensity formula.

Mamba’s generation reads a hidden state of size $O ( d_{\text{inner}} \cdot N \cdot L )$ bytes per sequence, where $d_{\text{inner}} = 2 d$ is the expanded channel dimension and $N$ is the state dimension (typically 16). This is fixed: it does not grow with $S$. At $S = 8 , 192$, the transformer’s KV cache dominates bandwidth; Mamba’s state is a rounding error beside the weight load.

The trade-off is recall fidelity. A fixed-size state is a lossy compression of history. At very long contextssay, multi-hop retrieval across a 128k-token document-a pure Mamba model will lose long-range references that a transformer would find cleanly. This is not a design flaw; it is a fundamental consequence of finite state capacity. The practical response is a hybrid architecture: use Mamba blocks for the vast majority of layers (cheap, constant-state generation) and insert sparse full-attention layers every 4-8 blocks as global checkpoints that can attend to anything in the prefix without compression loss. Early SSMs like S4 used fixed HiPPOinitialized matrices tuned for smooth continuous signals; they had no mechanism for content-based gating at all. Mamba’s selection mechanism is the critical upgrade.

A Worked Example

Consider a 7B Mamba model with $d = 4 , 096$, $d_{\text{inner}} = 8 , 192$, state dimension $N = 16$, $L = 32$ layers, $B = 8$ concurrent sequences, context $S = 8 , 192$, bf16 weights on an H100 (80 GB HBM, 3.35 TB/s bandwidth).

Transformer KV cache (equivalent architecture). Per layer per sequence:

Across 32 layers and $B = 8$: $128 \times 32 \times 8 = 32 , 768$ MB $= 32$ GB, plus 14 GB of model weights-46 GB total. This barely fits on one H100.

Mamba SSM state. Per layer per sequence:

Across 32 layers and $B = 8$: $256 \times 32 \times 8 = 65 , 536$ KB $= 64$ MB. Total inference memory: $14 , 064$ MB $\approx 14$ GB. The model fits with headroom for $B = 32$ or beyond.

Bandwidth per generation step.

• Transformer: reads 32 GB (KV) + 14 GB (weights) = 46 GB $\Rightarrow 46 / 3 . 35 \approx 13 . 7$ ms per token.

• Mamba: reads 0.064 GB (state) + 14 GB (weights) $\approx 14$ GB $\Rightarrow 14 / 3 . 35 \approx 4 . 2$ ms per token.

• Speedup: $13 . 7 / 4 . 2 \approx 3 . 3 \times$ from bandwidth reduction alone.

At $S = 32 , 768$ the transformer KV cache grows to 128 GB and no longer fits on a single H100. Mamba’s state remains at 64 MB regardless, and throughput continues to scale with $B$.

What You Actually Decide in Practice

• Pure Mamba vs. hybrid. Use pure Mamba only for workloads with short-to-medium context and no hard retrieval requirements-code completion under 2k tokens, for example. For long-document QA, agentic tool use, or RAG, insert one full-attention layer per 6-8 Mamba blocks. In my experience, placing full-attention layers non-uniformly (denser in the back half of the stack, around layers 20-32 of 32) gives the best quality-per-cost trade-off because long-range context matters more in deeper layers.

• State dimension $N$. The default of 16 is sufficient for typical language modeling. Increasing to 64 improves performance on long-context recall by a measurable margin at sequences exceeding 16k tokens, but quadruples the state memory and adds training compute. Worth it only if your deployment workload is consistently above 16k context.

• Selective scan kernels. Mamba’s parallel-scan training requires custom CUDA kernels. The inference path is a plain recurrence-no custom ops needed-but verify that your serving framework handles the recurrent forward pass without accidentally falling back to the scan path, which would be unnecessary overhead.

• Mixing with MoE. Hybrid Mamba-MoE models have shipped at production scale: a 52B MoE with 44 Mamba layers and 8 MoE-transformer layers demonstrated 40-60% lower per-token serving cost versus a dense transformer of equivalent quality. The Mamba layers handle local pattern matching; the sparse MoE transformer layers handle routing and long-range retrieval.

• Failure mode to watch for. Pure Mamba on retrieval-heavy workloads fails silently-the model produces confident, plausible-sounding answers that omit or corrupt facts from the far context. Add a full-attention layer or use retrieval augmentation before shipping to production if users will cite specific document content.

How This Shows Up in Interviews

• “Why is transformer generation memory-bound, and what architectural change would fix it?” Strong answer: Attention arithmetic intensity at generation is always $\approx 1$ because each sequence has its own KV cache-batching $B$ requests does not amortize it, since bytes transferred and FLOPs both scale with $B$ and cancel. The architectural fix is replacing the growing KV cache with a fixed-size recurrent state (SSM), making generation memory cost constant in context length.

• “What is associative recall, and why did early state-space models fail at it?” Strong answer: Associative recall is the task of retrieving a value given its key after many intervening tokens. Early SSMs use fixed linear recurrences-every token blends into the state identically regardless of content-so the query signal for a specific key gets diluted by unrelated tokens. Mamba’s input-dependent discretization ($\Delta_{t}$,

$B_{t}$, $C_{t}$ as functions of $x_{t}$) lets the model selectively amplify relevant bindings.

• Senior follow-up: “If you were deploying a 70B hybrid model with 32k context for an agentic workload, how would you design the layer composition?” Senior answer: Use 24 Mamba blocks and 8 full-attention layers in a 32-layer stack, placing full-attention layers at positions 4, 8, 16, 20, 24, 28- denser in the second half where long-range context is most relied upon. At $S = 32 , 768$ and $B = 4$, the full-attention KV cache for 8 layers is $8 \times 32768 \times 32 \times 128 \times 4 = 4$ GB; the Mamba state for 24 layers is 24 x 256 KB x 4 = 24 MB. Total inference memory stays under 20 GB, easily within one H100’s budget.

Key Takeaways

• Transformer generation is structurally memory-bound: attention arithmetic intensity is $\approx 1$ regardless of batch size because each sequence owns its KV cache and batching does not amortize the read.

• Mamba replaces the linearly growing KV cache with a fixed hidden state ($d_{\text{inner}} \times N$ per layer), making generation memory cost independent of context length.

• The critical innovation over S4 and earlier SSMs is the selective (input-dependent) state update, which enables content-based memory and passes associative recall benchmarks.

• At $S = 8 , 192$ and $B = 8$, a pure Mamba 7B model requires $\approx 14$ GB versus $\approx 46$ GB for an equivalent transformer-enabling $3 . 3 \times$ lower per-token latency from bandwidth alone.

• Production deployments use hybrid architectures: Mamba for local processing and cheap constant-state generation, with one sparse full-attention layer per 6-8 Mamba blocks for long-range retrieval.

17.1.2 Linear Attention as a Production Architecture

When an interviewer asks how you would serve a 400B-parameter MoE at low inter-token latency without quadratic memory, the answer that surfaces in real frontier-lab discussions is linear attention. Unlike SSMs, linear attention is a mathematically direct reformulation of the standard attention operator-no new parameterization, no custom scan kernels at inference, just a feature map that converts the quadratic softmax into a linear recurrence. The result is production models with hundreds of billions of parameters that serve at near-SSM latency while retaining transformer-class quality on most tasks. Knowing the kernel trick, the RNN dual, and where this architecture concretely breaks is the depth that separates mid-level from senior answers on architecture design interviews.

Standard causal attention at position $i$ computes:

The exponential makes this $O ( T^{2} )$ in sequence length and forces materializing a $T \times T$ attention matrix. Linear attention replaces the exponential kernel with a feature map $\phi : \mathbb{R}^{d_{\text{head}}} \rightarrow \mathbb{R}^{r}$ such that:

A common choice is $\phi ( x ) =$ ELU$( x ) + 1$, which ensures non-negative weights. With this substitution and using associativity of matrix multiplication:

where the state matrices update recurrently:

At generation step $i$, producing $y_{i}$ requires only $( S_{i - 1} , z_{i - 1} )$-no growing cache. The state $( S , z )$ per attention head is $r \times d_{\text{head}} + r$ values-fixed in size, completely independent of sequence length. During training the recurrence is parallelized via a parallel scan in $O ( T \log T )$ operations.

Linear attention converts softmax attention into an RNN with an $r \times d_{\text{head}}$ outer-product state, making generation $O ( 1 )$ in context length and replacing a growing KV cache with a fixed matrix load.

The throughput gain follows directly from the memory analysis. Softmax attention reads $O ( S \cdot n_{\text{heads}} \cdot d_{\text{head}} \cdot L )$ bytes per generation step, growing with every token emitted. Linear attention reads a fixed state of $O ( r \cdot d_{\text{head}} \cdot n_{\text{heads}} \cdot L )$ bytes per generation step regardless of how long the conversation has run. At context $S = 8 , 192$, that difference is substantial; at $S = 32 , 768$, it becomes the deciding factor in whether the model fits on a single accelerator at all.

The trade-off is sharpness of retrieval. Softmax attention is a sparse content router: at each step a handful of tokens receive nearly all the weight. Linear attention integrates all past keys and values into a single matrix $S_{t}$ via outer products. If two past tokens have similar $\phi ( k )$ directions, their values blend inseparably in $S_{t}$ and cannot be cleanly retrieved later. This means linear attention degrades on exact lookup tasks: fetching a specific fact from a long document, multi-hop reasoning, or string-copy probes. For local pattern matching, grammatical coherence, and generation tasks without hard retrieval demands, linear attention performs on par with softmax.

This trade-off is why production deployments are hybrids. MiniMax’s 456B-parameter MoE-a publicly documented production system-uses linear attention for the vast majority of layers with standard full-attention layers inserted at fixed intervals. The pattern is now industry-standard: linear and local-attention layers handle the bulk of computation cheaply; sparse full-attention layers handle long-range retrieval. Compared to Mamba, linear attention preserves the query-key-value interface of the transformer, making it straightforward to inherit pretrained weights or fine-tune an existing checkpoint toward more efficient inference without re-parameterizing the attention heads.

A Worked Example

A 7B hybrid model with $d = 4 , 096$, 32 heads, $d_{\text{head}} = 128$, feature-map dimension $r = d_{\text{head}} = 128$, 32 layers (24 linear attention, 8 full attention), bf16, $B = 8$, $S = 8 , 192$ on an H100 (80 GB HBM, 3.35 TB/s).

Full-attention KV cache (8 layers). Per layer per sequence:

For 8 layers and $B = 8$: $128 \times 8 \times 8 = 8 , 192$ MB $= 8$ GB.

Linear attention state (24 layers). Per head per layer: $r \times d_{\text{head}} \times 2$ bytes $= 128 \times 128 \times 2 = 32$ KB. Per layer (32 heads): 1 MB. For 24 layers and $B = 8$: $24 \times 1 \times 8 = 192$ MB.

Total inference memory: 14 GB (weights) + 8 GB (KV) + 0.2 GB (linear state) $\approx 22 . 2$ GB. The model fits with room to expand $B$ to 32 before hitting the 80 GB limit.

Compare to a pure transformer: 14 GB (weights) + 32 GB (all 32 layers KV at $B = 8$, $S = 8 , 192$) $= 46$ GB. Fits, but barely, with no headroom for larger batches.

Generation bandwidth per step:

• Hybrid: 14 GB + 8 GB + 0.2 GB = 22.2 GB $\Rightarrow 22 . 2 / 3 . 35 \approx 6 . 6$ ms per token.

• Pure transformer: 14 GB + 32 GB = 46 GB $\Rightarrow 46 / 3 . 35 \approx 13 . 7$ ms per token.

• Speedup: $13 . 7 / 6 . 6 \approx 2 . 1 \times$ in the memory-bound regime.

At $S = 32 , 768$, the pure-transformer KV cache grows to 128 GB and no longer fits on a single H100. The hybrid KV cache stays at 32 GB (only 8 layers), giving a total of $\approx 46$ GB-still comfortably within budget.

What You Actually Decide in Practice

• Hybrid layer ratio. A 3:1 ratio (three linear layers per one full-attention layer) recovers approximately 95% of full-transformer perplexity on standard language benchmarks while cutting KV cache to 25% of a full model’s cost. A 7:1 ratio is viable for low-retrieval workloads such as open-ended generation or summarization. Do not go beyond 8:1 on agentic or long-document workloads without running needle-in-a-haystack evaluations first.

• Feature map choice. ELU+1 is simple but can produce near-zero $z_{t}$ values early in a sequence, leading to division instability. Add a floor in your implementation: z = torch.clamp(z, min=1e-6). A more robust alternative is Gated Linear Attention (GLA), which adds an input-dependent forget gate-analogous to Mamba’s selection mechanism-and significantly improves performance on associative recall tasks at the cost of a slightly more complex recurrence.

• RoPE incompatibility. Standard RoPE rotates query and key vectors in a way that depends on the exponential attention structure; that structure disappears when you substitute the feature map. If your base model uses RoPE, you cannot apply linear attention directly to those layers without removing the positional encoding. In practice, restrict RoPE (or any rope-like encoding) to the full-attention layers and treat linear-attention layers as position-agnostic. They remain effective because local ordering cues are encoded in the token content itself.

• State size at large model width. At $d = 8 , 192$ (typical for a 70B model), the linear state per head per layer is $128 \times 128 \times 2 = 32$ KB, per layer 64 heads x 32 KB = 2 MB, across 56 linear layers at $B = 8$ gives $56 \times 2 \times 8 = 896$ MB. Manageable. But if you use a larger feature-map dimension $r = 4 \times d_{\text{head}} = 512$ for richer approximation, the per-head state quadruples-watch memory at $B > 16$.

• Training stability. Linear attention loss curves are typically noisier than softmax attention early in training, especially when $z_{t}$ is small. If initializing from a pretrained transformer, keep the first 4-8 layers as full attention during the initial fine-tuning phase and unfreeze the remaining linear layers gradually over the first 1k steps. Gradient clipping at 1.0 and a short linear warmup of 500 steps are usually sufficient to stabilize.

How This Shows Up in Interviews

• “Walk me through how linear attention eliminates the quadratic complexity of softmax attention.” Strong answer: Replace $\exp ( q^{\top} k )$ with $\phi ( q )^{\top} \phi ( k )$, then apply associativity to pull $\phi ( q_{i} )^{\top}$ outside the sum over $j$. The inner sum $S_{i} = \sum_{j \leq i} \phi ( k_{j} ) v_{j}^{\top}$ updates as a rank-1 outer-product addition each step and is $O ( r \cdot d_{\text{head}} )$ state-fixed, independent of sequence length. Training uses a parallel scan over this recurrence.

• “What does linear attention lose versus softmax attention, and when does that actually matter?” Strong answer: Softmax is a sparse content router that can concentrate almost all probability on a single token. Linear attention accumulates all history into a fixed matrix via outer products; if two past keys map to similar feature directions, their values blend inseparably. This matters for exact in-context lookup, multi-hop retrieval, and copy tasks. It does not matter for local coherence, syntax, or open-ended generation-which is why hybrids with 75-87.5% linear layers still achieve near-transformer perplexity.

• Senior follow-up: “How would you design a 70B hybrid model for an agentic workload with 64k context and hard factual lookup requirements?” Senior answer: Use 24 linear-attention layers and 8 full-attention layers in a 32-layer stack, with full-attention layers at positions 4, 8, 16, 20, 24, 28-denser toward the end where cross-document reasoning concentrates. At $S = 65 , 536$ and $B = 4$, the full-attention KV for 8 layers is $8 \times 65536 \times 64 \times 128 \times 4 \times 4 = 16$ GB; the linear state for 24 layers is $24 \times 64 \times 128 \times 128 \times 2 \times 4 = 384$ MB. Weights for a 70B model are $\approx 140$ GB in bf16, so this design requires tensor parallelism across 2-4 H100s; the KV cache stays under 20 GB per node, which is tractable.

Key Takeaways

• Linear attention replaces $\exp ( q^{\top} k )$ with a kernel feature map $\phi ( q )^{\top} \phi ( k )$, converting the growing KV cache into a fixed $r \times d_{\text{head}}$ recurrent state that is $O ( 1 )$ in context length.

• At $S = 8 , 192$ and $B = 8$, a 3:1 hybrid (24 linear, 8 full-attention layers) cuts inference bandwidth from 46 GB to 22 GB per step-a $2 . 1 \times$ latency improvement-and enables batch expansion that a pure transformer cannot afford.

• At $S = 32 , 768$ a pure transformer KV cache overflows a single H100; the hybrid fits comfortably by limiting the growing cache to 8 layers.

• Linear attention degrades on exact retrieval tasks; production architectures use a 3:1 to 7:1 hybrid ratio with full-attention layers handling long-range lookup. MiniMax’s 456B MoE is the public benchmark for this pattern.

• Standard RoPE is incompatible with the linear attention kernel; restrict position encodings to full-attention layers and treat linear-attention layers as position-agnostic.

17.2 Non-Autoregressive Generation

17.2.1 Diffusion Language Models

When a senior engineer at a frontier lab asks you how you would redesign the text generation stack if the autoregressive constraint were lifted, they are probing whether you understand the root cause of inference inefficiency rather than just its symptoms. Diffusion language models answer that question by abandoning sequential generation entirely, producing all output tokens simultaneously and then refining them across a fixed number of forward passes. This approach has moved from a research curiosity to a production-relevant architecture rapidly, and understanding its mechanics is now a legitimate interview topic at labs building the next generation of serving infrastructure.

Autoregressive models are forced to generate one token at a time because each token conditions on all preceding tokens. Diffusion models break this constraint by starting from a fully noisy or masked sequence of the target length and iteratively denoising it toward coherent text. At each refinement step $t$, the model receives the current noisy sequence $x_{t}$ and predicts a cleaner version $x_{t - 1}$, until $x_{0}$ is the final output after $K$ steps.

The decisive structural difference is that every refinement step processes the entire output sequence in a single forward pass. If you are generating a response of length $L$ tokens, the model sees all $L$ positions simultaneously at every step. This is structurally identical to the prefill phase of an autoregressive model, not its generation phase. There is no KV cache that grows with sequence length, because there is no sequential dependency to cache. Every step is a dense matrix multiply across the full sequence.

Formally, let $q ( x_{t} | x_{0} )$ be a fixed forward noising process (masking or Gaussian corruption) and $p_{\theta} ( x_{t - 1} | x_{t} )$ the learned denoising model. Training minimizes a denoising objective:

Inference draws $x_{K}$ from the prior and runs the chain $x_{K} \rightarrow x_{K - 1} \rightarrow \cdot \cdot \cdot \rightarrow x_{0}$ using a fixed number of steps, typically 10-50 in practice.

The payoff is that by paying the cost of a prefill-like forward pass $K$ times instead of paying the cost of a sequential memory-bottlenecked pass $L$ times, diffusion models trade a small number of expensive parallel steps for a large number of cheap-but-idle sequential steps.

The arithmetic intensity argument is decisive. During autoregressive generation, the effective arithmetic intensity per token is approximately 1 for the attention layers and $B$ for the MLP layers. At batch size 1, both collapse to the floor, and memory bandwidth - not compute - limits throughput. You spend most of every GPU cycle waiting for HBM reads.

In a diffusion model, each forward pass processes all $L$ output tokens at once. The arithmetic intensity of the MLP layers becomes $B \times L$. For a sequence of 1024 tokens and batch size 1, that is 1024 versus 1 - three orders of magnitude better GPU utilization per pass. The model is compute-limited for any reasonably long sequence, meaning you are doing useful arithmetic on every cycle of the accelerator.

The trade-off is quality. Autoregressive generation samples each token from its exact conditional distribution given the full prefix. Diffusion models make a weaker guarantee: the global distribution over output sequences is approximately correct after $K$ steps, but individual token predictions mid-chain are noisy. Early diffusion LMs struggled on tasks requiring precise factual recall or tight long-range consistency. On coding tasks, where correctness is mechanically verifiable, results are already competitive with similarly-sized AR models. The open question is whether more refinement steps or better training objectives will close the gap on harder tasks.

A Worked Example

Consider generating a 1024-token response using a 7B-parameter model on a single H100 (3.35 TB/s HBM bandwidth, 989 TFLOPS BF16).

Autoregressive baseline, batch size 1. Each generation step reads approximately the full parameter set: $7 \times 10^{9} \times 2$ bytes (BF16) $= 14$ GB. Latency per token $\approx 14$ GB $\div 3 . 35$ TB/s $\approx 4 . 2$ ms. Total for 1024 tokens: $4 . 2 \times 1024 \approx 4 . 3$ seconds. Throughput: approximately 238 tokens/second.

Diffusion model, 20 refinement steps. Each step is a full forward pass over 1024 tokens. FLOPs per pass:

Time per pass on H100: $1 . 43 \times 10^{13} \div 9 . 89 \times 10^{14} \approx 14 . 5$ ms. Total for 20 steps: $20 \times 14 . 5 = 290$ ms. Throughput: $1024 \div 0 . 29 \approx 3 , 500$ tokens/second.

That is roughly a $15 \times$ throughput improvement while generating the same number of tokens from a similarlysized model. The 20 refinement steps can be reduced to 10 for another $2 \times$ speedup at some quality cost, or raised to 50 for better quality at reduced throughput. The operating point is a continuous dial, not a fixed cost, which is qualitatively different from autoregressive generation where the step count equals the token count and is non-negotiable.

What You Actually Decide in Practice

• Refinement step count is a quality-latency dial, not a constant. Start at 20 steps as a reasonable default for production. Below 10, quality degrades sharply on long-form generation. Above 50, returns diminish and you approach the throughput of a small autoregressive model. Sweep this on your task distribution and set it per request tier, not globally.

• Output length must be committed at request time. Unlike AR models that stop at an EOS token, diffusion models require knowing the target sequence length upfront. Wasted padding positions cost compute proportionally. Build your serving layer to bucket requests by expected output length (powers of two: 128, 256, 512, 1024 work well) to keep average padding below $2 \times$.

• Architecture uses bidirectional attention. Because all positions are visible at every step, there is no causal mask. This means flash attention applies, prefix caching does not, and the attention FLOPs per step are $O ( L^{2} )$ rather than $O ( 1 )$ per new token. For very short outputs ($L < 64$), the quadratic cost erodes the throughput advantage rapidly.

• Quality gaps are real on recall-heavy tasks. In my experience testing diffusion LMs on tasks requiring exact entity lookup or precise numerical computation, quality drops noticeably compared to AR models of the same parameter count. Routing recall-heavy requests to an AR model and generation-heavy requests to a diffusion model is a production-viable strategy, but it requires a lightweight classifier at the request boundary.

• The draft model in speculative decoding is a natural fit. A small diffusion model can serve as the draft proposal generator in a speculative decoding setup, because it produces multiple tokens per forward pass natively. If the acceptance rate is reasonable, this combines the quality guarantee of the AR target model with the throughput of a diffusion draft.

How This Shows Up in Interviews

• “Autoregressive generation is memory-bottlenecked at batch size 1. What are the fundamental architectural alternatives?” Strong answer: contrast AR generation (arithmetic intensity $\approx 1$, memory-limited) with diffusion (all tokens in parallel, compute-limited, intensity $\approx B L$). Name the trade-off explicitly: throughput advantage is large for long outputs but shrinks for short ones; quality lags on recall tasks.

• “Walk me through how a diffusion language model generates text.” Strong answer: describe the iterative denoising chain from $x_{K}$ to $x_{0}$, emphasize that every step is a full-sequence forward pass (identical to prefill), explain the absence of a KV cache, and note that output length must be specified upfront. Give the arithmetic: 20 steps $\times$ prefill cost versus $L$ steps $\times$ generation cost.

• “If diffusion models are so much faster, why isn’t everyone using them?” Strong answer: the quality gap on factual recall and instruction following is real. Fixed output length at request time creates padding waste and serving complexity. The compute advantage only materializes at output lengths long enough to amortize per-step cost, roughly $L \geq 200$ tokens. Frontier labs are using them for high-throughput batch workloads, not interactive chat.

• Senior follow-up: “How would you use a diffusion model as the draft component in a speculative decoding pipeline serving a 70B AR target?” Senior answer: the diffusion draft generates $K$ token proposals in a single forward pass (no sequential sampling loop), which is faster than an AR draft generating $K$ tokens one-by-one. The AR target verifies all $K$ in parallel via prefill. Acceptance rate depends on distribution alignment; fine-tune the draft on the target model’s output distribution to improve it. The diffusion draft’s fixed-length constraint requires knowing the speculative lookahead window $K$ ahead of time, which is compatible with standard speculative decoding where $K$ is a fixed hyperparameter.

Key Takeaways

• Diffusion language models generate all output tokens in parallel at each refinement step, shifting inference from memory-limited (arithmetic intensity $\approx 1$) to compute-limited (intensity $\approx B L$).

• A 7B diffusion model with 20 refinement steps achieves approximately 3,500 tokens/second on an H100 versus approximately 238 tokens/second for a 7B autoregressive model at batch size 1 - roughly a $15 \times$ throughput advantage at $L = 1024$.

• The advantage shrinks for short outputs and disappears below roughly $L = 64$; build output-length bucketing into the serving scheduler to protect efficiency.

• Quality lags AR models on recall-heavy tasks; the refinement step count is the quality-latency dial to tune per task tier.

• There is no KV cache, which eliminates the most complex part of AR serving infrastructure but requires output length to be committed at request time.

17.2.2 Implications for Serving Cost

When inference costs are growing faster than revenue, the honest engineering question is not “which model is fastest in a benchmark?” but “what changes when I swap from an autoregressive backend to a non-autoregressive one, and which parts of my serving stack need to be redesigned?” Non-autoregressive generation does not just produce tokens faster - it changes the memory model, the scheduler contract, and the capacity planning calculus in ways that compound across a large deployment. Knowing these second-order effects is what separates a staff-level system design answer from a mid-level one.

The serving cost of an autoregressive model has two structurally different components: the model weights, which are static and shared across all requests, and the KV cache, which grows per-request and per-generatedtoken. Managing that growth requires PagedAttention-style virtual memory, copy-on-write prefix sharing, and eviction policies - all of which add engineering overhead and introduce fragmentation.

Non-autoregressive models collapse this structure. Because there is no sequential dependency across refinement steps, there is no KV cache to accumulate. Each step reads the model weights and the current noisy sequence, writes the refined sequence, and exits. Memory consumption per request is a flat $O ( L \cdot d )$ bytes where $L$ is the committed output length and $d$ is the model hidden dimension - constant across all $K$ refinement steps.

The total HBM footprint for a diffusion serving deployment is:

where $B$ is the batch size and $L$ is the fixed output length. Compare this to an AR deployment where $M_{\text{KV}}$ grows as the sequence extends and must be tracked dynamically. The diffusion model’s memory budget is fully computable before a single token is generated, which makes capacity planning a spreadsheet exercise rather than a runtime estimation problem.

Eliminating the KV cache turns the serving stack’s hardest dynamic-memory problem into a static allocation, at the cost of requiring output length to be committed at request time.

The economic argument is clearest when you look at cost per output token. Autoregressive generation is memory-bandwidth-limited: every token costs one full HBM sweep of the parameter set. For a 7B model in BF16, that is 14 GB of data movement per token. At 3.35 TB/s on an H100, the hardware-imposed minimum is 4.2 ms per token regardless of software optimization. No kernel tuning removes this floor.

Diffusion generation amortizes weight-sweep cost across all $L$ output tokens in a step. The effective bandwidth cost per output token across $K$ refinement steps is:

For $K = 20$ and $L = 512$, this is $20 / 512 \approx 0 . 04 \times$ the per-token cost of AR generation. Compute, not memory bandwidth, becomes the binding constraint - which is a far healthier operating regime because compute is easier to scale horizontally.

The catch is fixed output length. Autoregressive models can stream tokens and stop at any EOS, adapting response length to content at runtime. Diffusion models commit to $L$ upfront. Positions that go unused are wasted compute. In my experience, the serving system must either predict output length accurately or accept efficiency losses from padding; neither is free, but length bucketing at powers of two keeps average waste below $2 \times$.

A Worked Example

Consider 100 concurrent requests, each expecting up to 512 output tokens, served on a single H100 with 80 GB HBM.

Autoregressive backend (13B model, BF16):

• Weight memory: $13 \times 10^{9} \times 2 = 26$ GB.

• KV cache per sequence (40 layers, 32 KV heads, $d_{\text{head}} = 128$, 512 tokens): $512 \times 32 \times 128 \times 40 \times 2 \times 2 \approx 3 . 4$ GB.

• At batch size 16: $16 \times 3 . 4 + 26 = 80 . 4$ GB - already at the memory limit.

• Maximum practical batch size: 15-16 requests.

Diffusion backend (comparable model scale, BF16):

• Weight memory: 26 GB (same).

• Sequence buffer per request ($L = 512$, $d = 5 , 120$): $512 \times 5 , 120 \times 2 \approx 5$ MB.

• At batch size 100: $100 \times 5$ MB $+ 26$ GB $= 26 . 5$ GB total.

• 53 GB remain free for larger batches or additional model shards.

• Maximum practical batch size: limited by compute throughput, not memory - exactly the operating point you want.

The memory freed by eliminating the KV cache allows a batch size of 100 versus 16 on identical hardware, a $6 \times$ improvement that directly multiplies throughput before any algorithmic changes. Combined with the pertoken compute advantage from the previous subsection, the end-to-end throughput difference for long outputs is striking.

What You Actually Decide in Practice

• Output length prediction is a first-class infrastructure problem. Build a lightweight length predictor (a small classification head on a pooled prefix embedding works) that bins requests into length buckets before routing. Rounding up to the nearest power of two (128, 256, 512, 1024) keeps padding waste below $2 \times$ on most real distributions. Avoid variable-length padding within a batch; pad all requests in a batch to the same $L$ and process them together.

• Maintain separate queuing lanes per output-length bucket. Pack requests of similar expected length into batches; flush each lane when the batch is full or a timeout fires. This is simpler than AR continuous batching because there is no mid-sequence request admission: all requests in a batch start and finish together, so the scheduler contract is identical to training.

• Expose refinement step count as a per-request parameter. Low-stakes requests (short-form Q&A, autocompletion) can use 10-15 steps. High-stakes requests (long-form document generation, code for production) benefit from 30-50 steps. Implementing this as a tier enum in your request proto gives the product team a quality-cost lever without infrastructure changes.

• Route by task type, not just model size. Deploy AR backends for time-to-first-token-sensitive workloads (interactive chat, function calls, streaming responses) and diffusion backends for throughput-sensitive batch workloads (document summarization, bulk code generation, offline evaluation). A lightweight router that classifies on task type plus expected output length outperforms a single-model strategy on both cost and quality.

• Prefix caching does not apply; but prompt conditioning is cheaper. Because diffusion models use bidirectional attention, there is no causal prefix whose KV vectors can be cached across requests. The cost of conditioning on a long system prompt is paid again at every refinement step. If your serving traffic has a shared system prompt, this eliminates one of the major cost-saving mechanisms of modern AR serving systems. Design around it by keeping prompts short or by absorbing prompt information into the model at training time via instruction tuning.

How This Shows Up in Interviews

• “How does eliminating the KV cache change your serving infrastructure?” Strong answer: memory becomes static and fully predictable before generation begins; batch size is set by compute budget rather than memory budget; the scheduler no longer needs eviction policies or copy-on-write prefix sharing; but output length must be committed upfront, introducing a new length-prediction requirement.

• “What are the limits of the diffusion serving cost advantage?” Strong answer: three limits. First, short outputs ($L < 64$) do not amortize the per-step weight-sweep cost, shrinking or eliminating the advantage. Second, padding waste from poor length prediction can erode the memory savings. Third, prefix caching disappears, which hurts workloads with long shared system prompts. Name the cross-over point quantitatively: the diffusion backend breaks even with AR roughly when $K / L < 1$, i.e., the number of refinement steps is smaller than the output length.

• “How would you operate a mixed AR-diffusion cluster to minimize cost per quality point?” Strong answer: profile your request distribution by task type, expected output length, and quality sensitivity. Assign AR capacity to the interactive, recall-heavy, and short-output tail; assign diffusion capacity to the long-output, throughput-sensitive bulk. Size each backend to its SLA: AR optimized for P95 latency, diffusion optimized for tokens/second. Rebalance allocation weekly as the request mix shifts.

• Senior follow-up: “How does the capacity planning model change when the serving SLA is throughputper-dollar rather than latency?” Senior answer: AR throughput is bounded by memory bandwidth: tokens/second $\approx$ BW$/ ( 2 \times N_{\text{params}} )$, independent of batch size once memory is saturated. Diffusion throughput scales with compute: tokens/second $\approx$ TFLOPS $\times L / ( 2 \times N_{\text{params}} \times K )$, meaning it improves with longer outputs and fewer steps. The cost-per-token crossover is at $L = K$ (output length equals refinement step count); above that, diffusion dominates per dollar. Model this in your capacity plan and provision diffusion GPUs for the long-output tail of your distribution while using AR GPUs for the short-output head.

Key Takeaways

• Diffusion models have a flat, predictable per-request memory footprint ($O ( L \cdot d )$ bytes, constant across all refinement steps) versus an AR model’s linearly growing KV cache, enabling $\sim 6 \times$ larger batch sizes on identical hardware for 512-token outputs.

• The effective per-token bandwidth cost is $K / L$ times the AR cost; at $K = 20$, $L = 512$, this is $\sim 4 \%$ of the AR figure, shifting the bottleneck from memory bandwidth to compute.

• Output length must be committed at request time; build output-length bucketing (powers of two) into the scheduler to hold padding waste below $2 \times$.

• Prefix caching does not apply to diffusion models (bidirectional attention, no causal KV cache); workloads with long shared system prompts lose a major AR cost-saving mechanism.

• Route by task type and expected output length: diffusion for long-output batch workloads, AR for interactive or recall-heavy short-output workloads.