minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Introduce CudaDeviceMixin and CudaSRTPlatform (#24096)

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/platforms/test_platform_interface.py
+155 -2
View File
@@ -7,7 +7,10 @@ and the platform discovery / lazy initialization mechanism.
from unittest.mock import MagicMock, patch
import torch
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 (
CpuArchEnum,
DeviceCapability,
@@ -222,6 +225,111 @@ class TestSRTPlatform(CustomTestCase):
base = SRTPlatform()
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):
"""Tests for SRTPlatform method overrides via plugins."""
@@ -320,6 +428,16 @@ class TestResolvePlatformWithEnv(CustomTestCase):
class TestResolvePlatformAutoDiscover(CustomTestCase):
"""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.envs")
def test_single_plugin_activates(self, mock_envs, mock_load):
@@ -335,13 +453,48 @@ class TestResolvePlatformAutoDiscover(CustomTestCase):
self.assertEqual(result, mock_instance)
@patch("sglang.srt.platforms.load_plugins_by_group")
@patch("sglang.srt.platforms._is_cuda_available")
@patch("sglang.srt.platforms.envs")
def test_no_plugin_activates_fallback(self, mock_envs, mock_load):
"""When no plugin activates, return base SRTPlatform with warning."""
def test_no_plugin_activates_cuda_fallback(
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_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 = {}
result = _resolve_platform()
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.envs")