[Bugfix] Preserve BF16 residual semantics in fused MiniMax-H3 modulation (#6878)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+from vllm.triton_utils import HAS_TRITON
+
+from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
+    indexed_gate_rms_norm_scale_shift,
+    rms_norm_indexed_scale_shift,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, 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_preserves_bf16_residual_boundary() -> None:
+    torch.manual_seed(17)
+    rows, hidden_size, conditions = 11, 3072, 4
+    tensors = [
+        torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16),
+        torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16),
+        torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16),
+        torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16),
+        torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16),
+        torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16),
+    ]
+    residual, gate, branch, weight, shift, scale = tensors
+    indices = torch.arange(rows, device="cuda") % conditions
+    eps = 1e-6
+
+    residual_out, modulated_out = indexed_gate_rms_norm_scale_shift(
+        residual,
+        gate,
+        branch,
+        weight,
+        shift,
+        scale,
+        indices,
+        eps,
+    )
+    expected = rms_norm_indexed_scale_shift(
+        residual_out,
+        weight,
+        shift,
+        scale,
+        indices,
+        eps,
+    )
+
+    assert torch.equal(modulated_out, expected)

vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py

@@ -133,7 +133,8 @@ def _indexed_gate_rms_norm_scale_shift_kernel(
     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)
-    updated = residual + gate * branch
+    # Match the BF16 residual value consumed by the unfused RMSNorm path.
+    updated = (residual + gate * branch).to(tl.bfloat16).to(tl.float32)
     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)