From 6a3316dd1ec95a921118afa03ad0961b7cb70f0d Mon Sep 17 00:00:00 2001 From: DevashishLal-CB Date: Fri, 5 Jun 2026 16:55:50 -0700 Subject: [PATCH] [plugin] default device detection fixes for OOT platform plugins (#25337) Signed-off-by: Devashish Lal Co-authored-by: Devashish Lal --- python/sglang/srt/configs/device_config.py | 4 +++- python/sglang/srt/layers/rotary_embedding/base.py | 2 ++ .../model_executor/model_runner_kv_cache_mixin.py | 4 +--- python/sglang/srt/platforms/__init__.py | 2 +- python/sglang/srt/platforms/device_mixin.py | 4 ++-- python/sglang/srt/server_args.py | 7 +------ python/sglang/srt/utils/common.py | 14 +++++++------- 7 files changed, 17 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/configs/device_config.py b/python/sglang/srt/configs/device_config.py index 9836f935c..9d472af9c 100644 --- a/python/sglang/srt/configs/device_config.py +++ b/python/sglang/srt/configs/device_config.py @@ -3,6 +3,8 @@ from typing import Optional import torch +from sglang.srt.platforms import current_platform + logger = logging.getLogger(__name__) SUPPORTED_DEVICES = ["cuda", "xpu", "hpu", "cpu", "npu", "musa", "mps"] @@ -13,7 +15,7 @@ class DeviceConfig: gpu_id: Optional[int] def __init__(self, device: str = "cuda", gpu_id: int = -1) -> None: - if device in SUPPORTED_DEVICES: + if device in SUPPORTED_DEVICES or current_platform.is_out_of_tree(): self.device_type = device else: raise RuntimeError(f"Not supported device type: {device}") diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index ebfdb51f8..2928812b2 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/python/sglang/srt/layers/rotary_embedding/base.py @@ -8,6 +8,7 @@ import torch from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.srt.layers.utils import MultiPlatformOp +from sglang.srt.platforms import current_platform from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( cpu_has_amx_support, @@ -82,6 +83,7 @@ class RotaryEmbedding(MultiPlatformOp): and not (_is_npu) and not (_is_musa) and not (_is_mps) + and not (current_platform.is_out_of_tree()) ): # rotary_embedding from sglang.jit_kernel.rope and vllm._custom_ops has the same implementation. # TODO: Test on different devices and remove this conditional. diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index b3683bec2..7072103eb 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -36,6 +36,7 @@ from sglang.srt.mem_cache.memory_pool import ( ReqToTokenPool, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.srt.platforms import current_platform from sglang.srt.utils.common import ( get_available_gpu_memory, is_float4_e2m1fn_x2, @@ -382,9 +383,6 @@ class ModelRunnerKVCacheMixin: is_dsa_model = is_deepseek_dsa(self.model_config.hf_config) is_dsv4_model = is_deepseek_v4(self.model_config.hf_config) - # Out-of-tree platform plugin system — used by elif below - from sglang.srt.platforms import current_platform - self._validate_prefill_only_disable_kv_cache_pool_family( is_dsa_model, is_dsv4_model, current_platform ) diff --git a/python/sglang/srt/platforms/__init__.py b/python/sglang/srt/platforms/__init__.py index 9c68eefa5..c64a187a7 100644 --- a/python/sglang/srt/platforms/__init__.py +++ b/python/sglang/srt/platforms/__init__.py @@ -139,7 +139,7 @@ def _load_platform_class(qualname: str) -> type: return cls -def __getattr__(name: str): +def __getattr__(name: str) -> SRTPlatform: """Lazy initialization of current_platform on first access.""" if name == "current_platform": global _current_platform diff --git a/python/sglang/srt/platforms/device_mixin.py b/python/sglang/srt/platforms/device_mixin.py index c42c4b2bd..f5fa42ae6 100644 --- a/python/sglang/srt/platforms/device_mixin.py +++ b/python/sglang/srt/platforms/device_mixin.py @@ -155,8 +155,8 @@ class DeviceMixin: # ---- Device management ---- - def get_device(self, local_rank: int) -> "torch.device": - """[Planned] Return ``torch.device`` for the given local rank.""" + def get_device(self, device_id: int = 0) -> str: + """[Planned] Return ``torch.device`` for the given device id.""" raise NotImplementedError def set_device(self, device: "torch.device") -> None: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index be1b62874..82a3d31a2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -49,6 +49,7 @@ from sglang.srt.function_call.function_call_parser import FunctionCallParser from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.parser.reasoning_parser import ReasoningParser +from sglang.srt.platforms import current_platform from sglang.srt.utils.common import ( LORA_TARGET_ALL_MODULES, SUPPORTED_LORA_TARGET_MODULES, @@ -944,8 +945,6 @@ class ServerArgs: self._handle_xpu_backends() # Allow OOT platform plugins to apply server args defaults. - from sglang.srt.platforms import current_platform - current_platform.apply_server_args_defaults(self) # Handle piecewise CUDA graph. @@ -1359,8 +1358,6 @@ class ServerArgs: if is_hip() or is_npu() or is_cpu() or is_mps() or is_xpu(): self.disable_piecewise_cuda_graph = True # 5b. OOT platforms that don't support piecewise cuda graph - from sglang.srt.platforms import current_platform - if current_platform.is_out_of_tree(): if not current_platform.support_piecewise_cuda_graph(): self.disable_piecewise_cuda_graph = True @@ -2744,8 +2741,6 @@ class ServerArgs: 2.3 Otherwise, we will use triton backend. """ # OOT platforms provide their own default attention backend. - from sglang.srt.platforms import current_platform - if current_platform.is_out_of_tree(): return current_platform.get_default_attention_backend() diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index d90bfaef7..3606bc36a 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -94,6 +94,7 @@ from typing_extensions import Literal from sglang.srt.environ import envs from sglang.srt.observability.func_timer import enable_func_timer +from sglang.srt.platforms import current_platform from sglang.srt.utils.video_decoder import _BACKEND, VideoDecoderWrapper if TYPE_CHECKING: @@ -658,8 +659,6 @@ def get_available_gpu_memory( elif device == "mps": free_gpu_memory = psutil.virtual_memory().available else: - from sglang.srt.platforms import current_platform - if not current_platform.is_out_of_tree(): raise ValueError( f"Unsupported device type: {device!r}. " @@ -1893,8 +1892,6 @@ def get_mtgpu_memory_capacity(): def get_device_memory_capacity(device: str = None): # OOT platforms provide their own memory query via the platform class. - from sglang.srt.platforms import current_platform - if current_platform.is_out_of_tree(): mem_bytes = current_platform.get_device_total_memory() if mem_bytes: @@ -2079,7 +2076,12 @@ def get_device(device_id: Optional[int] = None) -> str: return "mps" return "mps:{}".format(device_id) - raise RuntimeError("No accelerator (CUDA, XPU, HPU, NPU, MUSA, MPS) is available.") + try: + return current_platform.get_device(device_id) + except Exception: + raise RuntimeError( + "No accelerator (CUDA, XPU, HPU, NPU, MUSA, MPS) or platform plugin is available." + ) @lru_cache(maxsize=1) @@ -2145,8 +2147,6 @@ def get_device_capability(device_id: int = 0) -> Tuple[int, int]: def get_compiler_backend(mode=None) -> str: # OOT platforms provide their own compile backend. - from sglang.srt.platforms import current_platform - if current_platform.is_out_of_tree(): return current_platform.get_compile_backend(mode)