From 5d3e6dc1b4c172202b080a8659f522deae76852b Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:16:45 +0800 Subject: [PATCH] [Bugfix][MiniMax-H3] Use Hopper-safe modulation precision (#7693) Signed-off-by: david6666666 <530634352@qq.com> --- .../common/ci_source_file_dependencies.yml | 1 + .../minimax_h3/test_minimax_h3_modulation.py | 53 +++++++++++++++++++ .../diffusion/layers/indexed_modulation.py | 49 +++++++++++++---- 3 files changed, 94 insertions(+), 9 deletions(-) diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index c101db48d..067af9fa6 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -100,6 +100,7 @@ x-hunyuan-image3: &hunyuan_image3 x-minimax-h3: &minimax_h3 - vllm_omni/diffusion/models/minimax_h3/ + - vllm_omni/diffusion/layers/indexed_modulation.py x-hunyuan-video: &hunyuan_video - vllm_omni/diffusion/models/hunyuan_video/ 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 7fde01354..fcca6f7ad 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/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) diff --git a/vllm_omni/diffusion/layers/indexed_modulation.py b/vllm_omni/diffusion/layers/indexed_modulation.py index 385ce1a1b..0f4f2e59e 100644 --- a/vllm_omni/diffusion/layers/indexed_modulation.py +++ b/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, )