[Platform] Route pin memory availability through current_platform (#28113)

Co-authored-by: N3u0ns <N3u0ns@users.noreply.github.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
N3ur0ns
2026-07-13 11:37:59 -07:00
committed by GitHub
co-authored by N3u0ns Alex Nails
parent afaa17a7f2
commit 11a82af5f8
7 changed files with 156 additions and 18 deletions
+6 -5
View File
@@ -371,6 +371,12 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
<td><strong>Active</strong></td> <td><strong>Active</strong></td>
<td>Get current peak memory usage in bytes</td> <td>Get current peak memory usage in bytes</td>
</tr> </tr>
<tr>
<td><code>is_pin_memory_available(device=None)</code></td>
<td><code>False</code></td>
<td><strong>Active</strong></td>
<td>Whether pinned host memory is available for a target device</td>
</tr>
<tr> <tr>
<td><code>get_torch_distributed_backend_str()</code></td> <td><code>get_torch_distributed_backend_str()</code></td>
<td><code>raise NotImplementedError</code></td> <td><code>raise NotImplementedError</code></td>
@@ -470,11 +476,6 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
<td><code>False</code></td> <td><code>False</code></td>
<td>Whether FP8 quantization is supported</td> <td>Whether FP8 quantization is supported</td>
</tr> </tr>
<tr>
<td><code>is_pin_memory_available()</code></td>
<td><code>True</code></td>
<td>Whether pinned memory is available</td>
</tr>
</tbody> </tbody>
</table> </table>
+3 -3
View File
@@ -125,9 +125,9 @@ class CpuSRTPlatform(CpuDeviceMixin, SRTPlatform):
supports_fp8 / support_cuda_graph / support_piecewise_cuda_graph keep the supports_fp8 / support_cuda_graph / support_piecewise_cuda_graph keep the
conservative SRTPlatform defaults (all False), so they are not repeated conservative SRTPlatform defaults (all False), so they are not repeated
here. Only is_pin_memory_available is overridden: the base defaults to here. is_pin_memory_available is repeated for explicitness: CPU has no GPU
True, but CPU has no GPU to pin host memory to. to pin host memory to.
""" """
def is_pin_memory_available(self) -> bool: def is_pin_memory_available(self, device=None) -> bool:
return False return False
+5
View File
@@ -52,6 +52,11 @@ class CudaDeviceMixin(DeviceMixin):
def get_available_memory(self, device_id: int = 0) -> tuple[int, int]: def get_available_memory(self, device_id: int = 0) -> tuple[int, int]:
return torch.cuda.mem_get_info(device_id) 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: def get_torch_distributed_backend_str(self) -> str:
return "nccl" return "nccl"
@@ -159,6 +159,10 @@ class DeviceMixin:
"""[Active] Get current peak memory usage in bytes.""" """[Active] Get current peak memory usage in bytes."""
raise NotImplementedError 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 # Planned methods — reserved interface. Core still uses hardcoded
# calls (e.g. torch.cuda.*). OOT implementations will NOT take # calls (e.g. torch.cuda.*). OOT implementations will NOT take
-4
View File
@@ -103,10 +103,6 @@ class SRTPlatform(DeviceMixin):
"""Whether this platform supports FP8 quantization.""" """Whether this platform supports FP8 quantization."""
return False 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: def support_cuda_graph(self) -> bool:
"""Whether this platform supports device graph capture and replay. """Whether this platform supports device graph capture and replay.
Controls CUDA graph (CudaGraphRunner) for the decode path. Controls CUDA graph (CudaGraphRunner) for the decode path.
+1 -5
View File
@@ -539,11 +539,7 @@ def get_available_gpu_memory(
def is_pin_memory_available(device=None) -> bool: def is_pin_memory_available(device=None) -> bool:
if not torch.cuda.is_available(): return current_platform.is_pin_memory_available(device)
return False
if device is not None and str(device) == "cpu":
return False
return True
def get_dispatch_device_backend(): def get_dispatch_device_backend():
@@ -19,6 +19,7 @@ from sglang.srt.platforms.device_mixin import (
PlatformEnum, PlatformEnum,
) )
from sglang.srt.platforms.interface import SRTPlatform 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.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -159,6 +160,11 @@ class TestDeviceMixin(CustomTestCase):
cuda = _make_device_mixin(PlatformEnum.CUDA, "cuda", "cuda") cuda = _make_device_mixin(PlatformEnum.CUDA, "cuda", "cuda")
self.assertFalse(cuda.is_out_of_tree()) 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") @patch("platform.machine")
def test_get_cpu_architecture(self, mock_machine): def test_get_cpu_architecture(self, mock_machine):
"""get_cpu_architecture maps common strings to CpuArchEnum.""" """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())
self.assertFalse(base.is_cuda_alike()) 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): class TestCudaDeviceMixin(CustomTestCase):
"""Tests for CUDA device operation defaults.""" """Tests for CUDA device operation defaults."""
@@ -211,6 +222,19 @@ class TestCudaDeviceMixin(CustomTestCase):
self.assertEqual(base.get_device_capability(1), DeviceCapability(9, 0)) self.assertEqual(base.get_device_capability(1), DeviceCapability(9, 0))
mock_get_device_capability.assert_called_once_with(1) 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.cuda.manual_seed_all")
@patch("torch.manual_seed") @patch("torch.manual_seed")
@patch("sglang.srt.platforms.device_mixin.np.random.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.supports_fp8())
self.assertFalse(base.support_cuda_graph()) self.assertFalse(base.support_cuda_graph())
self.assertFalse(base.support_piecewise_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())
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()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------