[diffusion] Flatten Wan VAE RMSNorm row addressing

[BBuf]

Summary

• address dense channels-last-3d Wan VAE activations as flat [pixel, channel] rows
• remove the per-program b/t/h/w div/mod chain and general stride reconstruction
• keep the normalization, dtype boundaries, SiLU path, launch geometry, and output layout unchanged
• add the dense-row guard, unit coverage, a registered benchmark, and kernel documentation

Why

Nsight Compute on the representative FastWan2.1 shape ([1, 96, 4, 480, 832], BF16 input / FP32 affine) shows 80.58% SM throughput but only 9.14% DRAM throughput and 93.93% achieved occupancy. Source counters attribute 120,771 / 196,793 samples to integer/index instructions. The hottest PCs are the reciprocal-based lowering of the row-to-b/t/h/w div/mod chain (I2F.RP, IABS, and MUFU.RCP).

Dense channels-last-3d already stores every pixel as a contiguous channel row, so those coordinates are unnecessary. stride(C) == 1 is checked explicitly because size-one channel tensors can satisfy the memory-format predicate with channel-first strides.

Performance

Single NVIDIA H200, cold-L2 registered benchmark:

Shape	main	PR	Speedup
FastWan2.1 BF16 [1,96,4,480,832]	1550.848 us	972.576 us	37.29%
FastWan2.2 FP32 [1,256,4,384,576]	921.184 us	553.760 us	39.89%

FastWan2.1 end-to-end ABBA, one H200, eager, 832x480x61, 3 steps, no compile, resident DiT:

Mode	main mean denoise / e2e	PR mean denoise / e2e	e2e change
lossless	0.8554 / 2.2129 s	0.8553 / 2.2151 s	-0.10% (control; kernel is not mounted)
quality=high	0.8513 / 1.9575 s	0.8533 / 1.8823 s	+3.84%

The quality-high full-stage trace reduces all 415 Wan RMSNorm+SiLU launches from 215.234 ms to 127.128 ms (40.94%). The largest C96 shape falls from 171.438 ms to 101.697 ms over 105 launches. Peak reserved memory stays at 31.402 GiB.

TurboWan2.1 T2V 1.3B provides a second model-level check at 832x480x81, 4 steps, quality=high, and no compile. Its main trace spends 286.083 ms across 545 Wan RMSNorm+SiLU launches (11.84% of GPU time). A clean main/PR/PR/main ABBA measured:

Model	main mean denoise / e2e	PR mean denoise / e2e	e2e change
TurboWan2.1 T2V 1.3B	1.0106 / 2.4235 s	1.0163 / 2.3325 s	+3.76%

Denoise is 0.57% slower, confirming that the request gain comes from the VAE path rather than timing drift in the DiT. Peak reserved memory remains 30.729 GiB.

The 14B TurboWan checkpoints provide resolution scaling checks with the same 545 launches:

Model	main mean denoise / e2e	PR mean denoise / e2e	kernel time	e2e change
14B 480P	6.6050 / 8.3559 s	6.6067 / 8.2559 s	289.647 -> 169.135 ms (-41.61%)	+1.20%
14B 720P	17.6763 / 21.2482 s	17.7044 / 21.0438 s	669.365 -> 389.957 ms (-41.74%)	+0.96%

Both ABBA brackets are byte-exact across main and PR. These longer DiT-heavy requests dilute the VAE gain, but independently confirm the kernel reduction at both production resolutions.

TurboWan2.2 I2V A14B adds a four-H200 CFG/Ulysses check at 1280x720x81. Main averaged 7.6857 / 12.1847 s and the PR averaged 7.6780 / 12.2241 s. Rank-0 recorded GPU time fell from 12,020.198 to 11,866.292 ms (-1.28%), but synchronized e2e regressed 0.32%, so this checkpoint is validation evidence rather than a claimed request speedup. All five outputs are byte-exact.

Output comparison

Prompt: A curious raccoon walks through a sunlit forest. Seed: 42.

main video · PR video · side-by-side video

All four FastWan quality-high outputs have the same SHA256: fcd22372b15c84dca9b7e18848a4642103bbe6f52dc7fbb3314e9d3e81bceb8a. All four TurboWan ABBA outputs have the same SHA256: 4eaef46a23e5e8c1e5a7a16941ba4c8b9080e2620ec2551848a11e0060270e46.

Validation

• pytest -q test/registered/kernels/ops/diffusion/test_norm.py -k wan_rmsnorm_silu: 4 passed
• pytest -q test/registered/kernels/ops/diffusion/test_model_fast_paths.py -k wan_vae_gate: 1 passed
• pre-commit on all changed files: passed after rebase onto current main
• all FastWan and TurboWan ABBA GPU-process boundaries were empty
• FastWan isolated cache cleanup: 29,228,743,029 bytes / 7 weight files -> 0
• TurboWan isolated cache cleanup: 16,777,680,195 bytes / 5 weight files -> 0
• TurboWan 14B 480P cache cleanup: baseline 42,521,216,678 bytes and ABBA 42,521,168,645 bytes / 7 weight files each -> 0
• TurboWan 14B 720P cache cleanup: baseline 42,519,295,635 bytes and ABBA 42,519,361,308 bytes / 7 weight files each -> 0
• TurboWan2.2 I2V A14B cache cleanup: baseline 71,103,355,450 bytes and ABBA 71,103,387,813 bytes / 10 weight files each -> 0

---

CI States

Latest PR Test (Base): ✅ Run #32644106273
Latest PR Test (Extra): ✅ Run #32661657125
Latest PR Test (AMD ROCm 7.2): ⏳ Run #32644106161

[BBuf]

/tag-and-rerun-ci extra

[BBuf]

/rerun-failed-ci

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/32644106273/job/97205704102?pr=35981