[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>

docs_new/docs/hardware-platforms/plugin.mdx

@@ -371,6 +371,12 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
       <td><strong>Active</strong></td>
       <td>Get current peak memory usage in bytes</td>
     </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>
       <td><code>get_torch_distributed_backend_str()</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>Whether FP8 quantization is supported</td>
     </tr>
-    <tr>
-      <td><code>is_pin_memory_available()</code></td>
-      <td><code>True</code></td>
-      <td>Whether pinned memory is available</td>
-    </tr>
   </tbody>
 </table>
 

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

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"
 

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

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.

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():

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()
 
 
 # ---------------------------------------------------------------------------