Inference: Fewer, smaller forwards

This is Part 3 of our Inference Engineering series. Read the other articles: Part 1: From noise to words · Part 2: Making every denoising pass fast.

Part 2 compressed the cost of one forward. That is only half of the throughput equation. Generation rate is

Kernels attack the denominator. The decoding algorithm attacks the numerator, and it also decides how wide each forward must stay. A 256-token canvas (generation window) that revises every position on every step pays for a wide forward even when most tokens are already settled. Everything below tries to commit more tokens per pass, warm the next canvas early, feed better conditioning into the next step, and stop paying attention over finished positions.

All Funnel numbers below use the Triton W4 expert stack from Part 2 on an NVIDIA DGX Spark (GB10). Every quantitative claim names its baseline.

Three generation paths

DiffusionGemma (Google’s open text-diffusion model) generates in fixed 256-token canvases. How those canvases are scheduled is what separates the paths.

Official Hugging Face sampler

Hugging Face model.generate() uses DiffusionGemma’s native sampler. One canvas is denoised at a time with entropy-bounded acceptance (commit tokens whose predictive uncertainty is low enough) and adaptive stopping. Part 2 kept this path fixed and only swapped BF16 (16-bit brain floating-point) experts for Triton W4A16. That gave a two-benchmark mean of 81.1 → 126.9 tok/s at 96/100 versus 94/100. It is not the Funnel result.

D2F pipeline

D2F overlaps denoising across multiple canvases. While canvas N is still noisy, canvas N+1 can already be decoding. Finished canvases are written into the encoder KV cache (stored attention keys and values). Active canvases share a block-causal multi-canvas decode. The project’s D2F LoRA (Low-Rank Adaptation) adapter was trained for this schedule and is merged before serving.

The pipe still recomputes a full active canvas while that canvas is live. It raises commits per forward, but it does not shrink the window the way Funnel does. On GSM8K with argmax renoising, draft spawn, k=2, and the Triton W4 stack, the measured pipe record is 92/100 at 196.5 tok/s (about 44.8 commits per forward). Same-session native baselines were roughly 95–97/100 near 92–95 tok/s. On HumanEval the pipe remains strong: 88/100 at 218 tok/s with a higher confidence bar and EOS-priority commit.

Funnel

Funnel keeps the same Triton W4 stack but changes how the active window is maintained. Tokens that have been finalised leave later decoder forwards. Their attention keys and values are promoted into a scattered KV cache (or arena: a preallocated buffer for finished K/V state). Later forwards therefore see a smaller active set. That is what “fewer, smaller forwards” means: fewer steps when commits are high, and cheaper steps once the window shrinks.

Funnel is the GSM8K speed path near 250 tok/s. Hard code still often prefers the D2F pipe’s full-canvas recomputation. The rest of this article walks every mechanism that makes those numbers possible.

Advancement 1: multi-canvas D2F scheduling

A single 256-token canvas forces a long serial chain of denoising steps before the next block can start. D2F keeps up to K canvases in flight. Canvas N+1 begins while N is still partially noisy. Between canvases the attention mask is block-causal: canvas i may attend to earlier canvases and to the prompt cache, but not to later canvases’ unfinished rows.

Each canvas carries a state:

• full: the front canvas. It may commit tokens. The front is always forced to full so generation never stalls on a non-committing passenger.

• semi: a later canvas that has graduated enough to commit, but is not yet the front.

• draft: a passenger that warms up via renoising but does not commit yet.

Graduation uses two thresholds. When the previous canvas’s decode fraction (share of positions already committed) reaches tau_add (default 0.5), a draft canvas becomes semi. When it reaches tau_act (default 0.95), that canvas becomes full. The LoRA was distilled for exactly this multi-canvas, block-causal student against a bidirectional teacher, so the schedule is not an inference hack bolted onto an unrelated checkpoint.

What it buys: more committed tokens per wall-clock second by overlapping work across canvases, without waiting for canvas N to finish before starting N+1.

Advancement 2: tau_conf (confidence commitment)

On every forward the model emits a distribution over the vocabulary at each uncommitted position. Confidence is the probability of the argmax (top prediction). A position commits only when

tau_conf is the main speed/accuracy knob. Lower values commit earlier and raise commits per forward (cpf); higher values wait longer and usually protect accuracy on harder text.

Typical serving choices:

• GSM8K Funnel records: τ_conf = 0.7

• HumanEval Funnel: τ_conf = 0.9 (code needs a higher bar)

• HumanEval pipe: τ_conf = 0.85 plus EOS-priority (below)

What it buys: the numerator of tok/s. Without early commitment, every position rides every forward until the whole canvas settles.

Advancement 3: argmax renoising

Rejected positions must be refilled before the next step. Classical diffusion would resample from a noise distribution. The record protocol instead writes the model’s current argmax back into those positions (pure argmax; a mix mode that randomly flips a fraction of positions exists but is not used for the published points).

That keeps the draft coherent between steps. Empirically, the argmax that is rejected at step t is often already the token the position eventually commits, so throwing it away and drawing uniform noise wastes useful structure.

A related escape hatch: if a canvas commits at most one token per forward for eight consecutive steps (Jacobi livelock), the next renoise uses a one-shot random restart, then argmax resumes. Without that, a canvas can spin for dozens of forwards on a stubborn tail.

What it buys: higher cpf and fewer wasted forwards, because the canvas stays near a coherent partial answer instead of being re-shuffled into noise.

Advancement 4: draft spawn

Draft spawn starts canvas N+1 early, before the previous canvas is ready to graduate under the normal tau_add rule. The spawn threshold is tau_add − 0.15. The new canvas enters as a non-committing draft: it receives argmax renoising and self-conditioning updates, but cannot freeze tokens until it graduates to semi or full.

The point is warmup. By the time canvas N approaches completion, canvas N+1 already carries a structured draft instead of fresh random tokens. The D2F pipe GSM8K record depends on this. Funnel multi-canvas runs use the same state machine.

What it buys: fewer cold starts on later canvases, which raises effective cpf across a multi-canvas answer.

Advancement 5: tau_eos (EOS-priority commit)

End-of-sequence (EOS) tokens are special. On code, a model trained on long tutorial answers can be weakly calibrated on when to stop, so junk tails appear after a correct solution. tau_eos gives an argmax-EOS a lower commit bar than ordinary tokens. If the model is pointing at EOS with confidence above tau_eos (often 0.5) but below tau_conf, the position still commits.

A stricter variant, EOS streak, requires two consecutive argmax-EOS decisions before the low bar applies. That costs speed and is optional.

Important caveat: EOS-priority helps this adapter on HumanEval and MBPP. It can truncate checkpoints trained on eval-shaped short answers. GSM8K Funnel records leave it off.

What it buys: shorter, cleaner code completions without waiting for ordinary tau_conf on the stop token.

Advancement 6: forced accept on the front canvas

If the front canvas is in full state, still has uncommitted positions, and no position cleared tau_conf on this step, the decoder force-commits the single highest-confidence remaining position. Generation cannot stall with zero commits forever.

This is a liveness rule, not an accuracy feature. Combined with the Jacobi random restart, it keeps worst-case forward counts bounded.

What it buys: no infinite stalls on a canvas that is almost done but never quite clears the bar.

Advancement 7: self-conditioning

Self-conditioning feeds the previous step’s denoising prediction into the next forward as an extra signal. Instead of seeing only the current noisy token ids, the decoder also sees a soft summary of what it predicted last time. DiffusionGemma implements this with a small conditioning module over previous logits (model scores over the vocabulary).

On Funnel, self-conditioning is on by default. A paired GSM8K A/B at τ_conf=0.7 measured:

So SC is not only a quality trick. On that run it also moved throughput up, consistent with slightly better commits per forward and slightly fewer wasted steps.

What it buys: a stronger reverse-step signal, which improves both accuracy and the numerator when the draft is more stable.

Advancement 8: top-k self-conditioning projection

Naively, self-conditioning softmaxes previous logits over the full vocabulary (hundreds of thousands of entries) and mixes them with the full embedding table. That is a large, mostly M-independent cost: a full-vocab fp32 softmax plus roughly a 1.5 GiB embedding read, every forward, even when the active window has shrunk to a handful of rows.

The fix is top-k projection. Keep only the top k logits (for example 64 or 256), renormalise that small simplex, and gather only those embedding rows. The soft embedding is then a tiny weighted sum instead of a full-vocab matmul. The decoder is called with the heavy full-vocab path disabled; a patched conditioning module injects the precomputed soft embeddings.

What it buys: self-conditioning stays affordable after Funnel has already shrunk the window, so the SC accuracy/speed benefit is not eaten by a vocabulary-wide gemm.

Advancement 9: Funnel’s shrinking window (scattered positions)

This is the geometric change that separates Funnel from the pipe.

In the pipe, an active canvas of size C is forwarded as C positions every step until retire. In Funnel, the forward window is only:

• positions that are still uncommitted, plus

• positions committed on the previous step that are waiting to promote their KV (the one-step-late rule below).

Committed tokens that have already been promoted no longer appear as query rows. Attention over finished text happens through the cache, not by recomputing those positions as live canvas tokens. Expert compute and attention both scale with the live window, so the forward gets cheaper as the answer settles.

Modelled ceilings for this scattered schedule sit around 2–3× versus always forwarding the full canvas, before counting multi-canvas overlap.

What it buys: lower seconds per forward as decode fraction rises. That is the “smaller forwards” half of the title.

Advancement 10: promote-one-step-late (correctness rule)

Capture and promote are easy to get wrong. At the forward where a token is about to commit, the input at that position may still be a draft value. Commitment happens after the forward, when confidence is read. If you stash K/V at commit decision time, you can cache the keys and values of the wrong token.

The rule: a newly committed token rides the window for one extra forward with its final token id as input. Only then is its K/V captured and promoted into the cache. Pending flags track that one-step delay.

RoPE (rotary position embeddings) is baked into keys at the token’s true absolute position, so cache entry order can be compacted; masks address columns by position, not by “left-to-right canvas index.”

What it buys: correctness. Without it, the scattered cache silently poisons later attention.

Advancement 11: scattered KV promotion and the arena

Promotion appends the captured K/V rows of finished tokens onto the persistent cache so later windows can attend to them read-only.

Two implementations:

1. Simple promote: concatenate captured rows onto each layer’s key/value tensors.

2. Arena: preallocate per-layer buffers. The live window writes in place at [alen : alen+q]. Attention reads the contiguous prefix [:alen+q]. Promotion left-compacts kept rows and advances alen. That removes per-step cat allocations and makes the layout graph-friendly.

Arena mode also stores each entry’s absolute position. Sliding-window attention layers can then band-mask entries that fall outside the 1024-token sliding window, which unlocks max_canvases=5 even when prompt plus all canvases would exceed the sliding window. Without arena position tracking, the harness clamps canvas count to stay inside the window.

What it buys: finished tokens leave the compute window but remain visible to attention; multi-canvas depth becomes practical.

Advancement 12: free front retire (vs encoder re-encode)

When the pipe finishes a canvas, it typically pays an encoder re-forward to write clean KV for that block. Funnel has already promoted K/V during decode, so retiring the front canvas is mostly bookkeeping: pop the finished canvas, append its tokens to the answer, recycle buffers.

An optional retire_refresh path re-encodes on retire for pipe-parity experiments. The speed path leaves refresh off and trusts promoted KV (encoder and decoder weights are tied, so captured decode KV matches what a clean encode would produce for the same token ids).

What it buys: one fewer heavy pass per finished canvas on the Funnel speed path.

Advancement 13: prefix attention

Profile work showed attention paying a large cost when additive float masks push scaled-dot-product attention off the flash path onto a slower math path, across all decoder layers.

The block-causal multi-canvas mask has special structure: segment i attends exactly to a contiguous KV prefix. Funnel can replace one big masked attention with per-segment mask-free attention over those prefix views, which stays on fast kernels and is mathematically the same when the band is a no-op. Sliding layers keep the masked path when the band actually clips.

What it buys: lower attention latency on multi-canvas forwards without changing the attention pattern.

Advancement 14: tau_decay and early spawn (straggler controls)

Two schedule extras target ugly tails:

• tau_decay: after a canvas has been active for more than 16 forwards, gradually lower tau_conf (floor 0.45). Stubborn positions eventually commit instead of burning dozens of forwards.

• early spawn: if the joint live window is already tiny (≤48 rows) and another canvas is still allowed, spawn immediately so those M-independent fixed costs (expert stream, lm head) ride on a fresh canvas’s work instead of a handful of leftover tokens.

On a GSM8K schedule sweep at τ_conf=0.7, base Funnel measured 45/50 at 253.8 tok/s; tau-decay and early-spawn arms sat at 45/50 near 257.5–257.6 tok/s. Same accuracy band, slightly higher speed.

What it buys: fewer pathological long tails without a measurable accuracy collapse on that sweep.

Advancement 15: commit-agree hysteresis

Optional agree mode commits only when confidence clears the bar and the argmax matches the previous step’s draft at that position (under argmax renoising, the input token is last step’s argmax for uncommitted rows). A confident first impression that flips on the next step is treated as unstable and is not frozen yet.

That costs roughly one forward per commit wave. On HumanEval it is an accuracy lever: the final Funnel result is 44/50, with the reported configurations running between 248.8 and 254.2 tok/s. Agree delays commitment when early confidence is unstable; it should be evaluated as an accuracy tradeoff, not assumed to be a speed improvement.

What it buys: accuracy on hard code by refusing one-step confidence spikes.

Advancement 16: supporting kernel pieces under Funnel

These are not new generation rules, but they sit under the Funnel measurements and matter for the denominator:

• fused RMSNorm (root-mean-square normalisation)

• fused QKV (query/key/value projection) with capture compatible with KV stash

• fused softmax-plus-argmax over the vocabulary for confidence

• optional fused QKV-post (RoPE, norms, arena write) in arena mode

• LoRA merged before serve (adapter matmuls otherwise tax every forward)

Part 2 already covered the expert W4 story. Part 3 assumes that stack and focuses on the algorithm above.

What the stack measures

GSM8K under Funnel (the ~250 tok/s band)

Compare to the D2F pipe on the same benchmark family: 92/100 at 196.5 tok/s. Funnel is faster; the pipe can still win accuracy on some sessions. Judge the band: n=50 accuracy moves by a couple of items run to run, and tok/s moves by tens under thermal drift on GB10.

HumanEval under Funnel (n=50, strong accuracy)

These sit in the same speed band as the GSM8K Funnel runs, with the higher confidence bar code needs. The D2F pipe’s final HumanEval point is 88/100 at 218 tok/s (with tau_eos). Funnel’s final score is 44/50. Both are 88% on their respective samples, but the different sample sizes mean this is not a matched accuracy comparison.

Serving split

• Routine maths and max-speed easy regimes → Funnel (arena, k=2, self-conditioning on, τ_conf around 0.7).

• Hard code → D2F pipe (higher τ_conf, tau_eos, draft spawn, full-canvas recomputation while a canvas is live).

Both share the Part 2 Triton W4 experts. The difference is the generation algorithm and the cache schedule.

Putting the series together

1. Part 1: discrete diffusion became a practical language-model shape.

2. Part 2: Triton W4A16 experts make official Hugging Face generation 1.56× faster by simple mean throughput on GSM8K and HumanEval, at 94/100 versus 96/100.

3. Part 3: D2F and Funnel raise commits per forward and shrink later forwards. Funnel plus that W4 stack reaches about 250 tok/s on GSM8K, and 44/50 near that speed on HumanEval n=50.

The remaining work is choosing the right generation path for the task, keeping confidence calibrated per domain, and measuring every claim against an explicit baseline.

References

1. Google DeepMind. (2026). DiffusionGemma model card.

2. Google DeepMind. (2026). Diffusion in Text Generation Explained.

3. Google DeepMind. (2026). DiffusionGemma model overview.