Inference: Making every denoising pass fast

This is Part 2 of our Inference Engineering series. Read the other articles: Part 1: From noise to words · Part 3: Fewer, smaller forwards.

Keep Hugging Face (an open-source machine-learning platform and model hub) model.generate() unchanged. Replace only DiffusionGemma’s BF16 (16-bit brain floating-point) experts with our Triton (a Python language for writing custom GPU kernels) W4A16 (4-bit weights, 16-bit activations) implementation. Across GSM8K (Grade School Math 8K) and HumanEval (a Python coding benchmark), simple mean throughput rises from 81.1 to 126.9 tok/s (tokens per second) while measured accuracy moves from 96/100 to 94/100.

Part 1 left DiffusionGemma (Google’s open text-diffusion model) with a concrete runtime and an unpaid bill. A 256-token canvas (generation window) can reconsider many positions in one forward, but that forward still has to move large expert matrices through a memory system that measured about 250 GB/s on this machine. Generation rate is still

This article attacks the denominator under the official sampler. Part 3 will change the algorithm that sets the numerator and the width of later forwards. Keep those experiments separate: the 1.56× figure below is not the Funnel speed band.

The experiment

Load the unmodified DiffusionGemma base once through Hugging Face Transformers (the library that loads and runs transformer models). Generate with official model.generate(). Then, without reloading the model or changing the sampler (the decoding procedure that chooses tokens), replace the expert weights and expert forward with our Triton W4A16 path. Ask the same questions again.

One model load. Arm 1 is pure BF16. The harness packs those same experts to W4A16 in place and runs arm 2. Both arms use:

• the same google/diffusiongemma-26B-A4B-it checkpoint, with arm 2 quantised in place;

• the same official Hugging Face model.generate() sampler;

• the same 50 GSM8K and 50 HumanEval prompts;

• the same random seed, reset before each benchmark;

• the same GB10 (DGX Spark GPU) machine and Python process.

There is no LoRA (Low-Rank Adaptation) adapter, no D2F pipeline, no Funnel schedule (a decoder that shrinks the active canvas as tokens finish), no custom confidence threshold, no fused RMSNorm (root-mean-square normalisation in one kernel), and no fused QKV (combined query-key-value projection) in this A/B. The only changed component is the routed experts: BF16 execution in the Hugging Face model versus W4A16 storage plus our grouped Triton expert kernels.

One limitation remains. BF16 always runs first and W4 second, so the run is not counterbalanced against thermal drift. Same-process design removes load variation, and both benchmarks move by about the same factor, but order still belongs in the report.

The result

Combined score is the unweighted sum of correct answers across 50 GSM8K and 50 HumanEval items, not an average of percentages.

The table rounds 126.85 to 126.9 tok/s. The 1.56× uses the unrounded mean.

This is close to quality-preserving on a 100-item sample, not mathematically lossless. HumanEval stays 47/50. GSM8K moves from 49/50 to 47/50. The quantised model solves 94 items where BF16 solves 96. W4 changes the weights and can change the denoising trajectory, so a larger evaluation would be needed for a tighter quality estimate.

Shorter answers do not explain the speedup. Average generated length rises slightly under W4, from 272 to 281 tokens on GSM8K and from 186 to 187 on HumanEval.

DiffusionGemma makes the expert swap matter. The model has 25.2B total parameters and 3.8B active per token. Each MoE (mixture-of-experts) layer routes a token through 8 experts from 128, plus one always-active shared expert. A 256-token denoising pass therefore moves a large, input-dependent set of expert matrices. The rest of this article explains why the A/B moves, what lower-level kernel numbers actually measure, and which extras failed once generation was complete.

How to read the remaining comparisons

Three evidence types appear below:

1. Hardware probe: a small synthetic program measures one property such as matrix throughput or memory bandwidth. It does not generate text.

2. Kernel microbenchmark: two implementations receive the same artificial tensors and run one operation, such as one MoE layer.

3. End-to-end decoder run: the model answers benchmark questions; we report accuracy and tok/s.

Every result names its baseline. Official Hugging Face generation means model.generate() with DiffusionGemma’s native sampler. Our custom decoder means one of two Python generation loops. D2F overlaps denoising across multiple canvases. Funnel also drops finished tokens from later forwards and caches their keys and values. Part 3 covers both algorithms.

A GPU kernel is one function launched on the GPU. Triton is a Python language for writing those kernels. A fused kernel combines work that would otherwise need several launches or intermediate tensors. Quantisation stores weights with fewer bits.

Start with the measured machine

Vendor specs are useful. A performance model should still use the box that ran the code. Standalone PyTorch probes on GB10 gave:

BF16 is a 16-bit floating-point format. FP8 is 8-bit floating-point. A TFLOP/s is a trillion floating-point operations per second. GB/s measures bytes moved per second. DRAM is the GPU’s main off-chip memory. Streaming multiprocessors (SMs) are the main compute blocks.

A roofline model asks whether an operation runs out of arithmetic or memory first. The ridge is measured compute over measured bandwidth:

Below about 388 useful FLOP per byte fetched, moving data is the likely limit.

For one expert with t token rows, input width d, and intermediate width m, a simplified intensity is

where b is bytes per stored weight. At t=1, one row still pays for a whole weight matrix. At t=256, many rows reuse that matrix and intensity rises.

In separate 256-row matrix probes, several non-expert BF16 projections reached roughly 82 to 93 per cent of the measured 250 GB/s. The custom four-bit expert-layer probe reached about 64 per cent (2.366 ms for the MoE W4 shape at M=256). A 20-forward bus-utilisation probe of the custom Funnel decoder found GPU kernels active during 92.4 per cent of wall time.

That does not prove the whole DiffusionGemma forward is compute-bound. It shows a busy GPU running bandwidth-sensitive kernels with little empty time between them. That shaped every optimisation that followed.

The first constraint was capacity

Decoder experts occupy about 85.1 GiB in BF16. That is only the expert slice. The full system also needs attention weights, a shared MLP (multi-layer perceptron), embeddings, the language-model head, activations, KV state, packing buffers, and any fine-tuning adapter.

W4A16 means:

• W4: each expert weight stored as a 4-bit integer;

• A16: activations stay 16-bit.

Against the same expert tensors in BF16, packing records

The packer quantises in groups: gate/up projections use groups of 128; the down projection uses groups of 64 because the intermediate width 704 is not divisible by 128. For a group g,

with integer qᵢ, scale s₍g₎, and zero point z₍g₎. Metadata means storage is not exactly one quarter of BF16, but it is roughly a 4× cut in expert weight bytes fetched each pass.

W4 is not lossless. A dequantised reference checks that the custom kernel matches rebuilt floats within a tolerance. End-to-end benchmarks check task accuracy separately. Neither check makes W4 identical to BF16.

Quantisation alone leaves a sparse execution problem

An MoE layer has many alternate feed-forward networks. A router scores them per token. DiffusionGemma picks 8 of 128, runs the token through those eight, weights the outputs, and adds them.

The work is irregular: experts get different row counts. A direct loop gathers rows, launches tiny gate/up and down multiplies, applies the activation, then scatters weighted outputs back. With a 256-token canvas and top-8 routing, an expert sees about 16 rows on average. Many tiny launches waste the GPU.

The grouped W4 kernel sorts assignments by expert, pads into regular blocks, and runs one grouped matrix across populated experts. Sparse routing stays; the work becomes more regular.

There was also a correctness trap after experts finished. Several outputs belong to one token and must be summed. An early path used index_add_ with GPU atomics. Atomic writes are individually safe, but their order is not fixed. Floating-point addition is order-sensitive. In iterative diffusion, a tiny confidence change can flip which token commits and send the rest of the canvas down another trajectory. The grouped kernels therefore never use index_add_ for the final combine. They reduce each token’s eight expert outputs in a fixed order.

The final best W4 MoE kernel

The production expert path is fused_moe_w4_v2.py (env MOE_V2=1): deeper fusion, activation folded into the first matrix epilogue, 32-row blocks, one fixed-order combine. Against the first grouped kernel on the same artificial MoE layer:

Fusion here means folding the nonlinear activation into the matrix kernel that already produced its inputs. Smaller 32-row blocks also waste less padding when an expert gets about 16 rows. Kernel comments from that sweep record about 209 GB/s on gate-up-plus-activation and about 158 GB/s on the down projection, under the measured 250 GB/s DRAM ceiling.

This is a synthetic one-layer microbenchmark with 256 artificial rows, random top-8 routing, and the same packed weights. It compares two custom grouped kernels. It is not Hugging Face model.generate(), and it is not a full D2F or Funnel decode.

Reworked output differed from the first implementation by 5.43×10⁻⁴ mean relative error and 1.05×10⁻³ at the largest relative deviation in the test. Running the reworked kernel twice on the same inputs produced zero bit mismatches. The 1.28× applies to one MoE layer only. Thirty layers, attention, the vocabulary head, cache handling, and the decoding algorithm sit outside it.

The small operations were not small in aggregate

Once experts improved, the profile showed work spread across hundreds of modules and a 262K vocabulary. The next changes cut launches and intermediate memory without changing the learned function.

RMSNorm

RMSNorm rescales each activation by its root-mean-square magnitude. The stock Hugging Face module spreads that across several PyTorch ops. Across the loaded encoder and decoder there are 665 RMSNorm modules. Our Triton replacement does each with one kernel. Baseline: stock Hugging Face/PyTorch RMSNorm in the same Transformers model. We do not quote a standalone end-to-end speedup for this swap. Numerical checks use a relative-error tolerance, not bitwise equality.

QKV projection

Attention builds query, key, and value for each token, usually with three linear projections. LoRA adds small trainable matrices beside frozen base weights. Leaving the adapter attached at inference costs extra matmuls. Merging folds those updates into the base weights once.

After merge, the three matrices can be concatenated and evaluated as one larger projection: three launches become one. Baseline: three separate projections in the same Hugging Face model object.

The speed figure below measures LoRA merging only, not QKV fusion. On the custom D2F pipeline at confidence 0.7 over 50 GSM8K questions, with W4 experts but without the later fused RMSNorm/QKV Funnel stack, leaving the adapter attached ran at 139.9 tok/s and merging it ran at 160.2 tok/s (+14.5%). That is not a Hugging Face model.generate() comparison. The arms needed separate full model loads, and GB10 speed drifts with temperature, so treat it as supporting evidence.

Softmax and argmax

The model emits a logit for every vocabulary item. Softmax turns those into probabilities; argmax picks the winner. Diffusion decoding needs both the winning token and its confidence at every canvas position. The naive path is:

probs = torch.softmax(logits.float(), dim=-1)
confidence, token = probs.max(dim=-1)

That materialises work over a 262,144-token vocabulary. For a 256-position canvas, that is more than 67 million logits.

The Triton kernel computes row max, normalisation sum, confidence, and argmax without storing the full probability matrix. Baseline: the PyTorch expression above on the same artificial logits. Argmax indices match exactly in the harness; confidence is accepted within 10⁻³ relative error. Part 3 explains why confidence feeds the commit rule. Exact indices do not imply bit-exact probabilities.

A forward profile of the custom Funnel decoder

Microbenchmarks find possible wins. End-to-end profiles show whether they matter.

This profile uses the custom Funnel decoder (Part 3). The model loads through Hugging Face Transformers with the D2F LoRA merged. Runtime uses W4 experts through the final grouped kernel, fused RMSNorm, fused QKV with key/value capture, and fused confidence-plus-argmax. It allows up to two active canvases, early preparation of the next canvas, and a preallocated KV arena. It does not use model.generate().

A three-question GSM8K smoke test covered 52 decoder forwards at 113.59 ms average wall time per forward:

Experts are the largest named decoder slice and nearly tie encoder/commit work, but neither owns the wall clock. Attention, the vocabulary head, cache mechanics, masks, and orchestration leave a large residual. That is why a 1.28× MoE-layer win does not multiply into a 1.28× generation win.

Amdahl’s law: if a component is fraction f of runtime and becomes s times faster,

With f≈ 29.79/113.59≈0.26, even making experts free would cap this profile near 1/(1-0.26)≈1.35. The next bottleneck would simply show up.

Three optimisations that did not survive the full system

Negative results matter because they map the real constraints.

CUDA graphs

A CUDA graph records a fixed GPU sequence and replays it with less CPU launch overhead. In an isolated full decoder forward through our patched Transformers model, eager took 143.13 ms and graph replay took 131.97 ms (1.08×). Inside the Funnel decoder in the same session, eager averaged 123.05 ms per forward and graph replay 125.79 ms (0.98×). The complete path did not improve. Neither arm is Hugging Face model.generate(); both are our custom forward.

Two-bit experts

W2 stores each expert weight in 2 bits. Inside the same Funnel family used for the dense A/B below, expert-only W4 scored 44/50 at 235.7 tok/s on 50 GSM8K questions. W2 scored 34/50 at 185.2 tok/s. That is W4 versus W2 inside Funnel, not versus native Hugging Face decoding. W2 lost accuracy and speed.

Dense W4 everywhere

Attention, the shared MLP, and the language-model head are dense: every active token uses them. In isolated 256-row tests, selected custom W4A16 dense kernels ran about 1.20× as fast as cuBLAS BF16 for QKV and 1.13× for a language-model-head tile. That is a matrix shape test, not decoding.

End to end under Funnel, expert-only W4 with BF16 dense weights scored 44/50 at 235.7 tok/s. Quantising attention, shared MLP, and the language-model head as well scored 46/50 at 204.1 tok/s. Different numerics mean different denoising trajectories, so this does not pin a precise slowdown. It also does not establish an end-to-end speed gain for dense W4. Expert W4 was essential; a good isolated dense kernel still has to survive conversion, scheduling, and cache behaviour.

What these measurements support

Claims with named baselines:

1. Official-sampler A/B: simple two-benchmark mean 81.1 → 126.9 tok/s (1.56×); combined score 96/100 → 94/100.

2. Hardware probes: about 97.1 TFLOP/s large BF16 matmul; about 250 GB/s streaming reads.

3. Expert storage: 85.1 GiB BF16 → about 21.3 GiB W4A16.

4. Final best W4 MoE kernel: 2.910 ms versus 3.739 ms for the first grouped kernel (1.28×) on the same artificial MoE layer.

5. Fused RMSNorm, QKV, and confidence-plus-argmax reduce launches or intermediate traffic; this article does not assign them unsupported end-to-end multipliers.

6. Custom D2F comparison: 160.2 tok/s with LoRA merged versus 139.9 tok/s attached, with the separate-load caveat.

7. Funnel profile: experts are the largest named decoder component; encoder work, attention, and glue still limit any one kernel.

8. No established full-system wins for CUDA graphs, W2 experts, or quantising every dense matrix.

One comparison is now justified: replacing only BF16 routed experts with Triton W4A16 experts makes official Hugging Face generation 1.56× faster by the simple mean of these two benchmark throughputs, at 94/100 versus 96/100. That claim does not hand the gain to RMSNorm, QKV, D2F, or Funnel.

That boundary is Part 3. Once one forward is cheaper, the larger opportunity is fewer forwards and smaller later forwards. D2F, argmax renoising, draft spawn, and Funnel change the generation algorithm. The Funnel GSM8K band near 250 tok/s belongs there: algorithm plus this kernel stack, not a W4-only Hugging Face model.generate() claim.

Kernels lower the denominator of

The generation algorithm sets the numerator, and it decides how wide the denominator must stay.

References

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

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

3. Williams, S., Waterman, A., and Patterson, D. (2009). Roofline: An Insightful Visual Performance Model for Multicore Architectures. Communications of the ACM.

4. NVIDIA. DGX Spark specifications.