From 4ddb22c4b2ed930d6565b2784736a0836ba78616 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 19 Aug 2026 18:13:23 +0800 Subject: [PATCH] [Hardware][MUSA] Enable MiniMax H3 fused kernels Signed-off-by: Xiaodong Ye --- .../minimax_h3/test_minimax_h3_modulation.py | 70 +++++++++++++++++++ .../attention/ops/minimax_h3_modulation.py | 34 ++++++--- vllm_omni/diffusion/layers/activation.py | 9 +++ .../minimax_h3/minimax_h3_transformer.py | 48 ++++++++----- vllm_omni/platforms/cuda/platform.py | 4 ++ vllm_omni/platforms/interface.py | 4 ++ vllm_omni/platforms/musa/platform.py | 4 ++ 7 files changed, 146 insertions(+), 27 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py 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 00000000000..ebe18ddb6df --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( + indexed_gate, + indexed_gate_rms_norm_scale_shift, + indexed_scale_shift_, + rms_norm_indexed_scale_shift, +) + + +def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + normalized = x.float() + variance = normalized.square().mean(-1, keepdim=True) + normalized = normalized * torch.rsqrt(variance + eps) + return (weight.float() * normalized).to(x.dtype) + + +def test_minimax_h3_modulation_cpu_fallback_with_strided_adaln() -> None: + """The fallback preserves H3's two-step BF16 rounding contract.""" + dtype = torch.bfloat16 + eps = 1e-6 + x = torch.tensor( + [[1.0, -2.0, 3.0, -4.0], [-5.0, 6.0, -7.0, 8.0]], + dtype=dtype, + ) + branch = torch.tensor( + [[0.5, -1.0, 1.5, -2.0], [2.5, -3.0, 3.5, -4.0]], + dtype=dtype, + ) + weight = torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=dtype) + indices = torch.tensor([0, 2], dtype=torch.int64) + + adaln = torch.arange(3 * 6 * x.shape[-1], dtype=torch.float32).reshape(3, 6 * x.shape[-1]).to(dtype) / 128 + shift, scale, gate, _, _, _ = adaln.chunk(6, dim=-1) + assert not shift.is_contiguous() + assert shift.stride() == (6 * x.shape[-1], 1) + + expected_scale_shift = (x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(dtype) + actual_scale_shift = indexed_scale_shift_(x.clone(), shift, scale, indices) + torch.testing.assert_close(actual_scale_shift, expected_scale_shift, atol=0, rtol=0) + + expected_gate = (x + gate.index_select(0, indices) * branch).to(dtype) + actual_gate = indexed_gate(x, gate, branch, indices) + torch.testing.assert_close(actual_gate, expected_gate, atol=0, rtol=0) + + normalized = _rms_norm(x, weight, eps) + expected_rms_mod = (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(dtype) + actual_rms_mod = rms_norm_indexed_scale_shift(x, weight, shift, scale, indices, eps) + torch.testing.assert_close(actual_rms_mod, expected_rms_mod, atol=0, rtol=0) + + expected_residual = expected_gate + normalized_residual = _rms_norm(expected_residual, weight, eps) + expected_residual_mod = ( + normalized_residual * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices) + ).to(dtype) + actual_residual, actual_residual_mod = indexed_gate_rms_norm_scale_shift( + x, + gate, + branch, + weight, + shift, + scale, + indices, + eps, + ) + torch.testing.assert_close(actual_residual, expected_residual, atol=0, rtol=0) + torch.testing.assert_close(actual_residual_mod, expected_residual_mod, atol=0, rtol=0) diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py index 53d3eb9b97e..5a8ded5bb0b 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py @@ -4,8 +4,18 @@ from __future__ import annotations import torch -import triton -import triton.language as tl +from vllm.triton_utils import HAS_TRITON, tl, triton + +from vllm_omni.platforms import current_omni_platform + + +def _can_use_triton(*tensors: torch.Tensor) -> bool: + if not HAS_TRITON or not current_omni_platform.supports_minimax_h3_triton_modulation(): + return False + device = tensors[0].device + return device.type == current_omni_platform.device_type and all( + tensor.device == device and tensor.ndim > 0 and tensor.stride(-1) == 1 for tensor in tensors + ) @triton.jit @@ -80,6 +90,7 @@ def _rms_norm_indexed_scale_shift_kernel( hidden_size: tl.constexpr, eps: tl.constexpr, stride_x_row, + stride_output_row, stride_shift_row, stride_scale_row, stride_indices, @@ -98,7 +109,7 @@ def _rms_norm_indexed_scale_shift_kernel( 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) tl.store( - output_ptr + row * hidden_size + columns, + output_ptr + row * stride_output_row + columns, normalized * (1.0 + scale) + shift, mask=mask, ) @@ -118,6 +129,8 @@ def _indexed_gate_rms_norm_scale_shift_kernel( hidden_size: tl.constexpr, eps: tl.constexpr, stride_residual_row, + stride_residual_out_row, + stride_modulated_out_row, stride_gate_row, stride_branch_row, stride_shift_row, @@ -134,7 +147,7 @@ def _indexed_gate_rms_norm_scale_shift_kernel( 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 - tl.store(residual_out_ptr + row * hidden_size + columns, updated, mask=mask) + tl.store(residual_out_ptr + row * stride_residual_out_row + 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 @@ -142,7 +155,7 @@ def _indexed_gate_rms_norm_scale_shift_kernel( 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) tl.store( - modulated_out_ptr + row * hidden_size + columns, + modulated_out_ptr + row * stride_modulated_out_row + columns, normalized * (1.0 + scale) + shift, mask=mask, ) @@ -155,7 +168,7 @@ def indexed_scale_shift_( indices: torch.Tensor, ) -> torch.Tensor: """Apply indexed scale/shift in-place to a disposable contiguous input.""" - if x.is_cpu: + if not _can_use_triton(x, shift, scale, indices): x.copy_((x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(x.dtype)) return x rows, hidden_size = x.shape @@ -185,7 +198,7 @@ def indexed_gate( indices: torch.Tensor, ) -> torch.Tensor: """Return ``x + gate[indices] * other`` without indexed temporaries.""" - if x.is_cpu: + if not _can_use_triton(x, gate, other, indices): return (x + gate.index_select(0, indices) * other).to(x.dtype) output = torch.empty_like(x) rows, hidden_size = x.shape @@ -218,7 +231,7 @@ def rms_norm_indexed_scale_shift( eps: float, ) -> torch.Tensor: """Fuse H3 RMSNorm with its indexed AdaLN affine transform.""" - if x.is_cpu: + if not _can_use_triton(x, weight, shift, scale, indices): input_dtype = x.dtype normalized = x.float() variance = normalized.pow(2).mean(-1, keepdim=True) @@ -238,6 +251,7 @@ def rms_norm_indexed_scale_shift( hidden_size, eps, x.stride(0), + output.stride(0), shift.stride(0), scale.stride(0), indices.stride(0), @@ -258,7 +272,7 @@ def indexed_gate_rms_norm_scale_shift( eps: float, ) -> tuple[torch.Tensor, torch.Tensor]: """Fuse gated residual, the following RMSNorm, and AdaLN affine.""" - if residual.is_cpu: + if not _can_use_triton(residual, gate, branch, weight, shift, scale, indices): input_dtype = residual.dtype residual_out = (residual + gate.index_select(0, indices) * branch).to(input_dtype) normalized = residual_out.float() @@ -286,6 +300,8 @@ def indexed_gate_rms_norm_scale_shift( hidden_size, eps, residual.stride(0), + residual_out.stride(0), + modulated_out.stride(0), gate.stride(0), branch.stride(0), shift.stride(0), diff --git a/vllm_omni/diffusion/layers/activation.py b/vllm_omni/diffusion/layers/activation.py index c9658cf6e97..41f04384fc8 100644 --- a/vllm_omni/diffusion/layers/activation.py +++ b/vllm_omni/diffusion/layers/activation.py @@ -30,6 +30,15 @@ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: self.op(out, x) return out + def forward_musa(self, x: torch.Tensor) -> torch.Tensor: + if x.device.type != "musa": + return self.forward_native(x) + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + self.op(out, x) + return out + def forward_npu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index af5f9b1dd70..31878de3d50 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -25,6 +25,7 @@ RowParallelLinear, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.triton_utils import HAS_TRITON from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout from vllm_omni.diffusion.attention.layer import Attention @@ -43,6 +44,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -746,6 +748,7 @@ def __init__( modality_num=MINIMAX_H3_ADALN_MODALITY_NUM, prefix=f"{prefix}.adaln_proj", ) + self.use_fused_modulation = HAS_TRITON and current_omni_platform.supports_minimax_h3_triton_modulation() def forward( self, @@ -777,14 +780,18 @@ def forward( ) = self.adaln_proj(t_emb) residual = x - h = rms_norm_indexed_scale_shift( - x, - self.norm1.weight, - shift_msa, - scale_msa, - combined_indices, - self.norm1.variance_epsilon, - ) + if self.use_fused_modulation: + h = rms_norm_indexed_scale_shift( + x, + self.norm1.weight, + shift_msa, + scale_msa, + combined_indices, + self.norm1.variance_epsilon, + ) + else: + h = self.norm1(x) + h = indexed_scale_shift_(h, shift_msa, scale_msa, combined_indices) h = self.attn( h, rope_table=rope_table, @@ -794,16 +801,21 @@ def forward( sp_seq_lens=sp_seq_lens, video_layout=video_layout, ) - x, h = indexed_gate_rms_norm_scale_shift( - residual, - gate_msa, - h, - self.norm2.weight, - shift_mlp, - scale_mlp, - combined_indices, - self.norm2.variance_epsilon, - ) + if self.use_fused_modulation: + x, h = indexed_gate_rms_norm_scale_shift( + residual, + gate_msa, + h, + self.norm2.weight, + shift_mlp, + scale_mlp, + combined_indices, + self.norm2.variance_epsilon, + ) + else: + x = indexed_gate(residual, gate_msa, h, combined_indices) + h = self.norm2(x) + h = indexed_scale_shift_(h, shift_mlp, scale_mlp, combined_indices) residual = x h = self.mlp(h) return indexed_gate(residual, gate_mlp, h, combined_indices) diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 4551126a7fe..7cccb5254bb 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -26,6 +26,10 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): _omni_enum = OmniPlatformEnum.CUDA + @classmethod + def supports_minimax_h3_triton_modulation(cls) -> bool: + return True + @classmethod def get_omni_ar_worker_cls(cls) -> str: return "vllm_omni.worker.gpu_ar_worker.GPUARWorker" diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 7759cb6d13d..ac85d273ea4 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -265,6 +265,10 @@ def create_autocast_context( def supports_cpu_offload(cls) -> bool: return True + @classmethod + def supports_minimax_h3_triton_modulation(cls) -> bool: + return False + @classmethod def supports_float64(cls) -> bool: return True diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index d667c9bf129..700a6f346cf 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -24,6 +24,10 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): _omni_enum = OmniPlatformEnum.MUSA + @classmethod + def supports_minimax_h3_triton_modulation(cls) -> bool: + return True + @classmethod def get_omni_ar_worker_cls(cls) -> str: return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"