[Bugfix][Diffusion] Align MiniMax-H3 Triton modulation operator precision with PyTorch reference (#7035)

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -190,3 +190,54 @@ def test_fused_modulation_preserves_bf16_residual_boundary() -> None:
     )
 
     assert torch.equal(modulated_out, expected)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_fused_modulation_matches_pytorch_reference() -> None:
+    torch.manual_seed(42)
+    rows, hidden_size, conditions = 128, 3072, 4
+    residual = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
+    gate = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    branch = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    indices = torch.arange(rows, device="cuda") % conditions
+    eps = 1e-6
+
+    # Native PyTorch reference
+    ref_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16)
+    ref_norm = ref_residual.float()
+    ref_var = ref_norm.pow(2).mean(-1, keepdim=True)
+    ref_norm = (weight.float() * (ref_norm * torch.rsqrt(ref_var + eps))).to(torch.bfloat16)
+    ref_modulated = (ref_norm * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(
+        torch.bfloat16
+    )
+
+    actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift(
+        residual, gate, branch, weight, shift, scale, indices, eps
+    )
+
+    torch.testing.assert_close(actual_residual, ref_residual, atol=5e-2, rtol=5e-2)
+    torch.testing.assert_close(actual_modulated, ref_modulated, atol=5e-2, rtol=5e-2)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_fused_modulation_exact_rounding_regression() -> None:
+    """Exact bit-level rounding regression: without enable_fp_fusion=False,
+    BF16 FMA skips intermediate rounding of gate * branch and differs from
+    PyTorch's unfused sequence."""
+    torch.manual_seed(42)
+    rows, hidden_size, conditions = 32, 1024, 2
+    residual = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
+    gate = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    branch = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
+    indices = torch.arange(rows, device="cuda") % conditions
+
+    expected_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16)
+    actual_residual = indexed_gate(residual, gate, branch, indices)
+    assert torch.equal(actual_residual, expected_residual)

vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py

@@ -18,6 +18,7 @@ def _iter_row_chunks(rows: int):
 
 def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) -> None:
     if device_type != "npu":
+        kwargs.setdefault("enable_fp_fusion", False)
         kernel[(rows,)](*args, 0, **kwargs)
         return
     for row_offset, chunk_rows in _iter_row_chunks(rows):
@@ -112,13 +113,14 @@ def _rms_norm_indexed_scale_shift_kernel(
     x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
     weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
     variance = tl.sum(x * x, axis=0) / hidden_size
-    normalized = x * tl.rsqrt(variance + eps) * weight
+    # Match the BF16 normalized value produced by unfused RMSNorm before AdaLN affine.
+    normalized = (x * tl.rsqrt(variance + eps) * weight).to(tl.bfloat16)
 
-    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
-    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0)
+    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
     tl.store(
         output_ptr + row * hidden_size + columns,
-        normalized * (1.0 + scale) + shift,
+        normalized * (scale + 1.0) + shift,
         mask=mask,
     )
 
@@ -150,21 +152,23 @@ def _indexed_gate_rms_norm_scale_shift_kernel(
     mask = columns < hidden_size
     index = tl.load(indices_ptr + row * stride_indices)
 
-    residual = tl.load(residual_ptr + row * stride_residual_row + columns, mask=mask, other=0.0).to(tl.float32)
-    gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32)
-    branch = tl.load(branch_ptr + row * stride_branch_row + columns, mask=mask, other=0.0).to(tl.float32)
+    residual = tl.load(residual_ptr + row * stride_residual_row + columns, mask=mask, other=0.0)
+    gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0)
+    branch = tl.load(branch_ptr + row * stride_branch_row + columns, mask=mask, other=0.0)
     # Match the BF16 residual value consumed by the unfused RMSNorm path.
-    updated = (residual + gate * branch).to(tl.bfloat16).to(tl.float32)
+    updated = residual + (gate * branch)
     tl.store(residual_out_ptr + row * hidden_size + columns, updated, mask=mask)
 
     weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
-    variance = tl.sum(updated * updated, axis=0) / hidden_size
-    normalized = updated * tl.rsqrt(variance + eps) * weight
-    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
-    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+    up_f32 = updated.to(tl.float32)
+    variance = tl.sum(up_f32 * up_f32, axis=0) / hidden_size
+    # Match the BF16 normalized value produced by unfused RMSNorm before AdaLN affine.
+    normalized = (up_f32 * tl.rsqrt(variance + eps) * weight).to(tl.bfloat16)
+    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0)
+    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
     tl.store(
         modulated_out_ptr + row * hidden_size + columns,
-        normalized * (1.0 + scale) + shift,
+        normalized * (scale + 1.0) + shift,
         mask=mask,
     )