[diffusion] ERNIE-Image bit-exact fused RMSNorm+scale/shift (H200 1024^2 e2e 15.63 -> 15.00 s, denoise -3.3%)

[BBuf]

Motivation

PR #33734 shipped the bit-exact residual-gate half of the ERNIE-Image adaLN block and explicitly dropped the other two fusions from closed PR #30170 — the CuTe-DSL fused_norm_scale_shift / fused_scale_residual_norm_scale_shift kernels — because they are not bit-exact: they compute norm(x)*(1+scale)+shift in one fp32 pass, and re-validating them behind a quality=high gate perturbed the 50-step trajectory to same-seed full-image PSNR 18.83 dB / SSIM 0.797 vs lossless, far below the PSNR>25 bar.

This PR is not a re-enable of those kernels. It reimplements both fusions as Triton kernels that are numerically re-engineered to be rounding-faithful to the eager chain, step for step, so they are bit-exact (torch.equal, full-image md5-identical) and can be wired on the default lossless path with no quality gate — the same standard #33734 used.

Per 1024x1024 image the block executes norm+modulate 2 sites x 36 blocks x 2 CFG x 50 steps = 7,200 times; eager spends 4-5 kernels per site (sgl_kernel.rmsnorm, 1+scale, mul, add, plus residual_gate_add_cuda at the second site), all bandwidth-bound on (1, 4216, 4096) bf16 activations.

What "rounding-faithful" means here

The eager RMSNorm.forward_cuda dispatches to sgl_kernel.rmsnorm -> flashinfer's CuTe-DSL RMSNormKernel. The Triton kernel replicates that kernel's exact arithmetic DAG for the contiguous bf16 regime ERNIE uses (H == 64 * threads_per_row, cluster_n == 1):

• per-"thread" fragment layout (cols = 8*tpr*b + 8*tx + v, v fastest) and its ordered sequential fp32 fadd chain (each x*x separately rounded via an opaque mul.rn.f32, no FMA contraction — MLIR vector.reduction semantics);
• warp reduction as the adjacent-pairs fold tree a shfl.bfly butterfly with offsets 1,2,4,8,16 produces (expressed with order-exact tl.reshape/tl.split folds), then the two warp sums added;
• rstd = rsqrt.approx.f32(sum_sq/H + eps) (inline PTX, matching cute.math.rsqrt fastmath);
• one final rounding for (bf16)(x * rstd * w), then the aten modulate chain's per-op bf16 boundaries (round(1+scale), round(y*(1+scale)), round(+shift)) reproduced with a bitcast round-to-nearest-even helper (same trick as indexed_modulation.py);
• the residual variant reproduces the eager pair round(gate*update), round(residual+...) — identical to the residual_gate_add_cuda fast path it subsumes at that site.

Because bit-exactness is a property of which rmsnorm implementation the platform dispatches to, the model-side wrappers verify torch.equal against the eager chain once at runtime on the first real input and permanently fall back to eager on any mismatch (e.g. FLASHINFER_USE_CUDA_NORM=1, or a future flashinfer numerics change). The fast path can never silently change outputs.

Modifications

• python/sglang/kernels/ops/diffusion/triton/rmsnorm_scale_shift_bitexact.py (new): the two Triton kernels + guards, registered as custom ops (torch.compile-safe fake impls).
• ernie_image.py: _ernie_norm_scale_shift / _ernie_gated_norm_scale_shift wrappers (first-call self-verification, one-time disable on exception/mismatch, same style as _ernie_residual_gate_add from [diffusion] ERNIE-Image bit-exact residual-gate fast path (H200 1024^2 e2e 16.17 -> 15.75 s) #33734) wired at both adaLN sites. The second site subsumes residual_gate_add_cuda + rmsnorm + modulate into one kernel and still returns the bit-identical residual stream.
• test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py (new): bit-exactness at the real ERNIE shape, a CFG-batched shape and the tpr=32 regime, plus an assertion that the fast paths were actually taken (not silently disabled).

Speed Tests (H200)

baidu/ERNIE-Image-Turbo, 1024x1024, 50 steps, bf16, eager, fixed prompt, seed=42, use_pe=False. Single process, 10 consecutive generations, first dropped, median of 9. Both arms on the same GPU, sequential.

Implementation	E2E wall (ms)	DenoisingStage (ms)
main run 1	15566.8	14825.3
main run 2	15695.1	14867.7
this PR	14999.6	14353.1
delta vs pooled main	-4.0%	-3.3%

Kernel microbenchmark at the real shape (1, 4216, 4096) bf16, 200 iters:

Site	eager	fused	speedup
norm*(1+scale)+shift (rmsnorm + 3 aten)	115.8 us	33.3 us	3.47x
res=residual+gate*upd; norm(res)*(1+scale)+shift (residual_gate_add_cuda + rmsnorm + 3 aten)	143.8 us	64.5 us	2.23x

Attribution: 36 blocks x (82.5 + 79.2) us saved x 2 CFG x 50 steps ~ 582 ms/image projected vs 493 ms measured (~85% materialized).

Accuracy Tests

Check	Result
Both kernels vs eager chain, torch.equal, (1,4216,4096)/(1,4096,4096)/(2,1140,4096)/(1,128,2048) bf16	True
Full 1024x1024 image, same prompt/seed, this PR vs main	md5 identical (324a06b56c4616aeaebd0d8d5017de9e on all four arms, incl. a second-GPU cross-check; also identical to the #33734-era main md5)
Determinism within each arm (9 images each)	single md5 per arm
Runtime self-verification	wired; disables fast path on any per-platform mismatch

Unit test (registered, base-b-kernel-unit):

3 passed

pre-commit run --files ...: all hooks pass.

Relation to #33734 / #30170

#30170 (closed) contained three fusions. #33734 shipped the residual-gate one (bit-exact as-is) and dropped the two norm fusions with the 18.83 dB evidence. This PR completes the set by re-deriving the two norm fusions with faithful numerics instead of re-gating the lossy CuTe-DSL implementations, which remain unused by ERNIE.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #31086934254
Latest PR Test (Extra): ❌ Run #31086936610

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/31086934254/job/92568691087?pr=33854