[Bugfix][MiniMax-H3] Use Hopper-safe modulation precision

Signed-off-by: david6666666 <530634352@qq.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -10,6 +10,7 @@ from vllm_omni.diffusion.layers.indexed_modulation import (
     _MAX_1D_GRID_SIZE,
     _iter_row_chunks,
     _launch_row_chunks,
+    _use_hopper_bf16_affine_semantics,
     indexed_gate,
     indexed_gate_rms_norm_scale_shift,
     indexed_scale_shift_,
@@ -190,3 +191,55 @@ 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(
+    torch.cuda.is_available() and torch.cuda.get_device_capability()[0] != 9,
+    reason="SM90-specific precision path",
+)
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_hopper_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
+
+    reference_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16)
+    normalized = reference_residual.float()
+    variance = normalized.pow(2).mean(-1, keepdim=True)
+    normalized = (weight.float() * (normalized * torch.rsqrt(variance + eps))).to(torch.bfloat16)
+    reference_modulated = (
+        normalized * (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, reference_residual, atol=5e-2, rtol=5e-2)
+    torch.testing.assert_close(actual_modulated, reference_modulated, atol=5e-2, rtol=5e-2)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_hopper_precision_dispatch_is_architecture_scoped(monkeypatch) -> None:
+    tensor = torch.empty(0, device="cuda")
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0))
+    assert _use_hopper_bf16_affine_semantics(tensor)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
+    assert not _use_hopper_bf16_affine_semantics(tensor)

vllm_omni/diffusion/layers/indexed_modulation.py

@@ -24,6 +24,18 @@ def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) ->
         kernel[(chunk_rows,)](*args, row_offset, **kwargs)
 
 
+def _use_hopper_bf16_affine_semantics(x: torch.Tensor) -> bool:
+    """Use the unfused AdaLN dtype boundaries on Hopper.
+
+    The fused FP32 modulation path is stable on Blackwell, but the SM90
+    kernel can accumulate enough extra precision in the repeated RMSNorm and
+    AdaLN operations to move the 50-step I2VA result below its CI threshold.
+    Keep the existing path on other architectures until their numerical
+    baselines are independently validated.
+    """
+    return x.device.type == "cuda" and torch.cuda.get_device_capability(x.device)[0] == 9
+
+
 @triton.jit
 def _indexed_scale_shift_kernel(
     output_ptr,
@@ -102,6 +114,7 @@ def _rms_norm_indexed_scale_shift_kernel(
     stride_scale_row,
     stride_indices,
     row_offset,
+    use_bf16_affine: tl.constexpr,
     block_n: tl.constexpr,
 ):
     row = tl.program_id(0) + row_offset
@@ -112,13 +125,20 @@ 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
-
-    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)
+    if use_bf16_affine:
+        # Match the unfused RMSNorm output consumed by AdaLN on Hopper.
+        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)
+        scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
+        output = normalized * (scale + 1.0) + shift
+    else:
+        normalized = x * 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)
+        output = normalized * (1.0 + scale) + shift
     tl.store(
         output_ptr + row * hidden_size + columns,
-        normalized * (1.0 + scale) + shift,
+        output,
         mask=mask,
     )
 
@@ -143,6 +163,7 @@ def _indexed_gate_rms_norm_scale_shift_kernel(
     stride_scale_row,
     stride_indices,
     row_offset,
+    use_bf16_affine: tl.constexpr,
     block_n: tl.constexpr,
 ):
     row = tl.program_id(0) + row_offset
@@ -159,12 +180,20 @@ def _indexed_gate_rms_norm_scale_shift_kernel(
 
     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)
+    if use_bf16_affine:
+        # Match the unfused RMSNorm output consumed by AdaLN on Hopper.
+        normalized = (updated * 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)
+        output = normalized * (scale + 1.0) + shift
+    else:
+        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)
+        output = normalized * (1.0 + scale) + shift
     tl.store(
         modulated_out_ptr + row * hidden_size + columns,
-        normalized * (1.0 + scale) + shift,
+        output,
         mask=mask,
     )
 
@@ -271,6 +300,7 @@ def rms_norm_indexed_scale_shift(
             shift.stride(0),
             scale.stride(0),
             indices.stride(0),
+            use_bf16_affine=_use_hopper_bf16_affine_semantics(x),
             block_n=triton.next_power_of_2(hidden_size),
             num_warps=8,
         )
@@ -324,6 +354,7 @@ def indexed_gate_rms_norm_scale_shift(
             shift.stride(0),
             scale.stride(0),
             indices.stride(0),
+            use_bf16_affine=_use_hopper_bf16_affine_semantics(residual),
             block_n=triton.next_power_of_2(hidden_size),
             num_warps=8,
         )