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

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

tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py

@@ -43,7 +43,7 @@ def _collective_worker(rank, rendezvous, backend):
                 assert value.item() == 2
 
             def encode_image(self, image):
-                assert rank == 0
+                self.collective()
                 assert image.size == (32, 32)
                 return torch.ones(1, 96)
 
@@ -115,23 +115,6 @@ def _collective_worker(rank, rendezvous, backend):
         torch.testing.assert_close(conditioning.audio_condition.cpu(), torch.full((160, 32), 3.0))
         assert conditioning.audio_condition_lengths == (80,)
 
-        image_media = (
-            MiniMaxH3EncoderMediaInput(
-                task="fl2va",
-                height=32,
-                width=32,
-                num_frames=96,
-                latent_t=29,
-                audio_t=160,
-                images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
-            )
-            if rank == 0
-            else None
-        )
-        image_conditioning = model._encode_local_media(image_media)
-        torch.testing.assert_close(image_conditioning.visual_condition.cpu(), torch.ones(1, 96))
-        assert image_conditioning.visual_condition_shapes == ((1, 2, 2),)
-
         model.video_vae.fail = True
         with pytest.raises(OmniClientError, match="codec failed on rank 1"):
             model._encode_local_media(media)

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)