[Perf] Minimax-H3 support fused RMSNorm and RoPE opt (#5801)

Signed-off-by: Fan <fan@FandeMacBook-Air.local>
Co-authored-by: Fan <fan@FandeMacBook-Air.local>

tests/diffusion/layers/test_norm.py

@@ -2,6 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers."""
 
+import sys
+import types
+
 import pytest
 import torch
 from pytest_mock import MockerFixture
@@ -299,6 +302,44 @@ def test_rmsnorm_numerical_correctness():
     torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
 
 
+def test_rmsnorm_dtype_matches_parameter_and_fp32_accumulation_contract():
+    """dtype controls gamma storage; reduced-precision inputs still use fp32 accumulation."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    eps = 1e-6
+    norm = RMSNorm(hidden_size, eps=eps, dtype=torch.bfloat16)
+    x = torch.randn(2, 4, hidden_size, dtype=torch.bfloat16)
+
+    assert norm.weight.dtype == torch.bfloat16
+
+    x_fp32 = x.float()
+    expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype)
+
+    torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
+
+
+def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.MonkeyPatch):
+    """The NPU fused op receives the bf16 checkpoint gamma used by H3."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    captured: dict[str, torch.Tensor] = {}
+
+    def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
+        captured["gamma"] = gamma
+        captured["epsilon"] = torch.tensor(epsilon)
+        return (x,)
+
+    monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm))
+    norm = RMSNorm(64, eps=1e-5, dtype=torch.bfloat16)
+    x = torch.randn(2, 4, 64, dtype=torch.bfloat16)
+
+    assert norm.forward_npu(x) is x
+    assert captured["gamma"] is norm.weight
+    assert captured["gamma"].dtype == torch.bfloat16
+    assert captured["epsilon"].item() == pytest.approx(1e-5)
+
+
 # ── RMSNorm compile-path regression tests ──
 
 

tests/diffusion/layers/test_rope_broadcast.py

@@ -10,10 +10,13 @@ forward_hip, and apply_rotary_emb_mindiesd paths.
 
 from __future__ import annotations
 
+import sys
+import types
+
 import pytest
 import torch
 
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rotary_emb_mindiesd
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -234,3 +237,33 @@ class TestRotaryEmbeddingNativeShapeRegression:
         output = rope.forward_native(x, cos, sin)
         assert output.shape == x.shape
         assert not torch.isnan(output).any()
+
+
+def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    """H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout."""
+    captured: dict[str, object] = {}
+
+    def rotary_position_embedding(x, cos, sin, **kwargs):
+        captured.update(x=x, cos=cos, sin=sin, kwargs=kwargs)
+        return x
+
+    monkeypatch.setitem(
+        sys.modules,
+        "mindiesd",
+        types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
+    )
+    x = torch.randn(11, 3, 96, dtype=torch.bfloat16)
+    cos = torch.randn(11, 96, dtype=torch.bfloat16)
+    sin = torch.randn(11, 96, dtype=torch.bfloat16)
+
+    actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)
+
+    assert actual is x
+    assert captured["x"] is x
+    assert captured["cos"] is cos
+    assert captured["sin"] is sin
+    assert captured["kwargs"] == {
+        "rotated_mode": "rotated_half",
+        "head_first": False,
+        "fused": True,
+    }

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -61,6 +61,33 @@ def test_packed_attention_is_a_regional_compile_boundary():
     assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
 
 
+def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
+    from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    attention = object.__new__(MiniMaxH3Attention)
+    nn.Module.__init__(attention)
+    attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
+    attention.rope._forward_method = attention.rope.forward_native
+
+    x = torch.randn(11, 3, 128, dtype=torch.bfloat16)
+    freqs_half = torch.randn(11, 48)
+    freqs = torch.cat((freqs_half, freqs_half), dim=-1)
+    actual = attention._apply_rope(x, freqs)
+
+    cos = torch.cos(freqs).to(x.dtype).unsqueeze(1)
+    sin = torch.sin(freqs).to(x.dtype).unsqueeze(1)
+    x_rot = x[..., :96]
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    expected_rot = x_rot * cos + torch.cat((-x2, x1), dim=-1) * sin
+    expected = torch.cat((expected_rot, x[..., 96:]), dim=-1)
+
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+    torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0)
+
+
 @pytest.mark.parametrize(
     ("tp_size", "message"),
     [

vllm_omni/diffusion/layers/norm.py

@@ -65,9 +65,9 @@ class LayerNorm(nn.LayerNorm, CustomOp):
 
 
 class RMSNorm(CustomOp):
-    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+    def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
         self.variance_epsilon = eps
 
     def _forward_fused(self, x: torch.Tensor) -> torch.Tensor:

vllm_omni/diffusion/layers/rope.py

@@ -89,10 +89,11 @@ class RotaryEmbedding(CustomOp):
            of 1st half and 2nd half (GPT-NeoX style).
     """
 
-    def __init__(self, is_neox_style: bool = False) -> None:
+    def __init__(self, is_neox_style: bool = False, half_head_dim: bool = True) -> None:
         super().__init__()
         self.is_neox_style = is_neox_style
         self.interleaved = not is_neox_style
+        self.half_head_dim = half_head_dim
         self.apply_rotary_emb_flash_attn = None
         self.has_mindie = False
         # ``find_spec("flash_attn")`` is True as long as *any* package publishes
@@ -110,6 +111,21 @@ class RotaryEmbedding(CustomOp):
         if find_spec("mindiesd") is not None:
             self.has_mindie = True
 
+    def _prepare_half_head_dim_cos_sin(
+        self,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Normalize shared cos/sin to the half-dim layout used by CUDA/native kernels."""
+        if cos.dim() == 3:
+            # All batch elements share the same rotary position encoding.
+            cos = cos[0]
+            sin = sin[0]
+        if not self.half_head_dim:
+            cos = cos[..., ::2] if self.interleaved else cos[..., : cos.shape[-1] // 2]
+            sin = sin[..., ::2] if self.interleaved else sin[..., : sin.shape[-1] // 2]
+        return cos, sin
+
     def forward_cuda(
         self,
         x: torch.Tensor,
@@ -118,10 +134,7 @@ class RotaryEmbedding(CustomOp):
     ) -> torch.Tensor:
         from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
 
-        if cos.dim() == 3:
-            # (B, S, D/2) -> (S, D/2)
-            cos = cos[0]
-            sin = sin[0]
+        cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin)
 
         x, squeezed = _ensure_batch_dim(x)
         output = apply_rotary_emb(
@@ -141,10 +154,7 @@ class RotaryEmbedding(CustomOp):
         if self.apply_rotary_emb_flash_attn is None:
             return self.forward_cuda(x, cos, sin)
 
-        if cos.dim() == 3:
-            # (B, S, D/2) -> (S, D/2)
-            cos = cos[0]
-            sin = sin[0]
+        cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin)
 
         x, squeezed = _ensure_batch_dim(x)
         output = self.apply_rotary_emb_flash_attn(
@@ -162,7 +172,7 @@ class RotaryEmbedding(CustomOp):
         sin: torch.Tensor,
     ) -> torch.Tensor:
         if self.has_mindie:
-            return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved)
+            return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved, self.half_head_dim)
         else:
             return self.forward_native(x, cos, sin)
 
@@ -191,9 +201,7 @@ class RotaryEmbedding(CustomOp):
         # All batch elements share the same rotary position encoding.
         # Strip the batch dim so the underlying op broadcasts over the batch,
         # consistent with forward_cuda / forward_hip / apply_rotary_emb_mindiesd.
-        if cos.dim() == 3:
-            cos = cos[0]
-            sin = sin[0]
+        cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin)
         return apply_rotary_emb_torch(
             x,
             cos,

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -33,6 +33,8 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
 )
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
@@ -168,16 +170,11 @@ def _reorder_grouped_qkv_to_qkv(
     )
 
 
-def _norm(size: int, *, eps: float, dtype: torch.dtype = _BF16_DTYPE) -> nn.RMSNorm:
+def _norm(size: int, *, eps: float, dtype: torch.dtype = _BF16_DTYPE) -> RMSNorm:
     # RMSNorm uses fp32 accumulation with bf16 inputs and outputs.
     # torch.nn.RMSNorm upcasts reduced-precision inputs for the variance
     # reduction, matching that accumulation semantic.
-    return nn.RMSNorm(size, eps=eps, dtype=dtype)
-
-
-def _rotate_half(x: torch.Tensor) -> torch.Tensor:
-    x1, x2 = torch.chunk(x, 2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
+    return RMSNorm(size, eps=eps, dtype=dtype)
 
 
 def _modulate_scale_shift(
@@ -230,20 +227,6 @@ class MiniMaxH3Rope(nn.Module):
         return torch.cat((half, half), dim=-1)  # [S, 96]
 
 
-def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
-    """Rotate the first rot_dim head dims; pass the rest through.
-
-    x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin
-    are cast to the activation dtype before the elementwise math.
-    """
-    rot_dim = freqs.shape[-1]
-    x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
-    cos = torch.cos(freqs).to(x.dtype).unsqueeze(1)  # [T, 1, rot_dim]
-    sin = torch.sin(freqs).to(x.dtype).unsqueeze(1)
-    x_rot = (x_rot * cos) + (_rotate_half(x_rot) * sin)
-    return torch.cat((x_rot, x_pass), dim=-1)
-
-
 class MiniMaxH3TimeEmbedder(nn.Module):
     def __init__(
         self,
@@ -349,6 +332,7 @@ class MiniMaxH3Attention(nn.Module):
         self.num_kv_heads = self.qkv_proj.num_kv_heads
         self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
+        self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
         self.out_proj = RowParallelLinear(
             inner_dim,
             arch.hidden_size,
@@ -367,6 +351,19 @@ class MiniMaxH3Attention(nn.Module):
             skip_sequence_parallel=skip_sequence_parallel,
         )
 
+    def _apply_rope(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+        """Rotate the first rot_dim head dims; pass the rest through.
+
+        x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin
+        are cast to the activation dtype before the elementwise math.
+        """
+        rot_dim = freqs.shape[-1]
+        x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
+        cos = torch.cos(freqs).to(x.dtype)  # [T, rot_dim]
+        sin = torch.sin(freqs).to(x.dtype)
+        x_rot = self.rope(x_rot, cos, sin)
+        return torch.cat((x_rot, x_pass), dim=-1)
+
     @torch.compiler.disable
     def _run_packed_attention(
         self,
@@ -450,8 +447,8 @@ class MiniMaxH3Attention(nn.Module):
         q = self.q_norm(q)
         k = self.k_norm(k)
         if rope_freqs is not None:
-            q = _apply_rope(q, rope_freqs)
-            k = _apply_rope(k, rope_freqs)
+            q = self._apply_rope(q, rope_freqs)
+            k = self._apply_rope(k, rope_freqs)
 
         # The packed layout uses a second document for alignment padding.
         # Local/Ulysses backends unpad it, while Ring keeps aligned rows for