Introduce CudaDeviceMixin and CudaSRTPlatform (#24096)

This commit is contained in:
Alex Nails
2026-05-15 10:59:02 -07:00
committed by GitHub
parent ee93795476
commit 4df42da658
10 changed files with 318 additions and 45 deletions
+3 -3
View File
@@ -203,6 +203,7 @@ from sglang.srt.observability.scheduler_metrics_mixin import (
) )
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.platforms import current_platform
from sglang.srt.plugins import load_plugins from sglang.srt.plugins import load_plugins
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
@@ -214,7 +215,6 @@ from sglang.srt.utils import (
broadcast_pyobj, broadcast_pyobj,
configure_gc_logger, configure_gc_logger,
configure_logger, configure_logger,
empty_device_cache,
freeze_gc, freeze_gc,
get_available_gpu_memory, get_available_gpu_memory,
get_bool_env_var, get_bool_env_var,
@@ -3454,7 +3454,7 @@ class Scheduler(
self.draft_worker.clear_cache_pool() self.draft_worker.clear_cache_pool()
if empty_cache: if empty_cache:
empty_device_cache(self.device_module) current_platform.empty_cache()
logger.info("Cache flushed successfully!") logger.info("Cache flushed successfully!")
success = True success = True
else: else:
@@ -3852,7 +3852,7 @@ class IdleSleeper:
and real_time() - self.last_empty_time > self.empty_cache_interval and real_time() - self.last_empty_time > self.empty_cache_interval
): ):
self.last_empty_time = real_time() self.last_empty_time = real_time()
empty_device_cache() current_platform.empty_cache()
def is_health_check_generate_req(recv_req): def is_health_check_generate_req(recv_req):
+12 -12
View File
@@ -406,7 +406,7 @@ class MambaPool:
return dst_index return dst_index
def get_cpu_copy(self, indices): def get_cpu_copy(self, indices):
torch.cuda.synchronize() current_platform.synchronize()
conv_cpu = [ conv_cpu = [
conv[:, indices].to("cpu", non_blocking=True) conv[:, indices].to("cpu", non_blocking=True)
for conv in self.mamba_cache.conv for conv in self.mamba_cache.conv
@@ -414,18 +414,18 @@ class MambaPool:
temporal_cpu = self.mamba_cache.temporal[:, indices].to( temporal_cpu = self.mamba_cache.temporal[:, indices].to(
"cpu", non_blocking=True "cpu", non_blocking=True
) )
torch.cuda.synchronize() current_platform.synchronize()
return conv_cpu, temporal_cpu return conv_cpu, temporal_cpu
def load_cpu_copy(self, mamba_cache_cpu, indices): def load_cpu_copy(self, mamba_cache_cpu, indices):
conv_cpu, temporal_cpu = mamba_cache_cpu conv_cpu, temporal_cpu = mamba_cache_cpu
torch.cuda.synchronize() current_platform.synchronize()
for i, conv in enumerate(self.mamba_cache.conv): for i, conv in enumerate(self.mamba_cache.conv):
conv[:, indices] = conv_cpu[i].to(conv.device, non_blocking=True) conv[:, indices] = conv_cpu[i].to(conv.device, non_blocking=True)
self.mamba_cache.temporal[:, indices] = temporal_cpu.to( self.mamba_cache.temporal[:, indices] = temporal_cpu.to(
self.mamba_cache.temporal.device, non_blocking=True self.mamba_cache.temporal.device, non_blocking=True
) )
torch.cuda.synchronize() current_platform.synchronize()
def get_contiguous_buf_infos(self): def get_contiguous_buf_infos(self):
""" """
@@ -982,7 +982,7 @@ class MHATokenToKVPool(KVCache):
return kv_data_ptrs, kv_data_lens, kv_item_lens return kv_data_ptrs, kv_data_lens, kv_item_lens
def get_cpu_copy(self, indices, mamba_indices=None): def get_cpu_copy(self, indices, mamba_indices=None):
torch.cuda.synchronize() current_platform.synchronize()
kv_cache_cpu = [] kv_cache_cpu = []
chunk_size = self.cpu_offloading_chunk_size chunk_size = self.cpu_offloading_chunk_size
for layer_id in range(self.layer_num): for layer_id in range(self.layer_num):
@@ -996,11 +996,11 @@ class MHATokenToKVPool(KVCache):
"cpu", non_blocking=True "cpu", non_blocking=True
) )
kv_cache_cpu[-1].append([k_cpu, v_cpu]) kv_cache_cpu[-1].append([k_cpu, v_cpu])
torch.cuda.synchronize() current_platform.synchronize()
return kv_cache_cpu return kv_cache_cpu
def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
torch.cuda.synchronize() current_platform.synchronize()
chunk_size = self.cpu_offloading_chunk_size chunk_size = self.cpu_offloading_chunk_size
for layer_id in range(self.layer_num): for layer_id in range(self.layer_num):
for i in range(0, len(indices), chunk_size): for i in range(0, len(indices), chunk_size):
@@ -1014,7 +1014,7 @@ class MHATokenToKVPool(KVCache):
v_chunk = v_cpu.to(self.v_buffer[0].device, non_blocking=True) v_chunk = v_cpu.to(self.v_buffer[0].device, non_blocking=True)
self.k_buffer[layer_id][chunk_indices] = k_chunk self.k_buffer[layer_id][chunk_indices] = k_chunk
self.v_buffer[layer_id][chunk_indices] = v_chunk self.v_buffer[layer_id][chunk_indices] = v_chunk
torch.cuda.synchronize() current_platform.synchronize()
def _get_key_buffer(self, layer_id: int): def _get_key_buffer(self, layer_id: int):
# for internal use of referencing # for internal use of referencing
@@ -1822,7 +1822,7 @@ class MLATokenToKVPool(KVCache):
return cache_k_nope, cache_k_rope return cache_k_nope, cache_k_rope
def get_cpu_copy(self, indices, mamba_indices=None): def get_cpu_copy(self, indices, mamba_indices=None):
torch.cuda.synchronize() current_platform.synchronize()
kv_cache_cpu = [] kv_cache_cpu = []
chunk_size = self.cpu_offloading_chunk_size chunk_size = self.cpu_offloading_chunk_size
for layer_id in range(self.layer_num): for layer_id in range(self.layer_num):
@@ -1833,11 +1833,11 @@ class MLATokenToKVPool(KVCache):
"cpu", non_blocking=True "cpu", non_blocking=True
) )
kv_cache_cpu[-1].append(kv_cpu) kv_cache_cpu[-1].append(kv_cpu)
torch.cuda.synchronize() current_platform.synchronize()
return kv_cache_cpu return kv_cache_cpu
def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
torch.cuda.synchronize() current_platform.synchronize()
chunk_size = self.cpu_offloading_chunk_size chunk_size = self.cpu_offloading_chunk_size
for layer_id in range(self.layer_num): for layer_id in range(self.layer_num):
for i in range(0, len(indices), chunk_size): for i in range(0, len(indices), chunk_size):
@@ -1846,7 +1846,7 @@ class MLATokenToKVPool(KVCache):
assert kv_cpu.shape[0] == len(chunk_indices) assert kv_cpu.shape[0] == len(chunk_indices)
kv_chunk = kv_cpu.to(self.kv_buffer[0].device, non_blocking=True) kv_chunk = kv_cpu.to(self.kv_buffer[0].device, non_blocking=True)
self.kv_buffer[layer_id][chunk_indices] = kv_chunk self.kv_buffer[layer_id][chunk_indices] = kv_chunk
torch.cuda.synchronize() current_platform.synchronize()
class MLATokenToKVPoolFP4(MLATokenToKVPool): class MLATokenToKVPoolFP4(MLATokenToKVPool):
@@ -1288,7 +1288,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# Single warmup all_reduce to initialize NCCL/RCCL communicator # Single warmup all_reduce to initialize NCCL/RCCL communicator
warmup_tensor = torch.zeros(1, device=torch.cuda.current_device()) warmup_tensor = torch.zeros(1, device=torch.cuda.current_device())
dist.all_reduce(warmup_tensor, group=tp_group_handle) dist.all_reduce(warmup_tensor, group=tp_group_handle)
torch.cuda.synchronize() current_platform.synchronize()
warmup_elapsed = time.perf_counter() - warmup_start warmup_elapsed = time.perf_counter() - warmup_start
logger.info( logger.info(
@@ -1847,7 +1847,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
f"group_rank={group_rank}, world_size={world_size}, group_name={group_name}, backend={backend}" f"group_rank={group_rank}, world_size={world_size}, group_name={group_name}, backend={backend}"
) )
torch.cuda.empty_cache() current_platform.empty_cache()
success = False success = False
message = "" message = ""
try: try:
@@ -1867,7 +1867,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
message = f"Failed to init group: {e}." message = f"Failed to init group: {e}."
logger.error(message) logger.error(message)
torch.cuda.empty_cache() current_platform.empty_cache()
return success, message return success, message
def send_weights_to_remote_instance( def send_weights_to_remote_instance(
@@ -1895,7 +1895,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
logger.error(message) logger.error(message)
return False, message return False, message
torch.cuda.empty_cache() current_platform.empty_cache()
success = False success = False
na = NetworkAddress(master_address, group_port) na = NetworkAddress(master_address, group_port)
message = "" message = ""
@@ -1915,7 +1915,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# destroy the process group after sending weights # destroy the process group after sending weights
del self._weights_send_group[group_name] del self._weights_send_group[group_name]
torch.distributed.distributed_c10d.destroy_process_group(send_group) torch.distributed.distributed_c10d.destroy_process_group(send_group)
torch.cuda.empty_cache() current_platform.empty_cache()
return success, message return success, message
def init_weights_update_group( def init_weights_update_group(
+6 -5
View File
@@ -106,6 +106,7 @@ from sglang.srt.model_loader.weight_utils import (
safetensors_weights_iterator, safetensors_weights_iterator,
set_runai_streamer_env, set_runai_streamer_env,
) )
from sglang.srt.platforms import current_platform
from sglang.srt.utils import ( from sglang.srt.utils import (
get_bool_env_var, get_bool_env_var,
get_device_capability, get_device_capability,
@@ -1221,7 +1222,7 @@ class QuantizedRLModelLoader(DefaultModelLoader):
del current_param_data del current_param_data
if is_last_update: if is_last_update:
gc.collect() gc.collect()
torch.cuda.empty_cache() current_platform.empty_cache()
logger.info("[QuantizedRL] Reload complete") logger.info("[QuantizedRL] Reload complete")
return updated_param_names, is_last_update return updated_param_names, is_last_update
@@ -1912,7 +1913,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
model.load_weights(qweight_iterator) model.load_weights(qweight_iterator)
torch.cuda.empty_cache() current_platform.empty_cache()
param_dict = dict(model.named_parameters()) param_dict = dict(model.named_parameters())
stacked_quant_state_dict: Dict[str, Dict[int, Any]] = {} stacked_quant_state_dict: Dict[str, Dict[int, Any]] = {}
@@ -2220,7 +2221,7 @@ class RemoteInstanceModelLoader(BaseModelLoader):
tp_rank=load_config.tp_rank, tp_rank=load_config.tp_rank,
instance_ip=instance_ip, instance_ip=instance_ip,
) )
torch.cuda.synchronize() current_platform.synchronize()
end_build_group_tic = time.time() end_build_group_tic = time.time()
logger.debug( logger.debug(
f"finish building group for remote instance, time used: {(end_build_group_tic - start_build_group_tic):.4f}s" f"finish building group for remote instance, time used: {(end_build_group_tic - start_build_group_tic):.4f}s"
@@ -2246,7 +2247,7 @@ class RemoteInstanceModelLoader(BaseModelLoader):
src=0, src=0,
group=client._model_update_group, group=client._model_update_group,
) )
torch.cuda.synchronize() current_platform.synchronize()
_post_load_weights(model) _post_load_weights(model)
end_get_weights_tic = time.time() end_get_weights_tic = time.time()
@@ -2257,7 +2258,7 @@ class RemoteInstanceModelLoader(BaseModelLoader):
torch.distributed.distributed_c10d.destroy_process_group( torch.distributed.distributed_c10d.destroy_process_group(
client._model_update_group client._model_update_group
) )
torch.cuda.empty_cache() current_platform.empty_cache()
def load_model_from_remote_instance_by_transfer_engine( def load_model_from_remote_instance_by_transfer_engine(
self, model, transfer_engine, seed_url, tp_rank self, model, transfer_engine, seed_url, tp_rank
+26 -2
View File
@@ -14,8 +14,12 @@ import logging
import pkgutil import pkgutil
from importlib.metadata import entry_points from importlib.metadata import entry_points
import torch
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.platforms.cuda import CudaSRTPlatform
from sglang.srt.platforms.interface import SRTPlatform from sglang.srt.platforms.interface import SRTPlatform
from sglang.srt.platforms.rocm import RocmSRTPlatform
from sglang.srt.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group from sglang.srt.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -23,6 +27,14 @@ logger = logging.getLogger(__name__)
_current_platform: SRTPlatform | None = None _current_platform: SRTPlatform | None = None
def _is_cuda_available() -> bool:
return bool(torch.cuda.is_available() and torch.version.hip is None)
def _is_rocm_available() -> bool:
return bool(torch.cuda.is_available() and torch.version.hip is not None)
def _resolve_platform() -> SRTPlatform: def _resolve_platform() -> SRTPlatform:
""" """
Discover and instantiate the active platform. Discover and instantiate the active platform.
@@ -39,7 +51,9 @@ def _resolve_platform() -> SRTPlatform:
SGLANG_PLATFORM unset (auto-discover): SGLANG_PLATFORM unset (auto-discover):
- Import and activate all discovered plugins - Import and activate all discovered plugins
- 0 activated → fallback base SRTPlatform - 0 activated + CUDA available → fallback CudaSRTPlatform
- 0 activated + ROCm available → fallback RocmSRTPlatform
- 0 activated + neither → fallback base SRTPlatform
- 1 activated → use it - 1 activated → use it
- N activated → RuntimeError (must set SGLANG_PLATFORM) - N activated → RuntimeError (must set SGLANG_PLATFORM)
@@ -90,7 +104,17 @@ def _resolve_platform() -> SRTPlatform:
logger.exception("Failed to activate platform plugin: %s", name) logger.exception("Failed to activate platform plugin: %s", name)
if len(activated) == 0: if len(activated) == 0:
logger.debug("No platform detected. Using base SRTPlatform with defaults.") if _is_cuda_available():
logger.debug(
"No platform plugin detected. Using CUDA SRTPlatform defaults."
)
return CudaSRTPlatform()
if _is_rocm_available():
logger.debug(
"No platform plugin detected. Using ROCm SRTPlatform defaults."
)
return RocmSRTPlatform()
logger.debug("No platform detected. Using base SRTPlatform.")
return SRTPlatform() return SRTPlatform()
if len(activated) == 1: if len(activated) == 1:
+75
View File
@@ -0,0 +1,75 @@
"""CUDA device operations for the SRT platform layer."""
from typing import Optional
import torch
from sglang.srt.platforms.device_mixin import (
DeviceCapability,
DeviceMixin,
PlatformEnum,
)
from sglang.srt.platforms.interface import SRTPlatform
class CudaDeviceMixin(DeviceMixin):
"""CUDA implementation of the shared device operations."""
_enum: PlatformEnum = PlatformEnum.CUDA
device_name: str = "cuda"
device_type: str = "cuda"
def get_device_total_memory(self, device_id: int = 0) -> int:
return int(torch.cuda.get_device_properties(device_id).total_memory)
def get_current_memory_usage(
self, device: Optional["torch.device"] = None
) -> float:
return float(torch.cuda.max_memory_allocated(device))
def get_device(self, local_rank: int) -> "torch.device":
return torch.device("cuda", local_rank)
def set_device(self, device: "torch.device") -> None:
torch.cuda.set_device(device)
def get_device_name(self, device_id: int = 0) -> str:
return str(torch.cuda.get_device_name(device_id))
def get_device_uuid(self, device_id: int = 0) -> str:
return str(torch.cuda.get_device_properties(device_id).uuid)
def get_device_capability(self, device_id: int = 0) -> DeviceCapability:
major, minor = torch.cuda.get_device_capability(device_id)
return DeviceCapability(major, minor)
def empty_cache(self) -> None:
torch.cuda.empty_cache()
def synchronize(self) -> None:
torch.cuda.synchronize()
def get_available_memory(self, device_id: int = 0) -> tuple[int, int]:
return torch.cuda.mem_get_info(device_id)
def get_torch_distributed_backend_str(self) -> str:
return "nccl"
@classmethod
def seed_everything(cls, seed: int | None = None) -> None:
if seed is not None:
super().seed_everything(seed)
torch.cuda.manual_seed_all(seed)
class CudaSRTPlatform(CudaDeviceMixin, SRTPlatform):
"""Default in-tree CUDA SRT platform."""
def supports_fp8(self) -> bool:
return True
def support_cuda_graph(self) -> bool:
return True
def support_piecewise_cuda_graph(self) -> bool:
return True
+4 -10
View File
@@ -26,10 +26,11 @@ Method status annotations:
""" """
import enum import enum
from typing import TYPE_CHECKING, NamedTuple, Optional import random
from typing import NamedTuple, Optional
if TYPE_CHECKING: import numpy as np
import torch import torch
class PlatformEnum(enum.Enum): class PlatformEnum(enum.Enum):
@@ -203,19 +204,12 @@ class DeviceMixin:
@classmethod @classmethod
def inference_mode(cls): def inference_mode(cls):
"""[Planned] Return inference mode context manager.""" """[Planned] Return inference mode context manager."""
import torch
return torch.inference_mode(mode=True) return torch.inference_mode(mode=True)
@classmethod @classmethod
def seed_everything(cls, seed: int | None = None) -> None: def seed_everything(cls, seed: int | None = None) -> None:
"""[Planned] Set random seeds for reproducibility across all libraries.""" """[Planned] Set random seeds for reproducibility across all libraries."""
if seed is not None: if seed is not None:
import random
import numpy as np
import torch
random.seed(seed) random.seed(seed)
np.random.seed(seed) np.random.seed(seed)
torch.manual_seed(seed) torch.manual_seed(seed)
+1 -6
View File
@@ -7,16 +7,11 @@ and adds SRT-specific subsystem factory methods, capability flags, and
configuration lifecycle hooks. configuration lifecycle hooks.
Out-of-tree platforms register via setuptools entry_points under the Out-of-tree platforms register via setuptools entry_points under the
"sglang.platform_plugins" group and should subclass SRTPlatform. "sglang.srt.platforms" group and should subclass SRTPlatform.
""" """
from typing import TYPE_CHECKING
from sglang.srt.platforms.device_mixin import DeviceMixin, PlatformEnum from sglang.srt.platforms.device_mixin import DeviceMixin, PlatformEnum
if TYPE_CHECKING:
pass
# Re-export for convenience # Re-export for convenience
__all__ = ["SRTPlatform", "PlatformEnum"] __all__ = ["SRTPlatform", "PlatformEnum"]
+31
View File
@@ -0,0 +1,31 @@
"""ROCm device operations for the SRT platform layer.
PyTorch exposes ROCm through the same ``torch.cuda.*`` API surface as CUDA
(HIP is a binary shim, and ``torch.device("rocm")`` does not exist). So
``RocmDeviceMixin`` inherits all device ops from ``CudaDeviceMixin`` and
only overrides identity (``_enum``, ``device_name``).
"""
from sglang.srt.platforms.cuda import CudaDeviceMixin
from sglang.srt.platforms.device_mixin import PlatformEnum
from sglang.srt.platforms.interface import SRTPlatform
class RocmDeviceMixin(CudaDeviceMixin):
"""ROCm device ops — identical surface to CUDA via torch.cuda's HIP shim."""
_enum: PlatformEnum = PlatformEnum.ROCM
device_name: str = "rocm"
# device_type stays "cuda" — torch.device("cuda") is the only valid
# device-type string for HIP devices in PyTorch.
class RocmSRTPlatform(RocmDeviceMixin, SRTPlatform):
"""Default in-tree ROCm SRT platform.
Capability flags (supports_fp8, support_cuda_graph, support_piecewise_cuda_graph)
keep the conservative SRTPlatform defaults rather than mirroring CudaSRTPlatform.
They are currently only consulted in OOT branches gated on is_out_of_tree(),
so the defaults are behaviorally inert for the in-tree ROCm path. A follow-up
that migrates AMD-specific gating off legacy is_hip() should set these here.
"""
@@ -7,7 +7,10 @@ and the platform discovery / lazy initialization mechanism.
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
import torch
from sglang.srt.platforms import _load_platform_class, _resolve_platform from sglang.srt.platforms import _load_platform_class, _resolve_platform
from sglang.srt.platforms.cuda import CudaDeviceMixin, CudaSRTPlatform
from sglang.srt.platforms.device_mixin import ( from sglang.srt.platforms.device_mixin import (
CpuArchEnum, CpuArchEnum,
DeviceCapability, DeviceCapability,
@@ -222,6 +225,111 @@ class TestSRTPlatform(CustomTestCase):
base = SRTPlatform() base = SRTPlatform()
self.assertEqual(base.get_compile_backend(mode="npugraph_ex"), "inductor") self.assertEqual(base.get_compile_backend(mode="npugraph_ex"), "inductor")
def test_base_device_identity_stays_unspecified(self):
"""The abstract SRT base should not claim any concrete in-tree device."""
base = SRTPlatform()
self.assertFalse(base.is_cuda())
self.assertFalse(base.is_cuda_alike())
class TestCudaDeviceMixin(CustomTestCase):
"""Tests for CUDA device operation defaults."""
def test_default_get_device_returns_cuda_device(self):
base = CudaSRTPlatform()
self.assertEqual(base.get_device(2), torch.device("cuda", 2))
def test_cuda_platform_identity(self):
base = CudaSRTPlatform()
self.assertTrue(base.is_cuda())
self.assertTrue(base.is_cuda_alike())
self.assertIsInstance(base, CudaDeviceMixin)
@patch("torch.cuda.get_device_properties")
def test_default_get_device_total_memory_uses_cuda(
self, mock_get_device_properties
):
mock_get_device_properties.return_value.total_memory = 123
base = CudaSRTPlatform()
self.assertEqual(base.get_device_total_memory(1), 123)
mock_get_device_properties.assert_called_once_with(1)
@patch("torch.cuda.max_memory_allocated", return_value=456)
def test_default_get_current_memory_usage_uses_cuda(
self, mock_max_memory_allocated
):
base = CudaSRTPlatform()
device = torch.device("cuda", 1)
self.assertEqual(base.get_current_memory_usage(device), 456.0)
mock_max_memory_allocated.assert_called_once_with(device)
@patch("torch.cuda.set_device")
def test_default_set_device_uses_cuda(self, mock_set_device):
base = CudaSRTPlatform()
device = torch.device("cuda", 1)
base.set_device(device)
mock_set_device.assert_called_once_with(device)
@patch("torch.cuda.get_device_name", return_value="NVIDIA H100")
def test_default_get_device_name_uses_cuda(self, mock_get_device_name):
base = CudaSRTPlatform()
self.assertEqual(base.get_device_name(1), "NVIDIA H100")
mock_get_device_name.assert_called_once_with(1)
@patch("torch.cuda.get_device_properties")
def test_default_get_device_uuid_uses_cuda(self, mock_get_device_properties):
mock_get_device_properties.return_value.uuid = "1234"
base = CudaSRTPlatform()
self.assertEqual(base.get_device_uuid(1), "1234")
mock_get_device_properties.assert_called_once_with(1)
@patch("torch.cuda.get_device_capability", return_value=(9, 0))
def test_default_get_device_capability_uses_cuda(self, mock_get_device_capability):
base = CudaSRTPlatform()
self.assertEqual(base.get_device_capability(1), DeviceCapability(9, 0))
mock_get_device_capability.assert_called_once_with(1)
@patch("torch.cuda.empty_cache")
def test_default_empty_cache_uses_cuda(self, mock_empty_cache):
base = CudaSRTPlatform()
base.empty_cache()
mock_empty_cache.assert_called_once_with()
@patch("torch.cuda.synchronize")
def test_default_synchronize_uses_cuda(self, mock_synchronize):
base = CudaSRTPlatform()
base.synchronize()
mock_synchronize.assert_called_once_with()
@patch("torch.cuda.mem_get_info", return_value=(123, 456), create=True)
def test_default_get_available_memory_uses_cuda(self, mock_mem_get_info):
base = CudaSRTPlatform()
self.assertEqual(base.get_available_memory(1), (123, 456))
mock_mem_get_info.assert_called_once_with(1)
def test_default_distributed_backend_is_nccl(self):
base = CudaSRTPlatform()
self.assertEqual(base.get_torch_distributed_backend_str(), "nccl")
@patch("torch.cuda.manual_seed_all")
@patch("torch.manual_seed")
@patch("sglang.srt.platforms.device_mixin.np.random.seed")
@patch("sglang.srt.platforms.device_mixin.random.seed")
def test_default_seed_everything_seeds_cuda(
self, mock_random_seed, mock_np_seed, mock_torch_seed, mock_cuda_seed
):
CudaSRTPlatform.seed_everything(123)
mock_random_seed.assert_called_once_with(123)
mock_np_seed.assert_called_once_with(123)
mock_torch_seed.assert_called_once_with(123)
mock_cuda_seed.assert_called_once_with(123)
def test_cuda_srt_platform_capabilities(self):
base = CudaSRTPlatform()
self.assertTrue(base.supports_fp8())
self.assertTrue(base.support_cuda_graph())
self.assertTrue(base.support_piecewise_cuda_graph())
class TestSRTPlatformOverrides(CustomTestCase): class TestSRTPlatformOverrides(CustomTestCase):
"""Tests for SRTPlatform method overrides via plugins.""" """Tests for SRTPlatform method overrides via plugins."""
@@ -320,6 +428,16 @@ class TestResolvePlatformWithEnv(CustomTestCase):
class TestResolvePlatformAutoDiscover(CustomTestCase): class TestResolvePlatformAutoDiscover(CustomTestCase):
"""Tests for _resolve_platform auto-discovery when SGLANG_PLATFORM is not set.""" """Tests for _resolve_platform auto-discovery when SGLANG_PLATFORM is not set."""
@patch("sglang.srt.platforms.torch")
def test_is_cuda_available_excludes_rocm(self, mock_torch):
"""ROCm exposes torch.cuda, but should not use the CUDA platform identity."""
mock_torch.cuda.is_available.return_value = True
mock_torch.version.hip = "6.0"
import sglang.srt.platforms as plat_mod
self.assertFalse(plat_mod._is_cuda_available())
@patch("sglang.srt.platforms.load_plugins_by_group") @patch("sglang.srt.platforms.load_plugins_by_group")
@patch("sglang.srt.platforms.envs") @patch("sglang.srt.platforms.envs")
def test_single_plugin_activates(self, mock_envs, mock_load): def test_single_plugin_activates(self, mock_envs, mock_load):
@@ -335,13 +453,48 @@ class TestResolvePlatformAutoDiscover(CustomTestCase):
self.assertEqual(result, mock_instance) self.assertEqual(result, mock_instance)
@patch("sglang.srt.platforms.load_plugins_by_group") @patch("sglang.srt.platforms.load_plugins_by_group")
@patch("sglang.srt.platforms._is_cuda_available")
@patch("sglang.srt.platforms.envs") @patch("sglang.srt.platforms.envs")
def test_no_plugin_activates_fallback(self, mock_envs, mock_load): def test_no_plugin_activates_cuda_fallback(
"""When no plugin activates, return base SRTPlatform with warning.""" self, mock_envs, mock_is_cuda_available, mock_load
):
"""When CUDA is available and no plugin activates, return CUDA defaults."""
mock_envs.SGLANG_PLATFORM.get.return_value = "" mock_envs.SGLANG_PLATFORM.get.return_value = ""
mock_is_cuda_available.return_value = True
mock_load.return_value = {}
result = _resolve_platform()
self.assertIsInstance(result, CudaSRTPlatform)
@patch("sglang.srt.platforms.load_plugins_by_group")
@patch("sglang.srt.platforms._is_cuda_available")
@patch("sglang.srt.platforms.envs")
def test_no_plugin_no_cuda_activates_base_fallback(
self, mock_envs, mock_is_cuda_available, mock_load
):
"""When no plugin or CUDA is available, return the abstract base platform."""
mock_envs.SGLANG_PLATFORM.get.return_value = ""
mock_is_cuda_available.return_value = False
mock_load.return_value = {} mock_load.return_value = {}
result = _resolve_platform() result = _resolve_platform()
self.assertIsInstance(result, SRTPlatform) self.assertIsInstance(result, SRTPlatform)
self.assertNotIsInstance(result, CudaSRTPlatform)
@patch("sglang.srt.platforms.load_plugins_by_group")
@patch("sglang.srt.platforms.torch")
@patch("sglang.srt.platforms.envs")
def test_no_plugin_rocm_does_not_activate_cuda_fallback(
self, mock_envs, mock_torch, mock_load
):
"""ROCm exposes torch.cuda but must not use the CUDA fallback platform."""
mock_envs.SGLANG_PLATFORM.get.return_value = ""
mock_torch.cuda.is_available.return_value = True
mock_torch.version.hip = "6.0"
mock_load.return_value = {}
result = _resolve_platform()
self.assertIsInstance(result, SRTPlatform)
self.assertNotIsInstance(result, CudaSRTPlatform)
@patch("sglang.srt.platforms.load_plugins_by_group") @patch("sglang.srt.platforms.load_plugins_by_group")
@patch("sglang.srt.platforms.envs") @patch("sglang.srt.platforms.envs")