From fcca4611fa516e2955b787e417ded5a6ff8fcad9 Mon Sep 17 00:00:00 2001 From: Yinghai Lu Date: Tue, 16 Jun 2026 04:49:05 -0700 Subject: [PATCH] [CAR] Let custom allreduce support VMM based allocation (#27593) Co-authored-by: Lianmin Zheng --- python/sglang/jit_kernel/all_reduce.py | 3 + .../distributed/custom_all_reduce_base.cuh | 3 + .../distributed/custom_all_reduce_pull.cuh | 6 +- .../distributed/custom_all_reduce.cuh | 94 ++- .../custom_all_reduce_v2.py | 31 +- .../custom_all_reduce_vmm_utils.py | 579 ++++++++++++++++++ .../eagle/test_deepseek_v3_fp4_mtp_small.py | 4 + 7 files changed, 713 insertions(+), 7 deletions(-) create mode 100644 python/sglang/srt/distributed/device_communicators/custom_all_reduce_vmm_utils.py diff --git a/python/sglang/jit_kernel/all_reduce.py b/python/sglang/jit_kernel/all_reduce.py index 26167c698..05b31127a 100644 --- a/python/sglang/jit_kernel/all_reduce.py +++ b/python/sglang/jit_kernel/all_reduce.py @@ -71,6 +71,9 @@ if TYPE_CHECKING: def post_init(self, handles: List[CUSTOM_AR_HANDLE]) -> None: ... def register_inputs(self, handles: List[List[CUSTOM_AR_PAIR]]) -> None: ... def set_cuda_graph_capture(self, is_capturing: bool) -> None: ... + def get_graph_capture_bases( + self, + ) -> Tuple[List[Tuple[int, int]], List[List[int]], List[int]]: ... def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None: ... def all_reduce( self, input: torch.Tensor, algo: AllReduceAlgo diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh index dc5f5beea..00e265513 100644 --- a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh +++ b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh @@ -21,6 +21,9 @@ inline void register_custom_all_reduce() { .def("post_init", &Class::post_init) .def("register_inputs", &Class::register_inputs) .def("set_cuda_graph_capture", &Class::set_cuda_graph_capture) + .def("get_graph_capture_ptrs", &Class::get_graph_capture_ptrs) + .def("get_graph_capture_bases", &Class::get_graph_capture_bases) + .def("register_peer_mapped_inputs", &Class::register_peer_mapped_inputs) .def("free_ipc_handles", &Class::free_ipc_handles) .def("free_storage", &Class::free_storage) .def("configure_pull", &Class::configure_pull); diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh index e8837af4c..0dfc63ac5 100644 --- a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh +++ b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh @@ -169,6 +169,8 @@ struct CustomAllReducePull : public CustomAllReduceBase { const auto stream = LaunchKernel::resolve_device(device); auto launch = LaunchKernel{num_blocks, m_cta_size, stream}; launch.enable_pdl(kUsePDL); + const auto input_bytes = static_cast(sizeof(DType) * num_items); + RuntimeCheck(input_bytes <= m_pull_buffer_bytes, "Input is too large, num items: ", num_items); const auto check_capturing = [&] { if (!m_is_graph_capturing) return false; // override to avoid cudaRT call overhead cudaStreamCaptureStatus status; @@ -177,13 +179,11 @@ struct CustomAllReducePull : public CustomAllReduceBase { }; if (check_capturing()) { // no-op if not really capturing, we're in a dummy run - const auto data_ptr = allocate_graph_capture_input(input_ptr); + const auto data_ptr = allocate_graph_capture_input(input_ptr, input_bytes); /// NOTE: we assume when the graph is replayed, the data_ptr should be ready launch(kernel, data_ptr, params, ctrl); } else { // 1.copy the input to the buffer - const auto input_bytes = static_cast(sizeof(DType) * num_items); - RuntimeCheck(input_bytes <= m_pull_buffer_bytes, "Input is too large, num items: ", num_items); RuntimeDeviceCheck(cudaMemcpyAsync(buffer_ptr, input_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream)); // 2. launch the all reduce kernel const auto data_ptr = get_data_ptr(); // use default buffer diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh index 239fac71a..bffe36727 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -213,6 +214,95 @@ struct CustomAllReduceBase : public tvm::ffi::Object { m_is_graph_capturing = enabled; } + tvm::ffi::Array get_graph_capture_ptrs() { + tvm::ffi::Array result; + const auto new_count = registered_count() - m_cum_registered_count; + result.reserve(new_count); + for (const auto ptr : std::span(m_graph_capture_inputs).subspan(m_cum_registered_count)) { + result.push_back(reinterpret_cast(ptr)); + } + return result; + } + + using BaseInfo = tvm::ffi::Tuple; // (base_ptr, size) + + /// Returns (unique_bases, per_input_base_indices, per_input_offset). + /// unique_bases[i] = (base_ptr, alloc_size) for each unique allocation. + /// per_input_base_indices[j] = indices of VMM allocations covering input j. + /// per_input_offset[j] = byte offset from the first allocation base for input j. + tvm::ffi::Tuple, tvm::ffi::Array>, tvm::ffi::Array> + get_graph_capture_bases() { + const auto new_inputs = std::span(m_graph_capture_inputs).subspan(m_cum_registered_count); + const auto new_input_bytes = std::span(m_graph_capture_input_bytes).subspan(m_cum_registered_count); + std::unordered_map base_to_idx; + tvm::ffi::Array bases; + tvm::ffi::Array> input_indices; + tvm::ffi::Array offsets; + input_indices.reserve(new_inputs.size()); + offsets.reserve(new_inputs.size()); + RuntimeCheck(new_inputs.size() == new_input_bytes.size(), "graph input metadata mismatch"); + for (const auto input_idx : irange(new_inputs.size())) { + const auto ptr = new_inputs[input_idx]; + auto remaining = new_input_bytes[input_idx]; + RuntimeCheck(remaining > 0, "Invalid graph capture input size: ", remaining); + + auto cursor = reinterpret_cast(ptr); + CUdeviceptr first_base = 0; + tvm::ffi::Array chunks; + while (remaining > 0) { + CUdeviceptr base = 0; + size_t size = 0; + const auto r = cuMemGetAddressRange(&base, &size, cursor); + RuntimeCheck(r == CUDA_SUCCESS, "cuMemGetAddressRange failed: ", r); + if (first_base == 0) first_base = base; + const auto byte_offset = static_cast(cursor - base); + RuntimeCheck( + byte_offset >= 0 && static_cast(byte_offset) < size, + "graph capture input at ", + reinterpret_cast(ptr), + " is outside VMM allocation [base=", + base, + ", size=", + size, + "]"); + + auto [it, inserted] = base_to_idx.try_emplace(base, bases.size()); + if (inserted) { + bases.push_back(BaseInfo{static_cast(base), static_cast(size)}); + } + chunks.push_back(it->second); + + const auto available = static_cast(size) - byte_offset; + const auto advance = std::min(remaining, available); + RuntimeCheck(advance > 0, "Failed to advance VMM graph capture span"); + remaining -= advance; + cursor += advance; + } + input_indices.push_back(chunks); + offsets.push_back(reinterpret_cast(ptr) - first_base); + } + using Result = + tvm::ffi::Tuple, tvm::ffi::Array>, tvm::ffi::Array>; + return Result(bases, input_indices, offsets); + } + + void register_peer_mapped_inputs(tvm::ffi::Array> peer_ptrs_per_input) { + const auto new_count = registered_count() - m_cum_registered_count; + RuntimeCheck(int64_t(peer_ptrs_per_input.size()) == new_count, "peer_ptrs count mismatch"); + if (new_count == 0) return; + std::vector data(new_count); + for (const auto j : irange(new_count)) { + const auto& ptrs = peer_ptrs_per_input[j]; + RuntimeCheck(ptrs.size() == m_num_gpu, "peer count mismatch"); + for (const auto i : irange(m_num_gpu)) { + data[j].input[i] = reinterpret_cast(static_cast(ptrs[i])); + } + } + const auto dst_ptr = get_data_ptr(m_cum_registered_count); + m_cum_registered_count += new_count; + RuntimeDeviceCheck(cudaMemcpy(dst_ptr, data.data(), sizeof(AllReduceData) * new_count, cudaMemcpyHostToDevice)); + } + void free_ipc_handles() { for (const auto& pair : m_ipc_cache) { host::RuntimeDeviceCheck(cudaIpcCloseMemHandle(pair.second)); @@ -238,10 +328,11 @@ struct CustomAllReduceBase : public tvm::ffi::Object { } protected: - AllReduceData* allocate_graph_capture_input(void* data_ptr) { + AllReduceData* allocate_graph_capture_input(void* data_ptr, int64_t input_bytes) { const auto count = registered_count(); RuntimeCheck(count < m_graph_buffer_count, "Graph buffer overflow, increase `graph_buffer_count`!"); m_graph_capture_inputs.push_back(data_ptr); + m_graph_capture_input_bytes.push_back(input_bytes); return get_data_ptr(count); } AllReduceData* get_data_ptr(int64_t which = -1) { @@ -316,6 +407,7 @@ struct CustomAllReduceBase : public tvm::ffi::Object { std::optional m_push_ctrl; void* m_storage = nullptr; std::vector m_graph_capture_inputs; + std::vector m_graph_capture_input_bytes; std::vector m_peer_storage; std::unordered_map m_ipc_cache; }; 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 a55a5e1e5..03442e15e 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 @@ -12,6 +12,10 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import can_use_custom_all_reduce_with_nvlink, is_weak_contiguous, ) +from sglang.srt.distributed.device_communicators.custom_all_reduce_vmm_utils import ( + VmmGraphInputManager, + is_vmm_pointer, +) from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) @@ -67,6 +71,12 @@ class CustomAllReduceV2: max_pull_blocks=max_pull_blocks, max_push_blocks=max_push_blocks, ) + self._vmm_graph_input_manager = VmmGraphInputManager( + obj=self.obj, + group=self.group, + rank=self.rank, + world_size=self.world_size, + ) self._post_init_obj() self.disabled = False log_info_on_rank0(logger, "Custom allreduce v2 initialized successfully") @@ -97,10 +107,21 @@ class CustomAllReduceV2: yield finally: self.obj.set_cuda_graph_capture(False) - # cannot call when graph is capturing assert ( - torch.cuda.is_current_stream_capturing() == False + not torch.cuda.is_current_stream_capturing() ), "Cannot register graph inputs while capturing CUDA graph" + raw_ptrs = self.obj.get_graph_capture_ptrs() + if raw_ptrs and is_vmm_pointer(raw_ptrs[0]): + self._vmm_graph_input_manager.register_graph_inputs() + else: + self._register_graph_inputs_ipc() + + def _register_graph_inputs_ipc(self): + """Register graph capture inputs via cudaIpcGetMemHandle. + + This is the fast path for cudaMalloc-backed allocations. Fails + on VMM pointers (expandable_segments). + """ pairs = self.obj.share_graph_inputs() handles = [handle for _, handle in pairs] offsets = [offset for offset, _ in pairs] @@ -108,7 +129,9 @@ class CustomAllReduceV2: offsets_all = self._share_list(offsets) result = [list(zip(o, h)) for o, h in zip(offsets_all, handles_all)] self.obj.register_inputs(result) - log_info_on_rank0(logger, f"Registering {len(pairs)} cuda graph addresses") + log_info_on_rank0( + logger, f"Registered {len(pairs)} cuda graph addresses via IPC" + ) def should_custom_ar(self, inp: torch.Tensor) -> bool: """Check if the input tensor is suitable for custom all-reduce.""" @@ -134,6 +157,8 @@ class CustomAllReduceV2: def close(self): if not self.disabled and hasattr(self, "obj"): self.obj.free(self.group) + if hasattr(self, "_vmm_graph_input_manager"): + self._vmm_graph_input_manager.close() def _all_reduce(self, input: torch.Tensor) -> torch.Tensor: """Perform the actual all-reduce via JIT kernel.""" diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_vmm_utils.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_vmm_utils.py new file mode 100644 index 000000000..239f6de68 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_vmm_utils.py @@ -0,0 +1,579 @@ +import logging +import os +import struct +import time +from typing import Any, List, Optional + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from sglang.srt.utils import log_info_on_rank0 + +logger = logging.getLogger(__name__) + +_drv = None +_FD_HEADER_BYTES = 24 +_FD_SEND_TIMEOUT_S = 120.0 + + +def _get_cuda_driver(): + """Lazily import cuda.bindings.driver (cached after first call).""" + global _drv + if _drv is None: + from cuda.bindings import driver + + _drv = driver + return _drv + + +def _check_drv(result_tuple, label): + """Check a cuda.bindings driver call result and return the value.""" + if not isinstance(result_tuple, tuple): + result_tuple = (result_tuple,) + err = result_tuple[0] + drv = _get_cuda_driver() + if err != drv.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"{label}: {err}") + return result_tuple[1] if len(result_tuple) > 1 else None + + +def is_vmm_pointer(ptr: int) -> bool: + """Check if a device pointer is VMM-backed (cuMemCreate/cuMemMap). + + cuMemRetainAllocationHandle succeeds only on pointers from cuMemCreate; + it fails on cudaMalloc pointers. + """ + drv = _get_cuda_driver() + err, handle = drv.cuMemRetainAllocationHandle(ptr) + if err == drv.CUresult.CUDA_SUCCESS: + drv.cuMemRelease(handle) + return True + return False + + +def _send_fd(sock, fd: int, src_rank: int, base_idx: int) -> None: + import array + import socket + + fds = array.array("i", [int(fd)]) + header = struct.pack(" None: + self.obj = obj + self.group = group + self.rank = rank + self.world_size = world_size + self._peer_mappings = [] + + def register_graph_inputs(self): + """Register graph capture inputs via VMM handle exchange. + + VMM-compatible path for expandable_segments. The C++ side deduplicates + graph capture pointers into unique base allocations via cuMemGetAddressRange. + Python exports handles for each unique base, imports + cuMemMaps peer + allocations, then registers the peer VAs. FABRIC handles are preferred; + POSIX file descriptors are used when FABRIC is unavailable. + """ + drv = _get_cuda_driver() + FABRIC = drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + POSIX_FD = ( + drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ) + FABRIC_HANDLE_BYTES = 64 + MAX_VMM_BASES = 4096 + MAX_CHUNKS_PER_INPUT = 16 + + t0 = time.perf_counter() + + bases_info, input_chunk_indices, input_offsets = ( + self.obj.get_graph_capture_bases() + ) + if not bases_info: + return + new_count = len(input_chunk_indices) + num_bases = len(bases_info) + device_id = torch.cuda.current_device() + + if num_bases > MAX_VMM_BASES: + raise RuntimeError( + f"Too many VMM bases to share: {num_bases} > {MAX_VMM_BASES}" + ) + + local_fabric_handles: List[bytes] = [] + local_posix_fds: List[int] = [] + retained_handles = [] + try: + for base_ptr, _ in bases_info: + alloc_h = _check_drv( + drv.cuMemRetainAllocationHandle(base_ptr), + "cuMemRetainAllocationHandle", + ) + retained_handles.append(alloc_h) + + local_fabric_error: Optional[Exception] = None + try: + for alloc_h in retained_handles: + fabric_h = _check_drv( + drv.cuMemExportToShareableHandle(alloc_h, FABRIC, 0), + "cuMemExportToShareableHandle(FABRIC)", + ) + local_fabric_handles.append(bytes(fabric_h.data)) + local_fabric_ok = True + except Exception as e: + local_fabric_error = e + local_fabric_ok = False + local_fabric_handles = [] + logger.info( + "FABRIC handle export failed on rank %s; falling back to " + "POSIX fd transport: %s", + self.rank, + e, + ) + + use_fabric = self._all_ranks_ok(local_fabric_ok) + if not use_fabric: + local_posix_error: Optional[Exception] = None + try: + for alloc_h in retained_handles: + fd = _check_drv( + drv.cuMemExportToShareableHandle(alloc_h, POSIX_FD, 0), + "cuMemExportToShareableHandle(POSIX_FD)", + ) + local_posix_fds.append(int(fd)) + local_posix_ok = True + except Exception as e: + local_posix_error = e + local_posix_ok = False + for fd in local_posix_fds: + try: + os.close(fd) + except OSError: + pass + local_posix_fds = [] + + if not self._all_ranks_ok(local_posix_ok): + local_cause = local_posix_error or local_fabric_error + message = ( + "VMM graph input registration failed: FABRIC export " + "failed on at least one rank and POSIX fd export failed " + "on at least one rank" + ) + if local_cause is not None: + message += f"; local rank {self.rank} error: {local_cause}" + raise RuntimeError(message) from local_posix_error + + local_input_chunks = [ + [int(idx) for idx in indices] for indices in input_chunk_indices + ] + for chunks in local_input_chunks: + if len(chunks) > MAX_CHUNKS_PER_INPUT: + raise RuntimeError( + "Too many VMM chunks for graph input: " + f"{len(chunks)} > {MAX_CHUNKS_PER_INPUT}" + ) + + # All-gather base metadata and per-input VMM spans. A captured tensor + # can cross expandable-segment allocation boundaries, so peer mappings + # must preserve each input's contiguous virtual-address span. FABRIC + # handles are inline metadata; POSIX fds are exchanged separately via + # SCM_RIGHTS because fd integers are process-local. + header_struct = struct.Struct(" local VA + peer_span_va = {} # (rank, chunk_indices...) -> (local VA, peer base) + new_mappings = [] + + def import_peer_handle(peer_rank: int, base_idx: int, fabric_handle): + if use_fabric: + return _check_drv( + drv.cuMemImportFromShareableHandle(fabric_handle, FABRIC), + f"cuMemImportFromShareableHandle(rank={peer_rank})", + ) + fd = posix_peer_fds[(peer_rank, base_idx)] + dup_fd = os.dup(fd) + try: + return _check_drv( + drv.cuMemImportFromShareableHandle(dup_fd, POSIX_FD), + f"cuMemImportFromShareableHandle(rank={peer_rank}, POSIX_FD)", + ) + finally: + try: + os.close(dup_fd) + except OSError: + pass + + try: + for peer_rank in range(self.world_size): + if peer_rank == self.rank: + for idx, (bp, _) in enumerate(bases_info): + peer_base_va[(peer_rank, idx)] = int(bp) + continue + + peer_bases = all_base_payload[peer_rank] + for idx, (_, fb, alloc_size) in enumerate(peer_bases): + imp_h = import_peer_handle(peer_rank, idx, fb) + prop = _check_drv( + drv.cuMemGetAllocationPropertiesFromHandle(imp_h), + "cuMemGetAllocationPropertiesFromHandle", + ) + gran = _check_drv( + drv.cuMemGetAllocationGranularity( + prop, + drv.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ), + "cuMemGetAllocationGranularity", + ) + va = _check_drv( + drv.cuMemAddressReserve(alloc_size, int(gran), 0, 0), + "cuMemAddressReserve", + ) + _check_drv( + drv.cuMemMap(int(va), alloc_size, 0, imp_h, 0), + "cuMemMap", + ) + access = drv.CUmemAccessDesc() + access.location.type = ( + drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + ) + access.location.id = device_id + access.flags = ( + drv.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + ) + _check_drv( + drv.cuMemSetAccess(int(va), alloc_size, [access], 1), + "cuMemSetAccess", + ) + peer_base_va[(peer_rank, idx)] = int(va) + new_mappings.append((int(va), alloc_size, [(0, alloc_size)])) + _check_drv(drv.cuMemRelease(imp_h), "cuMemRelease(peer)") + + # Build per-input peer VA lists and register. + peer_ptrs = [] + for j in range(new_count): + ptrs_j = [] + for rank in range(self.world_size): + chunks = all_input_chunks[rank][j] + off = all_input_offsets[rank][j] + if len(chunks) == 1: + ptrs_j.append(peer_base_va[(rank, chunks[0])] + off) + continue + + span_key = (rank, *chunks) + if span_key not in peer_span_va: + peer_bases = all_base_payload[rank] + first_base = peer_bases[chunks[0]][0] + last_base, _, last_size = peer_bases[chunks[-1]] + span_size = ( + int(last_base) + int(last_size) - int(first_base) + ) + if rank == self.rank: + span_va = int(first_base) + else: + span_va = _check_drv( + drv.cuMemAddressReserve(span_size, 0, 0, 0), + "cuMemAddressReserve(span)", + ) + mapped_chunks = [] + for chunk_idx in chunks: + base_ptr, fb, alloc_size = peer_bases[chunk_idx] + rel = int(base_ptr) - int(first_base) + imp_h = import_peer_handle(rank, chunk_idx, fb) + _check_drv( + drv.cuMemMap( + int(span_va) + rel, + int(alloc_size), + 0, + imp_h, + 0, + ), + "cuMemMap(span)", + ) + access = drv.CUmemAccessDesc() + access.location.type = ( + drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + ) + access.location.id = device_id + access.flags = ( + drv.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + ) + _check_drv( + drv.cuMemSetAccess( + int(span_va) + rel, + int(alloc_size), + [access], + 1, + ), + "cuMemSetAccess(span)", + ) + mapped_chunks.append((rel, int(alloc_size))) + _check_drv( + drv.cuMemRelease(imp_h), "cuMemRelease(span)" + ) + new_mappings.append( + (int(span_va), span_size, mapped_chunks) + ) + peer_span_va[span_key] = (int(span_va), int(first_base)) + + span_va, _ = peer_span_va[span_key] + ptrs_j.append(span_va + off) + peer_ptrs.append(ptrs_j) + + self.obj.register_peer_mapped_inputs(peer_ptrs) + self._peer_mappings.extend(new_mappings) + except Exception: + self._release_peer_mappings(new_mappings) + raise + finally: + for fd in posix_peer_fds.values(): + os.close(fd) + + elapsed_ms = (time.perf_counter() - t0) * 1000 + transport = "FABRIC" if use_fabric else "POSIX fd" + log_info_on_rank0( + logger, + f"Registered {new_count} cuda graph addresses via " + f"{transport} handles ({num_bases} unique allocations) " + f"in {elapsed_ms:.1f} ms", + ) + finally: + for fd in local_posix_fds: + os.close(fd) + for h in retained_handles: + _check_drv(drv.cuMemRelease(h), "cuMemRelease(retained)") + + def close(self): + if not self._peer_mappings: + return + self._release_peer_mappings(self._peer_mappings) + + def _all_ranks_ok(self, ok: bool) -> bool: + flag = torch.tensor([1 if ok else 0], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.BAND, group=self.group) + return flag.item() == 1 + + def _exchange_posix_fds(self, local_fds: List[int], peer_base_counts: List[int]): + import socket + import tempfile + import threading + + sock_kind = getattr(socket, "SOCK_SEQPACKET", socket.SOCK_STREAM) + sock_dir = tempfile.mkdtemp(prefix="sgl_ar_fd_") + sock_path = os.path.join(sock_dir, f"rank_{self.rank}.sock") + server = socket.socket(socket.AF_UNIX, sock_kind) + server.settimeout(_FD_SEND_TIMEOUT_S) + received_fds = {} + errors = [] + + def recv_loop(): + try: + for _ in range(self.world_size - 1): + conn, _ = server.accept() + with conn: + conn.settimeout(_FD_SEND_TIMEOUT_S) + while True: + packet = _recv_fd(conn) + if packet is None: + break + src_rank, base_idx, fd = packet + key = (src_rank, base_idx) + if key in received_fds: + os.close(fd) + raise RuntimeError(f"duplicate fd for {key}") + received_fds[key] = fd + except BaseException as e: + errors.append(e) + + try: + server.bind(sock_path) + server.listen(self.world_size) + paths = [None] * self.world_size + dist.all_gather_object(paths, sock_path, group=self.group) + + thread = threading.Thread(target=recv_loop, daemon=True) + thread.start() + try: + for peer_rank, peer_path in enumerate(paths): + if peer_rank == self.rank: + continue + with socket.socket(socket.AF_UNIX, sock_kind) as sock: + sock.settimeout(_FD_SEND_TIMEOUT_S) + sock.connect(peer_path) + for base_idx, fd in enumerate(local_fds): + _send_fd(sock, fd, self.rank, base_idx) + finally: + thread.join(_FD_SEND_TIMEOUT_S) + + if thread.is_alive(): + raise RuntimeError("timed out waiting for POSIX fd exchange") + if errors: + raise RuntimeError("POSIX fd exchange receive failed") from errors[0] + + expected = { + (rank, base_idx) + for rank, count in enumerate(peer_base_counts) + if rank != self.rank + for base_idx in range(count) + } + missing = expected.difference(received_fds) + extra = set(received_fds).difference(expected) + if missing or extra: + for fd in received_fds.values(): + os.close(fd) + raise RuntimeError( + "POSIX fd exchange mismatch: " + f"missing={sorted(missing)[:8]}, extra={sorted(extra)[:8]}" + ) + return received_fds + finally: + server.close() + try: + os.unlink(sock_path) + except FileNotFoundError: + pass + try: + os.rmdir(sock_dir) + except OSError: + pass + + def _release_peer_mappings(self, mappings): + drv = _get_cuda_driver() + while mappings: + va, span_size, mapped_chunks = mappings.pop() + for rel, size in mapped_chunks: + _check_drv(drv.cuMemUnmap(int(va) + int(rel), int(size)), "cuMemUnmap") + _check_drv( + drv.cuMemAddressFree(int(va), int(span_size)), "cuMemAddressFree" + ) diff --git a/test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py b/test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py index 3f7d9a5ef..e28dcbc2c 100644 --- a/test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py +++ b/test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py @@ -48,11 +48,15 @@ class TestDeepseekV3FP4MTP(CustomTestCase): "--model-loader-extra-config", '{"enable_multithread_load": true,"num_threads": 64}', ] + env = { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + } cls.process = popen_launch_server( cls.model, cls.base_url, timeout=SERVER_LAUNCH_TIMEOUT, other_args=other_args, + env=env, ) @classmethod