chore: move cuda_vmm_utils.py under srt/utils/ (#36053)

python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py

@@ -19,11 +19,11 @@ import torch.distributed as dist
 import torch.multiprocessing as mp
 from typing_extensions import ParamSpec
 
-from sglang.srt.cuda_vmm_utils import _gpu_fabric_clique
 from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from sglang.srt.distributed.parallel_state import in_the_same_node_as
 from sglang.srt.environ import envs as sglang_envs
 from sglang.srt.utils import is_cuda, is_hip, is_musa
+from sglang.srt.utils.cuda_vmm_utils import _gpu_fabric_clique
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py

@@ -31,16 +31,16 @@ from sglang.kernels.ops.communication.all_reduce import (
     IPCManager,
     custom_all_reduce,
 )
-from sglang.srt.cuda_vmm_utils import (
-    VmmGraphInputManager,
-    compute_graph_capture_bases,
-    is_vmm_pointer,
-)
 from sglang.srt.distributed.parallel_state import in_the_same_node_as
 from sglang.srt.environ import envs
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
 )
+from sglang.srt.utils.cuda_vmm_utils import (
+    VmmGraphInputManager,
+    compute_graph_capture_bases,
+    is_vmm_pointer,
+)
 
 from .configs.custom_all_reduce_v2 import (
     get_all_reduce_config,

python/sglang/srt/layers/moe/dwdp/layout.py

@@ -8,7 +8,7 @@ from typing import Dict, List, Optional, Tuple
 
 import torch
 
-from sglang.srt.cuda_vmm_utils import align_down, align_up
+from sglang.srt.utils.cuda_vmm_utils import align_down, align_up
 
 # one (start, end_capped) expert range per peer DWDP rank
 PeerRanges = List[Tuple[int, int]]

python/sglang/srt/layers/moe/dwdp/page_pool.py

@@ -8,7 +8,7 @@ from typing import Dict, List, Optional
 
 from cuda.bindings import driver as cuda
 
-from sglang.srt.cuda_vmm_utils import (
+from sglang.srt.utils.cuda_vmm_utils import (
     VmmReservation,
     align_up,
     check_drv,

python/sglang/srt/layers/moe/dwdp/transport.py

@@ -11,7 +11,12 @@ import torch
 import torch.distributed as dist
 from cuda.bindings import driver as cuda
 
-from sglang.srt.cuda_vmm_utils import (
+from sglang.srt.layers.moe.dwdp.layout import (
+    DwdpExpertLayout,
+    LayerWeightSpecs,
+    MnnvlHandleSet,
+)
+from sglang.srt.utils.cuda_vmm_utils import (
     VmmReservation,
     align_down,
     align_up,
@@ -23,11 +28,6 @@ from sglang.srt.cuda_vmm_utils import (
     make_device_allocation_prop,
     tensor_from_pointer,
 )
-from sglang.srt.layers.moe.dwdp.layout import (
-    DwdpExpertLayout,
-    LayerWeightSpecs,
-    MnnvlHandleSet,
-)
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/layers/moe/dwdp/weight_buffer.py

@@ -8,12 +8,6 @@ from typing import Dict, List, Optional, Tuple
 
 import torch
 
-from sglang.srt.cuda_vmm_utils import (
-    VmmReservation,
-    get_device_granularity,
-    make_device_allocation_prop,
-    tensor_from_pointer,
-)
 from sglang.srt.layers.moe.dwdp.layout import (
     EdgeInfo,
     LayerWeightSpecs,
@@ -21,6 +15,12 @@ from sglang.srt.layers.moe.dwdp.layout import (
     PageAlignedLayout,
 )
 from sglang.srt.layers.moe.dwdp.page_pool import PagePool, compute_slot_sizes
+from sglang.srt.utils.cuda_vmm_utils import (
+    VmmReservation,
+    get_device_granularity,
+    make_device_allocation_prop,
+    tensor_from_pointer,
+)
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/mem_cache/kv_vmm_backing.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, List, Optional, Sequence
 
 import torch
 
-from sglang.srt.cuda_vmm_utils import (
+from sglang.srt.utils.cuda_vmm_utils import (
     BumpArenaStub,
     VmmReservation,
     align_up,

python/sglang/srt/model_executor/runner_utils/pool.py

@@ -24,10 +24,10 @@ from typing import Any, Iterator, Optional
 
 import torch
 
-from sglang.srt.cuda_vmm_utils import BumpArenaStub
 from sglang.srt.environ import envs
 from sglang.srt.runtime_context import get_resources
 from sglang.srt.utils import is_cuda
+from sglang.srt.utils.cuda_vmm_utils import BumpArenaStub
 
 logger = logging.getLogger(__name__)
 _active_graph_pool_user: Optional[str] = None

python/sglang/srt/multimodal/transport/memory_pool.py

@@ -21,7 +21,7 @@ def align_up(value: int, alignment: int) -> int:
 def _driver_modules():
     from cuda.bindings import driver as cuda
 
-    from sglang.srt.cuda_vmm_utils import check_drv
+    from sglang.srt.utils.cuda_vmm_utils import check_drv
 
     return cuda, check_drv
 

python/sglang/srt/utils/cuda_vmm_transport_utils.py

@@ -10,22 +10,6 @@ from dataclasses import dataclass
 
 import torch
 
-from sglang.srt.cuda_vmm_utils import (
-    _FD_SEND_TIMEOUT_S,
-    VmmReservation,
-    _get_cuda_driver,
-    _recv_fd,
-    _send_fd,
-    align_up,
-    allocation_handle_type_name,
-    check_drv,
-    get_allocation_granularity,
-    get_device_allocation_handle_type,
-    import_and_map_alloc,
-    make_device_allocation_prop,
-    release_mappings,
-    tensor_from_pointer,
-)
 from sglang.srt.managers.schedule_batch import (
     Modality,
     MultimodalDataItem,
@@ -43,6 +27,22 @@ from sglang.srt.utils.cuda_ipc_transport_utils import (
     CudaIpcTensorTransportProxy,
     get_mm_feature_pool_size_per_worker,
 )
+from sglang.srt.utils.cuda_vmm_utils import (
+    _FD_SEND_TIMEOUT_S,
+    VmmReservation,
+    _get_cuda_driver,
+    _recv_fd,
+    _send_fd,
+    align_up,
+    allocation_handle_type_name,
+    check_drv,
+    get_allocation_granularity,
+    get_device_allocation_handle_type,
+    import_and_map_alloc,
+    make_device_allocation_prop,
+    release_mappings,
+    tensor_from_pointer,
+)
 
 logger = logging.getLogger(__name__)
 

test/registered/unit/test_cuda_vmm_utils.py

@@ -21,8 +21,8 @@ import torch.distributed as dist
 from cuda.bindings import driver as drv
 
 from sglang.kernels.jit.utils import cache_once
-from sglang.srt import cuda_vmm_utils
-from sglang.srt.cuda_vmm_utils import (
+from sglang.srt.utils import cuda_vmm_utils
+from sglang.srt.utils.cuda_vmm_utils import (
     check_drv,
     exchange_posix_fds,
     export_shareable_handles,