[BugFix][NPU] MiniMax-H3 RoPE crash: add the missing batch dim before the mindiesd kernel (#5896)

Signed-off-by: Weiming Liao <liaowm5@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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)

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]: