[TRTLLM-12891][feat] Support KV cache connector for v2_kvcm and extend VSWA support

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp

@@ -1323,7 +1323,11 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
         .def("__bool__", [](kv::ScratchDesc const& self) { return static_cast<bool>(self); });
 
     nb::class_<kv::AttnLifeCycle>(m, "AttnLifeCycle")
-        .def(nb::init<std::optional<int>, int>(), nb::arg("window_size"), nb::arg("num_sink_blocks"))
+        .def(nb::init<std::optional<int>, int>(), nb::arg("window_size").none(), nb::arg("num_sink_blocks"))
+        // Sink tokens round up to whole blocks. Bound rather than repeated in Python so the
+        // connector's view of a life cycle is built by the same code as the allocator's.
+        .def_static("make", &kv::AttnLifeCycle::make, nb::arg("window_size").none(), nb::arg("num_sink_tokens").none(),
+            nb::arg("tokens_per_block"))
         .def_prop_ro("window_size", [](kv::AttnLifeCycle const& self) { return self.windowSize; })
         .def_ro("num_sink_blocks", &kv::AttnLifeCycle::numSinkBlocks)
         .def("get_stale_range", &kv::AttnLifeCycle::getStaleRange, nb::arg("history_length"),

docs/source/features/kvcache.md

@@ -139,7 +139,7 @@ layer group in layer-group ID order.
 If neither `avg_seq_len` nor an explicit `pool_ratio` is configured, hybrid
 Mamba models warn and fall back to half of `max_seq_len`, which can produce a
 suboptimal pool split. Exact explicit boundaries currently require
-`MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. Hybrid
+`MambaHybridCacheManagerV2` and `max_beam_width=1`. Hybrid
 Mamba models select V2 by default (see
 [Selecting the KV Cache Manager](#selecting-the-kv-cache-manager)); set
 `use_kv_cache_manager_v2` to `false` to select the V1 C++

examples/llm-api/llm_kv_cache_connector.py

@@ -93,6 +93,8 @@
 from tensorrt_llm import LLM, SamplingParams, logger
 from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import (
     KvCacheConnectorScheduler, KvCacheConnectorWorker, SchedulerOutput)
+from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import \
+    valid_page_slots
 from tensorrt_llm.bindings.internal.batch_manager import LlmRequest
 from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig, TorchLlmArgs
 
@@ -113,6 +115,11 @@ def __init__(self, llm_args: TorchLlmArgs):
         self.kv_cache_tensor = None
 
     def register_kv_caches(self, kv_cache_tensor: torch.Tensor):
+        # This is the only registration hook this connector needs. A cache that
+        # describes itself as a layout instead still arrives here, through
+        # `register_kv_cache_layout`'s default, as long as one tensor can
+        # describe it. See llm_kv_cache_connector_vswa.py for the case where it
+        # cannot -- one attention window size per layer group.
         assert self.kv_cache_tensor is None, "KV cache tensor already registered"
         self.kv_cache_tensor = kv_cache_tensor
 
@@ -180,23 +187,35 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput):
 
             pending_load = self.pending_loads[req.request_id]
 
+            # Ordinal -> page slot for the blocks that have a page. Blocks with
+            # none keep their ordinal in `block_ids` so that entry `i` always
+            # describes the same token range; they are dropped here so no
+            # transfer can be built against one.
+            slots = dict(valid_page_slots(block_ids))
+
             for file_path, block_pos in zip(
                     pending_load, range(num_computed_blocks, len(block_ids))):
-                metadata.load.append((file_path, block_ids[block_pos]))
+                slot = slots.get(block_pos)
+                if slot is None:
+                    continue
+                metadata.load.append((file_path, slot))
 
             # Break up the remainder of the token sequence into chunks.
             chunks = self._chunk_tokens(req.new_tokens)
 
             # For each chunk that isn't already on device, and isn't in our connector cache, we need to save it.
             for block_pos in range(num_computed_blocks + len(pending_load),
                                    len(block_ids)):
+                slot = slots.get(block_pos)
+                if slot is None:
+                    continue
                 if len(chunks[block_pos]) == self.block_size:
                     hashed_tokens = self._hash_tokens(chunks[block_pos],
                                                       req.cache_salt)
 
                     file_path = self._file_path(hashed_tokens)
 
-                    metadata.save.append((file_path, block_ids[block_pos]))
+                    metadata.save.append((file_path, slot))
 
         self.pending_loads = {}