[TRTLLM-12891][feat] Support the KV cache connector on KVCacheManagerV2

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

@@ -42,6 +42,19 @@ These methods run on the leader process and drive the connector's behavior.
 
 * **`update_state_after_alloc(self, request: LlmRequest, block_ids: list[int])`**
   * **Description**: a callback to update internal state after KV cache blocks have been allocated for the prefill.
+  * **Note**: on `KVCacheManagerV2` with chunked prefill, `block_ids` covers only the blocks allocated for the first chunk, because V2 allocates per chunk rather than for the whole prompt. The remaining blocks arrive as append-deltas in `RequestData.new_block_ids` on subsequent chunks. A connector that treats this callback as its only source of block ids will under-plan; drive off `build_connector_meta` instead.
+  * **Note**: on `KVCacheManagerV2` under sliding-window attention, a block that the window has already passed holds no page, and is reported as `-1` (`BAD_PAGE_INDEX`) **in place** rather than being dropped from the list. This keeps each entry aligned with its block ordinal, so entry `i` always describes prompt tokens `[i * tokens_per_block, (i+1) * tokens_per_block)` and an append-delta over successive calls stays valid. Connectors must skip `-1` entries rather than treating them as page slots. The same applies to `RequestData.new_block_ids` and to `cache_block_ids` in `request_finished`.
+
+* **`cancel_load(self, request: LlmRequest, start: int, end: int)`**
+  * **Description**: Optional, with a no-op default. Tells the connector that the runtime will not consume KV it offered from `get_num_new_matched_tokens` for prompt tokens `[start, end)`, so any ownership taken for that range can be released. Offsets are absolute prompt positions, on the same scale as `num_computed_tokens`.
+  * **When it fires**: only on `KVCacheManagerV2`, which asks during a speculative scheduling pass and resolves the answer later. Two things can happen in between, and both are reported here: the runtime may fail to allocate pages to cover the offer, in which case the request falls back to computing the prefix locally; or the request may be cancelled, time out or fail before it ever reaches a batch, in which case the whole offer is released. A third case -- the local cache overtaking part of the offer because another request committed the same prefix -- is handled by the same callback but cannot arise today, since a request's local match is fixed when its cache is created and only its own completed forward passes extend it.
+  * **Caveat**: best-effort. For a synchronous load nothing has been transferred yet, so cancelling is exact. For `is_async=True` the transfer necessarily started inside `get_num_new_matched_tokens`, so it may already be in flight.
+
+##### Serving a prefix on `KVCacheManagerV2`
+
+V1 answers `get_num_new_matched_tokens` from C++ while the block manager holds its radix-tree mutex, so the local match and the query are atomic and the answer is consumed immediately. V2 has no such mutex, and its scheduling pass is speculative -- a prepared request can still be dropped at the token budget, at resize, at multimodal alignment or at cross attention, and retried in a later iteration.
+
+The contract for connectors is unchanged, and in particular `get_num_new_matched_tokens` is still called **exactly once per request** on both managers -- a request that is asked and then deferred is not asked again when it comes back. What differs is that on V2 the runtime may resolve the answer in a later iteration than the one it asked in, and may by then be unable to honour part or all of it. That is what `cancel_load` reports.
 
 #### 2. Worker Interface (`KvCacheConnectorWorker`)
 

examples/llm-api/llm_kv_cache_connector.py

@@ -116,6 +116,21 @@ def register_kv_caches(self, kv_cache_tensor: torch.Tensor):
         assert self.kv_cache_tensor is None, "KV cache tensor already registered"
         self.kv_cache_tensor = kv_cache_tensor
 
+    def register_kv_cache_layout(self, layout):
+        # KVCacheManagerV2 describes its pools rather than handing over one
+        # tensor. This example targets the uniform case, where a layer group's
+        # buffers coalesce into a single whole-slot region; indexing that
+        # region by page slot is the direct analogue of indexing a pool by
+        # block id, so the load/save paths below need no change.
+        assert self.kv_cache_tensor is None, "KV cache tensor already registered"
+        if len(layout.groups) != 1 or len(layout.groups[0].regions) != 1:
+            raise NotImplementedError(
+                "This example connector handles a single layer group with a "
+                "single coalesced region (uniform attention, one pool). Got "
+                f"{len(layout.groups)} layer group(s) and "
+                f"{sum(len(g.regions) for g in layout.groups)} region(s).")
+        self.kv_cache_tensor = layout.groups[0].regions[0].as_tensor()
+
     def start_load_kv(self, stream: torch.cuda.Stream):
         # Do all loads synchronously, and blockwise.
         for path, block_id in self._metadata.load:

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -269,7 +269,7 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice:
         assert self._page_table is not None
         layer_groups = self._page_table.layer_groups
 
-        is_gen_only = req.is_generation_only_request()
+        is_gen_only = req.is_generation_only_request
         cached_per_lg = (
             adapter.get_cached_token_count_per_layer_group(req, layer_groups)
             if is_gen_only

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -653,9 +653,9 @@ def _validate_or_fallback_kv_cache_manager_v2(
         # also go through the V2-incompatible-feature gate below.
         if issubclass(kv_cache_manager_cls, KVCacheManagerV2):
             sparse_attn_config = model_config.sparse_attention_config
+            # The KV connector is supported on V2 via the pool layout
+            # registration path, so it no longer forces a fallback to V1.
             incompat: List[str] = []
-            if self._kv_connector_manager is not None:
-                incompat.append("kv_connector_manager")
             if self._max_beam_width is not None and self._max_beam_width > 1:
                 incompat.append("max_beam_width > 1")
             if incompat:
@@ -682,8 +682,8 @@ def _validate_or_fallback_kv_cache_manager_v2(
                     raise NotImplementedError(
                         "Hybrid Mamba cache managers do not support "
                         f"{incompat_str}; CppMambaHybridCacheManager does not "
-                        "provide a compatible fallback. Use max_beam_width=1 "
-                        "and disable the KV connector.")
+                        "provide a compatible fallback. Disable the listed "
+                        "features to run hybrid linear models.")
                 # Plain V2 (explicitly enabled or selected by a model preference):
                 # V2 was a preference, not a structural requirement, so we can
                 # safely fall back to V1.

tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py

@@ -54,6 +54,7 @@
 
 if TYPE_CHECKING:
     from ..resource_manager import KVCacheManager
+    from .kv_cache_layout import KvCacheLayout
 
 
 # Used to store data for a single inflight request.
@@ -85,6 +86,11 @@ class RequestData:
     # remote object id) MUST mix cache_salt into their identifiers,
     # otherwise blocks from a different salt could be incorrectly reused.
     cache_salt: Optional[str] = None
+    # New page slot indices keyed by layer group, populated under
+    # KVCacheManagerV2. Blocks with no page in a group -- the sliding-window
+    # case -- appear as BAD_PAGE_INDEX in place, keeping ordinals stable.
+    # Empty under the V1 manager, whose block IDs are a single flat space.
+    new_block_ids_by_layer_group: Dict[int, List[int]] = field(default_factory=dict)
 
 
 # A class to store some basic data regarding all inflight requests.
@@ -135,6 +141,30 @@ def register_kv_caches(self, kv_cache_tensor: torch.Tensor):
             kv_cache_tensor: The contiguous KV cache tensor.
         """
 
+    def register_kv_cache_layout(self, layout: "KvCacheLayout") -> None:
+        """
+        Register the KV cache pools described by ``layout``.
+
+        Called instead of ``register_kv_caches`` when the KV cache manager is
+        ``KVCacheManagerV2``, whose memory cannot be expressed as a single
+        tensor: there is one slot address space per pool and one page-index
+        space per layer group.
+
+        The bytes for page slot ``i`` of a region live at
+        ``region.base + region.stride * i`` for ``region.size`` bytes, or
+        equivalently at ``region.as_tensor()[i]``. Page indices arriving in
+        ``RequestData.new_block_ids`` are scoped to a layer group and index the
+        regions of that group.
+
+        Args:
+            layout: Description of the KV cache pools; see ``KvCacheLayout``.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not implement register_kv_cache_layout, so it "
+            "cannot run against KVCacheManagerV2. Implement it, or select the V1 KV "
+            "cache manager with kv_cache_config.use_kv_cache_manager_v2=False."
+        )
+
     @abstractmethod
     def start_load_kv(self, stream: torch.cuda.Stream):
         """
@@ -255,6 +285,30 @@ def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]):
             block_ids: The KV cacheblock IDs that were allocated.
         """
 
+    def cancel_load(self, request: LlmRequest, start: int, end: int):
+        """
+        Called when the runtime will not consume KV that was offered by
+        ``get_num_new_matched_tokens`` for prompt tokens ``[start, end)``.
+
+        Offsets are absolute prompt positions, on the same scale as
+        ``num_computed_tokens``. This is raised when the local cache overtook
+        part or all of an offer while the request was waiting to be scheduled,
+        or when the runtime could not allocate pages to cover it.
+
+        Best-effort: release any ownership taken in
+        ``get_num_new_matched_tokens`` for that range. For a synchronous load
+        nothing has been transferred yet, so this is exact. For an
+        asynchronous load the transfer necessarily began inside
+        ``get_num_new_matched_tokens`` (see ``take_scheduled_requests_pending_load``),
+        so it may already be in flight and cancelling is genuinely lossy.
+
+        Args:
+            request: The request whose offer is being handed back.
+            start: First prompt position that will not be consumed.
+            end: One past the last prompt position that will not be consumed.
+        """
+        return
+
     def wait_for_initialization(self):
         """
         Some connectors need to wait for some resources to be initialized.
@@ -316,21 +370,49 @@ def loading_ids(self) -> Set[int]:
 class KvCacheConnectorSchedulerOutputRequest:
     def __init__(self):
         self.block_ids = []
+        self.block_ids_by_layer_group: Dict[int, List[int]] = {}
         self.tokens = []
 
     def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManager"):
-        block_ids = kv_cache_manager.get_cache_indices(req)
+        from ..kv_cache_manager_v2 import KVCacheManagerV2
+
+        is_v2 = isinstance(kv_cache_manager, KVCacheManagerV2)
         tokens = req.get_tokens(0)
 
-        # Commit hashes for any blocks that have become full since the last call
-        # and read back the full cumulative chain. The C++ side sets each block's
-        # mBlockKey/mHash on first call, so subsequent calls become pure lookups.
-        block_hashes = kv_cache_manager.commit_and_get_block_hashes(req)
+        new_block_ids_by_layer_group: Dict[int, List[int]] = {}
+        if is_v2:
+            # Block hashes and retention priorities have no V2 accessor yet, so
+            # they are reported empty rather than guessed at.
+            block_hashes = []
+            indices_by_group = kv_cache_manager.get_page_indices_by_layer_group(req)
+            for layer_group_id, indices in indices_by_group.items():
+                seen = self.block_ids_by_layer_group.setdefault(layer_group_id, [])
+                new_ids = indices[len(seen) :]
+                seen.extend(new_ids)
+                new_block_ids_by_layer_group[layer_group_id] = new_ids
+            # With a single layer group -- every non-VSWA, non-hybrid model --
+            # ``new_block_ids`` carries that group's indices so connectors that
+            # do not reason about layer groups keep working. With several groups
+            # it is left empty and ``new_block_ids_by_layer_group`` is the only
+            # correct source.
+            new_block_ids = (
+                next(iter(new_block_ids_by_layer_group.values()))
+                if len(new_block_ids_by_layer_group) == 1
+                else []
+            )
+            self.block_ids.extend(new_block_ids)
+        else:
+            block_ids = kv_cache_manager.get_cache_indices(req)
 
-        new_block_ids = block_ids[len(self.block_ids) :]
-        new_tokens = tokens[len(self.tokens) :]
+            # Commit hashes for any blocks that have become full since the last call
+            # and read back the full cumulative chain. The C++ side sets each block's
+            # mBlockKey/mHash on first call, so subsequent calls become pure lookups.
+            block_hashes = kv_cache_manager.commit_and_get_block_hashes(req)
+
+            new_block_ids = block_ids[len(self.block_ids) :]
+            self.block_ids.extend(new_block_ids)
 
-        self.block_ids.extend(new_block_ids)
+        new_tokens = tokens[len(self.tokens) :]
         self.tokens.extend(new_tokens)
 
         if req.state in (
@@ -346,9 +428,13 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag
             )  # Specdec with draft tokens is not supported yet.
 
         # Get retention priority for each new block only if retention config is provided
-        # (for priority-based offload filtering)
+        # (for priority-based offload filtering). Priorities stay None on
+        # KVCacheManagerV2: it does not implement `KvCacheRetentionConfig` at
+        # all -- its per-page priority comes from `custom_priority_callback`,
+        # which KVCacheManagerV2 never overrides -- so every page carries the
+        # default and reporting it would misdescribe what the user asked for.
         priorities = None
-        if req.kv_cache_retention_config is not None:
+        if not is_v2 and req.kv_cache_retention_config is not None:
             priorities = [
                 kv_cache_manager.get_priority_by_block_id(block_id) for block_id in new_block_ids
             ]
@@ -362,6 +448,7 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag
             block_hashes=block_hashes,
             priorities=priorities,
             cache_salt=req.cache_salt,
+            new_block_ids_by_layer_group=new_block_ids_by_layer_group,
         )
 
 
@@ -463,7 +550,18 @@ def _run_on_leader(self, f: Callable[[], Any]) -> Any:
             res = None
         return mpi_broadcast(res, root=0)
 
-    def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> int:
+    def query_num_new_matched_tokens(
+        self, request: LlmRequest, num_computed_tokens: int
+    ) -> Tuple[int, bool]:
+        """Ask the connector how much of the prompt it can serve. No side effects.
+
+        This is the half of ``get_num_new_matched_tokens`` that must run at most
+        once per request, because the connector ABC promises exactly one query
+        per request and connectors take ownership of remote blocks in it.
+        Callers that cannot commit to consuming the answer in the same iteration
+        (KVCacheManagerV2, whose scheduling pass is speculative) query here and
+        commit later via ``commit_new_matched_tokens``.
+        """
         if request.is_generation_only_request:
             raise RuntimeError("Connector API is not supported for generation-only requests!")
 
@@ -474,6 +572,18 @@ def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: i
         if num_tokens == 0 and load_kv_async:
             raise RuntimeError("load_kv_async must be False when num_tokens is 0!")
 
+        return num_tokens, load_kv_async
+
+    def commit_new_matched_tokens(
+        self, request: LlmRequest, num_tokens: int, load_kv_async: bool
+    ) -> None:
+        """Register the runtime's side of an answered query.
+
+        Must run in the iteration the request is actually scheduled:
+        ``external_loads`` is consumed and cleared by every
+        ``build_scheduler_output``, and ``new_async_requests.loading`` is read
+        by that same call to suppress the request's ``RequestData``.
+        """
         # TODO(jthomson04): This part is a bit ugly.
         # When the connector indicates that a request will be loaded
         # asynchronously, we need to suspend its execution. This is
@@ -488,8 +598,21 @@ def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: i
 
         request.py_num_connector_matched_tokens = num_tokens
 
+    def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> int:
+        """Query and commit in one step.
+
+        This is the V1 entry point, called from C++ while the block manager
+        holds the radix-tree mutex, so the local match and the query are atomic
+        with respect to the tree and the answer can be committed immediately.
+        """
+        num_tokens, load_kv_async = self.query_num_new_matched_tokens(request, num_computed_tokens)
+        self.commit_new_matched_tokens(request, num_tokens, load_kv_async)
         return num_tokens
 
+    def cancel_load(self, request: LlmRequest, start: int, end: int) -> None:
+        if self.scheduler is not None:
+            self.scheduler.cancel_load(request, start, end)
+
     def should_add_sequence(self, request: LlmRequest) -> bool:
         req_id = request.request_id
         return req_id not in self.finished_async_loading_requests