Consolidate CUDA VMM allocation helpers (#34199)

test/registered/unit/multimodal/test_cuda_vmm_transport.py

@@ -442,7 +442,7 @@ class TestCudaVmmTransport(CustomTestCase):
     def test_posix_export_fd_closes_when_allocation_setup_fails(self):
         with (
             patch(
-                "sglang.srt.utils.cuda_vmm_transport_utils._tensor_from_pointer",
+                "sglang.srt.utils.cuda_vmm_transport_utils.tensor_from_pointer",
                 side_effect=RuntimeError("forced storage failure"),
             ),
             patch(

test/registered/unit/multimodal/test_gpu_feature_transport.py

@@ -18,48 +18,27 @@ class TestCudaVmmFeatureTransport(unittest.TestCase):
         pool.memory_pool = object()
         pool.use_fabric = True
         pool.shareable_handle = b"handle"
-        pool._pool_pointer = 123
-        pool._allocation_handle = 456
-        pool._allocation_mapped = True
-        pool.allocation_size = 4096
+        allocation = MagicMock()
+        allocation.close.side_effect = [
+            RuntimeError("forced allocation close failure"),
+            None,
+        ]
+        pool._allocation = allocation
         pool.device_index = 0
-        driver = MagicMock()
-        driver.cuMemUnmap.return_value = "unmap"
-        driver.cuMemAddressFree.return_value = "address_free"
-        driver.cuMemRelease.return_value = "release"
-        failed_once = False
-
-        def check_driver(result, _operation):
-            nonlocal failed_once
-            if result == "address_free" and not failed_once:
-                failed_once = True
-                raise RuntimeError("forced address-free failure")
-            return result
 
         with (
-            patch.object(vmm, "_get_cuda_driver", return_value=driver),
             patch.object(vmm.torch.cuda, "device", return_value=nullcontext()),
-            patch.object(vmm, "check_drv", side_effect=check_driver),
-            self.assertRaisesRegex(RuntimeError, "forced address-free failure"),
+            self.assertRaisesRegex(RuntimeError, "forced allocation close failure"),
         ):
             pool._release_allocation()
 
-        self.assertFalse(pool._allocation_mapped)
-        self.assertEqual(pool._pool_pointer, 123)
-        self.assertEqual(pool._allocation_handle, 456)
+        self.assertIs(pool._allocation, allocation)
 
-        with (
-            patch.object(vmm, "_get_cuda_driver", return_value=driver),
-            patch.object(vmm.torch.cuda, "device", return_value=nullcontext()),
-            patch.object(vmm, "check_drv", side_effect=lambda result, _: result),
-        ):
+        with patch.object(vmm.torch.cuda, "device", return_value=nullcontext()):
             pool._release_allocation()
 
-        self.assertIsNone(pool._pool_pointer)
-        self.assertIsNone(pool._allocation_handle)
-        self.assertEqual(driver.cuMemUnmap.call_count, 1)
-        self.assertEqual(driver.cuMemAddressFree.call_count, 2)
-        self.assertEqual(driver.cuMemRelease.call_count, 1)
+        self.assertIsNone(pool._allocation)
+        self.assertEqual(allocation.close.call_count, 2)
 
     def test_model_class_controls_cuda_vmm_opt_in(self):
         from sglang.srt.managers.tokenizer_manager import TokenizerManager

test/registered/unit/test_cuda_vmm_utils.py

@@ -21,11 +21,15 @@ import torch.distributed as dist
 from cuda.bindings import driver as drv
 
 from sglang.kernels.jit.utils import cache_once
-from sglang.srt.distributed.device_communicators.vmm_utils import (
+from sglang.srt import cuda_vmm_utils
+from sglang.srt.cuda_vmm_utils import (
     check_drv,
     exchange_posix_fds,
     export_shareable_handles,
+    get_allocation_granularity,
+    get_device_allocation_handle_type,
     import_and_map_alloc,
+    make_device_allocation_prop,
     make_rw_access_desc,
     map_chunk_into_span,
     release_mappings,
@@ -107,6 +111,67 @@ def _assert_region(va: int, expected: int, peer: int, chunk: int) -> None:
     )
 
 
+@pytest.mark.parametrize(
+    ("rejected", "expected"),
+    [
+        ((_FABRIC,), _POSIX_FD),
+        ((_FABRIC, _POSIX_FD), 0),
+    ],
+)
+def test_default_handle_type_fallback(monkeypatch, rejected, expected) -> None:
+    device_id = torch.cuda.current_device()
+    create = drv.cuMemCreate
+
+    def reject_selected(size, prop, flags):
+        if prop.requestedHandleTypes in rejected:
+            return (drv.CUresult.CUDA_ERROR_NOT_SUPPORTED, None)
+        return create(size, prop, flags)
+
+    get_device_allocation_handle_type.cache_clear()
+    monkeypatch.setattr(cuda_vmm_utils, "is_gpu_fabric_ready", lambda _device: True)
+    monkeypatch.setattr(drv, "cuMemCreate", reject_selected)
+    try:
+        selected = get_device_allocation_handle_type(device_id)
+        prop = make_device_allocation_prop(device_id)
+        assert selected == expected
+        assert prop.requestedHandleTypes == expected
+        assert prop.allocFlags.gpuDirectRDMACapable == 0
+
+        explicit = make_device_allocation_prop(
+            device_id,
+            handle_types=_FABRIC,
+            gpu_direct_rdma=True,
+        )
+        assert explicit.requestedHandleTypes == _FABRIC
+        assert explicit.allocFlags.gpuDirectRDMACapable == 1
+
+        non_exportable = make_device_allocation_prop(device_id, handle_types=None)
+        assert non_exportable.requestedHandleTypes == 0
+
+        explicit_none = make_device_allocation_prop(device_id, handle_types=0)
+        assert explicit_none.requestedHandleTypes == 0
+
+        with pytest.raises(ValueError, match="handle_types must be"):
+            make_device_allocation_prop(device_id, handle_types="fabric")
+        with pytest.raises(ValueError, match="invalid CUDA handle-type value"):
+            make_device_allocation_prop(device_id, handle_types=42)
+    finally:
+        get_device_allocation_handle_type.cache_clear()
+
+
+def test_granularity_defaults_to_recommended(monkeypatch) -> None:
+    prop = make_device_allocation_prop(0, handle_types=None)
+    seen = []
+
+    def granularity(_prop, flag):
+        seen.append(flag)
+        return (drv.CUresult.CUDA_SUCCESS, _ALLOC_BYTES)
+
+    monkeypatch.setattr(drv, "cuMemGetAllocationGranularity", granularity)
+    assert get_allocation_granularity(prop) == _ALLOC_BYTES
+    assert seen == [_RECOMMENDED]
+
+
 @pytest.mark.parametrize("n_chunks", [1, 3])
 @pytest.mark.parametrize("transport", ["posix", "fabric"])
 def test_handle_roundtrip(transport: str, n_chunks: int) -> None: