[Diffusion] Fuse the Wan VAE conv bias epilogue and tile the RMSNorm+SiLU kernel

[Dayuxiaoshui]

Motivation

Follow-up to #38020 and #38182. Two costs in the Wan 2.1 VAE decoder that kernel-level profiles hide, found by grouping a quality=high decode at the aten-op level (key_averages(group_by_input_shape=True)):

1. The conv bias is a separate, unvectorised kernel. PyTorch's cuDNN path runs cudnn_convolution and then output.add_(bias.view(1, C, 1, 1, 1)). On channels_last outputs aten does not vectorise that broadcast add: 1.5 TB/s on H200 for [1, 96, 4, 480, 832] bf16. It shows up only as aten::conv3d minus aten::cudnn_convolution: 797 launches, 278 ms of a 2847 ms lossless decode and 227 ms of a 1876 ms high decode on a shared GPU (~10% in both modes).
2. wan_rmsnorm_silu ran one pixel per program with 4 warps over 96-384 channels: 0.63 TB/s on the largest shape, 16% of the high-quality decode.

Modifications

New kernel, lossless (bit-exact, unconditional)

• kernels/ops/diffusion/modulate/conv_bias_epilogue_triton.py: x.dtype(float(x) + float(bias)), the exact arithmetic of aten's add_ (fp32 opmath, one rounding), optionally fused with the residual add as x.dtype(x.dtype(x + bias) + h), i.e. the two roundings of the eager conv(x) + h chain. 4.1 TB/s. bias may be fp32 on a half-precision x (autocast): it is cast to x.dtype first, which is what autocast does before the conv call. The output is empty_like(x), so it keeps the dense channels_last(_3d) strides like the in-place add_. Predicate requires C > 1 with canonical strides and never raises; the kernel validates and raises ValueError.
• WanCausalConv3d._conv_with_epilogue: runs the conv without bias and applies it through the kernel; falls back to the exact aten ops (out.add_(bias.to(out.dtype).view(...))).
• residual_block_forward: conv2's bias and the block's residual add become one pass (residual=h); conv1's bias is deferred into norm2 (skip_bias=True + WanRMS_norm.forward(x, conv_bias=...), where the eager norm adds it first with the same arithmetic).

Quality-gated

• wan_rmsnorm_silu tiles [ROWS, C] consecutive pixels per program (32 rows at C <= 128 down to 4 at C = 1024): 0.63 -> ~3 TB/s at C = 96. The 2D reduction tree moves ~2e-6 of the elements by one bf16 ulp relative to the previous kernel, while both differ from eager on ~20% of elements, so the gated contract is unchanged. Optional conv_bias= folds the deferred conv1 bias into the same pass, applied exactly like aten's add_ before the statistics.
• FusedWanRMSNormSiLU.forward(x, conv_bias=None): folded under the gate, otherwise applied by the bit-exact epilogue kernel before the eager chain.

Tests / docs

• New test_conv_bias_epilogue.py: aten oracle (x.clone().add_(bias), + h) for bf16/fp16/fp32 x fp32/half bias, 4D and 5D, stride identity; F.conv3d with bias vs conv-without-bias + epilogue; predicate rejections (NCDHW, wrong bias width/dtype, residual layout/dtype, empty, C == 1, requires_grad).
• test_norm.py: ragged multi-row tiles at the three tile configurations; conv_bias fold equals the kernel on the pre-biased tensor bit for bit.
• test_model_fast_paths.py: a three-chunk tiny WanDecoder3d decode through the feature cache, fused path vs pure aten (torch.equal); deferred bias through the gated wrapper with the gate off (exact) and on (bounded relative error).
• README selection matrix and existing-fast-paths.md updated.

Accuracy Tests

1x H200, Wan 2.1 VAE weights (Wan-AI/Wan2.1-I2V-14B-480P-Diffusers), 81 frames 480x832.

Lossless

• Standalone decode, bf16 weights: torch.equal vs main: True.
• Standalone decode, fp32 weights + bf16 autocast (the Wan T2V pipeline setting): torch.equal vs main: True.
• End-to-end sglang generate, TurboWan2.1-T2V-1.3B, 4 steps, seed 1, 6 prompts: all 6 mp4s md5-identical to main (and to the outputs of the two previous PRs).
• Peak decode memory unchanged.

Quality path

Standalone high vs main lossless on the same latent: PSNR 58.4 dB (bf16), 59.1 dB (fp32 + autocast), unchanged from before this PR; the multi-row norm kernel differs from the previous gated kernel on ~2e-6 of elements by one bf16 ulp.

Speed Tests and Profiling

Standalone VAE decode, torch.profiler GPU kernel self time, idle H200:

weights / autocast	path	before	this PR
bf16	lossless	1324 ms / 7850 kernels	1253 ms / 7599
bf16	high	1112 ms / 4006	935 ms / 3420
fp32 + bf16 autocast	lossless	1593 ms	1530 ms
fp32 + bf16 autocast	high	1215 ms	1056 ms

Per-kernel (bf16, high): _wan_rmsnorm_silu_kernel 177 -> 57 ms over 609 launches; the 797 aten bias add_ launches are replaced by 396 _conv_bias_epilogue_kernel launches totalling 22 ms (649 launches / 35 ms on the lossless path). Qwen-Image 1024x1024 high: 43.3 -> 39.4 ms (shared norm kernel), lossless unchanged and still bit-exact.

End-to-end, TurboWan2.1-T2V-1.3B, 4 steps, 81f 480x832, main from a git worktree of the base commit, 6 prompts, median of the 5 warm requests. The decode stage is ~60% of a warm request. This box is shared; runs whose DmdDenoisingStage (identical code on both sides) drifted were discarded, and the cleanest interleaved pair is reported:

stage	main	this PR
DecodingStage, high	932 ms	785 ms (-15.8%)
DecodingStage, lossless	1535 ms	1501 ms (-2.2%)
DmdDenoisingStage	~890 ms	~890 ms

Micro-benchmarks behind the design (idle H200, bf16 [1, 96, 4, 480, 832]): aten bias add_ 0.409 ms (1.5 TB/s), Triton bias add 0.150 ms (4.1 TB/s), aten add_ + residual add 0.621 ms vs fused 0.215 ms; norm 0.970 ms -> 0.201 ms.

Checklist

• Format your code according to the Format code with pre-commit. (ruff format / ruff check with the pre-commit pinned v0.15.1, isort, codespell: clean)
• Add unit tests according to the Run and add unit tests. (test_conv_bias_epilogue.py, test_norm.py, test_layout.py, test_model_fast_paths.py: 393 passed, 1 skipped)
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #34327993512
Latest PR Test (Extra): ❌ Run #34327993130
Latest PR Test (AMD ROCm 7.2): ❌ Run #34327993453