diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index f85031ac4..72de7319e 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -17,6 +17,7 @@ from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, get_compiler_backend, + get_device, is_cpu, is_cuda, is_hip, @@ -822,7 +823,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding): beta_slow: int = 1, mscale: float = 1, mscale_all_dim: float = 0, - device: Optional[str] = "cuda" if not _is_npu else "npu", + device: Optional[str] = None, ) -> None: self.scaling_factor = scaling_factor self.extrapolation_factor = extrapolation_factor @@ -839,7 +840,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding): self.sin_cached_total = None self.cos_cached = None self.sin_cached = None - self.device = device + self.device = device if device is not None else get_device() super().__init__( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype )