[Bugfix] Fix .../mrope.py::apply_interleaved_rope() when torch.compile is used in torch==2.13 (#52005)
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
bastefaniak committed
Aug 14, 2026 at 00:08 UTC
71b0da7c4ecca10a9f11f535fc370cf988ce12d7
2 files changed
+52
-5
tests/kernels/core/test_mrope.py
+45
@@ -6,6 +6,7 @@ import pytest
6
import torch
7
8
from vllm.model_executor.layers.rotary_embedding import get_rope
9
+from vllm.model_executor.layers.rotary_embedding.mrope import apply_interleaved_rope
10
from vllm.platforms import current_platform
11
from vllm.transformers_utils.config import get_config
12
from vllm.utils.torch_utils import set_random_seed
@@ -60,6 +61,50 @@ MODELS_TO_TEST = [
61
num_tokens_list = [11, 8192]
62
63
64
+def test_apply_interleaved_rope():
65
+ mrope_section = [3, 1, 1]
66
+ x = torch.tensor(
67
+ [
68
+ [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]],
69
+ [[10, 11, 12, 13, 14], [15, 16, 17, 18, 19]],
70
+ [[20, 21, 22, 23, 24], [25, 26, 27, 28, 29]],
71
+ ]
72
+ )
73
+
74
+ result = apply_interleaved_rope(x, mrope_section)
75
+
76
+ expected = torch.tensor([[0, 11, 22, 3, 4], [5, 16, 27, 8, 9]])
77
+ torch.testing.assert_close(result, expected, rtol=0, atol=0)
78
+
79
+
80
+@pytest.mark.skipif(
81
+ not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only test."
82
+)
83
+def test_apply_interleaved_rope_torch_compile():
84
+ mrope_section = [24, 20, 20]
85
+ num_tokens = 8192
86
+ rotary_dim = sum(mrope_section) * 2
87
+ cache = torch.randn(
88
+ 3,
89
+ num_tokens,
90
+ rotary_dim,
91
+ device=device,
92
+ dtype=torch.bfloat16,
93
+ )
94
+ x = cache[..., : rotary_dim // 2]
95
+
96
+ expected = apply_interleaved_rope(x, mrope_section)
97
+ compiled_fn = torch.compile(
98
+ apply_interleaved_rope,
99
+ backend="inductor",
100
+ fullgraph=True,
101
+ )
102
+
103
+ result = compiled_fn(x, mrope_section)
104
+
105
+ torch.testing.assert_close(result, expected, rtol=0, atol=0)
106
+
107
+
108
@pytest.mark.skipif(
109
not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."
110
)
vllm/model_executor/layers/rotary_embedding/mrope.py
+7
-5
@@ -226,12 +226,14 @@ def triton_mrope(
226
def apply_interleaved_rope(x: torch.Tensor, mrope_section: list[int]) -> torch.Tensor:
227
"""Apply interleaved MRoPE to 3D rotary embeddings.
228
Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
229
- interleaved [THTHWHTHW...TT], preserving frequency continuity.
229
+ interleaved [THWTHWTHW...TT], preserving frequency continuity.
230
"""
231
- x_t = x[0].clone()
232
- x_t[..., 1 : mrope_section[1] * 3 : 3] = x[1, ..., 1 : mrope_section[1] * 3 : 3]
233
- x_t[..., 2 : mrope_section[2] * 3 : 3] = x[2, ..., 2 : mrope_section[2] * 3 : 3]
234
- return x_t
231
+ channels = torch.arange(x.shape[-1], device=x.device)
232
+ is_height = (channels % 3 == 1) & (channels < mrope_section[1] * 3)
233
+ is_width = (channels % 3 == 2) & (channels < mrope_section[2] * 3)
234
+
235
+ result = torch.where(is_height, x[1], x[0])
236
+ return torch.where(is_width, x[2], result)
237
238
239
class MRotaryEmbedding(RotaryEmbeddingBase):