diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 6ef4f938078c..a85106a3b9be 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -2879,7 +2879,7 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec `--enable-kv-cache-sharding` - Shard at-rest KV cache storage at page granularity across the group that currently replicates it (attention-CP ranks for GQA models under prefill CP; attention-TP ranks for MLA models), multiplying the group's unique KV capacity. Prefill workers of PD disaggregation only. + Shard at-rest KV cache storage at page granularity across the group that currently replicates it (attention-CP ranks for GQA models under prefill CP; attention-TP ranks for MLA models), multiplying the group's unique KV capacity. Prefill workers of PD disaggregation only. With NCCL 2.31.2 or newer and supported NVIDIA CUDA hardware, eligible single-node shard groups request zero-SM copy-engine prefix all-gather on a dedicated NCCL communicator. Older NCCL versions keep the SM-kernel all-gather. This path needs cuMem-backed NCCL memory, so on single-node deployments the server sets NCCL_CUMEM_ENABLE=1 for every NCCL communicator of all server processes unless NCCL_CUMEM_ENABLE is already set; it does not under --enable-memory-saver or SGLANG_MOONCAKE_CUSTOM_MEM_POOL, which keep the SM-kernel all-gather. Multi-node deployments keep NCCL's setting (MNNVL fabrics default it to 1); export NCCL_CUMEM_ENABLE=1 there to use this path for shard groups within one node. Set SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER=0 to disable this path. Leave NCCL_CTA_POLICY unset: a global value overrides both the gather communicator's zero-SM policy and the default policy of forward collectives. `False` bool flag (set to enable) diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index 686f20f150c9..b8cd4588802c 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -949,6 +949,11 @@ SGLang supports various environment variables that can be used to configure its Synchronize sampled token ids across tensor-parallel ranks. false + + SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER + With --enable-kv-cache-sharding, register the prefix gather scratch of single-node shard groups as an NCCL symmetric window so NCCL 2.31.2+ runs the all-gather on copy engines (zero SMs). Older NCCL versions keep the SM-kernel all-gather. On single-node deployments this also makes the server set NCCL_CUMEM_ENABLE=1 for every NCCL communicator unless it is already set, except under --enable-memory-saver or SGLANG_MOONCAKE_CUSTOM_MEM_POOL; see --enable-kv-cache-sharding for the full conditions. Set to false to keep the SM-kernel all-gather and the default NCCL_CUMEM_ENABLE=0. + true + diff --git a/python/sglang/srt/distributed/device_communicators/pynccl.py b/python/sglang/srt/distributed/device_communicators/pynccl.py index 46aef5b94523..0c20d630281a 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl.py @@ -34,6 +34,7 @@ def __init__( device: Union[int, str, torch.device], library_path: Optional[str] = None, is_symmetric_memory_enabled: bool = False, + cta_policy: Optional[int] = None, ): """ Args: @@ -44,6 +45,9 @@ def __init__( library_path: the path to the NCCL library. If None, it will use the default library path. is_symmetric_memory_enabled: whether symmetric memory is enabled. + cta_policy: ncclConfig_t.CTAPolicy for this communicator, e.g. + NCCL_CTA_POLICY_ZERO; None keeps the NCCL default. + The global NCCL_CTA_POLICY environment variable overrides it. It is the caller's responsibility to make sure each communicator is bind to a unique device. """ @@ -110,11 +114,14 @@ def __init__( # `torch.cuda.device` is a context manager that changes the # current cuda device to the specified one with torch.cuda.device(device): - if is_symmetric_memory_enabled: - # When symmetric memory is enabled, disable internal - # NCCL cuda event synchronizations to improve performance. + if is_symmetric_memory_enabled or cta_policy is not None: config = ncclConfig_t.create() - config.graphUsageMode = 1 + if is_symmetric_memory_enabled: + # When symmetric memory is enabled, disable internal + # NCCL cuda event synchronizations to improve performance. + config.graphUsageMode = 1 + if cta_policy is not None: + config.CTAPolicy = cta_policy self.comm: ncclComm_t = self.nccl.ncclCommInitRankConfig( self.world_size, self.unique_id, self.rank, config ) diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index 1e048372dd61..143be7df0bfc 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py @@ -1,10 +1,13 @@ +from __future__ import annotations + import ctypes import logging import os import tempfile import traceback +import weakref from contextlib import nullcontext -from typing import Optional +from typing import TYPE_CHECKING, Optional, Tuple import torch @@ -17,8 +20,80 @@ from sglang.srt.runtime_context import get_exec from sglang.srt.utils.common import torch_release +if TYPE_CHECKING: + from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator + from sglang.srt.distributed.device_communicators.pynccl_wrapper import ncclWindow_t + after_2_8_0 = torch_release >= (2, 8) + +class _NcclMemOwner: + """Exposes one ncclMemAlloc buffer to torch, which keeps this object alive + as long as any view of the wrapping tensor exists.""" + + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +def _release_nccl_window_buffer( + comm: PyNcclCommunicator, ptr: int, window: Optional[ncclWindow_t] +) -> None: + # Queued communication or computation may still be touching the buffer. + torch.cuda.synchronize(comm.device) + try: + if window is not None: + comm.deregister_comm_window(window) + finally: + comm.nccl.ncclMemFree(ptr) + + +def allocate_symmetric_buffer( + comm: PyNcclCommunicator, nbytes: int +) -> Tuple[torch.Tensor, bool]: + """Zeroed uint8 tensor over ncclMemAlloc memory registered as a symmetric + window on ``comm`` (a collective call), and whether NCCL returned a window. + + Every rank must request the same size and use matching offsets for views. + This allocation is independent of the shared symmetric memory pool and + --enable-symm-mem. It is deregistered and freed when the last view dies. + The returned flag means a symmetric window only on NCCL >= 2.29.2. + """ + with torch.cuda.device(comm.device): + ptr = comm.nccl.ncclMemAlloc(nbytes) + window = None + release = None + try: + window = comm.register_comm_window_raw(ptr, nbytes) + # NCCL >= 2.29.2 returns NULL when the comm has no symmetric memory + # support, usually because NCCL runs without cuMem (ncclMemAlloc is + # then a cudaMalloc); 2.28.x returns a local-registration handle. + window = window if window.value else None + owner = _NcclMemOwner(ptr, nbytes) + release = weakref.finalize( + owner, _release_nccl_window_buffer, comm, ptr, window + ) + # The process teardown reclaims the memory; NCCL may already be gone. + release.atexit = False + buf = torch.as_tensor(owner) + buf.zero_() + except BaseException: + try: + if release is None: + _release_nccl_window_buffer(comm, ptr, window) + else: + # Calling the finalizer also disables later GC cleanup. + release() + except Exception: + logger.exception("Failed to release NCCL buffer after setup failure") + raise + return buf, window is not None + + # C++ source for the NCCL allocator plugin # Key design: # 1. nccl_alloc_plug: Allocates memory via ncclMemAlloc and TRACKS the segment @@ -238,7 +313,7 @@ def get_nccl_mem_pool() -> torch.cuda.MemPool: # The context whose segments are being registered with NCCL right now. The # registration has no timeout of its own, so the scheduler watchdog reads this # to name a stall inside it. -_registering_context: Optional["SymmetricMemoryContext"] = None +_registering_context: Optional[SymmetricMemoryContext] = None def describe_registration_in_progress() -> Optional[str]: diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py b/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py index d9afda4331a1..7a7b66de9d9e 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py @@ -88,6 +88,9 @@ def find_nccl_library() -> str: # from nccl.h.in. NCCL_CONFIG_VERSION = 23000 # NCCL_VERSION(2, 30, 0) +# ncclConfig_t.CTAPolicy: run eligible collectives without SMs (copy engines). +NCCL_CTA_POLICY_ZERO = 0x02 + class ncclConfig_t(ctypes.Structure): """Mirror of ``ncclConfig_t`` (``ncclConfig_v22800``) from nccl.h.in. @@ -406,6 +409,14 @@ class NCCLLibrary: ), # ncclResult_t ncclCommWindowDeregister(ncclComm_t comm, ncclWindow_t win); Function("ncclCommWindowDeregister", ncclResult_t, [ncclComm_t, ncclWindow_t]), + # ncclResult_t ncclMemAlloc(void** ptr, size_t size); + Function( + "ncclMemAlloc", + ncclResult_t, + [ctypes.POINTER(buffer_type), ctypes.c_size_t], + ), + # ncclResult_t ncclMemFree(void* ptr); + Function("ncclMemFree", ncclResult_t, [buffer_type]), ] # class attribute to store the mapping from the path to the library @@ -441,7 +452,9 @@ def __init__(self, so_file: Optional[str] = None): if so_file not in NCCLLibrary.path_to_dict_mapping: _funcs: Dict[str, Any] = {} - exported_functions = NCCLLibrary.exported_functions + # Copy: optional symbols of this library must not leak into the + # class list that another library path binds from. + exported_functions = list(NCCLLibrary.exported_functions) if hasattr(self.lib, "ncclCommWindowRegister"): exported_functions.extend(NCCLLibrary.exported_functions_symm_mem) for func in exported_functions: @@ -650,6 +663,14 @@ def ncclCommWindowRegister( def ncclCommWindowDeregister(self, comm: ncclComm_t, window: ncclWindow_t) -> None: self.NCCL_CHECK(self._funcs["ncclCommWindowDeregister"](comm, window)) + def ncclMemAlloc(self, size: int) -> int: + ptr = buffer_type() + self.NCCL_CHECK(self._funcs["ncclMemAlloc"](ctypes.byref(ptr), size)) + return ptr.value + + def ncclMemFree(self, ptr: int) -> None: + self.NCCL_CHECK(self._funcs["ncclMemFree"](buffer_type(ptr))) + def ncclGroupStart(self) -> None: self.NCCL_CHECK(self._funcs["ncclGroupStart"]()) diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 574c09b84977..97ddac831014 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1871,8 +1871,27 @@ def _set_envs_and_config(server_args: ServerArgs): if cfg.nnodes > 1 and is_mnnvl_fabric_device(): os.environ.setdefault("NCCL_CUMEM_ENABLE", "1") os.environ.setdefault("NCCL_MNNVL_ENABLE", "1") + # The KV-shard gather registers its scratch as an NCCL symmetric + # window, which requires cuMem-backed NCCL memory. The pool keeps torch + # memory under the memory saver or a custom mem pool. Multi-node + # deployments keep NCCL's setting (MNNVL fabrics default it on above). + kv_shard_symmetric_scratch = ( + cfg.enable_kv_cache_sharding + and cfg.nnodes == 1 + and envs.SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER.get() + and not cfg.enable_memory_saver + and envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() is None + ) if "NCCL_CUMEM_ENABLE" not in os.environ or cfg.enable_symm_mem: - os.environ["NCCL_CUMEM_ENABLE"] = str(int(cfg.enable_symm_mem)) + if kv_shard_symmetric_scratch and not cfg.enable_symm_mem: + logger.info( + "Setting NCCL_CUMEM_ENABLE=1 for every NCCL communicator so the " + "KV-shard prefix gather can use a symmetric window. Set " + "SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER=0 to keep it off." + ) + os.environ["NCCL_CUMEM_ENABLE"] = str( + int(cfg.enable_symm_mem or kv_shard_symmetric_scratch) + ) if ( "NCCL_NVLS_ENABLE" not in os.environ or cfg.enable_nccl_nvls diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4510e77d9cb2..4d7336e7e7b8 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -612,6 +612,11 @@ class Envs: # HND KV layout folds (page, head) into one paged index for per-kv-head sparse # page tables (DP attn); paged backends like trtllm_mha consume it directly. SGLANG_USE_HND_KVCACHE = EnvBool(False) + # --enable-kv-cache-sharding: register single-node prefix gather scratch + # as NCCL symmetric memory and request copy engines on its dedicated + # communicator. Leave NCCL_CTA_POLICY unset to preserve per-comm policies. + # The engine defaults NCCL cuMem on for eligible single-node deployments. + SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER = EnvBool(True) # Attention (aiter, ROCm): route NEXTN spec draft_extend (EAGLE-v2 KV # catch-up) through aiter unified_attention (GQA-packed + split-KV) instead diff --git a/python/sglang/srt/mem_cache/page_interleave_pool.py b/python/sglang/srt/mem_cache/page_interleave_pool.py index 8e09ce8ac454..a1b5cef4daa9 100644 --- a/python/sglang/srt/mem_cache/page_interleave_pool.py +++ b/python/sglang/srt/mem_cache/page_interleave_pool.py @@ -24,13 +24,20 @@ scratch slot laid out ``[prefix region | chunk region | trash page]``: - The prefix is NCCL-allgathered from the shard group into the slot one layer - ahead, on a dedicated side stream with a dedicated ``PyNcclCommunicator`` - ahead. Under rotated - owner-classed allocation the owners of a cached prefix's pages are exactly - cyclic, so per-rank owned counts differ by <= 1: each rank sends its own - prefix pages in local-page order, padded to ``ceil(prefix_pages / N)`` pages - with the (never referenced) trash page — a *regular* in-place allgather, - owner-major output, no reorder pass. + ahead, on a dedicated side stream with a dedicated ``PyNcclCommunicator``. + Under rotated owner-classed allocation the owners of a cached prefix's + pages are exactly cyclic, so per-rank owned counts differ by <= 1: each + rank sends its own prefix pages in local-page order, padded to + ``ceil(prefix_pages / N)`` pages with the (never referenced) trash page — + a *regular* in-place allgather, owner-major output, no reorder pass. +- Symmetric gather scratch (``SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER``, + NCCL >= 2.31.2, single-node shard groups): both slots are carved out of one + ``ncclMemAlloc`` buffer registered as a symmetric window. Every rank carves + the same offsets, because CE addresses a peer's buffer by the local offset + into the window. + The dedicated gather communicator requests ``NCCL_CTA_POLICY_ZERO`` + for eligible copy-engine transfers. Leave the global ``NCCL_CTA_POLICY`` + unset so forward-pass communicators retain their default policy. - The chunk region is staged locally on the compute stream where the write path already holds the full chunk (GQA CP allgathers K/V before the pool write; MLA latent is replicated across the shard group at write time). @@ -51,7 +58,8 @@ from __future__ import annotations import logging -from contextlib import nullcontext +import math +from contextlib import contextmanager, nullcontext from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -61,6 +69,15 @@ ) from sglang.kernels.ops.kvcache import dsa_index_k_cache from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + allocate_symmetric_buffer, +) +from sglang.srt.distributed.device_communicators.pynccl_wrapper import ( + NCCL_CTA_POLICY_ZERO, + NCCLLibrary, +) +from sglang.srt.distributed.parallel_state import in_the_same_node_as +from sglang.srt.environ import envs from sglang.srt.mem_cache import page_interleave from sglang.srt.mem_cache.index_key_cache import IndexKeyCache from sglang.srt.mem_cache.memory_pool import ( @@ -72,7 +89,8 @@ unwrap_write_loc, ) from sglang.srt.mem_cache.utils import get_mla_kv_buffer_triton -from sglang.srt.utils import ceil_div, get_bool_env_var +from sglang.srt.utils import ceil_div, get_bool_env_var, is_cuda +from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: from sglang.srt.distributed.parallel_state import GroupCoordinator @@ -80,6 +98,64 @@ logger = logging.getLogger(__name__) +# (shape, dtype) of one scratch tensor. +_TensorSpec = Tuple[Tuple[int, ...], torch.dtype] + +_NUM_SCRATCH_SLOTS = 2 +# NCCL_WIN_REQUIRED_ALIGNMENT in nccl.h. NCCL checks it only at the window +# base; starting every slot tensor on it is a conservative choice. +_SCRATCH_ALIGN_BYTES = 4096 +# NCCL 2.31.2 fixes internal CE synchronization buffer initialization and the +# two-rank multicast-team setup race (upstream commit 430afbc). Zeroing our +# scratch cannot initialize NCCL's internal buffers. +_NCCL_SYMMETRIC_SCRATCH_MIN_VERSION = 23102 + + +def _symmetric_scratch_supported( + *, memory_saver_enabled: bool, custom_mem_pool, shard_group: GroupCoordinator +) -> bool: + """Collective over ``shard_group``: every rank returns the same verdict.""" + # ncclMemAlloc memory sits outside both the memory-saver region and a + # custom mem pool, so either one keeps the torch-allocated scratch. + if ( + not envs.SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER.get() + or memory_saver_enabled + or custom_mem_pool is not None + or not is_cuda() + or NCCLLibrary().ncclGetRawVersion() < _NCCL_SYMMETRIC_SCRATCH_MIN_VERSION + ): + return False + # Across nodes NCCL would run the windowed allgather on its hierarchical + # path (inter-node RMA proxy puts); only the in-node CE path is verified. + return all(in_the_same_node_as(shard_group.cpu_group, source_rank=0)) + + +def _symmetric_scratch_layout( + specs: Dict[str, _TensorSpec], num_slots: int +) -> Tuple[List[Dict[str, int]], int]: + """Byte offset of every slot tensor in one buffer, and the buffer size. + + Depends only on ``specs``, so it is identical on every rank. + """ + offsets = [] + nbytes = 0 + for _ in range(num_slots): + slot_offsets = {} + for name, (shape, dtype) in specs.items(): + slot_offsets[name] = nbytes + nbytes += ceil_align( + math.prod(shape) * dtype.itemsize, _SCRATCH_ALIGN_BYTES + ) + offsets.append(slot_offsets) + return offsets, nbytes + + +def _view_bytes( + buf: torch.Tensor, offset: int, shape: Tuple[int, ...], dtype: torch.dtype +) -> torch.Tensor: + nbytes = math.prod(shape) * dtype.itemsize + return buf[offset : offset + nbytes].view(dtype).view(shape) + class _ScratchSlot: """One half of the double buffer: scratch tensors + ready event. @@ -100,8 +176,9 @@ class PageInterleaveKVPoolMixin: """Shard-generic state and mechanics; mixed into concrete pools below. Subclasses call ``_init_page_shard_state`` after the base pool has created - its buffers, and implement ``_scratch_tensor_specs`` (per-slot tensors) and - ``_gather_pairs`` (pool buffer -> scratch tensor pairs of one layer). + its buffers, and implement ``_scratch_tensor_specs`` (per-slot tensor + shapes and dtypes) and ``_gather_pairs`` (pool buffer -> scratch tensor + pairs of one layer). """ # ---- init --------------------------------------------------------------- @@ -127,10 +204,20 @@ def _init_page_shard_state( self.shard_size = spec.shard_size self.device_module = torch.get_device_module(self.device) + # Check the pool's storage device before probing its NCCL shard group. + self._use_symmetric_scratch = False + if torch.device(self.device).type == "cuda": + self._use_symmetric_scratch = _symmetric_scratch_supported( + memory_saver_enabled=self.memory_saver_adapter.enabled, + custom_mem_pool=self.custom_mem_pool, + shard_group=shard_group, + ) # Dedicated communicator + stream so the layer-ahead gathers never # interleave with the group's main collectives. self.kv_gather_comm: PyNcclCommunicator = PyNcclCommunicator( - group=shard_group.cpu_group, device=shard_group.device + group=shard_group.cpu_group, + device=shard_group.device, + cta_policy=NCCL_CTA_POLICY_ZERO if self._use_symmetric_scratch else None, ) self.kv_gather_stream = self.device_module.Stream() @@ -138,28 +225,23 @@ def _init_page_shard_state( scratch_rows = spec.scratch_rows self._chunk_base = spec.max_prefix_tokens self._trash_base = spec.max_prefix_tokens + spec.chunk_tokens - with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): - with ( - torch.cuda.use_mem_pool(self.custom_mem_pool) - if self.custom_mem_pool - else nullcontext() - ): - self._slots = [ - _ScratchSlot( - self._scratch_tensor_specs(scratch_rows), self.device_module - ) - for _ in range(2) - ] - # Per-batch absolute scratch-page index of every logical page. - # Unplanned pages initially target the trash page. Logical page - # ids run [0, N * (size/ps + 1)) — N per physical page of one - # rank's pool incl. the reserved padded page. - self._page_pos = torch.full( - (spec.shard_size * (self.size // spec.page_size + 2),), - self._trash_base // spec.page_size, - dtype=torch.int32, - device=self.device, - ) + slot_tensors, self._scratch_window_registered = self._alloc_scratch( + scratch_rows + ) + self._slots = [ + _ScratchSlot(tensors, self.device_module) for tensors in slot_tensors + ] + with self._torch_alloc_context(): + # Per-batch absolute scratch-page index of every logical page. + # Unplanned pages initially target the trash page. Logical page + # ids run [0, N * (size/ps + 1)) — N per physical page of one + # rank's pool incl. the reserved padded page. + self._page_pos = torch.full( + (spec.shard_size * (self.size // spec.page_size + 2),), + self._trash_base // spec.page_size, + dtype=torch.int32, + device=self.device, + ) self._epoch = 0 self._shard_extend_active = False @@ -182,7 +264,8 @@ def _init_page_shard_state( logger.info( "Page-interleave KV sharding enabled: shard_rank=%d shard_size=%d " "page_size=%d prefix_tokens=%d chunk_tokens=%d scratch_rows=%d x2 " - "(%.2f MiB per rank, charged against the KV budget). " + "(%.2f MiB per rank, charged against the KV budget), " + "NCCL symmetric scratch=%s. " "Batch admission is limited by this scratch capacity.", self.shard_rank, self.shard_size, @@ -191,8 +274,57 @@ def _init_page_shard_state( spec.chunk_tokens, scratch_rows, scratch_bytes / (1024 * 1024), + self._scratch_window_registered, ) + @contextmanager + def _torch_alloc_context(self): + with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + with ( + torch.cuda.use_mem_pool(self.custom_mem_pool) + if self.custom_mem_pool + else nullcontext() + ): + yield + + def _alloc_scratch(self, rows: int) -> Tuple[List[Dict[str, torch.Tensor]], bool]: + """Every slot's tensors, and whether they sit in a symmetric window.""" + specs = self._scratch_tensor_specs(rows) + if self._use_symmetric_scratch: + return self._alloc_symmetric_scratch(specs) + with self._torch_alloc_context(): + slots = [ + { + name: torch.zeros(shape, dtype=dtype, device=self.device) + for name, (shape, dtype) in specs.items() + } + for _ in range(_NUM_SCRATCH_SLOTS) + ] + return slots, False + + def _alloc_symmetric_scratch( + self, specs: Dict[str, _TensorSpec] + ) -> Tuple[List[Dict[str, torch.Tensor]], bool]: + offsets, nbytes = _symmetric_scratch_layout(specs, num_slots=_NUM_SCRATCH_SLOTS) + buf, registered = allocate_symmetric_buffer(self.kv_gather_comm, nbytes) + if not registered: + logger.warning( + "NCCL returned no symmetric window for the KV-shard gather " + "scratch (NCCL_CUMEM_ENABLE=0, NCCL_WIN_ENABLE=0, or ranks " + "without CUDA P2P?); the prefix allgather falls back to NCCL " + "SM kernels." + ) + slots = [ + { + name: _view_bytes( + buf, offset=slot_offsets[name], shape=shape, dtype=dtype + ) + for name, (shape, dtype) in specs.items() + } + for slot_offsets in offsets + ] + return slots, registered + def set_kv_buffer_prefix_valid(self, *args, **kwargs): raise NotImplementedError( "prefix-valid commit is unsupported under logical-page KV sharding " @@ -223,7 +355,7 @@ def register_layer_transfer_counter(self, layer_transfer_counter): # ---- subclass hooks ------------------------------------------------------- - def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]: + def _scratch_tensor_specs(self, rows: int) -> Dict[str, _TensorSpec]: raise NotImplementedError def _gather_pairs(self, local_layer: int) -> List[Tuple[torch.Tensor, str]]: @@ -503,18 +635,10 @@ def __init__( assert not self.post_capture_active self._init_page_shard_state(shard_spec, shard_group) - def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]: + def _scratch_tensor_specs(self, rows: int) -> Dict[str, _TensorSpec]: return { - "k": torch.zeros( - (rows, self.head_num, self.head_dim), - dtype=self.store_dtype, - device=self.device, - ), - "v": torch.zeros( - (rows, self.head_num, self.v_head_dim), - dtype=self.store_dtype, - device=self.device, - ), + "k": ((rows, self.head_num, self.head_dim), self.store_dtype), + "v": ((rows, self.head_num, self.v_head_dim), self.store_dtype), } def _gather_pairs(self, local_layer: int): @@ -612,14 +736,8 @@ def __init__( assert not self.use_dsa, "DSA models are not supported by KV sharding yet" self._init_page_shard_state(shard_spec, shard_group) - def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]: - return { - "kv": torch.zeros( - (rows, 1, self.kv_cache_dim), - dtype=self.store_dtype, - device=self.device, - ), - } + def _scratch_tensor_specs(self, rows: int) -> Dict[str, _TensorSpec]: + return {"kv": ((rows, 1, self.kv_cache_dim), self.store_dtype)} def _gather_pairs(self, local_layer: int): return [(self.kv_buffer[local_layer], "kv")] @@ -885,14 +1003,13 @@ def __init__( def _create_index_key_cache(self) -> IndexKeyCache: return _PageInterleaveIndexKeyCache(self, self.index_buf_size) - def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]: - tensors = super()._scratch_tensor_specs(rows) - tensors["index_k"] = torch.zeros( + def _scratch_tensor_specs(self, rows: int) -> Dict[str, _TensorSpec]: + specs = super()._scratch_tensor_specs(rows) + specs["index_k"] = ( self.index_key_cache._buffer_shape(rows // self.page_size), - dtype=self.index_k_with_scale_buffer_dtype, - device=self.device, + self.index_k_with_scale_buffer_dtype, ) - return tensors + return specs def _gather_pairs(self, local_layer: int): pairs = super()._gather_pairs(local_layer) diff --git a/test/registered/unit/mem_cache/test_page_interleave_gather_multi_gpu.py b/test/registered/unit/mem_cache/test_page_interleave_gather_multi_gpu.py index 72a5489766c5..571ca8e73e6a 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_gather_multi_gpu.py +++ b/test/registered/unit/mem_cache/test_page_interleave_gather_multi_gpu.py @@ -30,15 +30,22 @@ exact KV bytes, separate target/draft storage and scratch plans, fragmented pages, partial tail pages, both MLA write APIs, and local/nonzero draft layer IDs across batches. +4. Copy-engine gathers: symmetric-window ownership and release, zero NCCL + kernel launches, cross-node fallback, and comparisons against ordinary + NCCL on 2 and 4 ranks, including CUDA graph replay. -Registered as a base-b 2-gpu-large unit test; skips when fewer than 2 CUDA -devices are visible. Run directly with: +Registered for base-b 2-gpu-large and base-c 4-gpu-b200; multi-GPU tests skip +when too few CUDA devices are visible. Run directly with: CUDA_VISIBLE_DEVICES=0,1 python3 \ test/registered/unit/mem_cache/test_page_interleave_gather_multi_gpu.py """ +import ctypes +import gc import os +import tempfile import unittest +import unittest.mock from types import SimpleNamespace import torch @@ -48,20 +55,33 @@ init_distributed_environment, initialize_model_parallel, ) +from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + allocate_symmetric_buffer, +) +from sglang.srt.distributed.device_communicators.pynccl_wrapper import ( + NCCL_CTA_POLICY_ZERO, + NCCLLibrary, +) +from sglang.srt.environ import envs from sglang.srt.mem_cache.allocator.page_interleave import PageInterleavePoolAllocator from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator from sglang.srt.mem_cache.page_interleave import PageShardSpec, get_kv_shard_group from sglang.srt.mem_cache.page_interleave_pool import ( + _NCCL_SYMMETRIC_SCRATCH_MIN_VERSION, PageInterleaveMHATokenToKVPool, PageInterleaveMLATokenToKVPool, + _symmetric_scratch_layout, ) from sglang.srt.runtime_context import get_context, get_parallel, publish from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.utils.common import temp_set_env from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase, publish_build_topology register_cuda_ci(est_time=120, stage="base-b", runner_config="2-gpu-large") +register_cuda_ci(est_time=120, stage="base-c", runner_config="4-gpu-b200") _GATHER_WORLD = 2 _GATHER_LAYER_NUM = 4 @@ -263,6 +283,7 @@ def _run_mla(rank, world, port): torch.distributed.barrier() if rank == 0: print("PASS: MLA page-interleave shard (attn-TP axis)") + return pool, req_to_token, prefix_len, seq_len def _run_mha(rank, world, port): @@ -683,5 +704,286 @@ def test_mtp_mha_shard_over_attention_cp(self): mp.spawn(_run_mtp, args=(29814, "mha"), nprocs=_GATHER_WORLD, join=True) +@unittest.skipIf(not torch.cuda.is_available(), "needs a CUDA device") +class TestSymmetricBufferLifetime(CustomTestCase): + def test_zeroed_buffer_deregisters_before_free(self): + """The allocation needs only window registration and memory ownership.""" + released = [] + + def free(ptr): + released.append(("free", ptr)) + torch.cuda.caching_allocator_delete(ptr) + + comm = SimpleNamespace( + device=torch.device("cuda", torch.cuda.current_device()), + nccl=SimpleNamespace( + ncclMemAlloc=torch.cuda.caching_allocator_alloc, + ncclMemFree=free, + ), + register_comm_window_raw=lambda ptr, size: ctypes.c_void_p(ptr), + deregister_comm_window=lambda window: released.append( + ("deregister", window.value) + ), + ) + buf, registered = allocate_symmetric_buffer(comm, 8192) + self.assertTrue(registered) + self.assertEqual(int(buf.count_nonzero()), 0) + ptr = buf.data_ptr() + del buf + gc.collect() + self.assertEqual(released, [("deregister", ptr), ("free", ptr)]) + + +def _assert_bytes(got, expected, label): + got = got.contiguous().view(torch.uint8) + expected = expected.contiguous().view(torch.uint8) + assert got.shape == expected.shape, label + assert torch.equal(got, expected), f"{label}: cache bytes differ" + + +def _gather_gpu_activity(pool, req_to_token, prefix_len, seq_len): + """Names of the GPU kernels and copies one sharded batch's gathers issue.""" + torch.cuda.synchronize() + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA] + ) as prof: + pool.begin_shard_extend( + req_to_token, torch.tensor([0]), [prefix_len], [seq_len] + ) + for layer_id in range(_GATHER_LAYER_NUM): + pool.get_key_buffer(layer_id) + torch.cuda.synchronize() + return [ + event.name + for event in prof.events() + if event.device_type == torch.autograd.DeviceType.CUDA + ] + + +def _run_zero_sm_mla(rank, world, port): + pool, req_to_token, prefix_len, seq_len = _run_mla(rank, world, port) + assert pool._use_symmetric_scratch and pool._scratch_window_registered + + # NCCL picks copy engines per collective, so a windowed gather can still + # land on ring or symmetric SM kernels; only the launched work tells. + activity = _gather_gpu_activity(pool, req_to_token, prefix_len, seq_len) + nccl_kernels = sorted({name for name in activity if "nccl" in name.lower()}) + assert not nccl_kernels, f"gather ran on NCCL SM kernels: {nccl_kernels}" + assert any("PtoP" in name for name in activity), "no peer copy-engine copy" + + scratch_bytes = sum( + tensor.numel() * tensor.element_size() + for slot in pool._slots + for tensor in slot.tensors.values() + ) + view = pool._slots[1].tensors["kv"] + # Deregistration can allocate a CUDA stream whose driver storage persists. + # Account for that separately from the backing allocation being released. + comm = pool.kv_gather_comm + deregister = comm.deregister_comm_window + taken_by_deregister = [] + + def measured_deregister(window): + free_before = torch.cuda.mem_get_info()[0] + deregister(window) + taken = free_before - torch.cuda.mem_get_info()[0] + taken_by_deregister.append(max(taken, 0)) + + comm.deregister_comm_window = measured_deregister + torch.cuda.synchronize() + torch.distributed.barrier() + free_with_pool = torch.cuda.mem_get_info()[0] + del pool + gc.collect() + torch.distributed.barrier() + free_with_view = torch.cuda.mem_get_info()[0] + assert not taken_by_deregister, "window deregistered under a live view" + # A peer's window maps this buffer until that peer deregisters too. + del view + gc.collect() + torch.distributed.barrier() + free_released = torch.cuda.mem_get_info()[0] + assert len(taken_by_deregister) == 1, "window not deregistered exactly once" + assert free_with_view - free_with_pool < scratch_bytes, "freed under a live view" + freed = free_released - free_with_view + taken_by_deregister[0] + assert freed >= scratch_bytes, "window buffer leaked" + + +def _run_cross_node_mla(rank, world, port): + with unittest.mock.patch( + "sglang.srt.mem_cache.page_interleave_pool.in_the_same_node_as", + lambda pg, source_rank: [r == source_rank for r in range(world)], + ): + pool, req_to_token, prefix_len, seq_len = _run_mla(rank, world, port) + assert not pool._use_symmetric_scratch and not pool._scratch_window_registered + activity = _gather_gpu_activity(pool, req_to_token, prefix_len, seq_len) + assert any("nccl" in name.lower() for name in activity) + + +def _run_zero_sm_allgather_comparison(rank, world, init_method): + """Compare the pool's in-place window contract with ordinary NCCL.""" + torch.cuda.set_device(rank) + torch.distributed.init_process_group( + backend="gloo", rank=rank, world_size=world, init_method=init_method + ) + device = torch.device(f"cuda:{rank}") + group = torch.distributed.group.WORLD + zero_sm = PyNcclCommunicator(group, device, cta_policy=NCCL_CTA_POLICY_ZERO) + ordinary = PyNcclCommunicator(group, device) + assert zero_sm.available and ordinary.available + max_send_bytes = 1 << 20 + + for dtype in (torch.uint8, torch.bfloat16, torch.float32): + # Both alternating slots start inside the window, rather than at its + # base. The input is each rank's slice of the registered receive view, + # exactly as in _prefetch_layer; its offset also changes with payload. + offsets, nbytes = _symmetric_scratch_layout( + { + "padding": ((4096,), torch.uint8), + "output": ((world * max_send_bytes // dtype.itemsize,), dtype), + }, + num_slots=2, + ) + buffer, registered = allocate_symmetric_buffer(zero_sm, nbytes) + assert registered, "comparison requires a real symmetric window" + slots = [ + buffer[slot["output"] : slot["output"] + world * max_send_bytes].view(dtype) + for slot in offsets + ] + assert all(slot["output"] > 0 for slot in offsets) + + def compare(slot_index, send_bytes, generation, graph=None): + count = send_bytes // dtype.itemsize + output = slots[slot_index][: world * count] + send = output[rank * count : (rank + 1) * count] + output.zero_() + values = ( + (torch.arange(count, device=device, dtype=torch.int32) * 17) + + rank * 29 + + generation * 31 + ) % 251 + send.copy_(values) + reference_input = send.clone() + reference = torch.empty_like(output) + if graph is None: + with zero_sm.change_state(enable=True): + zero_sm.all_gather(output, send) + else: + graph.replay() + # Use the same loaded NCCL library for the ordinary baseline, + # with torch allocations and the default communicator policy. + with ordinary.change_state(enable=True): + ordinary.all_gather(reference, reference_input) + _assert_bytes( + output, + reference, + f"rank {rank}, {dtype}, {send_bytes} B, generation {generation}", + ) + + for send_bytes in (256, 16 << 10, max_send_bytes): + for generation in range(6): + compare(generation % 2, send_bytes, generation) + + # Profile only the CE call: the ordinary baseline deliberately issues + # an NCCL kernel and must not contaminate the zero-SM assertion. + output = slots[1] + count = max_send_bytes // dtype.itemsize + send = output[rank * count : (rank + 1) * count] + torch.cuda.synchronize() + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA] + ) as prof: + with zero_sm.change_state(enable=True): + zero_sm.all_gather(output, send) + torch.cuda.synchronize() + activity = [ + event.name + for event in prof.events() + if event.device_type == torch.autograd.DeviceType.CUDA + ] + assert not any("nccl" in name.lower() for name in activity), activity + assert any("PtoP" in name for name in activity), activity + + # Replay the same registered addresses with fresh contributions. A + # comparison only at capture time would miss stale peer data on replay. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + with zero_sm.change_state(enable=True): + zero_sm.all_gather(output, send) + for generation in (7, 8, 9): + compare(1, max_send_bytes, generation, graph=graph) + del graph, send, output + slots.clear() + del buffer + gc.collect() + torch.distributed.barrier() + + torch.cuda.synchronize() + torch.distributed.destroy_process_group() + + +@unittest.skipIf( + torch.cuda.device_count() < 2, "page-interleave gather needs 2 CUDA devices" +) +class TestPageInterleaveZeroSmGatherMultiGpu(CustomTestCase): + def _require_patched_nccl(self): + # Gate on the library used by PyNccl, which may differ from torch's. + if NCCLLibrary().ncclGetRawVersion() < _NCCL_SYMMETRIC_SCRATCH_MIN_VERSION: + self.skipTest("the symmetric gather scratch needs NCCL >= 2.31.2") + + def _compare_allgather(self, world): + if torch.cuda.device_count() < world: + self.skipTest(f"copy-engine allgather comparison needs {world} GPUs") + self._require_patched_nccl() + with ( + temp_set_env(NCCL_CUMEM_ENABLE="1", NCCL_CTA_POLICY=None), + tempfile.TemporaryDirectory(prefix="sglang-zero-sm-gather-") as tmp, + ): + mp.spawn( + _run_zero_sm_allgather_comparison, + args=(world, f"file://{tmp}/rendezvous"), + nprocs=world, + join=True, + ) + + def test_copy_engine_matches_ordinary_allgather_two_gpu(self): + self._compare_allgather(2) + + def test_copy_engine_matches_ordinary_allgather_four_gpu(self): + self._compare_allgather(4) + + def test_copy_engine_window_gather_and_release(self): + """The zero-SM pool registers its scratch as a symmetric window, still + assembles the canonical prefix, and frees the window buffer only once + the pool and every scratch view are gone.""" + self._require_patched_nccl() + # Clear the global override before spawning: only the dedicated + # gather communicator requests ZERO. cuMem enables window registration. + with ( + temp_set_env(NCCL_CUMEM_ENABLE="1", NCCL_CTA_POLICY=None), + envs.SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER.override(True), + ): + mp.spawn( + _run_zero_sm_mla, + args=(_GATHER_WORLD, 29811), + nprocs=_GATHER_WORLD, + join=True, + ) + + def test_cross_node_group_keeps_sm_gather(self): + """A shard group spanning nodes keeps the torch scratch and the SM + allgather, even with cuMem on.""" + with ( + temp_set_env(NCCL_CUMEM_ENABLE="1", NCCL_CTA_POLICY=None), + envs.SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER.override(True), + ): + mp.spawn( + _run_cross_node_mla, + args=(_GATHER_WORLD, 29811), + nprocs=_GATHER_WORLD, + join=True, + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index 2451c676ae46..eeed10d04ae0 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -37,16 +37,20 @@ class draws (owners exactly cyclic along a chain), least-full root 8. Real sharded MLA pools: N ranks of stock-size pools hold N times one pool's tokens, each token stored on exactly one rank. -The real NCCL MLA/MHA/MTP gathers over 2 GPUs live in +The real NCCL MLA/MHA/MTP gathers, zero-SM gather correctness and release, +and symmetric buffer lifetime tests live in ``test_page_interleave_gather_multi_gpu.py``. DSA indexer and attention coverage lives in ``test_page_sharded_dsa_attention.py``. """ +import ctypes +import gc +import math import shutil import unittest import unittest.mock from array import array -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack, contextmanager, nullcontext from types import SimpleNamespace import numpy as np @@ -55,6 +59,10 @@ class draws (owners exactly cyclic along a chain), least-full root from sglang.srt.disaggregation.common.conn import CommonKVSender from sglang.srt.disaggregation.utils import filter_kv_indices_for_shard_rank +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + _release_nccl_window_buffer, + allocate_symmetric_buffer, +) from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.srt.mem_cache import page_interleave @@ -85,9 +93,12 @@ class draws (owners exactly cyclic along a chain), least-full root make_page_shard_spec, ) from sglang.srt.mem_cache.page_interleave_pool import ( + _NCCL_SYMMETRIC_SCRATCH_MIN_VERSION, PageInterleaveKVPoolMixin, PageInterleaveMHATokenToKVPool, PageInterleaveMLATokenToKVPool, + _symmetric_scratch_layout, + _symmetric_scratch_supported, ) from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache.components import ComponentType @@ -216,12 +227,147 @@ def test_estimate_matches_actual_mla_and_gqa_scratch_tensors(self): if use_mla else PageInterleaveMHATokenToKVPool ) - tensors = pool_cls._scratch_tensor_specs(pool, spec.scratch_rows) + specs = pool_cls._scratch_tensor_specs(pool, spec.scratch_rows) actual = 2 * sum( - tensor.numel() * tensor.element_size() - for tensor in tensors.values() + math.prod(shape) * dtype.itemsize for shape, dtype in specs.values() ) self.assertEqual(estimated, actual) + # The zero-SM buffer adds at most one window alignment per tensor. + _, symmetric_bytes = _symmetric_scratch_layout(specs, num_slots=2) + self.assertGreaterEqual(symmetric_bytes, estimated) + self.assertLess(symmetric_bytes, estimated + 2 * len(specs) * 4096) + + +class TestZeroSmScratchLayout(CustomTestCase): + def test_requires_nccl_with_copy_engine_race_fix(self): + self.assertEqual(_NCCL_SYMMETRIC_SCRATCH_MIN_VERSION, 23102) + for version in (22902, 23007, 23101, 23102): + with ( + self.subTest(version=version), + envs.SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER.override(True), + unittest.mock.patch( + "sglang.srt.mem_cache.page_interleave_pool.is_cuda", + return_value=True, + ), + unittest.mock.patch( + "sglang.srt.mem_cache.page_interleave_pool.NCCLLibrary" + ) as library, + unittest.mock.patch( + "sglang.srt.mem_cache.page_interleave_pool.in_the_same_node_as", + return_value=[True, True], + ) as same_node, + ): + library.return_value.ncclGetRawVersion.return_value = version + self.assertEqual( + _symmetric_scratch_supported( + memory_saver_enabled=False, + custom_mem_pool=None, + shard_group=SimpleNamespace(cpu_group=object()), + ), + version >= _NCCL_SYMMETRIC_SCRATCH_MIN_VERSION, + ) + if version < _NCCL_SYMMETRIC_SCRATCH_MIN_VERSION: + same_node.assert_not_called() + + def test_slot_tensors_are_window_aligned_and_disjoint(self): + """Both slots share one symmetric window: every tensor must start on + NCCL's 4 KiB window alignment and no two tensors may overlap.""" + rows = 13 * 16 # neither row width below fills a whole 4 KiB block + specs = { + "kv": ((rows, 1, 576), torch.bfloat16), + "index_k": ((rows // 16, 16 * 132), torch.uint8), + } + offsets, nbytes = _symmetric_scratch_layout(specs, num_slots=2) + spans = sorted( + (slot[name], slot[name] + math.prod(shape) * dtype.itemsize) + for slot in offsets + for name, (shape, dtype) in specs.items() + ) + self.assertEqual(len(spans), 4) + self.assertTrue(all(start % 4096 == 0 for start, _ in spans)) + for (_, end), (next_start, _) in zip(spans, spans[1:]): + self.assertLessEqual(end, next_start) + self.assertLessEqual(spans[-1][1], nbytes) + + +class TestSymmetricBufferFailureCleanup(CustomTestCase): + @contextmanager + def _fixture(self): + comm = SimpleNamespace( + device=torch.device("cuda:0"), + nccl=SimpleNamespace( + ncclMemAlloc=unittest.mock.Mock(return_value=1234), + ncclMemFree=unittest.mock.Mock(), + ), + register_comm_window_raw=unittest.mock.Mock( + return_value=ctypes.c_void_p(5678) + ), + deregister_comm_window=unittest.mock.Mock(), + ) + with ( + unittest.mock.patch("torch.cuda.device", return_value=nullcontext()), + unittest.mock.patch("torch.cuda.synchronize"), + unittest.mock.patch("torch.as_tensor") as as_tensor, + ): + yield comm, as_tensor + + def test_registration_failure_frees_allocation(self): + with self._fixture() as (comm, as_tensor): + comm.register_comm_window_raw.side_effect = RuntimeError("register failed") + with self.assertRaisesRegex(RuntimeError, "register failed"): + allocate_symmetric_buffer(comm, 8192) + gc.collect() + as_tensor.assert_not_called() + comm.deregister_comm_window.assert_not_called() + comm.nccl.ncclMemFree.assert_called_once_with(1234) + + def test_tensor_setup_failures_release_once(self): + for operation in ("wrap", "zero"): + with ( + self.subTest(operation=operation), + self._fixture() as (comm, as_tensor), + ): + if operation == "wrap": + as_tensor.side_effect = RuntimeError("wrap failed") + else: + as_tensor.return_value.zero_.side_effect = RuntimeError( + "zero failed" + ) + with self.assertRaisesRegex(RuntimeError, f"{operation} failed"): + allocate_symmetric_buffer(comm, 8192) + # Drop Mock's saved owner reference as well: a registered + # finalizer must not run a second time after eager cleanup. + as_tensor.reset_mock() + gc.collect() + comm.deregister_comm_window.assert_called_once() + self.assertEqual( + comm.deregister_comm_window.call_args.args[0].value, 5678 + ) + comm.nccl.ncclMemFree.assert_called_once_with(1234) + + def test_cleanup_failure_preserves_setup_error_without_double_free(self): + with self._fixture() as (comm, as_tensor): + as_tensor.side_effect = RuntimeError("wrap failed") + comm.deregister_comm_window.side_effect = RuntimeError("deregister failed") + with ( + self.assertLogs( + "sglang.srt.distributed.device_communicators.pynccl_allocator", + level="ERROR", + ), + self.assertRaisesRegex(RuntimeError, "wrap failed"), + ): + allocate_symmetric_buffer(comm, 8192) + as_tensor.reset_mock() + gc.collect() + comm.deregister_comm_window.assert_called_once() + comm.nccl.ncclMemFree.assert_called_once_with(1234) + + def test_deregister_failure_still_frees_allocation(self): + with self._fixture() as (comm, _): + comm.deregister_comm_window.side_effect = RuntimeError("deregister failed") + with self.assertRaisesRegex(RuntimeError, "deregister failed"): + _release_nccl_window_buffer(comm, 1234, ctypes.c_void_p(5678)) + comm.nccl.ncclMemFree.assert_called_once_with(1234) def _make_mtp_configurator(use_mla=True):