diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py new file mode 100644 index 000000000..dbeeaa66a --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +from vllm.triton_utils import HAS_TRITON + +from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( + indexed_gate_rms_norm_scale_shift, + rms_norm_indexed_scale_shift, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, 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_preserves_bf16_residual_boundary() -> None: + torch.manual_seed(17) + rows, hidden_size, conditions = 11, 3072, 4 + tensors = [ + torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16), + torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16), + torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16), + torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16), + torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16), + torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16), + ] + residual, gate, branch, weight, shift, scale = tensors + indices = torch.arange(rows, device="cuda") % conditions + eps = 1e-6 + + residual_out, modulated_out = indexed_gate_rms_norm_scale_shift( + residual, + gate, + branch, + weight, + shift, + scale, + indices, + eps, + ) + expected = rms_norm_indexed_scale_shift( + residual_out, + weight, + shift, + scale, + indices, + eps, + ) + + assert torch.equal(modulated_out, expected) diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py index 08e9e0e4a..f75833831 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py @@ -133,7 +133,8 @@ def _indexed_gate_rms_norm_scale_shift_kernel( 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) - updated = residual + gate * branch + # Match the BF16 residual value consumed by the unfused RMSNorm path. + updated = (residual + gate * branch).to(tl.bfloat16).to(tl.float32) 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)