From 9141ec11289ee57332423e8bd609c23fb20a5082 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sun, 23 Aug 2026 01:03:49 -0700 Subject: [PATCH] chore: move cuda_vmm_utils.py under srt/utils/ Relocate python/sglang/srt/cuda_vmm_utils.py to python/sglang/srt/utils/cuda_vmm_utils.py, next to the related cuda_vmm_transport_utils.py / cuda_ipc_transport_utils.py, and update all import sites. Pure move: the file contents are unchanged. --- .../custom_all_reduce_utils.py | 2 +- .../custom_all_reduce_v2.py | 10 +++--- python/sglang/srt/layers/moe/dwdp/layout.py | 2 +- .../sglang/srt/layers/moe/dwdp/page_pool.py | 2 +- .../sglang/srt/layers/moe/dwdp/transport.py | 12 +++---- .../srt/layers/moe/dwdp/weight_buffer.py | 12 +++---- python/sglang/srt/mem_cache/kv_vmm_backing.py | 2 +- .../srt/model_executor/runner_utils/pool.py | 2 +- .../srt/multimodal/transport/memory_pool.py | 2 +- .../srt/utils/cuda_vmm_transport_utils.py | 32 +++++++++---------- .../sglang/srt/{ => utils}/cuda_vmm_utils.py | 0 test/registered/unit/test_cuda_vmm_utils.py | 4 +-- 12 files changed, 41 insertions(+), 41 deletions(-) rename python/sglang/srt/{ => utils}/cuda_vmm_utils.py (100%) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py index 9eceb081d1af..d42fe63a3357 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py @@ -19,11 +19,11 @@ 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__) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 77833836aec7..3011c9958503 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -31,16 +31,16 @@ 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, diff --git a/python/sglang/srt/layers/moe/dwdp/layout.py b/python/sglang/srt/layers/moe/dwdp/layout.py index 563cc3dfa9cb..ee4874dec1c4 100644 --- a/python/sglang/srt/layers/moe/dwdp/layout.py +++ b/python/sglang/srt/layers/moe/dwdp/layout.py @@ -8,7 +8,7 @@ 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]] diff --git a/python/sglang/srt/layers/moe/dwdp/page_pool.py b/python/sglang/srt/layers/moe/dwdp/page_pool.py index b7a2b727a1c2..d3bfef3b8dad 100644 --- a/python/sglang/srt/layers/moe/dwdp/page_pool.py +++ b/python/sglang/srt/layers/moe/dwdp/page_pool.py @@ -8,7 +8,7 @@ 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, diff --git a/python/sglang/srt/layers/moe/dwdp/transport.py b/python/sglang/srt/layers/moe/dwdp/transport.py index 9b66378867b3..e325160ec2bc 100644 --- a/python/sglang/srt/layers/moe/dwdp/transport.py +++ b/python/sglang/srt/layers/moe/dwdp/transport.py @@ -11,7 +11,12 @@ 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 @@ make_device_allocation_prop, tensor_from_pointer, ) -from sglang.srt.layers.moe.dwdp.layout import ( - DwdpExpertLayout, - LayerWeightSpecs, - MnnvlHandleSet, -) logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/layers/moe/dwdp/weight_buffer.py b/python/sglang/srt/layers/moe/dwdp/weight_buffer.py index 0771d67b005b..dd17bf4ac359 100644 --- a/python/sglang/srt/layers/moe/dwdp/weight_buffer.py +++ b/python/sglang/srt/layers/moe/dwdp/weight_buffer.py @@ -8,12 +8,6 @@ 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 @@ 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__) diff --git a/python/sglang/srt/mem_cache/kv_vmm_backing.py b/python/sglang/srt/mem_cache/kv_vmm_backing.py index 3f76a6179999..fd69ea3333bc 100644 --- a/python/sglang/srt/mem_cache/kv_vmm_backing.py +++ b/python/sglang/srt/mem_cache/kv_vmm_backing.py @@ -6,7 +6,7 @@ import torch -from sglang.srt.cuda_vmm_utils import ( +from sglang.srt.utils.cuda_vmm_utils import ( BumpArenaStub, VmmReservation, align_up, diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index 7dd64dc8fa9f..2ff115188c18 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/python/sglang/srt/model_executor/runner_utils/pool.py @@ -24,10 +24,10 @@ 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 diff --git a/python/sglang/srt/multimodal/transport/memory_pool.py b/python/sglang/srt/multimodal/transport/memory_pool.py index aa23a40e5fa7..3a02435e94fc 100644 --- a/python/sglang/srt/multimodal/transport/memory_pool.py +++ b/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 diff --git a/python/sglang/srt/utils/cuda_vmm_transport_utils.py b/python/sglang/srt/utils/cuda_vmm_transport_utils.py index d0ad2d012cb3..1202aca7751f 100644 --- a/python/sglang/srt/utils/cuda_vmm_transport_utils.py +++ b/python/sglang/srt/utils/cuda_vmm_transport_utils.py @@ -10,22 +10,6 @@ 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, @@ -42,6 +26,22 @@ 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__) diff --git a/python/sglang/srt/cuda_vmm_utils.py b/python/sglang/srt/utils/cuda_vmm_utils.py similarity index 100% rename from python/sglang/srt/cuda_vmm_utils.py rename to python/sglang/srt/utils/cuda_vmm_utils.py diff --git a/test/registered/unit/test_cuda_vmm_utils.py b/test/registered/unit/test_cuda_vmm_utils.py index 776898f53c70..27f13bd77464 100644 --- a/test/registered/unit/test_cuda_vmm_utils.py +++ b/test/registered/unit/test_cuda_vmm_utils.py @@ -21,8 +21,8 @@ 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,