[Hardware][MUSA] Enable MiniMax H3 fused kernels

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)

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,15 +147,15 @@ 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
     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)
     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),

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)
 

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)

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"

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

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"