[BeatSeat]

Purpose

Fixes #7035.

Resolves the deterministic CI regression where the MiniMax-H3 I2VA accuracy test failed with PSNR 33.627597 dB (below the 34.000000 dB threshold on Builds 14708 and 15074).

Background & Root Cause

Before the modulation operator fusion in #6281, MiniMax-H3 evaluated the gated residual update, RMSNorm, and AdaLN affine modulation as separate operations:

# 1. Unfused RMSNorm materializes output as BF16
h = self.norm1(x)

# 2. AdaLN affine transformation evaluates in BF16
def _modulate_scale_shift(x, shift, scale, indices, *, dtype):
    return (x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(dtype)

PR #6878 restored the BF16 boundary on the gated residual before RMSNorm (updated = (residual + gate * branch).to(bfloat16)). However, inside the fused Triton kernels in vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py:

1. normalized was left as an unrounded FP32 intermediate value after RMSNorm.
2. The subsequent AdaLN affine modulation normalized * (1.0 + scale) + shift was computed in FP32 instead of native BF16 arithmetic.

Even the CPU fallback in #6281 explicitly preserved these BF16 casts:

normalized = (weight.float() * normalized).to(input_dtype)
return (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(input_dtype)

MiniMax-H3 contains 50 DiT layers. Over a 50-step flow-matching diffusion trajectory, the modulation kernel is invoked 2,500 times. The unrounded FP32 intermediate and affine operations accumulate across the non-linear attention and MLP blocks, causing the deterministic ~0.39 dB degradation down to 33.627 dB.

Key Changes

• In _rms_norm_indexed_scale_shift_kernel: truncate normalized to BF16, and perform the AdaLN scale/shift affine arithmetic in native BF16 (normalized * (scale + 1.0) + shift).
• In _indexed_gate_rms_norm_scale_shift_kernel: ensure both residual updated and second normalized truncate to BF16, evaluating affine modulation in native BF16.
• In tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py: add test_fused_modulation_matches_pytorch_reference to test numerical parity against the PyTorch reference implementation.

Test Plan

• Unit tests:

  pytest -v -o addopts="" tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

• 50-step flow-matching diffusion trajectory A/B benchmark on A800 GPU.

vLLM Version: 0.7.x
vLLM-Omni Commit: 548f286

Test Result

1. Unit Tests

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py::test_row_chunks_respect_ascend_grid_limit[0-expected0] PASSED [ 11%]
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py::test_row_chunks_respect_ascend_grid_limit[65535-expected1] PASSED [ 22%]
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py::test_row_chunks_respect_ascend_grid_limit[65536-expected2] PASSED [ 33%]
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py::test_row_chunks_respect_ascend_grid_limit[69888-expected3] PASSED [ 44%]
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py::test_launch_row_chunks_is_npu_only[npu-expected_grids_and_offsets0] PASSED [ 55%]
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py::test_launch_row_chunks_is_npu_only[cuda-expected_grids_and_offsets1] PASSED [ 66%]
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py::test_fused_modulation_preserves_bf16_residual_boundary PASSED [ 88%]
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py::test_fused_modulation_matches_pytorch_reference PASSED [100%]
================== 8 passed, 1 skipped in 1.56s ===================

2. A/B Benchmark Results (A800 GPU)

Comparison between current main (Method B) and this PR's aligned Triton kernel (Method C) against the unfused PyTorch reference (Method A):

Single DiT Block

Metric	Method B (Current Main)	Method C (This PR)	Delta
PSNR vs Reference	50.6907 dB	52.0429 dB	+1.352 dB 🚀
Max Absolute Diff	0.035156	0.027344	-22.2%
Bit Mismatch Rate	20.93%	17.16%	-3.77%

50-Step Diffusion Trajectory (Flow Matching)

Denoising Step	Method B (Current Main)	Method C (This PR)	Improvement
Step 1	53.64 dB	54.85 dB	+1.21 dB
Step 10	49.30 dB	50.55 dB	+1.25 dB
Step 25	47.12 dB	48.38 dB	+1.26 dB
Step 40	44.88 dB	46.13 dB	+1.25 dB
Step 50 (Final Latent)	42.75 dB	43.96 dB	+1.21 dB 🚀
Final Bit Mismatch Rate	12.54%	10.69%	-1.85%

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/index.md.

Module owners: @david6666666 @Isotr0py @princepride

Routing: @david6666666 via module of the changed files, module named in the PR description, semantic router, model owner; @Isotr0py via module of the changed files, module named in the PR description, CODEOWNERS; @princepride via module of the changed files, module named in the PR description

@BeatSeat, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 52f375ec3923 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[lishunyang12]

Set enable_fp_fusion=False; BF16 FMA still skips PyTorch's intermediate rounding.

[lishunyang12]

Add an exact rounding regression that fails on the previous kernel.

[BeatSeat]

Thanks @lishunyang12 for the review!

We have pushed an update addressing both points:

1. Passed enable_fp_fusion=False for CUDA launches in _launch_row_chunks so Triton does not generate BF16 FMA that skips intermediate rounding on residual + (gate * branch).
2. Added test_fused_modulation_exact_rounding_regression to verify exact bit-level equality against the unfused PyTorch reference.

Regarding relation with other PRs:
We noticed @david6666666 has opened #7693 which incorporates this BF16 modulation precision fix under a Hopper-scoped guard (use_bf16_affine), alongside VAE QK launch precision adjustments. If the team prefers landing the modulation precision fix as part of #7693, we can close this PR or fold the enable_fp_fusion=False and regression tests into #7693; otherwise, this PR is ready as a standalone operator fix.