[None][perf] BREAKING: Enable zero-copy mode for Mooncake Store

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/cudaVirtMem.cpp

@@ -19,6 +19,7 @@
 #include "kv_cache_manager_v2/exceptions.h"
 #include "kv_cache_manager_v2/utils/math.h"
 #include "tensorrt_llm/common/assert.h"
+#include "tensorrt_llm/common/logger.h"
 
 namespace tensorrt_llm::batch_manager::kv_cache_manager_v2
 {
@@ -102,6 +103,15 @@ PooledPhysMemAllocator::PooledPhysMemAllocator(size_t physMemSize)
     {
         throw std::runtime_error("PooledPhysMemAllocator: no supported physical memory allocation property");
     }
+    if (!mProp.allocFlags.gpuDirectRDMACapable)
+    {
+        // Otherwise this surfaces only later, as an opaque RDMA registration failure.
+        TLLM_LOG_WARNING(
+            "PooledPhysMemAllocator: device %d refused gpuDirectRDMACapable for KV cache memory with handle type %d, "
+            "so the pools are allocated without it. nvidia_peermem (ibv_reg_mr) cannot register this memory for "
+            "GPUDirect RDMA; dma-buf registration still can.",
+            static_cast<int>(mDeviceId), static_cast<int>(mProp.requestedHandleTypes));
+    }
 }
 
 PooledPhysMemAllocator::~PooledPhysMemAllocator()

docs/source/features/kv-cache-connector.md

@@ -456,7 +456,7 @@ says only what the engine then does with the pool.
 | `both` | yes | yes | An aggregated server, or a context server in a disaggregated deployment |
 | `producer` | no | yes | A server that fills the pool for others |
 | `consumer` | yes | no | A server that only replays what others stored |
-| `capacity` | no | no | A generation server: it lends memory and moves no KV |
+| `capacity` | no | no | A generation server: it lends memory and needs no GPUDirect RDMA |
 
 `capacity` is the right role for a disaggregated generation server. Generated tokens are rarely a
 reused prefix, and prompt KV reaches decode over the cache transceiver rather than through the
@@ -494,9 +494,48 @@ sliding-window attention, and `sparse_attention_config` with an index-V cache un
 `sparse_disable_index_value=True`; it bypasses individual requests that carry a LoRA adapter or
 multimodal content without hashes, leaving the rest of the deployment served.
 
-### How pages reach the pool
+### GPUDirect RDMA
 
-Every page passes through a pinned host buffer. `stage_through_host` defaults to `true` and is the
-only supported value today, so `false` is ignored with a warning; registering the KV pools with
-Mooncake instead is planned for a later release. Staging costs a copy each way and needs no
-GPUDirect RDMA. A `capacity` server stages nothing, holding no KV in the pool.
+By default the connector registers the KV pools with Mooncake so the NIC reads and writes pages in
+GPU memory directly. A KV cache manager V2 pool is many `cuMemCreate` allocations mapped back to back,
+and no registration may cross two of them, so the connector registers one allocation at a time.
+Which registration path Mooncake then takes is decided by `WITH_NVIDIA_PEERMEM`, and each path has
+its own preconditions:
+
+| `WITH_NVIDIA_PEERMEM` | Mooncake calls | Requires | Typical failure |
+|---|---|---|---|
+| unset or nonzero (default) | `ibv_reg_mr` | the `nvidia_peermem` kernel module, and KV pools allocated with `gpuDirectRDMACapable=1` | `EFAULT` ("Bad address") |
+| `0` | `ibv_reg_dmabuf_mr` | dma-buf support in the driver and NIC stack | `EINVAL` |
+
+Under `ibv_reg_mr`, a missing module and memory without `gpuDirectRDMACapable` both fail with the
+same `EFAULT` as a range crossing two allocations, so the error code alone does not say which. The
+cache manager asks for `gpuDirectRDMACapable` and logs a warning when the platform refuses it and it
+falls back to memory without it; `nvidia_peermem` cannot register such memory, while dma-buf can.
+
+Hosts with the open kernel modules and no `nvidia_peermem`, common on GB200 and GB300, need the
+dma-buf path. When the protocol is `rdma`, `WITH_NVIDIA_PEERMEM` is unset, neither `nvidia_peermem`
+nor the legacy `nv_peer_mem` module is loaded, and every visible device reports dma-buf support, the
+connector sets `WITH_NVIDIA_PEERMEM=0` itself and logs a warning saying so. A value the deployment
+sets is never overridden.
+
+Mooncake reads `WITH_NVIDIA_PEERMEM` once per process, when the first Mooncake client starts, so the
+connector only chooses for that first client. A later client that registers the KV pools over RDMA
+fails at startup when the path Mooncake already read is not the one it needs, for example after a
+TCP client started first on a host without `nvidia_peermem`. In a process that starts more than one
+Mooncake client, export the variable before the first one.
+
+A failed registration fails startup with a report of what the driver says about each range. At
+debug log level (`TLLM_LOG_LEVEL=debug`) the same report is logged before registering.
+
+The pools are registered once, at startup, and stay registered until shutdown. That is only valid
+while the memory behind them stays mapped, which is why `enable_kv_pool_rebalance` is
+[ignored under a connector](#kv-cache-tiers-are-gpu-only-under-a-connector): a rebalance remaps
+pool memory, and every range would have to be unregistered before it is unmapped and registered
+again afterwards.
+
+### Without GPUDirect RDMA
+
+A host without `nvidia_peermem` usually needs only `WITH_NVIDIA_PEERMEM=0`, which the connector
+selects on its own where it can. When neither path is available, set `stage_through_host: true` to
+copy pages through a pinned host buffer instead, at the cost of a copy each way. A `capacity` server
+needs neither, since it registers no pages.

tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py

@@ -37,7 +37,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Set, Tuple
 
 import torch
 
@@ -59,6 +59,10 @@
     from .kv_cache_layout import KvCacheLayout
 
 
+#: See `KvCacheConnectorWorker.get_allocation_fence`.
+AllocationFence = Callable[[Iterable[int]], None]
+
+
 # Used to store data for a single inflight request.
 @dataclass
 class RequestData:
@@ -170,6 +174,7 @@ class KvCacheConnectorWorker(ABC):
     def __init__(self, llm_args: TorchLlmArgs):
         self._llm_args = llm_args
         self._metadata = None
+        self._allocation_fence: Optional[AllocationFence] = None
         super().__init__()
 
     def __init_subclass__(cls, **kwargs):
@@ -204,6 +209,20 @@ def get_connector_meta(self) -> object:
     def _clear_connector_meta(self):
         self._metadata = None
 
+    def bind_allocation_fence(self, fence: Optional[AllocationFence]):
+        self._allocation_fence = fence
+
+    def get_allocation_fence(self) -> Optional[AllocationFence]:
+        """A callable that blocks until the named requests' pages are free of earlier GPU work.
+
+        A writer no CUDA stream orders, such as the NIC, must call it with the
+        requests it is about to write into first.
+
+        Returns:
+            The fence, or None when the KV cache manager offers none.
+        """
+        return self._allocation_fence
+
     def register_forward_pass_callable(self) -> Callable:
         """
         This callable will be called at the end of the forward pass.
@@ -743,6 +762,7 @@ def __init__(
         self.finished_async_loading_requests = dict()
 
         self._scheduler_output = None
+        self._allocation_fence: Optional[AllocationFence] = None
         self.scheduler_output_manager = KvCacheConnectorSchedulerOutputManager()
         self.prefix_reservations_enabled = False
         self._prefix_capability = None
@@ -1229,6 +1249,9 @@ def build_scheduler_output(
         self._scheduler_output = self.scheduler_output_manager.build_scheduler_output(
             scheduled_batch, async_requests, kv_cache_manager
         )
+        # Every page of this batch is allocated by now; see
+        # `KvCacheConnectorWorker.get_allocation_fence`.
+        self._allocation_fence = getattr(kv_cache_manager, "wait_allocation_fences", None)
         self._scheduler_output.prefix_loads = [
             load
             for reservation_id, load in self._prefix_loads.items()
@@ -1287,6 +1310,8 @@ def handle_metadata(self) -> object:
         self._bound_scheduler_output = self._scheduler_output
         self._scheduler_output = None
 
+        self.worker.bind_allocation_fence(self._allocation_fence)
+        self._allocation_fence = None
         self.worker.bind_connector_meta(metadata)
 
     def request_finished(

tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py

@@ -23,19 +23,108 @@
 and parallel layout because `build_kv_cache_layout_v2` derives it from the
 allocator's own aggregation. `bytes_per_page` goes into the key namespace so a
 geometry change cannot be read as a valid page.
+
+A V2 GPU pool is a virtual reservation mapped with many `cuMemCreate`
+allocations, and GPUDirect cannot register a range that crosses two of them.
+Given an `AllocationLookup`, registrations and transfer buffers are cut at
+those boundaries.
 """
 
-from typing import Dict, List, Sequence, Tuple
+from bisect import bisect_right
+from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple
 
 from ..kv_cache_layout import KvCacheLayout, KvCacheRegion
 
 __all__ = ["PageAddressing"]
 
+#: Resolves an address to `(base, size)` of the allocation mapped there, or
+#: None when it cannot. `gpudirect.allocation_range` asks the driver.
+AllocationLookup = Callable[[int], Optional[Tuple[int, int]]]
+
+
+def merge_intervals(intervals: Iterable[Tuple[int, int]]) -> List[Tuple[int, int]]:
+    """Collapse `(start, end)` byte ranges into a minimal disjoint cover.
+
+    A range may not be registered twice, but several regions routinely live
+    inside one pool allocation: sliding-window layer groups share it, and a
+    non-uniform slot such as MiniMax-M3's index-K beside K/V splits one pool
+    into several regions. Merging first spares the caller that distinction.
+    """
+    ordered = sorted((int(start), int(end)) for start, end in intervals if end > start)
+    merged: List[Tuple[int, int]] = []
+    for start, end in ordered:
+        if merged and start <= merged[-1][1]:
+            previous_start, previous_end = merged[-1]
+            merged[-1] = (previous_start, max(previous_end, end))
+        else:
+            merged.append((start, end))
+    return merged
+
+
+def allocation_spans(
+    start: int, end: int, lookup: AllocationLookup
+) -> Optional[List[Tuple[int, int]]]:
+    """Cut `[start, end)` at the boundaries of the allocations it covers.
+
+    Args:
+        start: First byte.
+        end: One past the last byte.
+        lookup: Resolves an address to its allocation.
+
+    Returns:
+        Consecutive `(start, end)` pieces, one per allocation touched, or None
+        if any address cannot be resolved.
+    """
+    spans: List[Tuple[int, int]] = []
+    cursor = start
+    while cursor < end:
+        allocation = lookup(cursor)
+        if allocation is None:
+            return None
+        base, size = allocation
+        stop = min(base + size, end)
+        if not base <= cursor < stop:
+            # An answer not containing `cursor` would loop or skip bytes.
+            return None
+        spans.append((cursor, stop))
+        cursor = stop
+    return spans
+
+
+def split_at_cuts(start: int, size: int, cuts: Sequence[int]) -> List[Tuple[int, int]]:
+    """Cut `[start, start + size)` at every address in `cuts` strictly inside it.
+
+    Args:
+        start: First byte.
+        size: Length in bytes.
+        cuts: Ascending addresses at which an allocation begins.
+
+    Returns:
+        Consecutive `(address, size)` pieces covering the input exactly.
+    """
+    end = start + size
+    pieces: List[Tuple[int, int]] = []
+    address = start
+    index = bisect_right(cuts, start)
+    while index < len(cuts) and cuts[index] < end:
+        pieces.append((address, cuts[index] - address))
+        address = cuts[index]
+        index += 1
+    pieces.append((address, end - address))
+    return pieces
+
 
 class PageAddressing:
     """Resolves `(layer group, page index)` to the byte ranges of that page."""
 
-    def __init__(self, layout: KvCacheLayout):
+    def __init__(self, layout: KvCacheLayout, allocation_lookup: Optional[AllocationLookup] = None):
+        """Index `layout`, and with `allocation_lookup` its allocation boundaries.
+
+        Args:
+            layout: The KV cache layout to address.
+            allocation_lookup: Resolves an address to its allocation. Without
+                it every range is left whole, which suffices for host staging.
+        """
         self._layout = layout
         self._regions: Dict[int, Tuple[KvCacheRegion, ...]] = {}
         self._bytes_per_page: Dict[int, int] = {}
@@ -58,11 +147,42 @@ def __init__(self, layout: KvCacheLayout):
                 )
             self._num_slots[group.layer_group_id] = slot_counts.pop()
 
+        self._registration_ranges: List[Tuple[int, int]] = []
+        self._unresolved: List[Tuple[int, int]] = []
+        # Interior allocation boundaries, ascending, consulted per transfer.
+        self._cuts: List[int] = []
+        for start, end in self._pool_spans():
+            pieces = (
+                allocation_spans(start, end, allocation_lookup)
+                if allocation_lookup is not None
+                else None
+            )
+            if pieces is None:
+                if allocation_lookup is not None:
+                    self._unresolved.append((start, end))
+                pieces = [(start, end)]
+            self._registration_ranges.extend(pieces)
+            self._cuts.extend(piece_start for piece_start, _ in pieces[1:])
+
+    def _pool_spans(self) -> List[Tuple[int, int]]:
+        """Merged byte spans the regions occupy, from first slot to last."""
+        spans: List[Tuple[int, int]] = []
+        for regions in self._regions.values():
+            for region in regions:
+                span_end = region.base + region.stride * (region.num_slots - 1) + region.size
+                spans.append((region.base, span_end))
+        return merge_intervals(spans)
+
     @property
     def layout(self) -> KvCacheLayout:
         """The layout this addressing was built from."""
         return self._layout
 
+    @property
+    def unresolved_ranges(self) -> Tuple[Tuple[int, int], ...]:
+        """Spans whose allocations the lookup could not resolve, left whole."""
+        return tuple(self._unresolved)
+
     @property
     def layer_group_ids(self) -> Tuple[int, ...]:
         """Layer group ids covered, in layout order."""
@@ -84,6 +204,9 @@ def num_slots(self, layer_group_id: int) -> int:
     def buffers(self, layer_group_id: int, page_index: int) -> Tuple[List[int], List[int]]:
         """Addresses and sizes of one page, in the order they concatenate.
 
+        A region crossing an allocation boundary contributes one buffer per
+        allocation, in order.
+
         Args:
             layer_group_id: Layer group the page index is scoped to.
             page_index: Page slot index within that group.
@@ -98,10 +221,25 @@ def buffers(self, layer_group_id: int, page_index: int) -> Tuple[List[int], List
                 f"page index {page_index} out of range [0, {num_slots}) for layer "
                 f"group {layer_group_id}"
             )
-        addresses = [region.base + region.stride * page_index for region in regions]
-        sizes = [region.size for region in regions]
+        addresses: List[int] = []
+        sizes: List[int] = []
+        for region in regions:
+            start = region.base + region.stride * page_index
+            for address, size in split_at_cuts(start, region.size, self._cuts):
+                addresses.append(address)
+                sizes.append(size)
         return addresses, sizes
 
+    def registration_ranges(self) -> List[Tuple[int, int]]:
+        """Byte ranges to hand to `register_buffer`, one per allocation.
+
+        Each region spans from its first slot to the end of its last, so every
+        slot is registered. Spans are merged, so a shared pool is registered
+        once, then cut at allocation boundaries. A span the lookup could not
+        resolve is returned whole; see `unresolved_ranges`.
+        """
+        return list(self._registration_ranges)
+
     def describe(self) -> str:
         """A one-line summary for startup logs."""
         parts: Sequence[str] = [

tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py

@@ -31,8 +31,6 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Optional, Tuple
 
-from tensorrt_llm.logger import logger
-
 if TYPE_CHECKING:
     from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig, TorchLlmArgs
 
@@ -289,21 +287,12 @@ class MooncakeStoreConnectorConfig:
     #: RPC without bounding how much a request may transfer.
     transfer_batch_size: int = 64
     #: Pass pages through a pinned host buffer instead of registering the KV
-    #: pools with Mooncake. The only transfer path supported today, so a False
-    #: here is raised back to True.
-    stage_through_host: bool = True
+    #: pools with Mooncake. Costs a copy each way, but works without GPUDirect
+    #: RDMA, which registering device memory requires.
+    stage_through_host: bool = False
 
     def __post_init__(self) -> None:
         """Reject settings that would fail later, inside a transfer."""
-        if not self.stage_through_host:
-            if self.role.transfers:
-                logger.warning(
-                    "mooncake-store: ignoring stage_through_host=False. Pages "
-                    "pass through pinned host memory; registering the KV pools "
-                    "takes one registration per pool mapping, which this "
-                    "connector does not do yet."
-                )
-            object.__setattr__(self, "stage_through_host", True)
         if not self.master_server_address:
             raise ValueError("master_server_address is required")
         if self.local_buffer_size <= 0:
@@ -346,7 +335,7 @@ def from_file(path: str) -> "MooncakeStoreConnectorConfig":
             namespace=str(raw.get("namespace", DEFAULT_NAMESPACE)),
             model_key=raw.get("model_key") or None,
             transfer_batch_size=int(raw.get("transfer_batch_size", 64)),
-            stage_through_host=bool(raw.get("stage_through_host", True)),
+            stage_through_host=bool(raw.get("stage_through_host", False)),
         )
 
     @staticmethod