From a413d0ceb214664d721ac4154c6158a976941656 Mon Sep 17 00:00:00 2001 From: david6666666 <530634352@qq.com> Date: Mon, 21 Sep 2026 01:33:36 +0000 Subject: [PATCH 1/2] [Bugfix][MiniMax-H3] Disable keyframe TF32 on Hopper Signed-off-by: david6666666 <530634352@qq.com> --- .../minimax_h3/test_minimax_h3_contract.py | 70 ++++++++++++------- vllm_omni/diffusion/models/minimax_h3/vae.py | 15 ++-- 2 files changed, 56 insertions(+), 29 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index aca8bdb83e5..8cc3f1a9ea8 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/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]) diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index ed28029bbab..c614ea0f293 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/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 From 66c500a71c2ecf1db53eec58f558353a18110fdc Mon Sep 17 00:00:00 2001 From: david6666666 <530634352@qq.com> Date: Mon, 21 Sep 2026 02:00:06 +0000 Subject: [PATCH 2/2] [Bugfix][MiniMax-H3] Keep fused modulation intermediates in FP32 The keyframe TF32-only change still fails H100 I2VA accuracy. Remove the additional Hopper-only BF16 norm/affine rounding introduced by #7693 while retaining the BF16 gated-residual boundary. This restores the existing fused FP32 contract on Hopper and leaves other architectures unchanged. Add exact tests for both kernels at a BF16 halfway value: the premature cast drops the RMS epsilon and changes 1.2265625 to 1.234375. The old BF16 branch fails both tests; the restored FP32 path passes. Signed-off-by: david6666666 <530634352@qq.com> --- .../minimax_h3/test_minimax_h3_modulation.py | 71 ++++++------------- .../diffusion/layers/indexed_modulation.py | 49 +++---------- 2 files changed, 34 insertions(+), 86 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py index fcca6f7adc3..df3b2268983 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/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)) diff --git a/vllm_omni/diffusion/layers/indexed_modulation.py b/vllm_omni/diffusion/layers/indexed_modulation.py index 0f4f2e59e9e..9fd402dec88 100644 --- a/vllm_omni/diffusion/layers/indexed_modulation.py +++ b/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, )