From f2e3d2e25bebc0c7110ccb945d7dcf4be81bf6db Mon Sep 17 00:00:00 2001 From: fan2956 Date: Thu, 6 Aug 2026 11:39:27 +0800 Subject: [PATCH] [Perf] Minimax-H3 support fused RMSNorm and RoPE opt (#5801) Signed-off-by: Fan Co-authored-by: Fan --- tests/diffusion/layers/test_norm.py | 41 ++++++++++++++++++ tests/diffusion/layers/test_rope_broadcast.py | 35 ++++++++++++++- .../minimax_h3/test_minimax_h3_parallel.py | 27 ++++++++++++ vllm_omni/diffusion/layers/norm.py | 4 +- vllm_omni/diffusion/layers/rope.py | 34 +++++++++------ .../minimax_h3/minimax_h3_transformer.py | 43 +++++++++---------- 6 files changed, 145 insertions(+), 39 deletions(-) diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py index 694f8f8e9..71125684b 100644 --- a/tests/diffusion/layers/test_norm.py +++ b/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 ── diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py index 29fcf1804..e1238ba44 100644 --- a/tests/diffusion/layers/test_rope_broadcast.py +++ b/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, + } diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py index d6f089d0a..bcd074483 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py +++ b/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"), [ diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 51c627975..14ab677a0 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/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: diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 6a045fca2..5d934866b 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/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, 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 7bd397a15..c48bbea48 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/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