EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

take aiter get_rope back (#7521)

Browse Source
This commit is contained in:
valarLip
2025-06-25 17:03:33 +08:00
committed by GitHub
parent 4c6675c4fc
commit 802815e40b
2 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -19,12 +19,15 @@ from sglang.srt.utils import (
_is_cuda = is_cuda() _is_cuda = is_cuda()
_is_hip = is_hip() _is_hip = is_hip()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_is_npu = is_npu() _is_npu = is_npu()
_is_cpu_amx_available = cpu_has_amx_support() _is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu() _is_cpu = is_cpu()
if _is_cuda: if _is_cuda:
from sgl_kernel import apply_rope_with_cos_sin_cache_inplace from sgl_kernel import apply_rope_with_cos_sin_cache_inplace
if _use_aiter:
from aiter.rotary_embedding import get_rope as aiter_get_rope
if is_npu(): if is_npu():
import torch_npu import torch_npu
@@ -1428,7 +1431,8 @@ def get_rope_wrapper(
device: Optional[str] = None, device: Optional[str] = None,
): ):
if device != "cpu": if device != "cpu":
return get_rope( wrapper = aiter_get_rope if _use_aiter else get_rope
return wrapper(
head_size, head_size,
rotary_dim, rotary_dim,
max_position, max_position,

2
python/sglang/srt/models/deepseek_v2.py
View File

@@ -125,8 +125,6 @@ if _is_hip:
decode_attention_fwd_grouped_rope, decode_attention_fwd_grouped_rope,
) )
if _use_aiter:
from aiter.rotary_embedding import get_rope
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)