[Perf] KVShard: 0SM KV Cache Allgather for KV Cache Sharding

docs/docs/advanced_features/server_arguments.mdx

@@ -2879,7 +2879,7 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-kv-cache-sharding`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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 <code>NCCL_CUMEM_ENABLE=1</code> for every NCCL communicator of all server processes unless <code>NCCL_CUMEM_ENABLE</code> is already set; it does not under <code>--enable-memory-saver</code> or <code>SGLANG_MOONCAKE_CUSTOM_MEM_POOL</code>, which keep the SM-kernel all-gather. Multi-node deployments keep NCCL's setting (MNNVL fabrics default it to 1); export <code>NCCL_CUMEM_ENABLE=1</code> there to use this path for shard groups within one node. Set <code>SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER=0</code> to disable this path. Leave <code>NCCL_CTA_POLICY</code> unset: a global value overrides both the gather communicator's zero-SM policy and the default policy of forward collectives.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
     </tr>

docs/docs/references/environment_variables.mdx

@@ -949,6 +949,11 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Synchronize sampled token ids across tensor-parallel ranks.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>With <code>--enable-kv-cache-sharding</code>, 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 <code>NCCL_CUMEM_ENABLE=1</code> for every NCCL communicator unless it is already set, except under <code>--enable-memory-saver</code> or <code>SGLANG_MOONCAKE_CUSTOM_MEM_POOL</code>; see <code>--enable-kv-cache-sharding</code> for the full conditions. Set to <code>false</code> to keep the SM-kernel all-gather and the default <code>NCCL_CUMEM_ENABLE=0</code>.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
   </tbody>
 </table>
 

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
                 )

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]:

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"]())
 

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

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