[diffusion] Wan VAE RMSNorm+SiLU fusion behind quality=high (H200 FastWan2.2 e2e 9.611 -> 9.125 s)

[BBuf]

Branch: pr-wan-vae-norm-silu-quality-high, single commit rebased on main (now that #33451 is merged). It reuses the DecodingStage gate wiring from #33451 (_sgl_vae_fast_path_gate, duck-typed) unchanged; the gate class itself is a deliberate inline copy, so this PR imports nothing from the FLUX.2 fast-path module. 4 files, +472/−4.

Motivation

This revives #30171 ("[diffusion] Fuse Wan VAE RMSNorm SiLU"), which fused the Wan VAE RMSNorm -> SiLU chains with a Triton kernel and was closed with the stated reason "consistency regression on B200": the kernel is numerically faithful (fp32 statistics, WanRMS_norm eager op boundaries preserved per step) but not bit-exact vs aten — a different reduction and SiLU path — so it moved the CI golden outputs. That is a policy constraint, not a numerics bug, and the infrastructure that has landed since removes it:

• [diffusion] Restrict request-level quality to two validated tiers: lossless (default) and high #33453 restricted request-level quality to two validated tiers: "lossless" (default) and "high";
• [diffusion] FLUX.2 VAE decoder fast path behind quality=high (H200: 1024^2 97.6->29.2 ms, 2048^2 437.2->168.5 ms) #33451 added the per-VAE fast-path gate + DecodingStage wiring: quality == "high" enables installed fast paths for the duration of a decode, the "lossless" default runs the original module path bit-for-bit.

This PR puts the #30171 kernel behind that gate. The "lossless" default is byte-identical to main (verified below), so golden-output CI is unaffected; only requests that opt in with --quality high run the fused kernel.

The workload motivation is unchanged from #30171: FastWan2.2-TI2V-5B is a 3-step DMD video model, so VAE decode dominates the request (DecodingStage is 5.77 s of the 9.61 s request, 60%, at 720p x 81 frames), and the decoder's WanRMS_norm -> SiLU chains (norm1/norm2 in every residual block + the output head, 29 sites) each run as ~5 strided aten kernels over up-to-1.8 GB activations.

Modifications

1. Kernel — new python/sglang/kernels/ops/diffusion/triton/wan_rmsnorm_silu.py (+196; the #30171 kernel migrated to the kernels.ops namespace with three deliberate changes):

• One program reduces one (b, t, h, w) pixel's channel row (channels_last_3d: channel dim innermost, fully coalesced) and applies F.normalize(x, dim=1) * scale * gamma [+ bias] + SiLU: fp32 norm statistics, each step materialized at the same dtype boundary as WanRMS_norm.forward eager, SiLU computed in fp32.
• New vs [diffusion] Fuse Wan VAE RMSNorm SiLU #30171 (a): supports the autocast dtype split — half-precision activations with fp32 affine params, which is how the bf16-decode Wan2.1 family (vae_decode_precision="bf16", fp32 VAE weights) actually runs. aten promotes to fp32 at the * gamma boundary; the kernel reproduces that promotion and returns fp32. Without this the kernel would never engage on Wan2.1-style decodes (FastWan2.2 decodes in fp32, where the uniform-dtype path applies).
• New vs [diffusion] Fuse Wan VAE RMSNorm SiLU #30171 (b): returns None for unsupported inputs (non-CUDA, grad mode, unsupported dtype combination, non-channels_last_3d layout, C > 1024) instead of raising — callers must fall back (same contract as [diffusion] FLUX.2 VAE decoder fast path behind quality=high (H200: 1024^2 97.6->29.2 ms, 2048^2 437.2->168.5 ms) #33451's kernel).
• New vs [diffusion] Fuse Wan VAE RMSNorm SiLU #30171 (c): t/h/w sizes are runtime args instead of tl.constexpr, so the causal chunked Wan decode (t=1 head chunk + t=2/4 steady chunks per temporal stage) does not trigger per-chunk-shape recompiles.

2. Wan fast path + gate — new runtime/models/vaes/wan_vae_cuda_opt.py (+159), wired via the existing CudaPlatform.optimize_vae hook (+7/−4 in platforms/cuda.py):

• maybe_optimize_wan_vae wraps every decoder WanRMS_norm -> SiLU chain (residual-block norm1/norm2 + the decoder output head; 29 sites on the FastWan2.2 VAE) in a FusedWanRMSNormSiLU module dispatching on a per-VAE VaeFastPathGate published as _sgl_vae_fast_path_gate — a deliberate inline copy of the [diffusion] FLUX.2 VAE decoder fast path behind quality=high (H200: 1024^2 97.6->29.2 ms, 2048^2 437.2->168.5 ms) #33451 gate (DecodingStage duck-types the attribute), keeping the two VAE fast-path modules independent. Gate on runs the fused kernel (with the None fallback); gate off runs F.silu(self.norm(x)), bit-identical to the original norm + nn.SiLU pair. DecodingStage gate wiring is reused from [diffusion] FLUX.2 VAE decoder fast path behind quality=high (H200: 1024^2 97.6->29.2 ms, 2048^2 437.2->168.5 ms) #33451 unchanged (zero changes to decoding.py here).
• Install is all-or-nothing and fail-closed: without Triton, under spatial-parallel decode, or with any decoder residual block / output head not matching the expected WanRMS_norm/nn.SiLU structure, nothing is installed and every request runs the unmodified decoder. The encoder is untouched (the gate is decode-scoped).
• Eager decodes only: under --enable-torch-compile Inductor already fuses the norm+SiLU chain when it compiles the decode, and routing it through an opaque custom op instead measured slightly slower (compiled FastWan decode 4.90 -> 4.98 s). The wrapper therefore steps aside via torch.compiler.is_compiling() and compiled decodes keep the Inductor path for both quality tiers — verified at parity below.

3. Gate sharing & parameter FQNs — FusedWanRMSNormSiLU registers gamma/bias directly on the wrapper (parameter names stay ...norm1.gamma, so by-name weight transfer and strict state_dict loads are unaffected) and inlines the bit-identical off-path op chain, mirroring FLUX.2's FusedGroupNormSiLU; VaeFastPathGate/GATE_ATTR are imported from flux2_vae_cuda_opt instead of duplicated (plus docstring trims there, no behavior change).

4. Tests — new test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py (+70, 3 cases): kernel numerics vs the eager chain in both production precision regimes (FastWan2.2 fp32-uniform; Wan2.1 bf16 x + fp32 affine with the aten fp32 promotion), each with and without bias; wrapper gate dispatch (gate off torch.equal vs the original module chain, gate on routes to the fused kernel). 3 passed on H200.

Applicability

The e2e win is proportional to VAE decode's share of the request, so it is workload-dependent by construction:

• Few-step / DMD video models (the target): FastWan2.2-TI2V-5B spends ~60% of an eager request in decode -> the fusion is clearly visible end-to-end (table below).
• Many-step models: a 50-step request is denoise-dominated, so the same 0.66 s decode delta dilutes to ~1–2% e2e (50-step CFG denoise at the measured 0.315 s/step for the same 5B DiT at 720p is ~31.5 s vs the 0.66 s decode delta). A live 50-step Wan2.1-T2V-1.3B A/B could not be run: its multi-shard UMT5 text encoder currently fails to load on main (unrelated qkv_proj weight-mapping regression), which blocks the whole pipeline; the Wan2.1 decode-level win is still measured directly in the harness table below (-11.9%).
• Compiled decodes (--enable-torch-compile): intentionally unaffected in both tiers (see above).

Speedup

All numbers H200 (single GPU, GPU-idle box checks around each run), this branch, seed 42. Three levels:

Kernel microbenchmark (real decoder norm shapes captured from production decodes via forward hooks; us per call, median of 13 rounds x 20 calls; aten = the eager F.silu(F.normalize(x, 1) * scale * gamma) chain at production dtypes):

workload (captured shape)	x dtype	aten us	fused us	speedup
fastwan 720p (1, 1024, 2, 96, 144)	fp32	324.6	73.6	4.41x
fastwan 720p (1, 1024, 1, 48, 72)	fp32	47.6	34.1	1.40x
fastwan 720p (1, 512, 4, 192, 288)	fp32	1224.2	273.0	4.48x
fastwan 720p (1, 512, 4, 384, 576)	fp32	4838.5	1082.8	4.47x
fastwan 720p (1, 256, 4, 384, 576)	fp32	2456.1	905.1	2.71x
wan21 480p (1, 384, 2, 120, 208)	bf16 (+fp32 affine)	222.4	69.5	3.20x
wan21 480p (1, 192, 4, 240, 416)	bf16 (+fp32 affine)	863.1	433.0	1.99x
wan21 480p (1, 96, 4, 480, 832)	bf16 (+fp32 affine)	1862.8	1510.5	1.23x

Component harness (frozen decode-input latents — FastWan captured from a production 720p x 81f TI2V request after scale_and_shift + preprocess_decoding, Wan2.1 fixed-seed synthetic at the production 832x480 x 81f latent shape; production decode replicated: fp32 decode for FastWan, bf16-autocast for Wan2.1; steady-state mean of 5):

workload	main baseline	lossless (this branch)	quality=high (this branch)	decode delta
FastWan2.2-TI2V-5B 720p x 81f	5918 ms	5910 ms (byte-identical)	5458 ms	-7.6%
Wan2.1-T2V-1.3B 480p x 81f	1833 ms	1839 ms (byte-identical)	1620 ms	-11.9%

"byte-identical" = sha256 of the full fp32 [0,1] output tensor equals the main baseline decode of the same latents.

End-to-end latency (#33451 protocol: one warmed DiffGenerator process serving sequential identical requests, per-request wall clock, first request dropped, median of 5; FastWan preset prompt/image, 1280x720, 81 frames):

configuration	default (lossless)	--quality high	delta
FastWan2.2-TI2V-5B, eager (default serving)	9.611 s (9.594–9.659)	9.125 s (9.123–9.248)	-0.486 s (-5.1%)
FastWan2.2-TI2V-5B, --enable-torch-compile	8.969 s (8.952–19.5*)	8.948 s (8.919–15.8*)	parity (by design)
50-step Wan model (dilution case)	—	—	est. ~-1–2% (see Applicability)

* max includes one late-compile-tail request; decode-stage times are flat at 4.90 s in both compiled modes (and were 4.98–5.07 s before the is_compiling bypass was added, i.e. the bypass turns a small compiled regression into exact parity).

Stage cross-check for the eager rows: DecodingStage 5.77 -> 5.11 s (-11.5%) while DmdDenoisingStage is flat (0.930 vs 0.932 s) — the entire e2e delta is decode, as expected.

Accuracy

• Lossless default is bit-exact end-to-end: gate-off harness decodes on this branch are byte-identical (sha256) to the main baseline for both the FastWan (fp32) and Wan2.1 (bf16-autocast) workloads, and gate-off timing is unregressed (5910 vs 5918 ms / 1839 vs 1833 ms). Golden-output CI is unaffected.
• fp32 ground-truth check (Wan2.1 latents, full decoder, fp32 no-autocast decode as truth): fast-path bf16 error is statistically identical to the baseline bf16 error — mean abs 0.001003 vs 0.000959, p99 0.004017 vs 0.003849 (on the [0,1] pixel scale). For the fp32 FastWan decode the baseline is the fp32 reference: fast-path deviation is mean abs 0.000011, p99 0.000079, max 0.003.
• e2e quality (maintainer bar for quality=high: PSNR > 25 dB vs the lossless output, same seed and prompt; per-frame metrics over all 81 frames decoded from the saved videos):

prompt	PSNR mean (dB)	PSNR min (dB)	SSIM mean	SSIM min	bar	PASS
cat walking (preset)	39.73	37.71	0.9671	0.9612	PSNR > 25	yes
red panda in snow	39.08	37.81	0.9661	0.9593	PSNR > 25	yes
waves on rocky shore	37.67	36.06	0.9598	0.9561	PSNR > 25	yes

All frames sit 11+ dB above the bar, in line with #30171's own 37.6–43 dB frame range. 4 side-by-side lossless-vs-high sample frame pairs: /tmp/prW_imgs/.

• Unit tests: 3 cases (see Modifications), 3 passed on H200.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests.
• Provide accuracy and speed benchmark results.
• Follow the SGLang code style guidance.

Sample frames (left: lossless, right: --quality high)

frame	side-by-side
cat prompt, frame 0
cat prompt, frame 40
red panda prompt, frame 40
ocean waves prompt, frame 40

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #30988802268
Latest PR Test (Extra): ✅ Run #31007345865

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[BBuf]

Fixed the multimodal-gen-component-accuracy failures: the wrapper nested the whole WanRMS_norm as a submodule, renaming ...norm1.gamma → ...norm1.norm.gamma (29 params), which failed the by-name weight-transfer check (81/110 and 167/196 < 98%). The wrapper now registers gamma/bias directly (FQNs unchanged — verified named_parameters() identical before/after install plus a strict state_dict round-trip on a full WanDecoder3d) and inlines the bit-identical off-path op chain, same pattern as FLUX.2's FusedGroupNormSiLU. Also deduped VaeFastPathGate by importing it from flux2_vae_cuda_opt and trimmed the tests (110 → 70 lines).

🤖 Generated with Claude Code

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/30988802268/job/92250050802?pr=33546