From 11a82af5f8371fcc1e3cb3cdf9d8706aaade0b20 Mon Sep 17 00:00:00 2001 From: N3ur0ns Date: Tue, 14 Jul 2026 02:37:59 +0800 Subject: [PATCH] [Platform] Route pin memory availability through current_platform (#28113) Co-authored-by: N3u0ns Co-authored-by: Alex Nails --- docs_new/docs/hardware-platforms/plugin.mdx | 11 +- python/sglang/srt/platforms/cpu.py | 6 +- python/sglang/srt/platforms/cuda.py | 5 + python/sglang/srt/platforms/device_mixin.py | 4 + python/sglang/srt/platforms/interface.py | 4 - python/sglang/srt/utils/common.py | 6 +- .../unit/platforms/test_platform_interface.py | 138 +++++++++++++++++- 7 files changed, 156 insertions(+), 18 deletions(-) diff --git a/docs_new/docs/hardware-platforms/plugin.mdx b/docs_new/docs/hardware-platforms/plugin.mdx index 0debdc29b..cc63cee75 100644 --- a/docs_new/docs/hardware-platforms/plugin.mdx +++ b/docs_new/docs/hardware-platforms/plugin.mdx @@ -371,6 +371,12 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat Active Get current peak memory usage in bytes + + is_pin_memory_available(device=None) + False + Active + Whether pinned host memory is available for a target device + get_torch_distributed_backend_str() raise NotImplementedError @@ -470,11 +476,6 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat False Whether FP8 quantization is supported - - is_pin_memory_available() - True - Whether pinned memory is available - diff --git a/python/sglang/srt/platforms/cpu.py b/python/sglang/srt/platforms/cpu.py index 46fd0251a..9c28ae676 100644 --- a/python/sglang/srt/platforms/cpu.py +++ b/python/sglang/srt/platforms/cpu.py @@ -125,9 +125,9 @@ class CpuSRTPlatform(CpuDeviceMixin, SRTPlatform): supports_fp8 / support_cuda_graph / support_piecewise_cuda_graph keep the conservative SRTPlatform defaults (all False), so they are not repeated - here. Only is_pin_memory_available is overridden: the base defaults to - True, but CPU has no GPU to pin host memory to. + here. is_pin_memory_available is repeated for explicitness: CPU has no GPU + to pin host memory to. """ - def is_pin_memory_available(self) -> bool: + def is_pin_memory_available(self, device=None) -> bool: return False diff --git a/python/sglang/srt/platforms/cuda.py b/python/sglang/srt/platforms/cuda.py index 5800886d2..d1e9f5e24 100644 --- a/python/sglang/srt/platforms/cuda.py +++ b/python/sglang/srt/platforms/cuda.py @@ -52,6 +52,11 @@ class CudaDeviceMixin(DeviceMixin): def get_available_memory(self, device_id: int = 0) -> tuple[int, int]: return torch.cuda.mem_get_info(device_id) + def is_pin_memory_available(self, device=None) -> bool: + if device is not None and str(device) == "cpu": + return False + return True + def get_torch_distributed_backend_str(self) -> str: return "nccl" diff --git a/python/sglang/srt/platforms/device_mixin.py b/python/sglang/srt/platforms/device_mixin.py index a87523aa7..5781b0191 100644 --- a/python/sglang/srt/platforms/device_mixin.py +++ b/python/sglang/srt/platforms/device_mixin.py @@ -159,6 +159,10 @@ class DeviceMixin: """[Active] Get current peak memory usage in bytes.""" raise NotImplementedError + def is_pin_memory_available(self, device=None) -> bool: + """[Active] Whether pinned host memory is available for a target device.""" + return False + # ------------------------------------------------------------------ # Planned methods — reserved interface. Core still uses hardcoded # calls (e.g. torch.cuda.*). OOT implementations will NOT take diff --git a/python/sglang/srt/platforms/interface.py b/python/sglang/srt/platforms/interface.py index 9a541900d..e95aa4c30 100644 --- a/python/sglang/srt/platforms/interface.py +++ b/python/sglang/srt/platforms/interface.py @@ -103,10 +103,6 @@ class SRTPlatform(DeviceMixin): """Whether this platform supports FP8 quantization.""" return False - def is_pin_memory_available(self) -> bool: - """Whether pinned memory is available on this platform.""" - return True - def support_cuda_graph(self) -> bool: """Whether this platform supports device graph capture and replay. Controls CUDA graph (CudaGraphRunner) for the decode path. diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 2725308bd..abee19bcf 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -539,11 +539,7 @@ def get_available_gpu_memory( def is_pin_memory_available(device=None) -> bool: - if not torch.cuda.is_available(): - return False - if device is not None and str(device) == "cpu": - return False - return True + return current_platform.is_pin_memory_available(device) def get_dispatch_device_backend(): diff --git a/test/registered/unit/platforms/test_platform_interface.py b/test/registered/unit/platforms/test_platform_interface.py index 9f7e0260c..50809184f 100644 --- a/test/registered/unit/platforms/test_platform_interface.py +++ b/test/registered/unit/platforms/test_platform_interface.py @@ -19,6 +19,7 @@ from sglang.srt.platforms.device_mixin import ( PlatformEnum, ) from sglang.srt.platforms.interface import SRTPlatform +from sglang.srt.platforms.rocm import RocmSRTPlatform from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -159,6 +160,11 @@ class TestDeviceMixin(CustomTestCase): cuda = _make_device_mixin(PlatformEnum.CUDA, "cuda", "cuda") self.assertFalse(cuda.is_out_of_tree()) + def test_pin_memory_default_is_conservative(self): + mixin = _make_device_mixin(PlatformEnum.OOT, "custom", "custom") + self.assertFalse(mixin.is_pin_memory_available()) + self.assertFalse(mixin.is_pin_memory_available(device="cpu")) + @patch("platform.machine") def test_get_cpu_architecture(self, mock_machine): """get_cpu_architecture maps common strings to CpuArchEnum.""" @@ -197,6 +203,11 @@ class TestSRTPlatform(CustomTestCase): self.assertFalse(base.is_cuda()) self.assertFalse(base.is_cuda_alike()) + def test_base_pin_memory_default_is_conservative(self): + base = SRTPlatform() + self.assertFalse(base.is_pin_memory_available()) + self.assertFalse(base.is_pin_memory_available(device="cpu")) + class TestCudaDeviceMixin(CustomTestCase): """Tests for CUDA device operation defaults.""" @@ -211,6 +222,19 @@ class TestCudaDeviceMixin(CustomTestCase): self.assertEqual(base.get_device_capability(1), DeviceCapability(9, 0)) mock_get_device_capability.assert_called_once_with(1) + def test_pin_memory_available_for_cuda_targets(self): + base = CudaSRTPlatform() + self.assertTrue(base.is_pin_memory_available()) + self.assertTrue(base.is_pin_memory_available(device="cuda")) + self.assertTrue(base.is_pin_memory_available(device=torch.device("cuda", 0))) + self.assertFalse(base.is_pin_memory_available(device="cpu")) + + def test_rocm_inherits_cuda_pin_memory_behavior(self): + base = RocmSRTPlatform() + self.assertTrue(base.is_pin_memory_available()) + self.assertTrue(base.is_pin_memory_available(device="cuda")) + self.assertFalse(base.is_pin_memory_available(device="cpu")) + @patch("torch.cuda.manual_seed_all") @patch("torch.manual_seed") @patch("sglang.srt.platforms.device_mixin.np.random.seed") @@ -291,8 +315,120 @@ class TestCpuDeviceMixin(CustomTestCase): self.assertFalse(base.supports_fp8()) self.assertFalse(base.support_cuda_graph()) self.assertFalse(base.support_piecewise_cuda_graph()) - # Override of the SRTPlatform default (True) — no GPU to pin to. + # CPU has no GPU to pin host memory to. self.assertFalse(base.is_pin_memory_available()) + self.assertFalse(base.is_pin_memory_available(device="cpu")) + + +class TestPinMemoryAvailability(CustomTestCase): + """Tests for common pin-memory helper dispatch through platforms.""" + + def test_srt_platform_does_not_shadow_device_mixin_pin_memory_override(self): + class M(DeviceMixin): + def is_pin_memory_available(self, device=None): + return device == "custom" + + class P(SRTPlatform, M): + pass + + self.assertTrue(P().is_pin_memory_available(device="custom")) + + def test_device_mixin_can_precede_srt_platform_for_pin_memory_override(self): + class M(DeviceMixin): + def is_pin_memory_available(self, device=None): + return device == "custom" + + class P(M, SRTPlatform): + pass + + self.assertTrue(P().is_pin_memory_available(device="custom")) + + def test_common_wrapper_dispatches_to_current_platform_with_device(self): + from sglang.srt.utils import common + + class P(SRTPlatform): + _enum = PlatformEnum.OOT + device_name = "custom" + device_type = "custom" + + def __init__(self): + self.calls = [] + + def is_pin_memory_available(self, device=None): + self.calls.append(device) + return True + + platform = P() + device = torch.device("cuda", 0) + with patch.object(common, "current_platform", platform): + self.assertTrue(common.is_pin_memory_available(device)) + + self.assertEqual(platform.calls, [device]) + + def test_common_wrapper_dispatches_to_current_platform_without_device(self): + from sglang.srt.utils import common + + class P(SRTPlatform): + _enum = PlatformEnum.OOT + device_name = "custom" + device_type = "custom" + + def __init__(self): + self.calls = [] + + def is_pin_memory_available(self, device=None): + self.calls.append(device) + return True + + platform = P() + with patch.object(common, "current_platform", platform): + self.assertTrue(common.is_pin_memory_available()) + + self.assertEqual(platform.calls, [None]) + + def test_oot_platform_override_true_is_used(self): + from sglang.srt.utils import common + + class P(SRTPlatform): + _enum = PlatformEnum.OOT + device_name = "custom" + device_type = "custom" + + def is_pin_memory_available(self, device=None): + return True + + with patch.object(common, "current_platform", P()): + self.assertTrue(common.is_pin_memory_available()) + + def test_oot_platform_override_false_is_used(self): + from sglang.srt.utils import common + + class P(SRTPlatform): + _enum = PlatformEnum.OOT + device_name = "custom" + device_type = "custom" + + def is_pin_memory_available(self, device=None): + return False + + with patch.object(common, "current_platform", P()): + self.assertFalse(common.is_pin_memory_available()) + + def test_oot_platform_without_override_uses_conservative_default(self): + from sglang.srt.utils import common + + class P(SRTPlatform): + _enum = PlatformEnum.OOT + device_name = "custom" + device_type = "custom" + + with ( + patch.object(common, "current_platform", P()), + patch("torch.cuda.is_available", return_value=True) as mock_cuda_available, + ): + self.assertFalse(common.is_pin_memory_available()) + + mock_cuda_available.assert_not_called() # ---------------------------------------------------------------------------