From 592637015ffccec514401b700548768232d9590c Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Sun, 9 Aug 2026 22:08:43 +0800 Subject: [PATCH] [BugFix][NPU] MiniMax-H3 RoPE crash: add the missing batch dim before the mindiesd kernel (#5896) Signed-off-by: Weiming Liao Co-authored-by: Hongsheng Liu --- tests/diffusion/layers/test_rope_broadcast.py | 34 ++++++++++++++++--- vllm_omni/diffusion/layers/rope.py | 7 ++-- 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py index e1238ba44..284049347 100644 --- a/tests/diffusion/layers/test_rope_broadcast.py +++ b/tests/diffusion/layers/test_rope_broadcast.py @@ -244,7 +244,7 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey captured: dict[str, object] = {} def rotary_position_embedding(x, cos, sin, **kwargs): - captured.update(x=x, cos=cos, sin=sin, kwargs=kwargs) + captured.update(cos=cos, sin=sin, kwargs=kwargs) return x monkeypatch.setitem( @@ -256,10 +256,8 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey 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) + 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"] == { @@ -267,3 +265,31 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey "head_first": False, "fused": True, } + + +def test_mindie_rope_restores_3d_layout(monkeypatch: pytest.MonkeyPatch) -> None: + """mindiesd rotary_position_embedding only accepts 4D ``x``.""" + captured: dict[str, torch.Tensor] = {} + + def rotary_position_embedding(x, cos, sin, **kwargs): + # Mirror mindiesd's own input guard. + assert x.dim() == 4, f"mindiesd requires 4D x, got {x.dim()}" + captured.update(x=x) + return x + + monkeypatch.setitem( + sys.modules, + "mindiesd", + types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding), + ) + + S, H, D = 11, 3, 96 + x = torch.randn(S, H, D, dtype=torch.bfloat16) + cos = torch.randn(S, D, dtype=torch.bfloat16) + sin = torch.randn(S, D, dtype=torch.bfloat16) + + actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False) + + assert captured["x"].shape == (1, S, H, D) + assert actual.shape == (S, H, D) + torch.testing.assert_close(actual, x, atol=0, rtol=0) diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 5d934866b..633dbd275 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -46,6 +46,8 @@ def apply_rotary_emb_mindiesd( ) -> torch.Tensor: from mindiesd import rotary_position_embedding + x, squeezed = _ensure_batch_dim(x) + if cos.dim() == 3: # (B, S, D/2) -> (S, D/2) cos = cos[0] @@ -57,13 +59,14 @@ def apply_rotary_emb_mindiesd( seqlen = cos.shape[0] sin = sin.unsqueeze(0).unsqueeze(2).unsqueeze(-1).expand(-1, -1, -1, -1, 2).reshape(1, seqlen, 1, -1) cos = cos.unsqueeze(0).unsqueeze(2).unsqueeze(-1).expand(-1, -1, -1, -1, 2).reshape(1, seqlen, 1, -1) - return rotary_position_embedding(x, cos, sin, rotated_mode="rotated_interleaved", head_first=False, fused=True) + out = rotary_position_embedding(x, cos, sin, rotated_mode="rotated_interleaved", head_first=False, fused=True) else: if half_head_dim: seqlen = cos.shape[0] sin = sin.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) cos = cos.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) - return rotary_position_embedding(x, cos, sin, rotated_mode="rotated_half", head_first=False, fused=True) + out = rotary_position_embedding(x, cos, sin, rotated_mode="rotated_half", head_first=False, fused=True) + return _restore_batch_dim(out, squeezed) def _ensure_batch_dim(x: torch.Tensor) -> tuple[torch.Tensor, bool]: