From 548f286dded14eefd42f5334c96cb8f30717c581 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Tue, 15 Sep 2026 14:29:59 +0800 Subject: [PATCH 1/2] fix(diffusion): align MiniMax-H3 Triton modulation operator precision with PyTorch reference - Truncate RMSNorm output to BF16 before AdaLN affine modulation in Triton kernels - Perform AdaLN scale/shift affine arithmetic in native BF16 to match unfused PyTorch implementation - Ensure indexed gate branch modulation preserves BF16 boundary - Resolves PSNR degradation across multi-step diffusion trajectories (improves final latent PSNR by +1.21 dB) Signed-off-by: BeatSeat --- .../minimax_h3/test_minimax_h3_modulation.py | 32 +++++++++++++++++++ .../attention/ops/minimax_h3_modulation.py | 29 +++++++++-------- 2 files changed, 48 insertions(+), 13 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 e10fbe56a9a..88cd1ad74a3 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -190,3 +190,35 @@ 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(not HAS_TRITON, reason="Triton required") +def test_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 + + # Native PyTorch reference + ref_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16) + ref_norm = ref_residual.float() + ref_var = ref_norm.pow(2).mean(-1, keepdim=True) + ref_norm = (weight.float() * (ref_norm * torch.rsqrt(ref_var + eps))).to(torch.bfloat16) + ref_modulated = (ref_norm * (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, ref_residual, atol=5e-2, rtol=5e-2) + torch.testing.assert_close(actual_modulated, ref_modulated, atol=5e-2, rtol=5e-2) diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py index 20327b7e945..c28c8b2d337 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py @@ -112,13 +112,14 @@ 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 + # Match the BF16 normalized value produced by unfused RMSNorm before AdaLN affine. + 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).to(tl.float32) - scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32) + 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) tl.store( output_ptr + row * hidden_size + columns, - normalized * (1.0 + scale) + shift, + normalized * (scale + 1.0) + shift, mask=mask, ) @@ -150,21 +151,23 @@ def _indexed_gate_rms_norm_scale_shift_kernel( mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) - residual = tl.load(residual_ptr + row * stride_residual_row + columns, mask=mask, other=0.0).to(tl.float32) - gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32) - branch = tl.load(branch_ptr + row * stride_branch_row + columns, mask=mask, other=0.0).to(tl.float32) + residual = tl.load(residual_ptr + row * stride_residual_row + columns, mask=mask, other=0.0) + gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0) + branch = tl.load(branch_ptr + row * stride_branch_row + columns, mask=mask, other=0.0) # Match the BF16 residual value consumed by the unfused RMSNorm path. - updated = (residual + gate * branch).to(tl.bfloat16).to(tl.float32) + updated = residual + (gate * branch) tl.store(residual_out_ptr + row * hidden_size + columns, updated, mask=mask) 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) + up_f32 = updated.to(tl.float32) + variance = tl.sum(up_f32 * up_f32, axis=0) / hidden_size + # Match the BF16 normalized value produced by unfused RMSNorm before AdaLN affine. + normalized = (up_f32 * 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) tl.store( modulated_out_ptr + row * hidden_size + columns, - normalized * (1.0 + scale) + shift, + normalized * (scale + 1.0) + shift, mask=mask, ) From 52f375ec3923469aad307126174dbac74342fc7d Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 18 Sep 2026 14:55:07 +0800 Subject: [PATCH 2/2] fix(diffusion): disable FP fusion on CUDA modulation and add exact rounding regression - Pass enable_fp_fusion=False by default for CUDA kernel launches in _launch_row_chunks to prevent BF16 FMA from skipping intermediate rounding on gate * branch - Add test_fused_modulation_exact_rounding_regression covering exact bit-level rounding against unfused PyTorch reference Signed-off-by: BeatSeat --- .../minimax_h3/test_minimax_h3_modulation.py | 19 +++++++++++++++++++ .../attention/ops/minimax_h3_modulation.py | 1 + 2 files changed, 20 insertions(+) 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 88cd1ad74a3..3c9b56e72cd 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -222,3 +222,22 @@ def test_fused_modulation_matches_pytorch_reference() -> None: torch.testing.assert_close(actual_residual, ref_residual, atol=5e-2, rtol=5e-2) torch.testing.assert_close(actual_modulated, ref_modulated, atol=5e-2, rtol=5e-2) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +def test_fused_modulation_exact_rounding_regression() -> None: + """Exact bit-level rounding regression: without enable_fp_fusion=False, + BF16 FMA skips intermediate rounding of gate * branch and differs from + PyTorch's unfused sequence.""" + torch.manual_seed(42) + rows, hidden_size, conditions = 32, 1024, 2 + 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) + indices = torch.arange(rows, device="cuda") % conditions + + expected_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16) + actual_residual = indexed_gate(residual, gate, branch, indices) + assert torch.equal(actual_residual, expected_residual) diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py index c28c8b2d337..949e3a9299d 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py @@ -18,6 +18,7 @@ def _iter_row_chunks(rows: int): def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) -> None: if device_type != "npu": + kwargs.setdefault("enable_fp_fusion", False) kernel[(rows,)](*args, 0, **kwargs) return for row_offset, chunk_rows in _iter_row_chunks(rows):