[XPU] Wire up SYCL apply_rotary_emb kernel in ApplyRotaryEmb (#55721)
Signed-off-by: Michal Ganczarenko <michal.ganczarenko@intel.com> Signed-off-by: Michał Ganczarenko <michal.ganczarenko@intel.com> Co-authored-by: Claude <noreply@anthropic.com>
Michał Ganczarenko committed
Sep 23, 2026 at 05:43 UTC
73c59d365b9e1e722fbea5e897e8efb646e3e513
3 files changed
+78
tests/kernels/core/test_apply_rotary_emb.py
+41
@@ -184,3 +184,44 @@ def test_rotary_embedding_dispatch(
184
- forward_cuda/forward methods should call ApplyRotaryEmb.forward()
185
"""
186
run_dispatch_test(test_case, device)
187
+
188
+
189
+@pytest.mark.skipif(not current_platform.is_xpu(), reason="XPU only test.")
190
+@pytest.mark.parametrize(
191
+ "num_tokens,num_heads,head_size,rot_dim,is_neox_style",
192
+ [
193
+ (128, 32, 128, 128, True), # typical text shape
194
+ (4096, 16, 80, 80, False), # GPT-J style (interleaved)
195
+ (4096, 32, 128, 64, True), # partial rotary (rot_dim < head_size)
196
+ ],
197
+ ids=["text", "gptj_interleaved", "partial_rotary"],
198
+)
199
+def test_apply_rotary_emb_xpu_matches_native(
200
+ num_tokens: int,
201
+ num_heads: int,
202
+ head_size: int,
203
+ rot_dim: int,
204
+ is_neox_style: bool,
205
+):
206
+ """ApplyRotaryEmb.forward_xpu (vllm_xpu_kernels SYCL op) must match
207
+ forward_native (the CustomOp default fallback) exactly."""
208
+ from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
209
+
210
+ vllm_config = VllmConfig(
211
+ compilation_config=CompilationConfig(custom_ops=["all", "+apply_rotary_emb"])
212
+ )
213
+ get_cached_compilation_config.cache_clear()
214
+
215
+ with set_current_vllm_config(vllm_config):
216
+ op = ApplyRotaryEmb(enforce_enable=True, is_neox_style=is_neox_style)
217
+
218
+ x = torch.randn(
219
+ num_tokens, num_heads, head_size, device="xpu", dtype=torch.bfloat16
220
+ )
221
+ x = x[..., :rot_dim].contiguous()
222
+ cos = torch.randn(num_tokens, rot_dim // 2, device="xpu", dtype=torch.bfloat16)
223
+ sin = torch.randn(num_tokens, rot_dim // 2, device="xpu", dtype=torch.bfloat16)
224
+
225
+ expected = op.forward_native(x, cos, sin)
226
+ actual = op.forward_xpu(x, cos, sin)
227
+ torch.testing.assert_close(actual, expected, atol=0, rtol=0)
vllm/_xpu_ops.py
+25
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
6
7
import torch
8
from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func
9
+from vllm_xpu_kernels.rotary import apply_rotary_emb
10
11
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
12
from vllm.logger import init_logger
@@ -249,6 +250,24 @@ def _xpu_ops_deepseek_scaling_rope_fake(
250
return query, key
251
252
253
+def _xpu_apply_rotary_emb_impl(
254
+ x: torch.Tensor,
255
+ cos: torch.Tensor,
256
+ sin: torch.Tensor,
257
+ is_neox_style: bool,
258
+) -> torch.Tensor:
259
+ return apply_rotary_emb(x, cos, sin, is_neox_style)
260
+
261
+
262
+def _xpu_apply_rotary_emb_fake(
263
+ x: torch.Tensor,
264
+ cos: torch.Tensor,
265
+ sin: torch.Tensor,
266
+ is_neox_style: bool,
267
+) -> torch.Tensor:
268
+ return torch.empty_like(x)
269
+
270
+
271
def _xpu_fp8_bmm_impl(
272
a: torch.Tensor,
273
b: torch.Tensor,
@@ -1257,6 +1276,12 @@ class xpu_ops:
1276
dispatch_key=current_platform.dispatch_key,
1277
)
1278
1279
+ direct_register_custom_op(
1280
+ op_name="xpu_apply_rotary_emb",
1281
+ op_func=_xpu_apply_rotary_emb_impl,
1282
+ fake_impl=_xpu_apply_rotary_emb_fake,
1283
+ )
1284
+
1285
direct_register_custom_op(
1286
op_name="xpu_mxfp8_quantize",
1287
op_func=_xpu_mxfp8_quantize_impl,
vllm/model_executor/layers/rotary_embedding/common.py
+12
@@ -237,6 +237,18 @@ class ApplyRotaryEmb(CustomOp):
237
output = self._post_process(output, origin_shape, origin_dtype)
238
return output
239
240
+ def forward_xpu(
241
+ self,
242
+ x: torch.Tensor,
243
+ cos: torch.Tensor,
244
+ sin: torch.Tensor,
245
+ ) -> torch.Tensor:
246
+ import vllm._xpu_ops # noqa: F401 registers torch.ops.vllm.xpu_apply_rotary_emb
247
+
248
+ x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin)
249
+ output = torch.ops.vllm.xpu_apply_rotary_emb(x, cos, sin, self.is_neox_style)
250
+ return self._post_process(output, origin_shape, origin_dtype)
251
+
252
def forward_hip(
253
self,
254
x: torch.Tensor,