From 815abce2bad2cee2f4e762f83852a108af8d44c0 Mon Sep 17 00:00:00 2001 From: Ting FU Date: Tue, 7 Jul 2026 17:54:39 +0800 Subject: [PATCH] [BugFix] fix possible shape mismatch when using ROPE (#4655) Signed-off-by: Semmer --- tests/diffusion/layers/test_rope_broadcast.py | 236 ++++++++++++++++++ vllm_omni/diffusion/layers/rope.py | 6 + 2 files changed, 242 insertions(+) create mode 100644 tests/diffusion/layers/test_rope_broadcast.py diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py new file mode 100644 index 000000000..29fcf1804 --- /dev/null +++ b/tests/diffusion/layers/test_rope_broadcast.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Shape broadcasting regression tests for RotaryEmbedding.forward_native. + +Covers the batch-dimension shape mismatch fix: forward_native now strips the +batch dim from 3D cos/sin before applying RoPE, consistent with forward_cuda, +forward_hip, and apply_rotary_emb_mindiesd paths. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.diffusion.layers.rope import RotaryEmbedding + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_inputs( + batch: int, + seq_len: int, + num_heads: int, + head_dim: int, + cos_batch: int | None = None, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Generate x, cos, sin tensors for RoPE testing. + + Args: + batch: Batch dimension of x. + seq_len: Sequence length. + num_heads: Number of attention heads. + head_dim: Head dimension (must be even). + cos_batch: Batch dimension of cos/sin. Defaults to batch if None. + dtype: Data type for the tensors. + """ + if cos_batch is None: + cos_batch = batch + + torch.manual_seed(42) + x = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype) + cos = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype) + sin = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype) + return x, cos, sin + + +class TestRotaryEmbeddingNativeShapeRegression: + """Regression tests for forward_native shape handling (#8e297d81).""" + + @pytest.fixture + def rope_neox(self) -> RotaryEmbedding: + """NeoX-style (non-interleaved) RoPE.""" + return RotaryEmbedding(is_neox_style=True) + + @pytest.fixture + def rope_gptj(self) -> RotaryEmbedding: + """GPT-J style (interleaved) RoPE.""" + return RotaryEmbedding(is_neox_style=False) + + # ── 2D cos/sin (no batch dim) ────────────────────────────────────── + + def test_2d_cos_sin_unchanged(self, rope_neox: RotaryEmbedding) -> None: + """2D cos/sin [S, D/2] pass through unchanged (no-op path).""" + x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=2) + # Manually drop batch dim to create 2D cos/sin + cos_2d = cos[0] + sin_2d = sin[0] + assert cos_2d.dim() == 2 + output = rope_neox.forward_native(x, cos_2d, sin_2d) + assert output.shape == x.shape + assert not torch.isnan(output).any() + assert not torch.isinf(output).any() + + # ── 3D cos/sin — batch-dim stripping ────────────────────────────── + + @pytest.mark.parametrize("cos_batch", [1, 2, 4]) + def test_3d_cos_sin_no_error(self, rope_neox: RotaryEmbedding, cos_batch: int) -> None: + """3D cos/sin with batch dim should not raise errors.""" + x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch) + output = rope_neox.forward_native(x, cos, sin) + assert output.shape == x.shape + + @pytest.mark.parametrize("interleaved", [False, True]) + def test_3d_cos_identical_to_2d(self, interleaved: bool) -> None: + """3D cos/sin with identical batch slices == 2D cos/sin result.""" + rope = RotaryEmbedding(is_neox_style=not interleaved) + B, S, H, D = 2, 16, 8, 64 + x = torch.randn(B, S, H, D) + cos_2d = torch.randn(S, D // 2) + sin_2d = torch.randn(S, D // 2) + + # 3D variant: expand the same 2D data across batch + cos_3d = cos_2d.unsqueeze(0).expand(B, -1, -1) + sin_3d = sin_2d.unsqueeze(0).expand(B, -1, -1) + + out_2d = rope.forward_native(x, cos_2d, sin_2d) + out_3d = rope.forward_native(x, cos_3d, sin_3d) + + torch.testing.assert_close(out_3d, out_2d, atol=0, rtol=0) + + # ── Batch size mismatch — the regression scenario ───────────────── + + @pytest.mark.parametrize( + "x_batch,cos_batch", + [ + (2, 1), # cos has fewer batch elements + (1, 2), # cos has more batch elements (takes [0]) + (4, 1), # large x batch, single cos + ], + ) + def test_batch_size_mismatch_regression( + self, + rope_neox: RotaryEmbedding, + x_batch: int, + cos_batch: int, + ) -> None: + """forward_native works when cos batch dim differs from x batch dim. + + This is the exact regression scenario — before the fix, a 3D cos/sin + with a different batch size than x would cause a shape mismatch error + inside apply_rotary_emb_torch during broadcast. + """ + x, cos, sin = _make_inputs(batch=x_batch, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch) + output = rope_neox.forward_native(x, cos, sin) + assert output.shape == x.shape + assert output.dtype == x.dtype + + # ── Batch consistency — all elements get same rotation ──────────── + + def test_2d_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None: + """With 2D cos, identical x slices → identical output slices.""" + B, S, H, D = 4, 16, 8, 64 + x_template = torch.randn(1, S, H, D) + x = x_template.expand(B, -1, -1, -1) # all batch elements identical + cos = torch.randn(S, D // 2) + sin = torch.randn(S, D // 2) + + output = rope_neox.forward_native(x, cos, sin) + + # All output batch slices must be identical + for i in range(1, B): + torch.testing.assert_close(output[i], output[0], atol=0, rtol=0) + + def test_3d_identical_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None: + """With 3D cos (identical slices), identical x → identical output.""" + B, S, H, D = 4, 16, 8, 64 + x_template = torch.randn(1, S, H, D) + x = x_template.expand(B, -1, -1, -1) + cos = torch.randn(1, S, D // 2).expand(B, -1, -1) + sin = torch.randn(1, S, D // 2).expand(B, -1, -1) + + output = rope_neox.forward_native(x, cos, sin) + + for i in range(1, B): + torch.testing.assert_close(output[i], output[0], atol=0, rtol=0) + + # ── Output shape / dtype invariants ─────────────────────────────── + + def test_output_shape_preserved(self, rope_neox: RotaryEmbedding) -> None: + """Output shape must match input shape for all cos/sin dimensionalities.""" + B, S, H, D = 2, 32, 4, 128 + x = torch.randn(B, S, H, D) + + # 2D cos/sin + cos_2d = torch.randn(S, D // 2) + sin_2d = torch.randn(S, D // 2) + assert rope_neox.forward_native(x, cos_2d, sin_2d).shape == (B, S, H, D) + + # 3D cos/sin, matching batch + cos_3d = torch.randn(B, S, D // 2) + sin_3d = torch.randn(B, S, D // 2) + assert rope_neox.forward_native(x, cos_3d, sin_3d).shape == (B, S, H, D) + + # 3D cos/sin, different batch + cos_3d_1 = torch.randn(1, S, D // 2) + sin_3d_1 = torch.randn(1, S, D // 2) + assert rope_neox.forward_native(x, cos_3d_1, sin_3d_1).shape == (B, S, H, D) + + def test_output_dtype_preserved(self, rope_neox: RotaryEmbedding) -> None: + """Output dtype must match input dtype.""" + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, dtype=dtype) + output = rope_neox.forward_native(x, cos, sin) + assert output.dtype == dtype, f"dtype mismatch: {output.dtype} != {dtype}" + + def test_output_no_nan_inf( + self, + rope_neox: RotaryEmbedding, + rope_gptj: RotaryEmbedding, + ) -> None: + """Output must be finite for both RoPE styles.""" + for rope in (rope_neox, rope_gptj): + x, cos, sin = _make_inputs(batch=2, seq_len=32, num_heads=4, head_dim=64) + output = rope.forward_native(x, cos, sin) + assert not torch.isnan(output).any(), f"NaN in {rope}" + assert not torch.isinf(output).any(), f"Inf in {rope}" + + def test_output_contiguous(self, rope_neox: RotaryEmbedding) -> None: + """Output should be contiguous.""" + x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64) + output = rope_neox.forward_native(x, cos, sin) + assert output.is_contiguous() + + # ── CFG-style doubled batch ─────────────────────────────────────── + + def test_cfg_style_doubled_batch(self, rope_neox: RotaryEmbedding) -> None: + """CFG (classifier-free guidance) doubles batch — cos may have + original batch size while x has 2× batch.""" + B, S, H, D = 2, 16, 8, 64 + # Simulate CFG: x has 2× batch (cond + uncond concatenated) + x = torch.randn(B * 2, S, H, D) + # cos/sin only have the original batch + cos = torch.randn(B, S, D // 2) + sin = torch.randn(B, S, D // 2) + + output = rope_neox.forward_native(x, cos, sin) + assert output.shape == (B * 2, S, H, D) + + # ── RotaryEmbeddingWan forward_native is unchanged ──────────────── + + def test_wan_forward_native_unchanged(self) -> None: + """RotaryEmbeddingWan.forward_native still uses its own implementation + (unflatten + stack + flatten) and should not crash with 3D cos/sin.""" + from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan + + rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) + B, S, H, D = 2, 16, 8, 64 + x = torch.randn(B, S, H, D) + cos = torch.randn(B, S, 1, D // 2) + sin = torch.randn(B, S, 1, D // 2) + + output = rope.forward_native(x, cos, sin) + assert output.shape == x.shape + assert not torch.isnan(output).any() diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 4cb714465..6a045fca2 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -188,6 +188,12 @@ class RotaryEmbedding(CustomOp): cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: + # 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] return apply_rotary_emb_torch( x, cos,