Fusing RMSNorm and FP8 Quantization Under torch.compile

FP8 serving with dynamic activation quantization has a small tax hidden in every decoder layer: the norm before each GEMM writes a bf16 activation to memory, and a second kernel immediately reads it back, reduces it, and quantizes it to FP8. Fusing residual-add, RMSNorm, and per-token FP8 quantization into one kernel that quantizes straight from fp32 registers eliminates that memory round trip and, more interestingly, one rounding step.

Shipping this on our production stack turned out to be an exercise in torch.compile archaeology: the natural design for the fusion cannot run at all under fullgraph piecewise CUDA graphs, for two independent reasons. This article covers the numerics (why removing a rounding step made benchmark accuracy go up), the two compile-time incompatibilities and the redesign that fixes them, how we probe for the silent failure mode fusions like this invite, and an A/B result that looks too good, with the decomposition showing why you shouldn't believe the headline number either.

TL;DR

• The unfused pipeline quantizes from bf16-rounded values; the fused kernel quantizes from fp32 registers. That deletes 5.6 KB/token of HBM traffic per site (−18%), one kernel launch, and one of two rounding steps.
• Under fullgraph torch.compile, the naive integration dies twice: JIT kernel objects that dynamo cannot trace (fixed with torch.library custom ops), and a producer→consumer hand-off via Python attributes on a tensor, a hard setattr graph break (fixed with an identity-matched single-slot hand-off chosen by empirically sweeping which mechanisms dynamo will trace).
• Every failure mode of the slot degrades to the unfused path: the win is lost, correctness is not. That property is what makes the design safe to ship.
• The fusion is numerics-changing (2–3% of FP8 codes differ), so byte-equality gating is the wrong tool. We used a three-stage silent-failure probe plus end-task accuracy: GSM8K improved in two independent runs (0.308 → 0.400 at 500 questions, ≈3σ), consistent with the removed rounding step.
• The A/B shows +36% throughput, but the decomposition attributes most of it to a speculative accept-length shift specific to the synthetic workload. The defensible production expectation is −5% mean latency plus the fusion savings.

Where norm+quant sits, and what fusing actually changes

Our production Gemma 4 26B-A4B deployment serves compressed-tensors FP8 (W8A8, dynamic): weights carry static per-channel scales, activations are quantized per token at every linear layer. In each of the 30 decoder layers, exactly two norm outputs feed FP8 GEMMs: the input layernorm into the QKV projection, and the pre-feedforward layernorm into the dense MLP. For an activation row $x\in {\mathbb{R}}^d$, e4m3 quantization with a per-token scale is

The subtlety is what gets quantized. The unfused pipeline rounds twice (norm output to bf16, then bf16 to e4m3), and both the mantissas and the absmax scale inherit the first rounding. The fused kernel holds the normed values in fp32 registers through both phases and rounds once:

On the memory side, the eliminated term is the bf16 re-read between the two kernels:

Why it couldn't boot: two ways to break fullgraph compile

Our stack runs the whole extend/verify and decode path dynamo-traced, inductor-compiled, and CUDA-graph-captured (piecewise, split only at attention). In that pipeline a dynamo graph break is not a slow path. It is fatal at boot: the warmup compile raises and the scheduler dies. The natural design for this fusion breaks fullgraph tracing twice.

Incompatibility 1: opaque JIT kernel objects. The fused kernels are JIT modules whose call objects dynamo cannot trace (Unsupported method call ... class 'Function'). The fix is standard but essential: register both kernel wrappers as torch.library.custom_op with register_fake meta functions, declaring the residual mutation. That makes each kernel an opaque single node that dynamo, inductor, and graph capture all handle correctly.

Incompatibility 2: the tensor-attribute hand-off. The harder problem is architectural. The producing norm must hand the quantized activation to the consuming linear layer through module boundaries the fusion doesn't control. The obvious mechanism is attaching attributes to the norm's bf16 output tensor, out._fp8_data = q, which under dynamo is setattr(Tensor, str, Tensor): a hard graph break, three to six per layer. In break-tolerant compile the feature limps along in eager sub-frames; under fullgraph it cannot boot.

Rather than guess at a replacement, we swept the candidate hand-off mechanisms under fullgraph compile on the real checkpoint shapes:

The chosen design is a single module-global slot with identity-matched pop semantics: the flagged norm publishes (anchor, q, s) where the anchor is its bf16 output tensor; the linear layer consumes the pair only if its input is that anchor (Python object identity), clearing the slot either way. Three properties make it correct. First, publish and consume are adjacent in trace order (norm then GEMM, nothing between), so dynamo resolves the dict read at trace time and bakes the consumption branch into the graph; nothing needs to survive graph-partition boundaries at runtime. Second, identity anchoring means any view, clone, or unrelated tensor simply misses and falls back to re-quantizing: every failure mode degrades to the unfused path, losing the win but never correctness. Third, only norms constructed with the fusion flag publish, which excludes our speculative-decoding draft model, which shares the decoder-layer class but runs unquantized bf16 and must not pay for quantization nobody consumes.

Gating a numerics-changing fusion

For our SWA radix islands, the correctness gate was byte-equality at temperature 0, because that change was semantically neutral. This one is not: fused and unfused FP8 codes differ on 2–3% of elements by design. Byte-equality is the wrong tool, and the dangerous failure mode changes accordingly. It is not a crash but a silent no-op: if the consumption branch didn't get baked into the graph, the linear layer quietly re-quantizes and the feature delivers nothing while reporting nothing.

The gate therefore builds a positive proof chain. Stage one is a poisoned-fallback trace proof: compile norm-plus-linear with the fallback quantization kernels replaced by a poison that returns zeros. If the compiled output equals the eager output and is nonzero, the fused hand-off is provably the path in the graph; a control norm without the flag must produce zeros, proving the poison is live. This stage runs in minutes without a server, and it caught both incompatibilities above on its first two runs. Stage two boots the verbatim production stack off and on and demands the on-boot be healthy under a burst, with outputs expected to differ (the fusion is not bitwise). Stage three checks the JIT build artifact exists after the on-boot and not after the off-boot. Alongside: kernel-level numerics bounds (exact e4m3 code agreement 96.7–97.8% against an independently rounded fp32 reference, per-token scale error under 0.3%) and end-task accuracy below.

Accuracy went up, and the sign is not a surprise

An accuracy improvement from a performance fusion deserves suspicion, but here the mechanism points the right way: the unfused path's double rounding is a small systematic degradation, and the 2–3% of borderline codes that change all move toward the fp32 ground truth. Two independent runs agreeing in direction, one at ≈3σ, is consistent with that story. We'd still describe the honest claim as "no regression, plausibly a small gain," not as a feature.

The A/B, and why we don't believe our own headline

Our first benchmark pair ran the two arms on different GPUs and showed +37%, which is not credible for a norm fusion. The corrected A/B runs both arms serially on one machine with an identically seeded workload, and still shows a large number:

Here is the fine print. On the GSM8K workload, accept length was flat (2.08 off vs 1.99 on). The +0.94 shift appears on random word-salad prompts, whose degenerate continuations draft easily; small numerics changes plausibly tip generation between more- and less-draftable patterns. The shift replicated across two independent runs, so it is stable on this workload, but we do not forecast +36% on production traffic. The defensible expectation is the −5% mean E2E plus the raw fusion savings, with the accept-length shift filed as workload-dependent upside. When a fusion of two small kernels appears to buy a third more throughput, the correct response is to decompose the gain, not to publish it.

What we'd tell anyone shipping this

• The two fixes are prerequisites for any torch-compiled stack: custom-op registration for the kernels, and a fullgraph-safe hand-off (our slot, or an equivalent). Without them the feature either kills boot or silently does nothing, depending on your compile mode.
• The slot is deliberately fragile in a checkable way. A future refactor that inserts a view or clone between a flagged norm and its GEMM silently reverts that site to unfused. The poisoned-fallback probe is the regression guard; it runs in minutes and belongs next to any edit of the layer's forward.
• A numerics-changing flag contaminates other gates. Any temperature-0 replay or byte-equality check elsewhere in the stack must now pin this flag on both sides, or it fails for reasons that have nothing to do with what it tests.
• Coverage is 2 of 3 quantized GEMM inputs per layer: the MoE expert path quantizes inside the fused-MoE kernel and needs its own analogous fusion; that's the natural follow-up.

FAQ

Doesn't quantizing "more accurately" just mean different, not better? The fp32-register values are the ground truth both paths approximate; the unfused path adds an intermediate rounding before the final one. Removing it strictly reduces approximation error per element. Whether that moves end-task metrics is empirical; here it did, in the direction the mechanism predicts.

Why not just use torch.compile to fuse the norm and quant automatically? Inductor can fuse elementwise chains, but this kernel's two-phase structure (a CTA reduction for rstd, then a second reduction for the absmax scale, with the vector held in registers and an in-place residual write) is a hand-scheduled pattern that pattern-matching does not currently produce, and the hand-off to the GEMM's quantized input is a dataflow rewiring, not a fusion.

Does the fused kernel change the residual stream? No: the residual update matched the fp32 reference exactly in unit tests (0.000 difference), and the bf16 normed output is written identically. Only the FP8 codes and scales the GEMM consumes differ.

What hardware does this need? e4m3 stores require SM89+ (Ada/Hopper). Our stack runs it on H100; it is irrelevant on pre-FP8 GPUs.

Why is a +0.9 accept-length shift worth +30% throughput? With speculative decoding, tokens per verify step scale with accept length (τ); at τ ≈ 8–9, one extra accepted draft per step is roughly a tenth more decode throughput, and the effect compounds with the shorter benchmark tail. Our benchmark article discusses why accept length is the most workload-sensitive metric in the stack.

Free tools from Netra

All of these run in your browser. No signup, and your data is not uploaded to us.

• LLM VRAM Calculator: size the GPU for any model, precision and context length.
• Fine-Tuning Memory Calculator: the training-time memory budget.
• Token Counter: measure prompts under the GPT, Claude and Llama tokenizers.
• Free OCR: turn scans and PDFs into LLM-ready text.
• Web to Markdown: strip a page to clean Markdown before it reaches a prompt.
• Limbus: local-first image segmentation.
• sam3.c: SAM3 inference in pure C.

References

• torch.library custom ops: the registration mechanism that makes opaque kernels compile-safe.
• Micikevicius et al., FP8 Formats for Deep Learning: the e4m3/e5m2 formats and their rounding behavior.
• Radix Islands: How Eager SWA Eviction Doubled Our Serving Concurrency: the companion optimization, whose byte-equality gate is the contrast case for the methodology here.
• Up to 4x faster than vLLM: the serving benchmark methodology.