[Bugfix][MiniMax-H3] Fix Hopper keyframe and modulation precision

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -3,7 +3,7 @@
 
 import json
 import sys
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
 from typing import Any, TypeVar
@@ -2837,41 +2837,63 @@ def test_dit_encoder_selection_keeps_vae_resident():
     component.offload_to_cpu.assert_not_called()
 
 
-def test_keyframe_encode_pins_and_restores_cudnn_settings():
+@pytest.mark.parametrize(
+    ("capability", "allow_tf32"),
+    [((8, 6), True), ((9, 0), False), ((10, 0), True), ((10, 3), True), (None, True)],
+)
+@pytest.mark.parametrize("initial_tf32", [False, True])
+@pytest.mark.parametrize("fail_encode", [False, True])
+def test_keyframe_encode_pins_and_restores_cudnn_settings(
+    monkeypatch, capability, allow_tf32, initial_tf32, fail_encode
+):
+    from vllm.platforms.interface import DeviceCapability
+
     from vllm_omni.diffusion.models.minimax_h3.vae import (
         _minimax_h3_keyframe_encode_context,
     )
+    from vllm_omni.platforms import current_omni_platform
+
+    def get_capability(device_id):
+        assert device_id == 1
+        return DeviceCapability(*capability) if capability is not None else None
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(current_omni_platform, "get_device_capability", get_capability)
 
     cudnn = torch.backends.cudnn
-    original = (
-        cudnn.enabled,
-        cudnn.benchmark,
-        cudnn.deterministic,
-        cudnn.allow_tf32,
-    )
-    try:
-        cudnn.enabled = False
-        cudnn.benchmark = True
-        cudnn.deterministic = False
-        cudnn.allow_tf32 = False
-
-        with _minimax_h3_keyframe_encode_context(torch.device("cuda")):
+    with cudnn.flags(enabled=False, benchmark=True, deterministic=False, allow_tf32=initial_tf32):
+        error_context = pytest.raises(RuntimeError, match="encode failed") if fail_encode else nullcontext()
+        with error_context, _minimax_h3_keyframe_encode_context(torch.device("cuda:1")):
             assert cudnn.enabled
             assert not cudnn.benchmark
             assert cudnn.deterministic
-            assert cudnn.allow_tf32
+            assert cudnn.allow_tf32 is allow_tf32
+            if fail_encode:
+                raise RuntimeError("encode failed")
 
         assert not cudnn.enabled
         assert cudnn.benchmark
         assert not cudnn.deterministic
-        assert not cudnn.allow_tf32
-    finally:
-        (
-            cudnn.enabled,
-            cudnn.benchmark,
-            cudnn.deterministic,
-            cudnn.allow_tf32,
-        ) = original
+        assert cudnn.allow_tf32 is initial_tf32
+
+
+def test_keyframe_encode_cpu_keeps_cudnn_settings(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.vae import (
+        _minimax_h3_keyframe_encode_context,
+    )
+    from vllm_omni.platforms import current_omni_platform
+
+    def unexpected_capability_query(device_id):
+        pytest.fail("CPU keyframe encode must not query CUDA capability")
+
+    monkeypatch.setattr(current_omni_platform, "get_device_capability", unexpected_capability_query)
+    cudnn = torch.backends.cudnn
+    with cudnn.flags(enabled=False, benchmark=True, deterministic=False, allow_tf32=False):
+        with _minimax_h3_keyframe_encode_context(torch.device("cpu")):
+            assert not cudnn.enabled
+            assert cudnn.benchmark
+            assert not cudnn.deterministic
+            assert not cudnn.allow_tf32
 
 
 @pytest.mark.parametrize("fail_encode", [False, True])

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -10,7 +10,6 @@
     _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_,
@@ -194,52 +193,28 @@ def test_fused_modulation_preserves_bf16_residual_boundary() -> None:
 
 
 @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 current_omni_platform.is_cuda(), reason="CUDA required")
 @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
+@pytest.mark.parametrize("fused_gate", [False, True])
+def test_fused_modulation_keeps_fp32_affine_intermediates(fused_gate: bool) -> None:
+    device = current_omni_platform.get_torch_device()
+    rows, hidden_size = 2, 5376
+    x = torch.ones(rows, hidden_size, device=device, dtype=torch.bfloat16)
+    weight = torch.full((hidden_size,), 1.125, device=device, dtype=x.dtype)
+    shift = torch.zeros(1, hidden_size, device=device, dtype=x.dtype)
+    scale = torch.full_like(shift, 0.09375)
+    indices = torch.zeros(rows, device=device, dtype=torch.int64)
     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)
+    if fused_gate:
+        residual, actual = indexed_gate_rms_norm_scale_shift(
+            x, torch.zeros_like(shift), x, weight, shift, scale, indices, eps
+        )
+        assert torch.equal(residual, x)
+    else:
+        actual = rms_norm_indexed_scale_shift(x, weight, shift, scale, indices, eps)
+
+    # Without the RMS epsilon, 1.125 * 1.09375 = 1.23046875 is exactly
+    # halfway between two BF16 values. FP32 normalization keeps the epsilon
+    # and rounds down to 1.2265625. Premature BF16 rounding drops the epsilon
+    # and instead rounds the tie up to 1.234375, changing every output row.
+    assert torch.equal(actual, torch.full_like(x, 1.2265625))

vllm_omni/diffusion/layers/indexed_modulation.py

@@ -24,18 +24,6 @@ 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,
@@ -114,7 +102,6 @@ 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
@@ -125,17 +112,13 @@ 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
-    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
+    # Keep RMSNorm and AdaLN intermediates in FP32 on every architecture.
+    # Rounding the normalized value to BF16 before the affine transform
+    # changes the fused computation, including values at rounding boundaries.
+    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,
         output,
@@ -163,7 +146,6 @@ 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
@@ -180,17 +162,10 @@ 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
-    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
+    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,
         output,
@@ -300,7 +275,6 @@ 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,
         )
@@ -354,7 +328,6 @@ 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,
         )

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -31,6 +31,7 @@
     BoundedAllocatorCache,
     PinnedModuleStager,
 )
+from vllm_omni.platforms import current_omni_platform
 
 from .chunked_decode import decode_h3_chunks
 from .ops import install_h3_vae_optimizations
@@ -56,15 +57,19 @@ def _minimax_h3_keyframe_encode_context(
         yield
         return
 
-    # The official keyframe latent uses cuDNN's TF32 convolution path. The
-    # default non-deterministic algorithm can select numerically different
-    # reductions on H100s, and the difference is amplified by the denoiser.
-    # Pin both the algorithm and math mode for this sensitive encode only.
+    # Deterministic cuDNN algorithms can still use reduced-precision TF32.
+    # On Hopper, TF32 can shift the keyframe latent enough to push I2VA below
+    # the reference accuracy thresholds. Use FP32 convolutions for this
+    # sensitive encode only; preserve the existing math mode elsewhere.
+    allow_tf32 = True
+    if current_omni_platform.is_cuda():
+        capability = current_omni_platform.get_device_capability(device.index or 0)
+        allow_tf32 = capability is None or capability.major != 9
     with torch.backends.cudnn.flags(
         enabled=True,
         benchmark=False,
         deterministic=True,
-        allow_tf32=True,
+        allow_tf32=allow_tf32,
     ):
         yield