[Diffusion] MiniMax-H3 fuse the video VAE decoder's RMSNorm and QK RoPE

[niehen6174]

Motivation

Each block of the H3 video VAE decoder runs eager RMSNorm, a scaled residual add, and per-head Q/K RMSNorm plus partial NeoX RoPE. The decoder is a 36-block ViT and runs on every H3 request. Decode is still one tile per call.

Modifications

• h3_vae_rmsnorm, h3_vae_scale_add_rmsnorm, h3_vae_qk_rmsnorm_rope (kernels/ops/diffusion/norm/h3_vae_fused_norm.py): FP32 RMSNorm, the scaled residual add fused into the next RMSNorm, and affine-free Q/K RMSNorm plus partial NeoX RoPE (head 64 / rotary 48). A contract mismatch stays on the eager path. MINIMAX_H3_VAE_DECODER_FUSED_NORM=1 by default.
• The decoder broadcasts one rotary table across the tiles in a batch, so the Q/K kernel also runs for B > 1.

Accuracy

Fused matches eager to a max absolute error of about 5e-5 on a 28-tile output. test_h3_vae_fused_norm.py checks each kernel. test_fused_decoder_matches_unfused_forward checks the decoder and that the fused kernels actually ran.

Speed

One H200, fp16 autocast, released decoder shape, one tile per call. A 1344×768, 10 s, 24 fps clip is 14 temporal units, 28 tiles each. The unit row is measured. The 10 s row is that unit times 14: same decoder, no extra run.

	main	This PR	Change
One temporal unit	656.3 ms	589.5 ms	-66.8 ms (-10.2%)
10 s VAE decode, 1×H200	9.19 s	8.25 s	-0.94 s (-10.2%)

4×H200, FastH3 4-step, ulysses 4, VSA sparsity 0.9. One request is 1344×768, 10.125 s, 24 fps, 5 sampling steps, seed 1001. Each number is the mean of 8 requests after one warmup. The decoding stage includes the audio VAE.

	main	This PR	Change
End to end	11.00 s	10.75 s	-0.25 s (-2.3%)
VAE decode	2.50 s	2.27 s	-0.23 s (-9.4%)

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. test_h3_vae_fused_norm.py, test_fused_decoder_matches_unfused_forward, test_batched_neox_rope_matches_per_tile_rows.
• Update documentation according to Write documentations. Kernels are registered in kernels/ops/diffusion/__init__.py. The README selection matrix is not updated yet.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): 🚫 Run #37256954420
Latest PR Test (Extra): ❌ Run #37256954215
Latest PR Test (AMD ROCm 10): ❌ Run #37256954368

[mickqian]

/tag-and-rerun-ci

[niehen6174]

/rerun-failed-ci