diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index f931098be..9e30c6f9e 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -19,12 +19,15 @@ from sglang.srt.utils import ( _is_cuda = is_cuda() _is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_npu = is_npu() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() if _is_cuda: 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(): import torch_npu @@ -1428,7 +1431,8 @@ def get_rope_wrapper( device: Optional[str] = None, ): if device != "cpu": - return get_rope( + wrapper = aiter_get_rope if _use_aiter else get_rope + return wrapper( head_size, rotary_dim, max_position, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 10079227d..c6617858e 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -125,8 +125,6 @@ if _is_hip: decode_attention_fwd_grouped_rope, ) -if _use_aiter: - from aiter.rotary_embedding import get_rope logger = logging.getLogger(__name__)