[Intel GPU]Add sycl mrope pass for xpu device (#27646)
This commit is contained in:
@@ -19,10 +19,17 @@ from sglang.srt.layers.rotary_embedding.yarn import (
|
|||||||
yarn_linear_ramp_mask,
|
yarn_linear_ramp_mask,
|
||||||
)
|
)
|
||||||
from sglang.srt.server_args import get_global_server_args
|
from sglang.srt.server_args import get_global_server_args
|
||||||
from sglang.srt.utils import cpu_has_amx_support, is_cuda, is_npu, support_triton
|
from sglang.srt.utils import (
|
||||||
|
cpu_has_amx_support,
|
||||||
|
is_cuda,
|
||||||
|
is_npu,
|
||||||
|
is_xpu,
|
||||||
|
support_triton,
|
||||||
|
)
|
||||||
|
|
||||||
_is_cuda = is_cuda()
|
_is_cuda = is_cuda()
|
||||||
_is_npu = is_npu()
|
_is_npu = is_npu()
|
||||||
|
_is_xpu = is_xpu()
|
||||||
_is_cpu_amx_available = cpu_has_amx_support()
|
_is_cpu_amx_available = cpu_has_amx_support()
|
||||||
|
|
||||||
if _is_cuda:
|
if _is_cuda:
|
||||||
@@ -31,6 +38,8 @@ if _is_cuda:
|
|||||||
if _is_npu:
|
if _is_npu:
|
||||||
import torch_npu
|
import torch_npu
|
||||||
|
|
||||||
|
if _is_xpu:
|
||||||
|
from sgl_kernel import multimodal_rotary_embedding
|
||||||
|
|
||||||
import triton
|
import triton
|
||||||
import triton.language as tl
|
import triton.language as tl
|
||||||
@@ -377,7 +386,20 @@ class MRotaryEmbedding(RotaryEmbedding):
|
|||||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||||
assert positions.ndim in (1, 2)
|
assert positions.ndim in (1, 2)
|
||||||
if positions.ndim == 2 and self.mrope_section:
|
if positions.ndim == 2 and self.mrope_section:
|
||||||
return self.forward_triton(positions, query, key)
|
multimodal_rotary_embedding(
|
||||||
|
query,
|
||||||
|
key,
|
||||||
|
self.cos_sin_cache,
|
||||||
|
positions,
|
||||||
|
self.mrope_section,
|
||||||
|
self.head_size,
|
||||||
|
self.rotary_dim,
|
||||||
|
self.mrope_interleaved,
|
||||||
|
self.mrope_interleaved_glm,
|
||||||
|
self.is_neox_style,
|
||||||
|
self.axis_map,
|
||||||
|
)
|
||||||
|
return query, key
|
||||||
return self.forward_native(positions, query, key, fused_set_kv_buffer_arg)
|
return self.forward_native(positions, query, key, fused_set_kv_buffer_arg)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
|||||||
Reference in New Issue
Block a user