@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3bd9c49158
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 11, 10:00 UTC

[CustomOp] Extract ApplyRotaryEmb as CustomOp and unify the dispatch logic (#29873)

Signed-off-by: shen-shanshan <467638484@qq.com> Co-authored-by: gcanlin <canlinguosdu@gmail.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

Shanshan Shen committed Dec 16, 2025 at 11:08 UTC 3bd9c491583d45ae9f3f24b10e99626f502014b4
14 files changed +553 -280
tests/kernels/core/test_apply_rotary_emb.py new
+203
@@ -0,0 +1,203 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Tests for ApplyRotaryEmb CustomOp dispatch behavior.
5 +
6 +This test ensures that RotaryEmbedding classes correctly call the appropriate
7 +ApplyRotaryEmb methods based on the calling context:
8 +
9 +1. RotaryEmbedding.forward_native() -> ApplyRotaryEmb.forward_native()
10 +2. RotaryEmbedding.forward_cuda() -> ApplyRotaryEmb.forward() (auto-dispatch)
11 +3. RotaryEmbedding.forward_hip() -> ApplyRotaryEmb.forward() (auto-dispatch)
12 +"""
13 +
14 +from dataclasses import dataclass
15 +
16 +import pytest
17 +import torch
18 +
19 +from vllm.config import (
20 + CompilationConfig,
21 + VllmConfig,
22 + get_cached_compilation_config,
23 + set_current_vllm_config,
24 +)
25 +from vllm.platforms import current_platform
26 +
27 +CUDA_DEVICES = ["cuda:0"]
28 +
29 +
30 +@dataclass
31 +class RotaryEmbeddingTestCase:
32 + """Test case configuration for RotaryEmbedding dispatch tests."""
33 +
34 + name: str
35 + rope_class: type
36 + rope_kwargs: dict
37 + method_name: str # forward_native, forward_cuda, forward
38 + positions_shape: tuple # (num_tokens,) or (3, num_tokens) or (4, num_tokens)
39 + expect_forward_native: bool # Should call ApplyRotaryEmb.forward_native()
40 + expect_forward: bool # Should call ApplyRotaryEmb.forward()
41 +
42 +
43 +def get_test_cases() -> list[RotaryEmbeddingTestCase]:
44 + """Generate test cases for all RotaryEmbedding classes."""
45 + from vllm.model_executor.layers.rotary_embedding.ernie45_vl_rope import (
46 + Ernie4_5_VLRotaryEmbedding,
47 + )
48 + from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
49 + from vllm.model_executor.layers.rotary_embedding.xdrope import XDRotaryEmbedding
50 +
51 + common_kwargs = {
52 + "head_size": 128,
53 + "rotary_dim": 128,
54 + "max_position_embeddings": 4096,
55 + "base": 10000,
56 + "is_neox_style": True,
57 + "dtype": torch.bfloat16,
58 + }
59 +
60 + return [
61 + # MRotaryEmbedding tests
62 + RotaryEmbeddingTestCase(
63 + name="MRotaryEmbedding.forward_native",
64 + rope_class=MRotaryEmbedding,
65 + rope_kwargs={**common_kwargs, "mrope_section": [16, 24, 24]},
66 + method_name="forward_native",
67 + positions_shape=(3, 32), # 2D for multimodal
68 + expect_forward_native=True,
69 + expect_forward=False,
70 + ),
71 + RotaryEmbeddingTestCase(
72 + name="MRotaryEmbedding.forward_cuda_1d",
73 + rope_class=MRotaryEmbedding,
74 + rope_kwargs={**common_kwargs, "mrope_section": [16, 24, 24]},
75 + method_name="forward_cuda",
76 + positions_shape=(32,), # 1D triggers apply_rotary_emb path
77 + expect_forward_native=False,
78 + expect_forward=True,
79 + ),
80 + # XDRotaryEmbedding tests
81 + RotaryEmbeddingTestCase(
82 + name="XDRotaryEmbedding.forward",
83 + rope_class=XDRotaryEmbedding,
84 + rope_kwargs={
85 + **common_kwargs,
86 + "scaling_alpha": 1.0,
87 + "xdrope_section": [16, 16, 16, 16],
88 + },
89 + method_name="forward",
90 + positions_shape=(4, 32), # 4D for P/W/H/T
91 + expect_forward_native=False,
92 + expect_forward=True,
93 + ),
94 + # Ernie4_5_VLRotaryEmbedding tests
95 + RotaryEmbeddingTestCase(
96 + name="Ernie4_5_VLRotaryEmbedding.forward_native",
97 + rope_class=Ernie4_5_VLRotaryEmbedding,
98 + rope_kwargs={**common_kwargs, "mrope_section": [22, 22, 20]},
99 + method_name="forward_native",
100 + positions_shape=(3, 32), # 2D for multimodal
101 + expect_forward_native=True,
102 + expect_forward=False,
103 + ),
104 + ]
105 +
106 +
107 +def run_dispatch_test(
108 + test_case: RotaryEmbeddingTestCase,
109 + device: str,
110 +):
111 + """Run a dispatch test for a RotaryEmbedding class."""
112 + vllm_config = VllmConfig(
113 + compilation_config=CompilationConfig(custom_ops=["all", "+apply_rotary_emb"])
114 + )
115 + get_cached_compilation_config.cache_clear()
116 +
117 + with set_current_vllm_config(vllm_config):
118 + rope = test_case.rope_class(**test_case.rope_kwargs).to(device=device)
119 +
120 + apply_rotary_emb = rope.apply_rotary_emb
121 +
122 + # Verify custom op is enabled
123 + if test_case.expect_forward_native:
124 + assert (
125 + apply_rotary_emb._forward_method != apply_rotary_emb.forward_native
126 + ), "Test setup error: ApplyRotaryEmb custom op should be enabled"
127 +
128 + # Setup call tracking
129 + call_tracker = {"forward_native_called": False, "forward_called": False}
130 + original_forward_native = apply_rotary_emb.forward_native
131 + original_forward = apply_rotary_emb.forward
132 +
133 + def tracked_forward_native(*args, **kwargs):
134 + call_tracker["forward_native_called"] = True
135 + return original_forward_native(*args, **kwargs)
136 +
137 + def tracked_forward(*args, **kwargs):
138 + call_tracker["forward_called"] = True
139 + return original_forward(*args, **kwargs)
140 +
141 + apply_rotary_emb.forward_native = tracked_forward_native
142 + apply_rotary_emb.forward = tracked_forward
143 +
144 + try:
145 + num_tokens = test_case.positions_shape[-1]
146 + num_q_heads = 8
147 + num_kv_heads = 2
148 + head_size = test_case.rope_kwargs["head_size"]
149 + max_position = test_case.rope_kwargs["max_position_embeddings"]
150 +
151 + positions = torch.randint(
152 + 0, max_position // 4, test_case.positions_shape, device=device
153 + )
154 + query = torch.randn(
155 + num_tokens, num_q_heads * head_size, dtype=torch.bfloat16, device=device
156 + )
157 + key = torch.randn(
158 + num_tokens,
159 + num_kv_heads * head_size,
160 + dtype=torch.bfloat16,
161 + device=device,
162 + )
163 +
164 + # Call the method under test
165 + method = getattr(rope, test_case.method_name)
166 + method(positions, query.clone(), key.clone())
167 +
168 + # Verify expectations
169 + if test_case.expect_forward_native:
170 + assert call_tracker["forward_native_called"], (
171 + f"{test_case.name} should call ApplyRotaryEmb.forward_native()"
172 + )
173 + if not test_case.expect_forward:
174 + assert not call_tracker["forward_called"], (
175 + f"{test_case.name} should NOT call ApplyRotaryEmb.forward(). "
176 + "Bug: when +apply_rotary_emb is enabled, forward_native() "
177 + "incorrectly dispatches to CUDA/HIP kernels."
178 + )
179 + if test_case.expect_forward:
180 + assert call_tracker["forward_called"], (
181 + f"{test_case.name} should call ApplyRotaryEmb.forward()"
182 + )
183 + finally:
184 + apply_rotary_emb.forward_native = original_forward_native
185 + apply_rotary_emb.forward = original_forward
186 +
187 +
188 +@pytest.mark.skipif(
189 + not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."
190 +)
191 +@pytest.mark.parametrize("test_case", get_test_cases(), ids=lambda tc: tc.name)
192 +@pytest.mark.parametrize("device", CUDA_DEVICES)
193 +def test_rotary_embedding_dispatch(
194 + test_case: RotaryEmbeddingTestCase,
195 + device: str,
196 +):
197 + """
198 + Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method.
199 +
200 + - forward_native methods should call ApplyRotaryEmb.forward_native()
201 + - forward_cuda/forward methods should call ApplyRotaryEmb.forward()
202 + """
203 + run_dispatch_test(test_case, device)
vllm/model_executor/layers/rotary_embedding/base.py
+17 -3
@@ -7,7 +7,7 @@ import torch
7 from vllm._aiter_ops import rocm_aiter_ops
8 from vllm.model_executor.custom_op import CustomOp
9
10 -from .common import apply_rotary_emb_torch
10 +from .common import ApplyRotaryEmb
11
12
13 @CustomOp.register("rotary_embedding")
@@ -49,6 +49,10 @@ class RotaryEmbeddingBase(CustomOp):
49 rocm_aiter_ops.is_triton_rotary_embed_enabled()
50 )
51
52 + self.apply_rotary_emb = ApplyRotaryEmb(
53 + is_neox_style=self.is_neox_style,
54 + )
55 +
56 def _compute_inv_freq(self, base: float) -> torch.Tensor:
57 """Compute the inverse frequency."""
58 # NOTE(woosuk): To exactly match the HF implementation, we need to
@@ -123,7 +127,12 @@ class RotaryEmbedding(RotaryEmbeddingBase):
127 query = query.view(num_tokens, -1, head_size)
128 query_rot = query[..., :rotary_dim]
129 query_pass = query[..., rotary_dim:]
126 - query_rot = apply_rotary_emb_torch(query_rot, cos, sin, is_neox_style)
130 + query_rot = ApplyRotaryEmb.forward_static(
131 + query_rot,
132 + cos,
133 + sin,
134 + is_neox_style,
135 + )
136 query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
137
138 # key may be None in some cases, e.g. cross-layer KV sharing
@@ -132,7 +141,12 @@ class RotaryEmbedding(RotaryEmbeddingBase):
141 key = key.view(num_tokens, -1, head_size)
142 key_rot = key[..., :rotary_dim]
143 key_pass = key[..., rotary_dim:]
135 - key_rot = apply_rotary_emb_torch(key_rot, cos, sin, is_neox_style)
144 + key_rot = ApplyRotaryEmb.forward_static(
145 + key_rot,
146 + cos,
147 + sin,
148 + is_neox_style,
149 + )
150 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
151 return query, key
152
vllm/model_executor/layers/rotary_embedding/common.py
+153 -71
@@ -2,19 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import math
5 -from collections.abc import Callable
6 -from functools import cache
5 from importlib.util import find_spec
6
7 import torch
8
9 from vllm.logger import init_logger
12 -from vllm.platforms import current_platform
10 +from vllm.model_executor.custom_op import CustomOp
11 from vllm.utils.torch_utils import direct_register_custom_op
12
15 -if current_platform.is_cuda():
16 - from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
17 -
13 logger = init_logger(__name__)
14
15
@@ -32,71 +27,6 @@ def rotate_gptj(x: torch.Tensor) -> torch.Tensor:
27 return x.flatten(-2)
28
29
35 -def apply_rotary_emb_torch(
36 - x: torch.Tensor,
37 - cos: torch.Tensor,
38 - sin: torch.Tensor,
39 - is_neox_style: bool,
40 -) -> torch.Tensor:
41 - cos = cos.unsqueeze(-2).to(x.dtype)
42 - sin = sin.unsqueeze(-2).to(x.dtype)
43 - if is_neox_style:
44 - x1, x2 = torch.chunk(x, 2, dim=-1)
45 - else:
46 - x1 = x[..., ::2]
47 - x2 = x[..., 1::2]
48 - o1 = x1 * cos - x2 * sin
49 - o2 = x2 * cos + x1 * sin
50 - if is_neox_style:
51 - return torch.cat((o1, o2), dim=-1)
52 - else:
53 - return torch.stack((o1, o2), dim=-1).flatten(-2)
54 -
55 -
56 -def apply_rotary_emb_dispatch(
57 - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, is_neox_style: bool
58 -) -> torch.Tensor:
59 - """
60 - Args:
61 - x: [num_tokens, num_heads, head_size]
62 - cos: [num_tokens, head_size // 2]
63 - sin: [num_tokens, head_size // 2]
64 - is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
65 - positional embeddings.
66 - """
67 - if current_platform.is_cuda():
68 - return apply_rotary_emb(x.unsqueeze(0), cos, sin, not is_neox_style).squeeze(0)
69 - else:
70 - return apply_rotary_emb_torch(x, cos, sin, is_neox_style)
71 -
72 -
73 -@cache
74 -def dispatch_rotary_emb_function(
75 - default: Callable[..., torch.Tensor] | None = None,
76 -) -> Callable[..., torch.Tensor]:
77 - if current_platform.is_cuda():
78 - return apply_rotary_emb
79 -
80 - # if torch compile is not enabled
81 - # use rotary embedding function from flash_attn package
82 - # otherwise use the naive pytorch embedding implementation
83 - # is faster when torch compile is enabled.
84 - if current_platform.is_rocm() and not torch.compiler.is_compiling():
85 - if find_spec("flash_attn") is not None:
86 - from flash_attn.ops.triton.rotary import apply_rotary
87 -
88 - return apply_rotary
89 - else:
90 - logger.warning(
91 - "flash_attn is not installed. Falling back to PyTorch "
92 - "implementation for rotary embeddings."
93 - )
94 - if default is not None:
95 - return default
96 -
97 - return apply_rotary_emb_torch
98 -
99 -
30 # yarn functions
31 # Inverse dim formula to find dim based on number of rotations
32 def yarn_find_correction_dim(
@@ -186,3 +116,155 @@ direct_register_custom_op(
116 mutates_args=["query", "key"], # These tensors are modified in-place
117 fake_impl=_flashinfer_rotary_embedding_fake,
118 )
119 +
120 +
121 +@CustomOp.register("apply_rotary_emb")
122 +class ApplyRotaryEmb(CustomOp):
123 + def __init__(
124 + self,
125 + enforce_enable: bool = False,
126 + is_neox_style: bool = True,
127 + enable_fp32_compute: bool = False,
128 + ) -> None:
129 + super().__init__(enforce_enable)
130 + self.is_neox_style = is_neox_style
131 + self.enable_fp32_compute = enable_fp32_compute
132 +
133 + self.apply_rotary_emb_flash_attn = None
134 + if find_spec("flash_attn") is not None:
135 + from flash_attn.ops.triton.rotary import apply_rotary
136 +
137 + self.apply_rotary_emb_flash_attn = apply_rotary
138 +
139 + @staticmethod
140 + def forward_static(
141 + x: torch.Tensor,
142 + cos: torch.Tensor,
143 + sin: torch.Tensor,
144 + is_neox_style: bool = True,
145 + enable_fp32_compute: bool = False,
146 + ) -> torch.Tensor:
147 + """
148 + Args:
149 + x: [batch_size (optional), seq_len, num_heads, head_size]
150 + cos: [seq_len, head_size // 2]
151 + sin: [seq_len, head_size // 2]
152 + is_neox_style: Whether to use the Neox-style or GPT-J-style.
153 + enable_fp32_compute: Temporarily convert x, cos, sin to FP32 dtype
154 + for higher accuracy.
155 + """
156 + origin_dtype = x.dtype
157 + if enable_fp32_compute:
158 + x = x.float()
159 +
160 + cos = cos.unsqueeze(-2).to(x.dtype)
161 + sin = sin.unsqueeze(-2).to(x.dtype)
162 +
163 + if is_neox_style:
164 + x1, x2 = torch.chunk(x, 2, dim=-1)
165 + else:
166 + x1 = x[..., ::2]
167 + x2 = x[..., 1::2]
168 +
169 + o1 = x1 * cos - x2 * sin
170 + o2 = x2 * cos + x1 * sin
171 +
172 + if is_neox_style:
173 + output = torch.cat((o1, o2), dim=-1)
174 + else:
175 + output = torch.stack((o1, o2), dim=-1).flatten(-2)
176 +
177 + if enable_fp32_compute:
178 + output = output.to(origin_dtype)
179 + return output
180 +
181 + def forward_native(
182 + self,
183 + x: torch.Tensor,
184 + cos: torch.Tensor,
185 + sin: torch.Tensor,
186 + ) -> torch.Tensor:
187 + output = self.forward_static(
188 + x, cos, sin, self.is_neox_style, self.enable_fp32_compute
189 + )
190 + return output
191 +
192 + def forward_cuda(
193 + self,
194 + x: torch.Tensor,
195 + cos: torch.Tensor,
196 + sin: torch.Tensor,
197 + ) -> torch.Tensor:
198 + from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
199 +
200 + origin_dtype = x.dtype
201 + if self.enable_fp32_compute:
202 + x = x.float()
203 + cos = cos.float()
204 + sin = sin.float()
205 +
206 + origin_shape = x.shape
207 + if len(origin_shape) == 3:
208 + # x: [seq_len, num_heads, head_size]
209 + x = x.unsqueeze(0)
210 +
211 + """
212 + Arguments of apply_rotary_emb() in vllm_flash_attn:
213 + x: [batch_size, seq_len, nheads, headdim]
214 + cos, sin: [seqlen_rotary, rotary_dim / 2]
215 + interleaved: defalut as False (Neox-style).
216 + ...
217 + """
218 + interleaved = not self.is_neox_style
219 + output = apply_rotary_emb(x, cos, sin, interleaved)
220 +
221 + if len(origin_shape) == 3:
222 + output = output.squeeze(0)
223 + if self.enable_fp32_compute:
224 + output = output.to(origin_dtype)
225 + return output
226 +
227 + def forward_hip(
228 + self,
229 + x: torch.Tensor,
230 + cos: torch.Tensor,
231 + sin: torch.Tensor,
232 + ) -> torch.Tensor:
233 + if self.apply_rotary_emb_flash_attn is not None:
234 + origin_dtype = x.dtype
235 + if self.enable_fp32_compute:
236 + x = x.float()
237 + cos = cos.float()
238 + sin = sin.float()
239 +
240 + origin_shape = x.shape
241 + if len(origin_shape) == 3:
242 + # x: [seq_len, num_heads, head_size]
243 + x = x.unsqueeze(0)
244 +
245 + """
246 + Arguments of apply_rotary() in flash_attn:
247 + x: [batch_size, seq_len, nheads, headdim]
248 + cos, sin: [seqlen_rotary, rotary_dim / 2]
249 + interleaved: defalut as False (Neox-style).
250 + ...
251 + """
252 + interleaved = not self.is_neox_style
253 + output = self.apply_rotary_emb_flash_attn(
254 + x, cos, sin, interleaved=interleaved
255 + ).type_as(x)
256 +
257 + if len(origin_shape) == 3:
258 + output = output.squeeze(0)
259 + if self.enable_fp32_compute:
260 + output = output.to(origin_dtype)
261 + else:
262 + # Falling back to PyTorch native implementation.
263 + output = self.forward_native(x, cos, sin)
264 +
265 + return output
266 +
267 + def extra_repr(self) -> str:
268 + s = f"is_neox_style={self.is_neox_style}"
269 + s += f"enable_fp32_compute={self.enable_fp32_compute}"
270 + return s
vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py
+10 -3
@@ -4,7 +4,6 @@
4
5 import torch
6
7 -from .common import apply_rotary_emb_dispatch
7 from .mrope import MRotaryEmbedding
8
9
@@ -55,14 +54,22 @@ class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
54 query = query.view(num_tokens, -1, self.head_size)
55 query_rot = query[..., : self.rotary_dim]
56 query_pass = query[..., self.rotary_dim :]
58 - query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin, self.is_neox_style)
57 + query_rot = self.apply_rotary_emb.forward_native(
58 + query_rot,
59 + cos,
60 + sin,
61 + )
62 query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
63
64 key_shape = key.shape
65 key = key.view(num_tokens, -1, self.head_size)
66 key_rot = key[..., : self.rotary_dim]
67 key_pass = key[..., self.rotary_dim :]
65 - key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin, self.is_neox_style)
68 + key_rot = self.apply_rotary_emb.forward_native(
69 + key_rot,
70 + cos,
71 + sin,
72 + )
73 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
74 return query, key
75
vllm/model_executor/layers/rotary_embedding/mrope.py
+20 -5
@@ -8,7 +8,6 @@ import torch
8 from vllm.triton_utils import tl, triton
9
10 from .base import RotaryEmbeddingBase
11 -from .common import apply_rotary_emb_dispatch
11 from .yarn_scaling_rope import YaRNScalingRotaryEmbedding, yarn_get_mscale
12
13
@@ -301,14 +300,22 @@ class MRotaryEmbedding(RotaryEmbeddingBase):
300 query = query.view(num_tokens, -1, self.head_size)
301 query_rot = query[..., : self.rotary_dim]
302 query_pass = query[..., self.rotary_dim :]
304 - query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin, self.is_neox_style)
303 + query_rot = self.apply_rotary_emb.forward_native(
304 + query_rot,
305 + cos,
306 + sin,
307 + )
308 query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
309
310 key_shape = key.shape
311 key = key.view(num_tokens, -1, self.head_size)
312 key_rot = key[..., : self.rotary_dim]
313 key_pass = key[..., self.rotary_dim :]
311 - key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin, self.is_neox_style)
314 + key_rot = self.apply_rotary_emb.forward_native(
315 + key_rot,
316 + cos,
317 + sin,
318 + )
319 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
320 return query, key
321
@@ -347,13 +354,21 @@ class MRotaryEmbedding(RotaryEmbeddingBase):
354 query = query.view(num_tokens, -1, self.head_size)
355 query_rot = query[..., : self.rotary_dim]
356 query_pass = query[..., self.rotary_dim :]
350 - query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin, self.is_neox_style)
357 + query_rot = self.apply_rotary_emb(
358 + query_rot,
359 + cos,
360 + sin,
361 + )
362 query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
363
364 key = key.view(num_tokens, -1, self.head_size)
365 key_rot = key[..., : self.rotary_dim]
366 key_pass = key[..., self.rotary_dim :]
356 - key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin, self.is_neox_style)
367 + key_rot = self.apply_rotary_emb(
368 + key_rot,
369 + cos,
370 + sin,
371 + )
372 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
373 return query, key
374
vllm/model_executor/layers/rotary_embedding/xdrope.py
+62 -4
@@ -4,7 +4,6 @@
4 import numpy as np
5 import torch
6
7 -from .common import apply_rotary_emb_dispatch
7 from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding
8
9
@@ -36,7 +35,7 @@ class XDRotaryEmbedding(DynamicNTKAlphaRotaryEmbedding):
35 dtype,
36 )
37
39 - def forward(
38 + def forward_native(
39 self,
40 positions: torch.Tensor,
41 query: torch.Tensor,
@@ -68,14 +67,73 @@ class XDRotaryEmbedding(DynamicNTKAlphaRotaryEmbedding):
67 query = query.view(num_tokens, -1, self.head_size)
68 query_rot = query[..., : self.rotary_dim]
69 query_pass = query[..., self.rotary_dim :]
71 - query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin, self.is_neox_style)
70 + query_rot = self.apply_rotary_emb.forward_native(
71 + query_rot,
72 + cos,
73 + sin,
74 + )
75 + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
76 +
77 + key_shape = key.shape
78 + key = key.view(num_tokens, -1, self.head_size)
79 + key_rot = key[..., : self.rotary_dim]
80 + key_pass = key[..., self.rotary_dim :]
81 + key_rot = self.apply_rotary_emb.forward_native(
82 + key_rot,
83 + cos,
84 + sin,
85 + )
86 + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
87 + return query, key
88 +
89 + def forward_cuda(
90 + self,
91 + positions: torch.Tensor,
92 + query: torch.Tensor,
93 + key: torch.Tensor | None = None,
94 + offsets: torch.Tensor | None = None,
95 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
96 + """PyTorch-native implementation equivalent to forward().
97 +
98 + Args:
99 + positions:
100 + [4, num_tokens] (P/W/H/T positions with multimodal inputs)
101 + query: [num_tokens, num_heads * head_size]
102 + key: [num_tokens, num_kv_heads * head_size]
103 + """
104 + assert positions.ndim == 2
105 + assert key is not None
106 +
107 + num_tokens = positions.shape[-1]
108 + cos_sin = self.cos_sin_cache[positions]
109 + cos, sin = cos_sin.chunk(2, dim=-1)
110 + cos = torch.cat(
111 + [m[i] for i, m in enumerate(cos.split(self.xdrope_section, dim=-1))], dim=-1
112 + )
113 + sin = torch.cat(
114 + [m[i] for i, m in enumerate(sin.split(self.xdrope_section, dim=-1))], dim=-1
115 + )
116 +
117 + query_shape = query.shape
118 + query = query.view(num_tokens, -1, self.head_size)
119 + query_rot = query[..., : self.rotary_dim]
120 + query_pass = query[..., self.rotary_dim :]
121 + query_rot = self.apply_rotary_emb(
122 + query_rot,
123 + cos,
124 + sin,
125 + )
126 query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
127
128 key_shape = key.shape
129 key = key.view(num_tokens, -1, self.head_size)
130 key_rot = key[..., : self.rotary_dim]
131 key_pass = key[..., self.rotary_dim :]
78 - key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin, self.is_neox_style)
132 + key_rot = self.apply_rotary_emb(
133 + key_rot,
134 + cos,
135 + sin,
136 + )
137 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
138 return query, key
139
vllm/model_executor/models/dots_ocr.py
+13 -27
@@ -29,6 +29,9 @@ from vllm.model_executor.layers.linear import (
29 RowParallelLinear,
30 )
31 from vllm.model_executor.layers.quantization import QuantizationConfig
32 +from vllm.model_executor.layers.rotary_embedding.common import (
33 + ApplyRotaryEmb,
34 +)
35 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36 from vllm.model_executor.models.interfaces import (
37 MultiModalEmbeddings,
@@ -158,32 +161,6 @@ class DotsOCRProcessingInfo(Qwen2VLProcessingInfo):
161 return processor
162
163
161 -def rotate_half(x):
162 - """Rotates half the hidden dims of the input."""
163 - x1 = x[..., : x.shape[-1] // 2]
164 - x2 = x[..., x.shape[-1] // 2 :]
165 - return torch.cat((-x2, x1), dim=-1)
166 -
167 -
168 -def apply_rotary_pos_emb_vision(
169 - tensor: torch.Tensor, freqs: torch.Tensor
170 -) -> torch.Tensor:
171 - orig_dtype = tensor.dtype
172 - tensor = tensor.float()
173 -
174 - cos = freqs.cos()
175 - sin = freqs.sin()
176 -
177 - cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
178 - sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
179 -
180 - output = (tensor * cos) + (rotate_half(tensor) * sin)
181 -
182 - output = output.to(orig_dtype)
183 -
184 - return output
185 -
186 -
164 class VisionRotaryEmbedding(nn.Module):
165 def __init__(self, dim: int, theta: float = 10000.0) -> None:
166 super().__init__()
@@ -298,6 +275,11 @@ class DotsVisionAttention(nn.Module):
275 prefix=f"{prefix}.attn",
276 )
277
278 + self.apply_rotary_emb = ApplyRotaryEmb(
279 + enforce_enable=True,
280 + enable_fp32_compute=True,
281 + )
282 +
283 def forward(
284 self,
285 hidden_states: torch.Tensor,
@@ -318,7 +300,11 @@ class DotsVisionAttention(nn.Module):
300
301 if rotary_pos_emb is not None:
302 qk_concat = torch.cat([q, k], dim=0)
321 - qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
303 + qk_rotated = self.apply_rotary_emb(
304 + qk_concat,
305 + rotary_pos_emb.cos(),
306 + rotary_pos_emb.sin(),
307 + )
308 q, k = torch.chunk(qk_rotated, 2, dim=0)
309
310 context_layer = self.attn(
vllm/model_executor/models/ernie45_vl.py
+14 -49
@@ -33,7 +33,7 @@ import numpy as np
33 import torch
34 import torch.nn as nn
35 import torch.nn.functional as F
36 -from einops import rearrange, repeat
36 +from einops import rearrange
37 from transformers import BatchFeature
38
39 from vllm.attention.backends.registry import AttentionBackendEnum
@@ -53,6 +53,9 @@ from vllm.model_executor.layers.linear import (
53 RowParallelLinear,
54 )
55 from vllm.model_executor.layers.quantization import QuantizationConfig
56 +from vllm.model_executor.layers.rotary_embedding.common import (
57 + ApplyRotaryEmb,
58 +)
59 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
60 from vllm.multimodal import MULTIMODAL_REGISTRY
61 from vllm.multimodal.inputs import (
@@ -69,7 +72,6 @@ from vllm.multimodal.processing import (
72 PromptUpdate,
73 )
74 from vllm.multimodal.profiling import BaseDummyInputsBuilder
72 -from vllm.platforms import current_platform
75 from vllm.sequence import IntermediateTensors
76 from vllm.utils.tensor_schema import TensorSchema, TensorShape
77
@@ -89,52 +91,6 @@ logger = init_logger(__name__)
91 # === Vision Transformer === #
92
93
92 -def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
93 - if not interleaved:
94 - x1, x2 = x.chunk(2, dim=-1)
95 - return torch.cat((-x2, x1), dim=-1)
96 - else:
97 - x1, x2 = x[..., ::2], x[..., 1::2]
98 - return rearrange(
99 - torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2
100 - )
101 -
102 -
103 -def apply_rotary_emb_torch(
104 - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
105 -) -> torch.Tensor:
106 - """
107 - x: (batch_size, seqlen, nheads, headdim)
108 - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
109 - """
110 - ro_dim = cos.shape[-1] * 2
111 - assert ro_dim <= x.shape[-1]
112 - cos = repeat(
113 - cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
114 - )
115 - sin = repeat(
116 - sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
117 - )
118 - return torch.cat(
119 - [
120 - x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin,
121 - x[..., ro_dim:],
122 - ],
123 - dim=-1,
124 - )
125 -
126 -
127 -def apply_rotary_pos_emb_vision(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
128 - t_ = t.float()
129 - cos = freqs.cos()
130 - sin = freqs.sin()
131 - apply_rotary_emb = apply_rotary_emb_torch
132 - if current_platform.is_cuda():
133 - from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
134 - output = apply_rotary_emb(t_, cos, sin).type_as(t)
135 - return output
136 -
137 -
94 def all_gather_interleave(local_tensor, hidden_size: int, tp_size: int):
95 """All-gather the input tensor interleavely across model parallel group."""
96 import torch.distributed as dist
@@ -200,6 +156,11 @@ class Ernie4_5_VisionAttention(nn.Module):
156 prefix=f"{prefix}.attn",
157 )
158
159 + self.apply_rotary_emb = ApplyRotaryEmb(
160 + enforce_enable=True,
161 + enable_fp32_compute=True,
162 + )
163 +
164 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
165 # [s, b, 3 * head * head_dim]
166 seq_len, bs, _ = qkv.shape
@@ -244,7 +205,11 @@ class Ernie4_5_VisionAttention(nn.Module):
205 q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v))
206 if rotary_pos_emb is not None:
207 qk_concat = torch.cat([q, k], dim=0)
247 - qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
208 + qk_rotated = self.apply_rotary_emb(
209 + qk_concat,
210 + rotary_pos_emb.cos(),
211 + rotary_pos_emb.sin(),
212 + )
213 q, k = torch.chunk(qk_rotated, 2, dim=0)
214
215 output = self.attn(
vllm/model_executor/models/glm4_1v.py
+10 -3
@@ -65,6 +65,9 @@ from vllm.model_executor.layers.linear import (
65 )
66 from vllm.model_executor.layers.quantization import QuantizationConfig
67 from vllm.model_executor.layers.rotary_embedding import get_rope
68 +from vllm.model_executor.layers.rotary_embedding.common import (
69 + ApplyRotaryEmb,
70 +)
71 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
72 from vllm.model_executor.models.module_mapping import MultiModelKeys
73 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -95,7 +98,7 @@ from .interfaces import (
98 SupportsMultiModal,
99 SupportsPP,
100 )
98 -from .qwen2_vl import _create_qwen2vl_field_factory, apply_rotary_pos_emb_vision
101 +from .qwen2_vl import _create_qwen2vl_field_factory
102 from .utils import (
103 AutoWeightsLoader,
104 WeightsMapper,
@@ -304,6 +307,8 @@ class Glm4vVisionAttention(nn.Module):
307 multimodal_config=multimodal_config,
308 )
309
310 + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)
311 +
312 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
313 # [s, b, 3 * head * head_dim]
314 seq_len, bs, _ = qkv.shape
@@ -339,8 +344,10 @@ class Glm4vVisionAttention(nn.Module):
344 if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
345 # [2 * b, s, heads, head_dim]
346 qk_concat = torch.cat([q, k], dim=0)
342 - qk_rotated = apply_rotary_pos_emb_vision(
343 - qk_concat, rotary_pos_emb_cos, rotary_pos_emb_sin
347 + qk_rotated = self.apply_rotary_emb(
348 + qk_concat,
349 + rotary_pos_emb_cos,
350 + rotary_pos_emb_sin,
351 )
352 q, k = torch.chunk(qk_rotated, 2, dim=0)
353
vllm/model_executor/models/keye.py
+9 -13
@@ -30,6 +30,9 @@ from vllm.model_executor.layers.linear import (
30 RowParallelLinear,
31 )
32 from vllm.model_executor.layers.quantization import QuantizationConfig
33 +from vllm.model_executor.layers.rotary_embedding.common import (
34 + ApplyRotaryEmb,
35 +)
36 from vllm.model_executor.model_loader.weight_utils import (
37 default_weight_loader,
38 maybe_remap_kv_scale_name,
@@ -59,7 +62,6 @@ from vllm.multimodal.processing import (
62 PromptUpdate,
63 )
64 from vllm.multimodal.profiling import BaseDummyInputsBuilder
62 -from vllm.platforms import current_platform
65 from vllm.sequence import IntermediateTensors
66 from vllm.utils.tensor_schema import TensorSchema, TensorShape
67
@@ -341,20 +343,14 @@ def apply_rotary_pos_emb_flashatt(
343 cos = cos.chunk(2, dim=-1)[0].contiguous()
344 sin = sin.chunk(2, dim=-1)[0].contiguous()
345
344 - if current_platform.is_cuda():
345 - from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
346 - elif current_platform.is_rocm():
347 - from flash_attn.ops.triton.rotary import apply_rotary as apply_rotary_emb
348 - else:
349 - # For other platforms, use PyTorch fallback
350 - from vllm.model_executor.layers.rotary_embedding.common import (
351 - apply_rotary_emb_torch,
352 - )
346 + apply_rotary_emb = ApplyRotaryEmb(
347 + enforce_enable=True,
348 + enable_fp32_compute=True,
349 + )
350
354 - apply_rotary_emb = partial(apply_rotary_emb_torch, is_neox_style=True)
351 + q_embed = apply_rotary_emb(q, cos, sin)
352 + k_embed = apply_rotary_emb(k, cos, sin)
353
356 - q_embed = apply_rotary_emb(q.float(), cos.float(), sin.float()).type_as(q)
357 - k_embed = apply_rotary_emb(k.float(), cos.float(), sin.float()).type_as(k)
354 return q_embed, k_embed
355
356
vllm/model_executor/models/paddleocr_vl.py
+11 -44
@@ -22,7 +22,7 @@ from typing import Annotated, Literal
22 import numpy as np
23 import torch
24 import torch.nn as nn
25 -from einops import rearrange, repeat
25 +from einops import rearrange
26 from transformers import BatchFeature, PretrainedConfig
27 from transformers.activations import GELUActivation
28 from transformers.modeling_outputs import (
@@ -47,7 +47,7 @@ from vllm.model_executor.layers.linear import (
47 )
48 from vllm.model_executor.layers.quantization import QuantizationConfig
49 from vllm.model_executor.layers.rotary_embedding.common import (
50 - dispatch_rotary_emb_function,
50 + ApplyRotaryEmb,
51 )
52 from vllm.model_executor.model_loader.weight_utils import (
53 default_weight_loader,
@@ -130,47 +130,6 @@ def smart_resize(
130 return h_bar, w_bar
131
132
133 -def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
134 - if not interleaved:
135 - x1, x2 = x.chunk(2, dim=-1)
136 - return torch.cat((-x2, x1), dim=-1)
137 - x1, x2 = x[..., ::2], x[..., 1::2]
138 - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2)
139 -
140 -
141 -def apply_rotary_emb_torch(
142 - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
143 -) -> torch.Tensor:
144 - """
145 - x: (batch_size, seqlen, nheads, headdim)
146 - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
147 - """
148 - ro_dim = cos.shape[-1] * 2
149 - assert ro_dim <= x.shape[-1]
150 - cos = repeat(
151 - cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
152 - )
153 - sin = repeat(
154 - sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
155 - )
156 - return torch.cat(
157 - [
158 - x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin,
159 - x[..., ro_dim:],
160 - ],
161 - dim=-1,
162 - )
163 -
164 -
165 -def apply_rotary_pos_emb_vision(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
166 - rotary_emb_function = dispatch_rotary_emb_function(default=apply_rotary_emb_torch)
167 - t_ = t.float()
168 - cos = freqs.cos()
169 - sin = freqs.sin()
170 - output = rotary_emb_function(t_, cos, sin).type_as(t)
171 - return output
172 -
173 -
133 class PaddleOCRVLProcessingInfo(BaseProcessingInfo):
134 def get_hf_config(self):
135 return self.ctx.get_hf_config()
@@ -609,6 +568,10 @@ class SiglipAttention(nn.Module):
568 multimodal_config=multimodal_config,
569 prefix=f"{prefix}.attn",
570 )
571 + self.apply_rotary_emb = ApplyRotaryEmb(
572 + enforce_enable=True,
573 + enable_fp32_compute=True,
574 + )
575
576 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
577 seq_len, bs, _ = qkv.shape
@@ -651,7 +614,11 @@ class SiglipAttention(nn.Module):
614
615 if rotary_pos_emb is not None:
616 qk_concat = torch.cat([q, k], dim=0)
654 - qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
617 + qk_rotated = self.apply_rotary_emb(
618 + qk_concat,
619 + rotary_pos_emb.cos(),
620 + rotary_pos_emb.sin(),
621 + )
622 q, k = torch.chunk(qk_rotated, 2, dim=0)
623
624 context_layer = self.attn(
vllm/model_executor/models/qwen2_5_vl.py
+9 -3
@@ -60,6 +60,9 @@ from vllm.model_executor.layers.linear import (
60 )
61 from vllm.model_executor.layers.quantization import QuantizationConfig
62 from vllm.model_executor.layers.rotary_embedding import get_rope
63 +from vllm.model_executor.layers.rotary_embedding.common import (
64 + ApplyRotaryEmb,
65 +)
66 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
67 from vllm.model_executor.models.module_mapping import MultiModelKeys
68 from vllm.model_executor.models.vision import should_torch_compile_mm_vit
@@ -95,7 +98,6 @@ from .qwen2_vl import Qwen2VLDummyInputsBuilder as Qwen2_5_VLDummyInputsBuilder
98 from .qwen2_vl import (
99 Qwen2VLMultiModalProcessor,
100 Qwen2VLProcessingInfo,
98 - apply_rotary_pos_emb_vision,
101 )
102 from .utils import (
103 AutoWeightsLoader,
@@ -353,6 +355,8 @@ class Qwen2_5_VisionAttention(nn.Module):
355 multimodal_config=multimodal_config,
356 )
357
358 + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)
359 +
360 def forward(
361 self,
362 x: torch.Tensor,
@@ -378,8 +382,10 @@ class Qwen2_5_VisionAttention(nn.Module):
382 qk_reshaped = einops.rearrange(
383 qk, "b s two head head_dim -> (two b) s head head_dim", two=2
384 )
381 - qk_rotated = apply_rotary_pos_emb_vision(
382 - qk_reshaped, cos=rotary_pos_emb_cos, sin=rotary_pos_emb_sin
385 + qk_rotated = self.apply_rotary_emb(
386 + qk_reshaped,
387 + rotary_pos_emb_cos,
388 + rotary_pos_emb_sin,
389 )
390 qk_rotated = qk_rotated.view(
391 2,
vllm/model_executor/models/qwen2_vl.py
+7 -14
@@ -59,8 +59,7 @@ from vllm.model_executor.layers.linear import (
59 from vllm.model_executor.layers.quantization import QuantizationConfig
60 from vllm.model_executor.layers.rotary_embedding import get_rope
61 from vllm.model_executor.layers.rotary_embedding.common import (
62 - apply_rotary_emb_torch,
63 - dispatch_rotary_emb_function,
62 + ApplyRotaryEmb,
63 )
64 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
65 from vllm.model_executor.models.module_mapping import MultiModelKeys
@@ -280,16 +279,6 @@ class Qwen2VisionMLP(nn.Module):
279 return x
280
281
283 -def apply_rotary_pos_emb_vision(
284 - t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
285 -) -> torch.Tensor:
286 - rotary_emb_function = dispatch_rotary_emb_function(
287 - default=partial(apply_rotary_emb_torch, is_neox_style=True)
288 - )
289 - output = rotary_emb_function(t, cos, sin).type_as(t)
290 - return output
291 -
292 -
282 class Qwen2VisionAttention(nn.Module):
283 def __init__(
284 self,
@@ -341,6 +330,8 @@ class Qwen2VisionAttention(nn.Module):
330 multimodal_config=multimodal_config,
331 )
332
333 + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)
334 +
335 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
336 # [s, b, 3 * head * head_dim]
337 seq_len, bs, _ = qkv.shape
@@ -387,8 +378,10 @@ class Qwen2VisionAttention(nn.Module):
378
379 # [2 * b, s, heads, head_dim]
380 qk_concat = torch.cat([q, k], dim=0)
390 - qk_rotated = apply_rotary_pos_emb_vision(
391 - qk_concat, rotary_pos_emb_cos, rotary_pos_emb_sin
381 + qk_rotated = self.apply_rotary_emb(
382 + qk_concat,
383 + rotary_pos_emb_cos,
384 + rotary_pos_emb_sin,
385 )
386 q, k = torch.chunk(qk_rotated, 2, dim=0)
387
vllm/model_executor/models/siglip2navit.py
+15 -41
@@ -6,7 +6,6 @@ within a vision language model."""
6 from collections.abc import Iterable
7
8 import torch
9 -from einops import rearrange, repeat
9 from torch import nn
10 from torch.nn import functional as F
11 from transformers import Siglip2VisionConfig
@@ -26,6 +25,9 @@ from vllm.model_executor.layers.linear import (
25 RowParallelLinear,
26 )
27 from vllm.model_executor.layers.quantization import QuantizationConfig
28 +from vllm.model_executor.layers.rotary_embedding.common import (
29 + ApplyRotaryEmb,
30 +)
31 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
32 from vllm.platforms import current_platform
33
@@ -146,40 +148,6 @@ class Siglip2VisionEmbeddings(nn.Module):
148 return patch_embeds
149
150
149 -# copy from flash_attn/layers/rotary.py
150 -def rotate_half(x, interleaved=False):
151 - if not interleaved:
152 - x1, x2 = x.chunk(2, dim=-1)
153 - return torch.cat((-x2, x1), dim=-1)
154 - else:
155 - x1, x2 = x[..., ::2], x[..., 1::2]
156 - return rearrange(
157 - torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2
158 - )
159 -
160 -
161 -def apply_rotary_emb_torch(x, cos, sin, interleaved=False):
162 - """
163 - x: (batch_size, seqlen, nheads, headdim)
164 - cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
165 - """
166 - ro_dim = cos.shape[-1] * 2
167 - assert ro_dim <= x.shape[-1]
168 - cos = repeat(
169 - cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
170 - )
171 - sin = repeat(
172 - sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
173 - )
174 - return torch.cat(
175 - [
176 - x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin,
177 - x[..., ro_dim:],
178 - ],
179 - dim=-1,
180 - )
181 -
182 -
151 def apply_rotary_pos_emb(
152 q: torch.Tensor,
153 k: torch.Tensor,
@@ -189,14 +157,20 @@ def apply_rotary_pos_emb(
157 ) -> tuple[torch.Tensor, torch.Tensor]:
158 cos = cos.chunk(2, dim=-1)[0].contiguous()
159 sin = sin.chunk(2, dim=-1)[0].contiguous()
192 - if is_flash_attn_backend and current_platform.is_cuda():
193 - from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
160
195 - apply_rotary_emb_func = apply_rotary_emb
161 + apply_rotary_emb = ApplyRotaryEmb(
162 + enforce_enable=True,
163 + enable_fp32_compute=True,
164 + )
165 +
166 + if is_flash_attn_backend and not current_platform.is_cuda():
167 + apply_rotary_emb_func = apply_rotary_emb.forward_cuda
168 else:
197 - apply_rotary_emb_func = apply_rotary_emb_torch
198 - q_embed = apply_rotary_emb_func(q.float(), cos.float(), sin.float()).type_as(q)
199 - k_embed = apply_rotary_emb_func(k.float(), cos.float(), sin.float()).type_as(k)
169 + apply_rotary_emb_func = apply_rotary_emb.forward_native
170 +
171 + q_embed = apply_rotary_emb_func(q, cos, sin)
172 + k_embed = apply_rotary_emb_func(k, cos, sin)
173 +
174 return q_embed, k_embed
175
176