From caa25a6d4a895cd7b29329b6e12919a33f19b3fa Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Sun, 6 Sep 2026 00:36:14 +0800 Subject: [PATCH] [None][feat] Support the KV cache connector on KVCacheManagerV2, including VSWA KVCacheManager (V1) answers a connector's two questions -- where the KV pages are, and which page a request's block occupies -- with one flat index space over a single primary pool. KVCacheManagerV2 has one slot address space per pool and one page-index space per layer group, so it answers them differently. V1 also rejects a connector outright on a VSWA model, and never consults the connector's prefix contribution during scheduling. Describe the pools rather than hand over one tensor. register_kv_cache_layout receives a KvCacheLayout of byte ranges per layer group, assembled from V2's own layout API so coalescing comes from the allocator. Its default reconstructs the single-pool tensor and forwards to register_kv_caches, so a connector written against V1 runs unchanged wherever one tensor describes the cache. Report page indices per layer group, positionally aligned. A block with no page in a group holds BAD_PAGE_INDEX in place, so entry i keeps describing tokens [i * tokens_per_block, (i+1) * tokens_per_block) and an append-delta stays valid. valid_page_slots and KvCacheRegion.slot_tensor keep such an entry from becoming a device address, the same two-layer discipline the disaggregated path runs. update_state_after_alloc and request_finished gain *_by_layer_group forms that default to the flat ones for a single group. Implementing a per-layer-group form also satisfies the abstract flat method it replaces, so a connector written only for VSWA carries no dead stubs. Bind AttnLifeCycle::make so the out-of-window mask is built from the same life cycle the allocator uses rather than a Python copy of the rounding rule. Serve a connector-supplied prefix from KVCacheManagerV2.prepare_resources, downstream of every stage that can still drop a request, so an asked request always reaches request_finished. V2 allocates per context chunk, so an offer reaching past the current chunk is served as far as the allocation can grow and the rest is computed locally. Reject at bring-up what cannot be honoured: a host or disk cache tier, and a connector prefix without block reuse. Report block_hashes and priorities as empty on V2 rather than guessing, and warn when a retention config has no effect. Signed-off-by: Yueh-Ting Chen --- .../batch_manager/kvCacheManagerV2.cpp | 6 +- docs/source/features/kv-cache-connector.md | 214 +- docs/source/features/kvcache.md | 2 +- examples/llm-api/llm_kv_cache_connector.py | 23 +- .../llm-api/llm_kv_cache_connector_vswa.py | 359 +++ tensorrt_llm/_torch/pyexecutor/_util.py | 6 +- .../connectors/kv_cache_connector.py | 448 +++- .../pyexecutor/connectors/kv_cache_layout.py | 335 +++ .../kv_cache/kv_cache_manager_v2.py | 359 ++- tensorrt_llm/_torch/pyexecutor/llm_request.py | 12 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 165 +- .../_torch/pyexecutor/py_executor_creator.py | 37 +- .../_torch/pyexecutor/resource_manager.py | 10 +- .../pyexecutor/scheduler/scheduler_v2.py | 42 +- .../llmapi/data/kv_connector_vswa_prompt.txt | 43 + .../defs/llmapi/test_llm_api_connector.py | 2128 ++++++++++++++++- .../integration/test_lists/test-db/l0_a10.yml | 101 +- .../kv_cache/test_kv_cache_manager_v2.py | 3 + .../test_kv_cache_v2_first_new_block_probe.py | 11 +- .../kv_cache/test_kv_cache_v2_scheduler.py | 19 +- .../kv_cache/test_kv_pool_rebalance.py | 9 + .../kv_cache/test_mamba_cache_manager.py | 18 +- .../_torch/executor/test_kv_cache_layout.py | 908 +++++++ .../executor/test_kv_connector_v2_prefix.py | 796 ++++++ ...est_kv_connector_v2_prefix_real_manager.py | 636 +++++ .../executor/test_send_kv_async_split.py | 4 +- .../executor/test_token_budget_fallback.py | 8 +- .../multi_gpu/test_kv_pool_rebalance_tp.py | 1 + tests/unittest/_torch/test_connector.py | 432 +++- .../test_kv_cache_stats_behavior.py | 1 + 30 files changed, 6948 insertions(+), 188 deletions(-) create mode 100644 examples/llm-api/llm_kv_cache_connector_vswa.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py create mode 100644 tests/integration/defs/llmapi/data/kv_connector_vswa_prompt.txt create mode 100644 tests/unittest/_torch/executor/test_kv_cache_layout.py create mode 100644 tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py create mode 100644 tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 52662f3ced2c..8572d2d1c567 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/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(self); }); nb::class_(m, "AttnLifeCycle") - .def(nb::init, int>(), nb::arg("window_size"), nb::arg("num_sink_blocks")) + .def(nb::init, 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"), diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index dc16209c31ad..902d6844f53a 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -39,9 +39,217 @@ These methods run on the leader process and drive the connector's behavior. * **`request_finished(self, request: LlmRequest, cache_block_ids: list[int]) -> bool`** * **Description**: Called when a request completes generation. * **Returns**: A boolean indicating if an asynchronous save operation is underway. If `True`, the system waits for the operation to complete before releasing the KV cache blocks. + * **Note**: under sliding-window attention `cache_block_ids` covers the live window, not the whole prompt. See [What a connector can persist under a sliding window](#what-a-connector-can-persist-under-a-sliding-window). * **`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**: with chunked prefill, `block_ids` covers only the blocks allocated for the first chunk. The remaining blocks arrive as append-deltas in `RequestData.new_block_ids` on later calls to `build_connector_meta`, on the entries under `scheduler_output.cached_requests`. A connector that treats this callback as its only source of block ids will under-plan. Read both lists: + + ```python + def build_connector_meta(self, scheduler_output): + for req in scheduler_output.new_requests: # the first chunk + self._plan(req.request_id, req.new_block_ids) + for req in scheduler_output.cached_requests: # every later chunk + self._plan(req.request_id, req.new_block_ids) + ``` + + Both example connectors walk `new_requests` only, so neither one demonstrates this. + +* **`update_state_after_alloc_by_layer_group(self, request: LlmRequest, block_ids_by_layer_group: list[list[int]])`** +* **`request_finished_by_layer_group(self, request: LlmRequest, cache_block_ids_by_layer_group: list[list[int]]) -> bool`** + * **Description**: the per-layer-group forms of the two callbacks above, indexed by layer group id. Entry `[g][i]` is the page slot of block ordinal `i` in layer group `g`. + * **When they are called**: whenever the KV cache reports page indices per layer group. A page index is scoped to its group — one group per attention window size — so a cache with more than one group can be described no other way, and the flat `block_ids` / `cache_block_ids` are empty there. With a single layer group the flat lists carry that group's indices as well, so a connector that implements only the flat forms keeps working on those models. + * **Which form to implement**: implement exactly one complete set. + + | Set | Models it covers | + |---|---| + | per-layer-group | every model, VSWA included | + | flat | non-VSWA, non-hybrid only | + + A set is complete when both of its methods are defined. Mixing the two — one method from each — is rejected during executor bring-up, naming the method that is missing. An existing flat connector needs no change on the models it already covers: with a single layer group the base per-layer-group implementation folds back to the flat call. Hybrid / linear-attention models are not enabled for the connector yet; the per-layer-group set is the shape their cache will need. See [Running under VSWA](#running-under-vswa). + +##### Running under VSWA + +Under variable sliding-window attention the KV cache allocates one pool per attention window size, and a page index only means something inside its own layer group. A single tensor and a single flat block list cannot describe that, so three methods have to be implemented together: + +| Method | Replaces | +|---|---| +| `KvCacheConnectorWorker.register_kv_cache_layout` | `register_kv_caches` | +| `KvCacheConnectorScheduler.update_state_after_alloc_by_layer_group` | `update_state_after_alloc` | +| `KvCacheConnectorScheduler.request_finished_by_layer_group` | `request_finished` | + +Implementing the per-layer-group form of a pair is enough — the flat method it replaces does not also have to be defined. + +All three are checked during executor bring-up, before any request is admitted. `register_kv_cache_layout` refuses there, naming the group and region counts it could not describe, and the two scheduler methods are checked alongside it. Nothing is deferred to the first request, so a partial implementation costs a start-up failure rather than one after the model is loaded. + +`examples/llm-api/llm_kv_cache_connector_vswa.py` is a worked connector for this case. A VSWA connector has to do five things: + +1. **Address pages per group.** `layout.groups[g].regions[r]` gives the byte ranges; `region.slot_tensor(i)` is page slot `i` *of that group*. `layout.group_of_layer(layer_id)` maps a model layer back to its group, which is what the per-layer `wait_for_layer_load` / `save_kv_layer` hooks need. +2. **Read the per-group block lists.** `RequestData.new_block_ids_by_layer_group[g]` carries the page slots; the flat `new_block_ids` is empty. +3. **Carry the layer group in the cache key, and in every transfer target.** This one is a correctness requirement, not a convenience. The same token range exists in *every* layer group holding **different** KV, so a key derived from the token sequence alone collides across groups and one group's bytes will overwrite another's — then be loaded back into the wrong group. Mix `layer_group_id` (or the window size, or the layer set) into the identifier, and carry `(layer_group_id, page_slot)` rather than `page_slot` alone as the transfer target. +4. **Filter out-of-window blocks through `valid_page_slots`**, and size the store for the window rather than the prompt. See below. +5. **Serve a block only when every group holds it.** A full-attention group keeps the whole prompt while a sliding group keeps only its window, so the prefix that can be served back is bounded by the smallest window. Stop the lookup at the first block ordinal any group misses. + +##### `KvCacheLayout` reference + +```python +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import KvCacheLayout +``` + +The type is passed to `register_kv_cache_layout`; importing it is only needed for a type annotation. + +A layout describes the byte ranges that repeat once per page slot. It describes ranges rather than +implying them, which is what lets one type cover MLA (a pool simply has no `value` buffer), block +scales, sliding-window attention and hybrid models without any of them being a special case. + +| Attribute | Meaning | +|---|---| +| `layout.tokens_per_block` | Tokens covered by one page. | +| `layout.dtype` | Element type of the KV data, for a typed view over a region. | +| `layout.groups` | The layer groups, as `KvCacheLayerGroupLayout`. | +| `layout.group(layer_group_id)` | One group by id. Raises `KeyError` if absent. | +| `layout.group_of_layer(layer_id)` | The group owning a model layer — what the per-layer hooks route on. | +| `layout.as_single_pool_tensor()` | The `[num_blocks, num_layers, kv_factor, block_size]` view a single-pool cache hands `register_kv_caches`, or `None` when the cache cannot be described that way. This is what the default `register_kv_cache_layout` calls. | + +Each `KvCacheLayerGroupLayout`: + +| Attribute | Meaning | +|---|---| +| `group.layer_group_id` | The index page slots are scoped to. Dense, starting at 0. | +| `group.window_size` | Attention window for the group, or `None` for full attention. | +| `group.layer_ids` | Global model layer indices in the group — the same index space `wait_for_layer_load` and `save_kv_layer` receive. | +| `group.regions` | The `KvCacheRegion`s making up one page of this group. | +| `group.bytes_per_page` | Total bytes the group occupies for one page slot. | + +Each `KvCacheRegion` is a contiguous byte range that repeats once per page slot: + +| Attribute | Meaning | +|---|---| +| `region.base` | Device address of page slot 0. | +| `region.size` | Bytes the region covers within one slot. | +| `region.stride` | Distance between consecutive slots. | +| `region.num_slots` | Number of page slots. | +| `region.buffers` | The `(layer_id, role, expansion)` tuples the region covers, in memory order. `role` is the cache manager's own name, e.g. `"key"` / `"value"`. | +| `region.address_of(slot)` | `base + stride * slot`. Raises `IndexError` outside `[0, num_slots)`. | +| `region.as_tensor(dtype=torch.uint8)` | A strided `[num_slots, size // itemsize]` view; row `i` is page slot `i`. Accepts any subscript, including `-1`. | +| `region.slot_tensor(slot_id, dtype=torch.uint8)` | The bytes of one page slot, raising `IndexError` outside `[0, num_slots)`. The guarded form of `as_tensor(dtype)[slot_id]`. | + +`size` is not necessarily `stride`: a region covers one run of adjacent buffers within a slot, and a +slot may hold several runs. For a model with uniform layer shapes the buffers coalesce into a single +region spanning the whole slot, which is the whole-page transfer. A group with more than one region +must be addressed region by region. + +The addresses are device addresses, and they stay valid because every cache tier below GPU is +rejected at bring-up while a connector is attached. See [KV cache tiers](#kv-cache-tiers-are-gpu-only-under-a-connector). + +##### KV cache tiers are GPU-only under a connector + +A connector registers device addresses and holds them across iterations. Evicting a page to another +tier reassigns its GPU slot underneath the connector, so with a connector attached: + +* Setting `KvCacheConfig.host_cache_size` or `KvCacheConfig.disk_cache_size` above zero **fails at + bring-up**, with a message naming both settings. +* Leaving `host_cache_size` unset **drops the host tier** rather than failing, with a log line. That + tier is provisioned automatically only to give the `MAX_UTILIZATION` scheduler somewhere to spill + to via suspend/resume, which a connector run does not use. +* `enable_kv_pool_rebalance` is **ignored** — startup and inference continue, and the rebalance + simply never runs. Rebalance suspends every active request and runs a defragmenting migration that + reassigns the same page slots a tier eviction would. + +The practical consequence is that a KV-exhausted connector deployment has no secondary tier to +fall back on. The remedies are `kv_cache_config.max_tokens`, +`kv_cache_config.free_gpu_memory_fraction`, or lowering `max_num_tokens` to hand memory back to the +KV pool; the scheduler's exhaustion error says so directly when a connector is attached. + +##### Block reuse alongside the connector + +* Specify `KvCacheConfig.enable_block_reuse=True` alongside a connector. Without it the connector's + prefix is never honoured: either the combination is rejected at start-up, or the lookup, the reads + and the device copies are performed and discarded, at no correctness cost but at full latency + cost. +* The start-up check reads the value the cache resolved, not the one you passed: some quantization + algorithms, some SM versions and hybrid linear models turn block reuse off on their own, so this + error can appear without the flag being set anywhere in your configuration. + +##### A page slot must not be reassigned underneath the connector + +The connector holds page indices across iterations, and `RequestData` reports only the pages appended +since the last call. Anything that hands a slot the connector already knows about to a different +request therefore goes unreported and corrupts the next transfer against it. Three configurations do +that, and each is rejected at bring-up. + +| Configuration | Mechanism | +|---|---| +| Speculative decoding | Rejected draft tokens shrink a request's page list, and the freed slot goes to whichever request allocates next. The connector is never told the tail block moved. | +| A capacity scheduler policy other than `GUARANTEED_NO_EVICT` | A destroyed-and-replayed request comes back on different pages, and the connector's per-request block delta is then measured against pages that were freed with it. | +| A host or disk cache tier | Tier eviction reassigns the GPU slot. See [KV cache tiers are GPU-only under a connector](#kv-cache-tiers-are-gpu-only-under-a-connector). | + +The exact set the runtime refuses depends on your cache configuration; the bring-up error is +authoritative. Speculative decoding is unsupported with a connector wherever it is not refused — +the same page-list shrink happens there. + +##### `RequestData` fields that may not be populated + +Two `RequestData` fields can be reported empty depending on the cache configuration. A connector must tolerate both. + +| Field | When empty | Consequence | +|---|---|---| +| `block_hashes` | `[]` | No block-hash accessor exists on this path. Nothing in the runtime reads the field, and neither example connector uses it — both hash the token sequence themselves. A connector that keys its external store on `block_hashes` gets no key and therefore no hits and no saves; it does not mis-address a transfer. | +| `priorities` | `None` | `KvCacheRetentionConfig` is not honoured, so every page carries the default priority. A warning is logged the first time a request carrying a retention config is reported. The gap is wider than the connector: the retention config has no effect either way in that configuration. | + +Both are gaps to be closed rather than intended behaviour. + +##### Blocks with no page + +A block that has no page in a layer group is reported as `-1` (`BAD_PAGE_INDEX`) **in place**, not dropped from the list. This keeps each entry aligned with its block ordinal, so entry `i` always describes tokens `[i * tokens_per_block, (i+1) * tokens_per_block)` and an append-delta over successive calls stays valid. + +That alignment is also why the list is not safe to index with directly: `-1` is a valid Python and PyTorch subscript, so it resolves to the *last* page slot of the pool rather than raising — a transfer against another request's KV. Two API points keep a page index from reaching device memory unchecked. + +| | | +|---|---| +| `valid_page_slots(page_indices)` | Yields `(block_ordinal, page_slot)` for the entries that address a page. The ordinal is preserved, so the token range a page covers is still recoverable. | +| `region.slot_tensor(slot_id)` | The bytes of one page slot, raising `IndexError` on a slot outside `[0, num_slots)`. | + +Build transfer targets with `valid_page_slots` and address them with `slot_tensor`. This covers `block_ids`, `cache_block_ids`, `RequestData.new_block_ids`, and both `*_by_layer_group` forms. + +```python +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import valid_page_slots + +for ordinal, slot in valid_page_slots(cache_block_ids): + tokens = all_tokens[ordinal * tokens_per_block:(ordinal + 1) * tokens_per_block] + store.put(self._key(tokens), region.slot_tensor(slot)) +``` + +Alignment is not the same as stability. Under a sliding window an entry that was reported with a page reads back as `-1` once the window passes its block, and the delta — which carries only the ordinals appended since the last call — does not restate it. Save a block when its tokens complete rather than deferring: a completed block is far inside the window for any usable window size, whereas a deferred save can reach a slot the cache has already reclaimed. + +##### What a connector can persist under a sliding window + +Under sliding-window attention, a connector can persist **at most `window_size` tokens per sequence**, not `prompt_len`. + +The KV cache manager reclaims a block's page once the window has moved past it, so by the time `request_finished` runs there is no readable KV for anything older than the last `window_size` tokens. Those ordinals report no page (see [Blocks with no page](#blocks-with-no-page)), and the page slots offered to save from cover the live window only. A prefix-caching connector on such a model therefore caches a tail rather than a prefix, and the prefix it can serve back on a later request is bounded the same way. + +This is a property of the cache, not of the connector: the blocks are gone whether or not a connector is attached. The same bound applies to the KV cache transceiver, which drops the same range before sending. + +##### Serving a prefix + +A connector that implements only the flat callbacks and `register_kv_caches` runs unchanged on any model with a single *non-sliding* attention window. Where the cache describes itself as pools rather than one tensor it calls `register_kv_cache_layout` instead — but that method's default reconstructs the single-pool tensor, in the same `[num_blocks, num_layers, kv_factor, block_size]` shape and KV dtype, and forwards it to `register_kv_caches`. The same applies to the two block-id callbacks: their per-layer-group forms default to the flat ones when there is a single layer group. + +Variable sliding-window attention is the case where that stops working, because the cache then allocates one pool per window size and a page index is scoped to a layer group. See [Running under VSWA](#running-under-vswa). + +A model whose layers all share one sliding window stays a single layer group, so the flat callbacks still apply and such a connector is not refused. What differs is that the callbacks cover the **live window only**. Blocks the window has passed report `-1` (`BAD_PAGE_INDEX`) in place — the list stays aligned to block ordinals, so entry `i` still describes tokens `[i * tokens_per_block, (i+1) * tokens_per_block)`, but the up-front blocks carry no page and are not available to load into or save from. Filter with `valid_page_slots`, described in [Blocks with no page](#blocks-with-no-page); without it a `-1` resolves to the last page slot of the pool. A warning naming the window size is logged at start-up when a flat-only connector is attached to such a model. + +`get_num_new_matched_tokens` is asked once the batch for the upcoming forward pass is final. A request that is asked is therefore a request that runs, and the connector can take ownership of remote blocks in the query and release it in `request_finished`. + +Two things are worth knowing when tuning a deployment. + +* **The runtime may honour less than you offer.** With chunked prefill the cache is allocated per context chunk, which is what bounds its memory, so an offer reaching past the current chunk requires the runtime to grow the allocation and that can fail under pressure. The runtime then serves the part it can cover and computes the rest locally. The amount actually served is what `RequestData.computed_position` reflects; the unserved remainder needs no action from the connector beyond its usual `request_finished` cleanup. +* **The query is not part of the scheduler's budget.** The scheduler sizes a request's chunk as if the connector will serve nothing, so a served prefix reduces the work in the forward pass but does not free budget for another request in the same iteration. + +Specify `enable_block_reuse=True` alongside the connector for any of this to run; see [Block reuse alongside the connector](#block-reuse-alongside-the-connector). + +`get_num_new_matched_tokens` is called **at most once per KV allocation**. This is the precise form of the "once per request" rule: if a request's KV cache is destroyed and the request is replayed -- which `MAX_UTILIZATION` does under memory pressure -- the replay asks again, because the pages the first answer described are gone. + +**Deployment note.** Under a connector, a workload that was token-bound becomes KV-bound: the connector removes forward-pass tokens but its prefix still occupies GPU pages. Lowering `max_num_tokens` to hand memory back to the KV pool is usually the right adjustment, the opposite of the guidance for a connector-free deployment. #### 2. Worker Interface (`KvCacheConnectorWorker`) @@ -49,7 +257,11 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **`register_kv_caches(self, kv_cache_tensor: torch.Tensor)`** * **Description**: Called at initialization. Provides the worker with the GPU KV cache tensors. - * **Arguments**: `kv_cache_tensor` is the underlying storage tensor for the KV cache. + * **Arguments**: `kv_cache_tensor` is the underlying storage tensor for the KV cache, shaped `[num_blocks, num_layers, kv_factor, block_size]`. Row `block_id` is that block's KV for every layer. Dimension 1 is indexed by model layer in ascending order, so the `layer_idx` passed to `wait_for_layer_load` and `save_kv_layer` indexes it directly — no mapping is supplied, and none is needed. + +* **`register_kv_cache_layout(self, layout: KvCacheLayout)`** + * **Description**: Called at initialization *instead of* `register_kv_caches` when the cache describes itself as pools rather than one tensor. `KvCacheLayout` gives byte ranges per layer group: `layout.groups[g].regions[r]`, where the data for page slot `i` is at `region.base + region.stride * i` for `region.size` bytes, or `region.slot_tensor(i, dtype)`. Full attribute reference: [`KvCacheLayout` reference](#kvcachelayout-reference). + * **Default**: reconstructs the single-pool tensor and forwards it to `register_kv_caches`, so a connector that does not override this needs no changes for any single-window model. It raises when the cache cannot be described as one tensor — several layer groups (VSWA), or several regions (block scales, layers of differing size). * **`start_load_kv(self, stream: torch.cuda.Stream)`** * **Description**: Initiates the loading of KV blocks from the external source into the GPU memory. diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 23a85fc4b03a..4a51dc855348 100644 --- a/docs/source/features/kvcache.md +++ b/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++ diff --git a/examples/llm-api/llm_kv_cache_connector.py b/examples/llm-api/llm_kv_cache_connector.py index 882478993e5a..73e847fa3425 100644 --- a/examples/llm-api/llm_kv_cache_connector.py +++ b/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,9 +187,18 @@ 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) @@ -190,13 +206,16 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput): # 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 = {} diff --git a/examples/llm-api/llm_kv_cache_connector_vswa.py b/examples/llm-api/llm_kv_cache_connector_vswa.py new file mode 100644 index 000000000000..4ae42f83a555 --- /dev/null +++ b/examples/llm-api/llm_kv_cache_connector_vswa.py @@ -0,0 +1,359 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +### :title KV Cache Connector under VSWA +### :order 6 +### :section Customization +"""KV connector for a cache with one attention window size per layer group. + +`llm_kv_cache_connector.py` is the connector to start from. It works whenever a +single tensor can describe the whole KV cache, which is every model with one +attention window size. This one covers what changes when that stops being true +-- variable sliding-window attention (VSWA), where the cache allocates one pool +per window size. + +Five things differ, and each is marked `VSWA:` below. + +1. Pages are addressed per layer group. A page index is scoped to a group, so + the flat `block_ids` list does not exist and `register_kv_caches` is never + called; `register_kv_cache_layout` is the entry point. +2. Every block-id callback switches to its `*_by_layer_group` form. +3. A cache key must include the layer group. The same token range lives in every + group holding *different* KV, so a key derived from tokens alone collides + across groups and one group's bytes overwrite another's. +4. A sliding group offers only its live window to save. Blocks the window has + passed report no page and hold no readable KV, so `valid_page_slots` drops + them and a connector persists at most `window_size` tokens per sequence for + such a group -- not `prompt_len`. Size the store for that. +5. A block is only servable when *every* group holds it. The full-attention + group keeps the whole prompt while the sliding group keeps a tail, so the + prefix this connector can serve back is bounded by the smallest window. The + lookup below intersects across groups and stops at the first ordinal any + group misses. + +The cache key covers one block's tokens, as in the flat example. That assumes a +block's tokens determine its KV, which two prompts sharing a block but not the +prefix before it break: the second reads back KV computed under the first one's +prefix. A production connector chains the prefix into the key. + +Run with a VSWA model, for example: + + python llm_kv_cache_connector_vswa.py --model \ + --max-attention-window 1024 1024 1024 1024 1024 32768 +""" + +import hashlib +import os +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import click +import torch + +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 ( + KvCacheRegion, + valid_page_slots, +) +from tensorrt_llm.bindings.internal.batch_manager import LlmRequest +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KvCacheConnectorConfig, TorchLlmArgs + +CONNECTOR_CACHE_FOLDER_KEY = "CONNECTOR_CACHE_FOLDER" + + +@dataclass +class VswaConnectorMetadata: + # (path, layer_group_id, page_slot) -- the group is part of every entry + # because a page slot only means something inside its own group. + load: List[Tuple[str, int, int]] = field(default_factory=list) + save: List[Tuple[str, int, int]] = field(default_factory=list) + + +class VswaKvCacheConnectorWorker(KvCacheConnectorWorker): + def __init__(self, llm_args: TorchLlmArgs): + super().__init__(llm_args) + # VSWA (1): one region per layer group, not one tensor for the whole + # cache. The region is kept rather than a `[num_slots, bytes]` view of + # it, because `slot_tensor` checks the page slot before addressing it + # and a plain view accepts any subscript. + self.group_regions: Dict[int, KvCacheRegion] = {} + self.layer_to_group: Dict[int, int] = {} + + def register_kv_cache_layout(self, layout) -> None: + # VSWA (1): a group's buffers may coalesce into several regions, so a + # region is addressed per (group, region). `region.slot_tensor(i)` is the + # bytes of page slot `i` of that group. + for group in layout.groups: + if len(group.regions) != 1: + raise NotImplementedError( + f"layer group {group.layer_group_id} has " + f"{len(group.regions)} regions; this example handles one. " + "Address `group.regions` individually to support more." + ) + self.group_regions[group.layer_group_id] = group.regions[0] + for layer_id in group.layer_ids: + self.layer_to_group[layer_id] = group.layer_group_id + logger.info( + f"layer group {group.layer_group_id}: window={group.window_size}, " + f"{len(group.layer_ids)} layers, {group.bytes_per_page} bytes per page" + ) + + def start_load_kv(self, stream: torch.cuda.Stream): + for path, group_id, slot in self._metadata.load: + cpu_tensor = torch.load(path, map_location="cpu", weights_only=True) + self.group_regions[group_id].slot_tensor(slot).copy_(cpu_tensor, non_blocking=False) + + def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream): + # `layer_to_group` is what turns a per-layer hook into the group whose + # pages that layer reads. A connector that overlapped the transfer with + # compute would wait here only on the group this layer belongs to. + pass + + def save_kv_layer(self, layer_idx: int, stream: torch.cuda.Stream): + pass + + def wait_for_save(self, stream: torch.cuda.Stream): + stream.synchronize() + for path, group_id, slot in self._metadata.save: + if Path(path).exists(): + continue + torch.save(self.group_regions[group_id].slot_tensor(slot).cpu(), path) + + def get_finished( + self, finished_gen_req_ids: List[int], started_loading_req_ids: List[int] + ) -> Tuple[List[int], List[int]]: + return [], [] + + +class VswaKvCacheConnectorLeader(KvCacheConnectorScheduler): + def __init__(self, llm_args: TorchLlmArgs): + super().__init__(llm_args) + self.block_size = self._llm_args.kv_cache_config.tokens_per_block + # VSWA (5): the lookup has to check every group, so the leader needs the + # group count before the first request, and the configured window list + # is all it has that early. That list is not what the cache groups on: + # two configured windows at or above `max_seq_len` describe one group, + # not two. Count effective windows instead. + windows = self._llm_args.kv_cache_config.max_attention_window or [None] + max_seq_len = self._llm_args.max_seq_len + self.num_layer_groups = len( + {self._effective_window(window, max_seq_len) for window in windows} + ) + # request_id -> list of per-group file paths, one entry per matched block + # ordinal, in ordinal order starting at the first locally uncomputed one. + self.pending_loads: Dict[int, List[List[str]]] = {} + self.cache_folder = os.environ.get(CONNECTOR_CACHE_FOLDER_KEY, "./connector_cache") + os.makedirs(self.cache_folder, exist_ok=True) + + @staticmethod + def _effective_window(window: Optional[int], max_seq_len: Optional[int]) -> Optional[int]: + """The window a layer group is formed on, not the one that was configured. + + Mirrors `_resolve_v2_max_attention_window_vec`: clamp to `max_seq_len`, + then report full attention (`None`) for a window that reaches it. + `max_seq_len` is `None` when it was inferred from the model rather than + configured, and the count then falls back to the configured values -- + `build_connector_meta` rejects the request if that disagrees with the + cache. + """ + if window is None or window <= 0: + return None + if max_seq_len is None: + return int(window) + return None if int(window) >= int(max_seq_len) else int(window) + + # VSWA (3): the group id goes into the key. Without it, group 0 and group 1 + # hash the same tokens to the same file and overwrite each other's KV. + def _file_path(self, tokens: List[int], layer_group_id: int, salt: Optional[str]) -> str: + digest = hashlib.sha256(repr((tokens, layer_group_id, salt)).encode()).hexdigest() + return os.path.join(self.cache_folder, f"{digest}.pt") + + def _chunk_tokens(self, tokens: List[int]) -> List[List[int]]: + return [tokens[i : i + self.block_size] for i in range(0, len(tokens), self.block_size)] + + def get_num_new_matched_tokens( + self, request: LlmRequest, num_computed_tokens: int + ) -> Tuple[int, bool]: + self.pending_loads[request.request_id] = [] + + # Partial blocks are not stored, so a partial local match has nothing + # to append to. + if num_computed_tokens % self.block_size != 0: + return 0, False + + computed_blocks = num_computed_tokens // self.block_size + remaining = request.get_tokens(0)[computed_blocks * self.block_size :] + + for chunk in self._chunk_tokens(remaining): + if len(chunk) != self.block_size: + break + paths = [ + self._file_path(chunk, group_id, request.cache_salt) + for group_id in range(self.num_layer_groups) + ] + # VSWA (5): every group or none. A block the sliding group dropped + # is unservable even though the full-attention group still has it, + # because the sliding layers would then attend to KV that was never + # written. + if not all(Path(path).exists() for path in paths): + break + self.pending_loads[request.request_id].append(paths) + + matched = len(self.pending_loads[request.request_id]) * self.block_size + logger.info( + f"VSWA KV CONNECTOR: matched {matched} tokens " + f"({len(self.pending_loads[request.request_id])} blocks x " + f"{self.num_layer_groups} groups) for request {request.request_id}" + ) + return matched, False + + def update_state_after_alloc_by_layer_group( + self, request: LlmRequest, block_ids_by_layer_group: List[List[int]] + ) -> None: + # VSWA (2): the flat `update_state_after_alloc` is never called here. + pass + + def build_connector_meta(self, scheduler_output: SchedulerOutput): + # NOTE: This is a simplified implementation, and does not work with + # chunked prefill. A request appears in `new_requests` once, carrying + # the pages allocated so far; with chunked prefill the later chunks' + # pages arrive under `cached_requests`, whose entries this loop never + # reads. Blocks whose page is not yet allocated drop out of + # `valid_by_group` below and are never revisited, so those chunks go + # unsaved. + + metadata = VswaConnectorMetadata() + + for req in scheduler_output.new_requests: + pending_load = self.pending_loads.pop(req.request_id, []) + by_group = req.new_block_ids_by_layer_group + # VSWA (2): the flat list is empty with several groups. + if len(by_group) != self.num_layer_groups: + raise RuntimeError( + f"expected {self.num_layer_groups} layer groups from the " + f"cache, got {len(by_group)}. The window list this leader " + "derived its group count from does not describe the cache." + ) + + # `computed_position` excludes what the connector said it would + # serve, so this is where the locally computed prefix ends and the + # matched blocks begin. + num_computed_blocks = req.computed_position // self.block_size + + # VSWA (4): a sliding group holds pages for its live window only, so + # its list reports no page for the ordinals the window has passed. + # Those ordinals stay in place to keep entry `i` describing the same + # token range, and drop out here so no transfer targets them. + valid_by_group = [dict(valid_page_slots(slots)) for slots in by_group] + + for offset, paths in enumerate(pending_load): + ordinal = num_computed_blocks + offset + for group_id, path in enumerate(paths): + slot = valid_by_group[group_id].get(ordinal) + if slot is None: + continue + metadata.load.append((path, group_id, slot)) + + chunks = self._chunk_tokens(req.new_tokens) + for ordinal in range(num_computed_blocks + len(pending_load), len(chunks)): + if len(chunks[ordinal]) != self.block_size: + continue + for group_id, valid_slots in enumerate(valid_by_group): + slot = valid_slots.get(ordinal) + if slot is None: + continue + path = self._file_path(chunks[ordinal], group_id, req.cache_salt) + metadata.save.append((path, group_id, slot)) + + return metadata + + def request_finished_by_layer_group( + self, request: LlmRequest, cache_block_ids_by_layer_group: List[List[int]] + ) -> bool: + # VSWA (2) and (4): per group, and a sliding group's list covers its live + # window only -- everything older is -1 and holds no readable KV. + self.pending_loads.pop(request.request_id, None) + return False + + +def build_llm( + model: str, + max_attention_window: List[int], + max_seq_len: Optional[int] = None, + free_gpu_memory_fraction: float = 0.5, + use_kv_cache_manager_v2: "bool | str" = True, + enable_block_reuse: bool = True, +): + """An `LLM` wired to this connector. Shared by `main` and the e2e test.""" + connector_config = KvCacheConnectorConfig( + connector_module=__name__, + connector_scheduler_class="VswaKvCacheConnectorLeader", + connector_worker_class="VswaKvCacheConnectorWorker", + ) + return LLM( + model=model, + backend="pytorch", + cuda_graph_config=None, + disable_overlap_scheduler=True, + max_seq_len=max_seq_len, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=free_gpu_memory_fraction, + # One window size per layer, repeated cyclically. More than one + # distinct value is what makes the cache allocate a layer group + # per window -- and what this example exists to demonstrate. + max_attention_window=list(max_attention_window), + # VSWA needs the layout-describing registration path. "auto" is + # enough for a model that declares the preference itself, which + # Gemma-3 does. + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + enable_block_reuse=enable_block_reuse, + ), + kv_connector_config=connector_config, + ) + + +@click.command() +@click.option("--model", type=str, required=True) +@click.option("--max-attention-window", type=int, multiple=True, required=True) +def main(model: str, max_attention_window: Tuple[int, ...]): + with tempfile.TemporaryDirectory() as cache_folder: + os.environ[CONNECTOR_CACHE_FOLDER_KEY] = cache_folder + prompt = "The future of AI is" + params = SamplingParams(max_tokens=16, ignore_eos=True) + + # Cold: nothing is cached, so every full block is saved. + llm = build_llm(model, list(max_attention_window)) + try: + print("cold:", llm.generate([prompt], params)[0].outputs[0].text) + finally: + llm.shutdown() + + # Warm: a fresh instance shares only the disk cache, so anything served + # back came through the connector. + llm = build_llm(model, list(max_attention_window)) + try: + print("warm:", llm.generate([prompt], params)[0].outputs[0].text) + finally: + llm.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index cfcb6e849212..190c621785ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -741,9 +741,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 through the pool layout + # registration path, so it no longer forces a fallback. 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: @@ -771,7 +771,7 @@ def _validate_or_fallback_kv_cache_manager_v2( "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.") + "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. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 99da2a42265c..99bdff750e77 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -48,12 +48,22 @@ ) from tensorrt_llm.bindings.internal.batch_manager import LlmRequest from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.logger import logger from ..llm_request import get_draft_token_length from ..scheduler import ScheduledRequests if TYPE_CHECKING: from ..resource_manager import KVCacheManager + from .kv_cache_layout import KvCacheLayout + + +# `logger.warning_once` key for the KVCacheManagerV2 retention diagnostic +# emitted by `KvCacheConnectorSchedulerOutputRequest.update_and_build_data`. +# `log_once` marks a key as seen before consulting the log level +# (tensorrt_llm/logger.py:284-287), so the key survives a run that printed +# nothing; a test that wants to observe the warning has to clear it first. +V2_RETENTION_IGNORED_LOG_KEY = "kv_connector_v2_retention_config_ignored" # Used to store data for a single inflight request. @@ -85,6 +95,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 per layer group, indexed by layer group id. A block + # that is out of window, or has no page in that group, appears as + # BAD_PAGE_INDEX in place, keeping ordinals aligned to token ranges. Empty + # when block IDs are a single flat space, which `new_block_ids` describes. + new_block_ids_by_layer_group: List[List[int]] = field(default_factory=list) # A class to store some basic data regarding all inflight requests. @@ -98,12 +113,60 @@ class SchedulerOutput: cached_requests: List[RequestData] = field(default_factory=list) +def _flat_form_unavailable(flat: str, grouped: str) -> Callable: + """A stand-in for ``flat`` that names the form this connector implements.""" + + def unavailable(self, *args, **kwargs): + raise NotImplementedError( + f"{type(self).__name__} implements {grouped}, not {flat}. {flat} is " + "what a cache that describes itself as one flat list calls, rather " + "than one list per layer group; implement it too to run on such a " + "model." + ) + + unavailable.__name__ = flat + unavailable.__qualname__ = flat + return unavailable + + +def _satisfy_flat_abstracts(cls: type, base: type, pairs: Dict[str, str]) -> None: + """Let an override of the per-layer-group form stand in for the flat one. + + ``request_finished`` and ``request_finished_by_layer_group`` are two + spellings of one callback, and a connector implements whichever one its + cache reports in. Only the flat spelling is abstract, so without this a + VSWA-only connector would not instantiate at all, and would have to carry a + dead flat method purely to clear the abstract flag -- abstractness is + tracked per method name, so overriding the per-group form does not clear it. + + Called from ``__init_subclass__``, which runs before ``ABCMeta`` collects + ``__abstractmethods__``, so the injected method counts as an implementation. + A connector that implements neither form is left alone and still fails at + construction with a ``TypeError`` naming the flat method. + """ + for flat, grouped in pairs.items(): + if getattr(cls, grouped, None) is getattr(base, grouped): + # The per-layer-group form is the base default, so nothing stands in. + continue + if not getattr(getattr(cls, flat, None), "__isabstractmethod__", False): + continue + setattr(cls, flat, _flat_form_unavailable(flat, grouped)) + + class KvCacheConnectorWorker(ABC): def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args self._metadata = None super().__init__() + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + _satisfy_flat_abstracts( + cls, + KvCacheConnectorWorker, + {"register_kv_caches": "register_kv_cache_layout"}, + ) + def bind_connector_meta(self, metadata: object): self._metadata = metadata @@ -135,6 +198,40 @@ 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 cache has one slot + address space per pool and one page-index space per layer group, which + a single tensor cannot always describe. 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]``; the page indices in + ``RequestData.new_block_ids_by_layer_group[g]`` are scoped to layer + group ``g`` and index that group's regions. With a single layer group + they are also reported flat, in ``RequestData.new_block_ids``. + + Args: + layout: Description of the KV cache pools; see ``KvCacheLayout``. + + The default forwards a single-pool cache to ``register_kv_caches`` with + the tensor shape that path has always received, so a connector that does + not override this keeps working wherever one tensor describes the cache. + Override it for several layer groups (variable sliding-window + attention), several regions (block scales, or layers of differing size), + or to address regions directly. + """ + tensor = layout.as_single_pool_tensor() + if tensor is None: + raise NotImplementedError( + f"{type(self).__name__} does not implement " + "register_kv_cache_layout, and this KV cache cannot be described " + f"as one pool tensor: {len(layout.groups)} layer group(s), " + f"{sum(len(group.regions) for group in layout.groups)} region(s). " + "Implement register_kv_cache_layout to address the regions directly." + ) + self.register_kv_caches(tensor) + @abstractmethod def start_load_kv(self, stream: torch.cuda.Stream): """ @@ -201,6 +298,17 @@ def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args super().__init__() + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + _satisfy_flat_abstracts( + cls, + KvCacheConnectorScheduler, + { + "request_finished": "request_finished_by_layer_group", + "update_state_after_alloc": "update_state_after_alloc_by_layer_group", + }, + ) + @abstractmethod def build_connector_meta(self, scheduler_output: SchedulerOutput): """ @@ -237,6 +345,9 @@ def request_finished(self, request: LlmRequest, cache_block_ids: List[int]) -> b Args: request: The request that finished generating tokens. + cache_block_ids: Page slot indices to save from. Under a sliding + window this covers the live window only -- blocks the window + has passed hold no readable KV and appear as ``-1``. Returns: Whether the request is performing asynchronous saving operations. @@ -245,6 +356,33 @@ def request_finished(self, request: LlmRequest, cache_block_ids: List[int]) -> b (determined by ``get_finished`` on the workers). """ + def request_finished_by_layer_group( + self, request: LlmRequest, cache_block_ids_by_layer_group: List[List[int]] + ) -> bool: + """ + Per-layer-group form of ``request_finished``. + + Called instead of the flat form when the KV cache has more than one + layer group. Same return contract as ``request_finished``. Implement it + together with ``update_state_after_alloc_by_layer_group``; a connector + missing either one is rejected during executor bring-up for such a + model. + + Args: + request: The request that finished generating tokens. + cache_block_ids_by_layer_group: Page slot indices per layer group, + indexed by layer group id, with ``-1`` where a block has no page + in that group. + """ + if len(cache_block_ids_by_layer_group) != 1: + raise NotImplementedError( + f"{type(self).__name__} does not implement " + "request_finished_by_layer_group, and this KV cache has " + f"{len(cache_block_ids_by_layer_group)} layer groups, whose page " + "indices cannot share one list." + ) + return self.request_finished(request, cache_block_ids_by_layer_group[0]) + @abstractmethod def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): """ @@ -255,6 +393,35 @@ def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): block_ids: The KV cacheblock IDs that were allocated. """ + def update_state_after_alloc_by_layer_group( + self, request: LlmRequest, block_ids_by_layer_group: List[List[int]] + ) -> None: + """ + Per-layer-group form of ``update_state_after_alloc``. + + Called instead of the flat form when the KV cache has more than one + layer group, where a single list cannot describe the allocation: a page + index is scoped to a group, so indices from different groups collide. + Implementing this is what makes a connector usable under variable + sliding-window attention; a connector that implements only the flat form + is rejected during executor bring-up for such a model. + + Args: + request: The request that was allocated resources. + block_ids_by_layer_group: Page slot indices per layer group, indexed + by layer group id. Entry ``[g][i]`` describes block ordinal + ``i`` of group ``g``, and is ``-1`` (``BAD_PAGE_INDEX``) where + that block has no page in that group. + """ + if len(block_ids_by_layer_group) != 1: + raise NotImplementedError( + f"{type(self).__name__} does not implement " + "update_state_after_alloc_by_layer_group, and this KV cache has " + f"{len(block_ids_by_layer_group)} layer groups, whose page indices " + "cannot share one list." + ) + self.update_state_after_alloc(request, block_ids_by_layer_group[0]) + def wait_for_initialization(self): """ Some connectors need to wait for some resources to be initialized. @@ -316,21 +483,48 @@ def loading_ids(self) -> Set[int]: class KvCacheConnectorSchedulerOutputRequest: def __init__(self): self.block_ids = [] + self.block_ids_by_layer_group: List[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.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: List[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) + while len(self.block_ids_by_layer_group) < len(indices_by_group): + self.block_ids_by_layer_group.append([]) + for layer_group_id, indices in enumerate(indices_by_group): + seen = self.block_ids_by_layer_group[layer_group_id] + new_ids = indices[len(seen) :] + seen.extend(new_ids) + new_block_ids_by_layer_group.append(new_ids) + # A page index is scoped to a layer group, so several groups cannot + # be flattened into `new_block_ids`; with one group -- every + # non-VSWA, non-hybrid model -- that group's indices are the flat + # list, and with several the by-group field is the only source. + new_block_ids = ( + new_block_ids_by_layer_group[0] 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) - self.block_ids.extend(new_block_ids) + new_block_ids = block_ids[len(self.block_ids) :] + self.block_ids.extend(new_block_ids) + + new_tokens = tokens[len(self.tokens) :] self.tokens.extend(new_tokens) if req.state in ( @@ -345,13 +539,28 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag req ) # 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) + # Get retention priority for each new block only if retention config is + # provided (for priority-based offload filtering). Priorities stay None + # under KVCacheManagerV2, which honours no `KvCacheRetentionConfig`: + # every page carries the default there, so reporting a priority would + # misdescribe what the user asked for. Warn rather than report nothing + # in silence -- the user configured retention and it is not in effect. priorities = None if 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 - ] + if is_v2: + logger.warning_once( + "KvCacheRetentionConfig has no effect in this configuration: no " + "per-block retention priority is honoured, so RequestData.priorities is " + "reported as None and a connector cannot filter offloads by priority. Set " + "kv_cache_config.use_kv_cache_manager_v2=False to keep retention " + "priorities on the connector path.", + key=V2_RETENTION_IGNORED_LOG_KEY, + ) + else: + priorities = [ + kv_cache_manager.get_priority_by_block_id(block_id) + for block_id in new_block_ids + ] return RequestData( req.request_id, @@ -362,6 +571,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, ) @@ -411,6 +621,18 @@ def build_scheduler_output( def record_new_matched_tokens(self, request: LlmRequest, num_new_matched_tokens: int): self.external_loads[request.request_id] = num_new_matched_tokens + def reset_request(self, request_id: int) -> None: + """Drop the per-request deltas so a replay is reported as a new request. + + ``block_ids`` and ``tokens`` are cumulative, so every + ``update_and_build_data`` reports only what was appended since the last + call. Once an allocation is destroyed that delta is taken against pages + that no longer exist, and the replay lands in ``cached_requests``, which + a connector that walks only ``new_requests`` never reads. + """ + self.requests.pop(request_id, None) + self.external_loads.pop(request_id, None) + class KvCacheConnectorManager(KvCacheConnectorManagerCpp): """ @@ -463,10 +685,24 @@ 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: - if request.is_generation_only_request: - raise RuntimeError("Connector API is not supported for generation-only requests!") - + 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. + + The connector ABC promises one query per allocation, so a caller must + reach this at most once per request per allocation whatever it does with + the result. A caller that cannot consume the whole answer -- one + allocating per context chunk, where an offer can outrun the reserved + pages -- records only the part it honours via + ``commit_new_matched_tokens``. + + Generation-only requests are rejected by the caller, not here: the two + callers hold different objects. ``get_num_new_matched_tokens`` is handed + a C++-side ``LlmRequest`` where the flag is a property, while + ``KVCacheManagerV2`` holds the Python subclass where it is a method, so + no single spelling of the check reads correctly in both. + """ num_tokens, load_kv_async = self._run_on_leader( lambda: self.scheduler.get_num_new_matched_tokens(request, num_computed_tokens) ) @@ -474,6 +710,33 @@ 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. + + ``num_tokens`` is what the runtime will actually consume, which may be + less than the offer; the unconsumed tail is recomputed locally and the + connector releases the whole request at ``request_finished``. Must run + in the iteration the request is scheduled, because every + ``build_scheduler_output`` consumes and clears ``external_loads`` and + reads ``new_async_requests.loading``. + + ``load_kv_async`` passes through unmodified even on a short honour, and + even when nothing is honoured at all -- the one pair + ``query_num_new_matched_tokens`` rejects on the way in. A connector that + answered asynchronously has already started the transfer, because a + parked request is skipped by ``build_scheduler_output`` and so never + reaches ``start_load_kv``; dropping the flag would leave the request in + the batch and let prefill write those pages concurrently. So a connector + must tolerate a load whose result is discarded: the runtime parks the + request, waits for the load, then computes the range locally and + overwrites it. There is no way to say otherwise -- the ABC has no + ``cancel_load``, and a parked request carries no honoured count into + ``build_connector_meta``. + """ # 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 +751,112 @@ 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. + + 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. That + caller allocates for the whole prompt in the same call and so never + honours less than it was offered, which is what makes fusing the two + safe here. + + ``request`` arrives from C++, which casts it to the base-class + ``LlmRequest`` rather than the ``_torch`` subclass, so only what + ``nanobind/batch_manager/bindings.cpp`` binds is readable below. + """ + if request.is_generation_only_request: + raise RuntimeError("Connector API is not supported for generation-only requests!") + + 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 + _GROUPED_SCHEDULER_METHODS = ( + "update_state_after_alloc_by_layer_group", + "request_finished_by_layer_group", + ) + + def _missing_grouped_scheduler_methods(self) -> List[str]: + """Which per-layer-group forms this scheduler leaves at the base default. + + Collective: the answer is broadcast from the leader, so every rank has to + take part -- calling it on the leader alone hangs the others. + """ + return self._run_on_leader( + lambda: [ + name + for name in self._GROUPED_SCHEDULER_METHODS + # A scheduler that does not inherit `KvCacheConnectorScheduler` has no + # default to be holding, so absence here is not the flat form. + if getattr(type(self.scheduler), name, None) + is getattr(KvCacheConnectorScheduler, name) + ] + ) + + def reject_flat_only_scheduler(self, num_layer_groups: int) -> None: + """Refuse a scheduler that cannot describe a multi-group cache. + + A page index is scoped to a layer group, so with more than one group the + flat callbacks are never called and both per-layer-group forms have to + be implemented. Checked here rather than left to the base defaults, + which raise on the first scheduled request: by then the model is loaded, + and the two forms are checked independently, so a connector overriding + one and leaving the other flat would start up cleanly. + + The scheduler lives on the leader alone, so the answer is broadcast -- + every rank raises, rather than rank 0 failing while the others wait. + """ + if num_layer_groups <= 1: + return + missing = self._missing_grouped_scheduler_methods() + if not missing: + return + raise NotImplementedError( + f"This KV cache has {num_layer_groups} layer groups, whose page indices " + "cannot share one list, so the KV connector scheduler must implement " + f"{' and '.join(missing)}. Implement both per-layer-group forms; the " + "flat forms they replace are not called for this cache." + ) + + def warn_flat_scheduler_under_swa(self, window_size: Optional[int]) -> None: + """Warn when a flat-only scheduler meets a single sliding-window group. + + One layer group means the flat callbacks still carry that group's page + indices, so such a connector runs -- and is allowed to, because refusing + it would stop connectors that work today. What differs + is that a sliding window reports ``BAD_PAGE_INDEX`` for every block it has + passed, and ``-1`` is a valid subscript. The warning is what makes that + difference findable before it becomes a transfer against another + request's KV. + """ + if window_size is None: + return + # The probe is collective, so every rank runs it; only the rank holding + # the scheduler reports, or the message repeats once per rank. + missing = self._missing_grouped_scheduler_methods() + if not missing or self.scheduler is None: + return + logger.warning( + "The KV connector scheduler implements only the flat page-index callbacks, and " + f"this model has a single sliding attention window ({window_size} tokens). Blocks " + "the window has passed are reported as BAD_PAGE_INDEX (-1) in place, and indexing " + "a page list with -1 addresses the last page slot of the pool -- another request's " + "KV. Filter page indices through " + "tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout.valid_page_slots. See " + "docs/source/features/kv-cache-connector.md, 'Blocks with no page'." + ) + + def reset_request_state(self, request: LlmRequest) -> None: + """Tell the connector bookkeeping that this request's allocation died. + + Only a cache that can destructively pause and replay a live request + needs this; under ``GUARANTEED_NO_EVICT`` an allocation dies only when + the request finishes. ``KVCacheV2Scheduler`` coerces the policy to + ``MAX_UTILIZATION`` whatever was configured, so everything keyed to the + old allocation has to go with it. + """ + self.scheduler_output_manager.reset_request(request.request_id) + def should_add_sequence(self, request: LlmRequest) -> bool: req_id = request.request_id return req_id not in self.finished_async_loading_requests @@ -540,12 +907,20 @@ def handle_metadata(self) -> object: self.worker.bind_connector_meta(metadata) - def request_finished(self, req: LlmRequest, cache_block_ids: List[int]) -> bool: + def request_finished( + self, + req: LlmRequest, + cache_block_ids: List[int], + cache_block_ids_by_layer_group: Optional[List[List[int]]] = None, + ) -> bool: """ Called when a request is finished generating tokens. Args: req: The request that finished generating tokens. + cache_block_ids: Flat page slot indices to save from. + cache_block_ids_by_layer_group: The same, per layer group. Used when + the connector implements the per-layer-group API. Returns: Whether the request is performing asynchronous saving @@ -556,9 +931,21 @@ def request_finished(self, req: LlmRequest, cache_block_ids: List[int]) -> bool: if req.request_id in self.finished_async_loading_requests: del self.finished_async_loading_requests[req.request_id] - saving_async = self._run_on_leader( - lambda: self.scheduler.request_finished(req, cache_block_ids) - ) + # A caller with no per-group view of the cache leaves this unset, and the + # flat list is the whole description. Every per-group caller passes one + # list per layer group, empty lists included, so a request whose cache + # was already released still reaches the form the connector implements + # rather than the grouped default's "0 layer groups" refusal. + if cache_block_ids_by_layer_group: + saving_async = self._run_on_leader( + lambda: self.scheduler.request_finished_by_layer_group( + req, cache_block_ids_by_layer_group + ) + ) + else: + saving_async = self._run_on_leader( + lambda: self.scheduler.request_finished(req, cache_block_ids) + ) # This is similar to take_scheduled_requests_pending_load. # We need to update the request's state to indicate that it's still being used, but isn't schedulable. @@ -620,8 +1007,23 @@ def get_finished(self) -> List[LlmRequest]: # The execution loop will call _terminate_request on these requests. return list(all_finished.saving.values()) - def update_state_after_alloc(self, req: LlmRequest, block_ids: List[int]): - if self.scheduler is not None: + def update_state_after_alloc( + self, + req: LlmRequest, + block_ids: List[int], + block_ids_by_layer_group: Optional[List[List[int]]] = None, + ): + if self.scheduler is None: + return + # A cache that reports per layer group is reported that way, whatever + # the group count -- including one empty list per group for a request + # whose cache is already gone. The connector's own default folds a single + # group back to the flat form, so a connector written against that keeps + # receiving exactly what it always has. A caller with no per-group view + # leaves this unset and the flat list is the whole description. + if block_ids_by_layer_group: + self.scheduler.update_state_after_alloc_by_layer_group(req, block_ids_by_layer_group) + else: self.scheduler.update_state_after_alloc(req, block_ids) def set_scheduler_output(self, scheduler_output: SchedulerOutput): diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py new file mode 100644 index 000000000000..b0749a168390 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py @@ -0,0 +1,335 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""KV cache layout description handed to a KV connector. + +The cache allocates *pool groups*, each holding ``num_slots`` slots. A slot holds +the coalesced buffers of one *layer group* (a life cycle), and a buffer is keyed +by ``BufferId(layer_id, role)``. A connector therefore cannot be handed a single +pool tensor: there is one slot address space per pool, and one index space per +layer group. + +Instead it is handed :class:`KvCacheLayout` -- a description of the byte ranges +that repeat per page slot. The addressing contract is taken verbatim from +``AggregatedPageDesc``:: + + (base + stride * i + Range(0, size) for i in aggregated_page_indices) + +where ``i`` comes from ``_KVCache.get_aggregated_page_indices(layer_group_id)``. + +Because ranges are described rather than implied, this covers MLA (a pool simply +has no VALUE buffer), sliding-window attention and hybrid models (one layer group +per window size) without any of them being special cases. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, Iterable, Iterator, List, Optional, Tuple + +import torch + +from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor + +if TYPE_CHECKING: + from ..kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + +__all__ = [ + "KvCacheBufferRef", + "KvCacheLayerGroupLayout", + "KvCacheLayout", + "KvCacheRegion", + "build_kv_cache_layout_v2", + "valid_page_slots", +] + + +def valid_page_slots(page_indices: Iterable[int]) -> Iterator[Tuple[int, int]]: + """Yield ``(block_ordinal, page_slot)`` for the entries that address a page. + + A page-index list holds one entry per block ordinal, so entry ``i`` describes + tokens ``[i * tokens_per_block, (i + 1) * tokens_per_block)``. A block with no + page in the layer group -- past the end of a sliding window, or never + allocated there -- holds ``BAD_PAGE_INDEX`` in place so that alignment + survives, which makes the list unsafe to index with directly. + + Iterating here keeps the ordinal, which a caller needs to find the token range + a page covers, and drops the entries that address no page, which no transfer + may be built against. + """ + for ordinal, slot in enumerate(page_indices): + if slot >= 0: + yield ordinal, slot + + +@dataclass(frozen=True) +class KvCacheBufferRef: + """One ``(layer, role)`` buffer covered by a region, in memory order.""" + + #: Global model layer index -- the same index space the per-layer connector + #: hooks (``wait_for_layer_load`` / ``save_kv_layer``) receive. + layer_id: int + #: Native role name as the cache manager spells it, e.g. "key" / "value". + #: Deliberately not an enum: roles are an open vocabulary on the manager + #: side, and a connector should not need updating when one is added. + role: str + #: Page expansion factor for heterogeneous tokens-per-block layers. + expansion: int = 1 + + +@dataclass(frozen=True) +class KvCacheRegion: + """A contiguous byte range that repeats once per page slot. + + The data for page slot ``i`` lives at ``base + stride * i``, for ``size`` + bytes. ``size`` is not necessarily equal to ``stride``: a region covers one + run of adjacent buffers within a slot, and a slot may hold several runs. + """ + + base: int + size: int + stride: int + num_slots: int + buffers: Tuple[KvCacheBufferRef, ...] + + def address_of(self, slot_id: int) -> int: + """Device address of this region for ``slot_id``.""" + if not 0 <= slot_id < self.num_slots: + raise IndexError(f"slot_id {slot_id} out of range [0, {self.num_slots})") + return self.base + self.stride * slot_id + + def as_tensor(self, dtype: torch.dtype = torch.uint8) -> torch.Tensor: + """A strided ``[num_slots, size // itemsize]`` view; row ``i`` is slot ``i``. + + Defaults to ``uint8``. A region may span several roles whose element + types differ, and a connector that only moves bytes should not have to + care; callers that want a typed view can pass ``dtype`` explicitly. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + if self.size % itemsize or self.stride % itemsize: + raise ValueError( + f"region size {self.size} and stride {self.stride} must both be " + f"multiples of {dtype} itemsize {itemsize}" + ) + return convert_to_torch_tensor( + TensorWrapper( + self.base, + dtype, + shape=(self.num_slots, self.size // itemsize), + strides=(self.stride // itemsize, 1), + ) + ) + + def slot_tensor(self, slot_id: int, dtype: torch.dtype = torch.uint8) -> torch.Tensor: + """A ``[size // itemsize]`` view of one page slot. + + The guarded form of ``as_tensor(dtype)[slot_id]``. A page index that + addresses no page is ``BAD_PAGE_INDEX`` (-1), which the strided view + accepts as a subscript and resolves to its last row -- a live page + belonging to some other request. Addressing a slot through here raises + ``IndexError`` instead, so a page index reaches device memory only once + it has been checked. Pair it with ``valid_page_slots`` to filter a whole + list at the source. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + if self.size % itemsize: + raise ValueError( + f"region size {self.size} must be a multiple of {dtype} itemsize {itemsize}" + ) + return convert_to_torch_tensor( + TensorWrapper( + self.address_of(slot_id), + dtype, + shape=(self.size // itemsize,), + strides=(1,), + ) + ) + + +@dataclass(frozen=True) +class KvCacheLayerGroupLayout: + """One layer group -- the unit that page indices are scoped to.""" + + layer_group_id: int + #: Global model layer indices belonging to this group. + layer_ids: Tuple[int, ...] + #: Attention window for this group, or None for full attention. + window_size: Optional[int] + regions: Tuple[KvCacheRegion, ...] + + @property + def bytes_per_page(self) -> int: + """Total bytes this group occupies for a single page slot.""" + return sum(region.size for region in self.regions) + + +@dataclass(frozen=True) +class KvCacheLayout: + """What a connector is handed in place of a single KV cache pool tensor.""" + + tokens_per_block: int + groups: Tuple[KvCacheLayerGroupLayout, ...] + #: Element type of the KV data, for typed views over a region. + dtype: torch.dtype = torch.uint8 + + def group(self, layer_group_id: int) -> KvCacheLayerGroupLayout: + for group in self.groups: + if group.layer_group_id == layer_group_id: + return group + raise KeyError(f"no layer group {layer_group_id} in layout") + + def group_of_layer(self, layer_id: int) -> KvCacheLayerGroupLayout: + """The layer group owning a global model layer index.""" + for group in self.groups: + if layer_id in group.layer_ids: + return group + raise KeyError(f"layer {layer_id} is not covered by this layout") + + def as_single_pool_tensor(self) -> Optional[torch.Tensor]: + """A ``[num_slots, num_layers, kv_factor, block_size]`` view, or None. + + This is the shape a single-pool KV cache manager hands + ``register_kv_caches``, so a connector written against that signature + keeps working when the same cache is described as a layout. Returns None + when the cache cannot be described that way -- several layer groups, + several regions (block scales, or layers of differing size), or a page + expansion factor that breaks the uniform grid. + + Dimension 1 is the region's buffers in memory order, laid out layer-major + and ascending by construction: the storage config walks ``config.layers`` + in order and appends each layer's buffers to the coalesced buffer, then + assigns offsets by walking that list. So the dimension is indexed by + layer directly. The order is not re-derived here; if that construction + ever changes, this is what changes with it. + """ + if len(self.groups) != 1 or len(self.groups[0].regions) != 1: + return None + group = self.groups[0] + region = group.regions[0] + num_layers = len(group.layer_ids) + if not num_layers or len(region.buffers) % num_layers: + return None + kv_factor = len(region.buffers) // num_layers + if any(buffer.expansion != 1 for buffer in region.buffers): + return None + itemsize = torch.tensor([], dtype=self.dtype).element_size() + if region.size % (itemsize * len(region.buffers)): + return None + block_size = region.size // itemsize // len(region.buffers) + return region.as_tensor(self.dtype).unflatten(1, (num_layers, kv_factor, block_size)) + + +def _global_layer_ids(manager: "KVCacheManagerV2", local_layer_ids) -> List[int]: + """Map internal layer ids to global model layer indices. + + ``pp_layers`` is the local-to-global table the manager already keeps. Models + that map several internal layers onto one model layer (the sparse-attention + virtual-layer path) have no single global index per internal layer, so they + are rejected rather than silently mislabelled. + """ + if hasattr(manager, "_layer_attn_to_layer_id"): + raise NotImplementedError( + "KV connector layout is not supported for managers with virtual " + "attention layers (sparse attention): an internal layer does not map " + "to a single model layer index." + ) + pp_layers = manager.pp_layers + return [int(pp_layers[int(lid)]) for lid in local_layer_ids] + + +def _window_size(init_config, local_layer_id: int) -> Optional[int]: + layers = init_config.layers + if local_layer_id >= len(layers): + raise ValueError(f"no layer config for internal layer {local_layer_id}") + window = getattr(layers[local_layer_id], "window_size", None) + return None if window is None else int(window) + + +def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: + """Describe the cache's page regions for a KV connector. + + Every region is reported with ``desc.base`` verbatim and no tier predicate, + so the addresses are device addresses only because tiers below GPU are + rejected while a connector is attached + (``PyExecutor._reject_non_gpu_cache_tiers``). Reads only the manager's + public layout API -- ``layer_grouping``, ``all_buffer_ids``, + ``get_aggregated_pages``, ``pool_group_descs`` -- so it assumes nothing + about dimension order, kv factor, or pool count. + """ + impl = manager.impl + init_config = impl.init_config + + # A pool group's slot count applies to every layer group drawn from it. + # Note LayerGroupId and PoolGroupIndex are distinct index spaces; the + # variants of a pool group name the layer groups it backs. + slots_by_group: Dict[int, int] = {} + for pool_group in impl.pool_group_descs: + for variant in pool_group.slot_desc.variants: + slots_by_group[int(variant.layer_group_id)] = int(pool_group.num_slots) + + buffers_by_layer: Dict[int, List] = {} + for buffer_id in impl.all_buffer_ids: + buffers_by_layer.setdefault(int(buffer_id.layer_id), []).append(buffer_id) + + groups: List[KvCacheLayerGroupLayout] = [] + for layer_group_id, local_layer_ids in enumerate(impl.layer_grouping): + local_layer_ids = [int(lid) for lid in local_layer_ids] + if not local_layer_ids: + continue + + num_slots = slots_by_group[layer_group_id] + global_by_local = dict(zip(local_layer_ids, _global_layer_ids(manager, local_layer_ids))) + + buffer_ids = [b for lid in local_layer_ids for b in buffers_by_layer.get(lid, ())] + + regions: List[KvCacheRegion] = [] + for desc in impl.get_aggregated_pages(buffer_ids): + if int(desc.layer_group_id) != layer_group_id: + continue + regions.append( + KvCacheRegion( + base=int(desc.base), + size=int(desc.size), + stride=int(desc.stride), + num_slots=num_slots, + buffers=tuple( + KvCacheBufferRef( + layer_id=global_by_local[int(b.id.layer_id)], + role=str(b.id.role), + expansion=int(b.expansion), + ) + for b in desc.buffers + ), + ) + ) + + groups.append( + KvCacheLayerGroupLayout( + layer_group_id=layer_group_id, + layer_ids=tuple(global_by_local[lid] for lid in local_layer_ids), + window_size=_window_size(init_config, local_layer_ids[0]), + regions=tuple(regions), + ) + ) + + # A region may span roles whose element types differ; the KV dtype is the + # one a typed view over K/V data needs, and anything else stays uint8. + try: + dtype = binding_to_torch_dtype(manager.dtype) + except (AssertionError, KeyError, TypeError): + dtype = torch.uint8 + + return KvCacheLayout( + tokens_per_block=int(manager.tokens_per_block), + groups=tuple(groups), + dtype=dtype, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 746766070332..5cdc3d8923ab 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1098,6 +1098,27 @@ def _rewind_context_cursor(req: LlmRequest, reuse: int) -> None: req.context_current_position = reuse +def _settle_context_cursor(req: LlmRequest, reuse: int, tokens_per_block: int) -> None: + """Point a context cursor at a ``reuse``-token prefix, chunk included. + + The chunk is position-relative, so moving the position without it leaves the + pair describing two different ranges -- a request re-entering after a + narrowed chunk comes back with ``position + chunk`` short of ``prompt_len``, + silently chunked. Spanning to the end is the right floor: the chunked + scheduler overwrites it, and the non-chunked path wants exactly this value. + """ + # ``reuse`` is ``num_committed_tokens``: the locally matched prefix alone. A + # connector-served prefix is held -- pages, and the cache's ``history_length`` + # -- but never committed, because committing is for tokens this worker + # computed. Settling a re-entering first chunk on the commit depth alone + # therefore points it below KV the cache holds, and ``history_length`` cannot + # be lowered to agree. + reuse = max(reuse, req.py_connector_served_position) + _rewind_context_cursor(req, reuse) + req.set_prepopulated_prompt_len(reuse, tokens_per_block) + req.context_chunk_size = req.context_remaining_length + + class KVCacheManagerV2(BaseResourceManager): # Filled lazily by _cold_pool_group_membership(); the grouping is fixed after construction. # Declared on the class so it is present even when an instance is built without running __init__. @@ -1142,10 +1163,10 @@ def __init__( self.mapping = mapping self.dtype = dtype self.is_disagg = is_disagg + self.kv_connector_manager = kv_connector_manager + # Filled on first use; the layer grouping does not change after init. + self._connector_life_cycle_by_group: Optional[List[AttnLifeCycle]] = None - assert kv_connector_manager is None, ( - "kv_connector_manager is not supported for KVCacheManagerV2" - ) assert max_beam_width == 1, "max_beam_width must be 1 for KVCacheManagerV2" self.kv_cache_type = kv_cache_type @@ -1416,7 +1437,18 @@ def append_to_kv_heads_per_layer( logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB") cache_tiers: List[CacheTierConfig] = [GpuCacheTierConfig(quota=int(quota))] - if kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: + if kv_connector_manager is not None and kv_cache_config.host_cache_size is None: + # A connector holds device addresses for its pages, and evicting a + # page to another tier reassigns its GPU slot underneath it. The + # automatic host tier only exists to give MAX_UTILIZATION somewhere + # to spill to, which a connector run never uses, so drop it here; an + # explicitly configured host_cache_size is rejected at bring-up. + host_quota = 0 + logger.info( + "KV cache manager v2 host tier disabled: a KV connector is attached " + "and registers GPU page addresses that tier migration would invalidate." + ) + elif kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: host_quota = kv_cache_config.host_cache_size else: # The V2 MAX_UTILIZATION scheduler relies on suspend/resume to @@ -3451,6 +3483,13 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) # Disagg generation receives prompt KV from the context worker; # scratch blocks are only valid for local prefill chunks. kv_cache.enable_swa_scratch_reuse = False + elif self._connector_may_serve(req): + # Same reason, one step earlier: a connector writes real cache + # content into these blocks. Whether it will is not known until + # `prepare_resources` asks, and the flag has to be off before + # `resize_context` can take scratch slots, so it is cleared for + # every servable request rather than only the served ones. + kv_cache.enable_swa_scratch_reuse = False if not self._resume_and_restore(req.py_request_id, kv_cache): return None return kv_cache.num_committed_tokens @@ -3476,8 +3515,7 @@ def prepare_context(self, req: LlmRequest) -> bool: # First chunk only: num_committed_tokens holds at the initial prefix # until context end, so reapplying later would rewind the cursor. if req.is_first_context_chunk and self.enable_block_reuse: - _rewind_context_cursor(req, reused) - req.set_prepopulated_prompt_len(reused, self.tokens_per_block) + _settle_context_cursor(req, reused, self.tokens_per_block) return True def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: @@ -3544,8 +3582,7 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: if reused is None: return False if self.enable_block_reuse: - _rewind_context_cursor(req, reused) - req.set_prepopulated_prompt_len(reused, self.tokens_per_block) + _settle_context_cursor(req, reused, self.tokens_per_block) kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is None: @@ -3631,6 +3668,75 @@ def resume_request(self, req: LlmRequest) -> bool: # ---- prepare_resources ---- + def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: + """Attention life cycle per layer group. + + Every layer in a group shares one life cycle, which is how + ``impl.layer_grouping`` partitions layers, so the first layer describes + the group. A group with no window -- full attention, or a group whose + layers carry no window at all -- yields an empty stale range. + """ + if self._connector_life_cycle_by_group is None: + layers = self.kv_cache_manager_py_config.layers + self._connector_life_cycle_by_group = [ + AttnLifeCycle.make( + getattr(layer, "window_size", None), + getattr(layer, "num_sink_tokens", None), + self.tokens_per_block, + ) + for layer in (layers[int(layer_ids[0])] for layer_ids in self.impl.layer_grouping) + ] + return self._connector_life_cycle_by_group + + def _stale_block_range(self, layer_group_id: int, history_length: int) -> Tuple[int, int]: + """Block ordinals ``[start, end)`` that ``layer_group_id`` no longer reads. + + The cache's own life cycle decides this, so the connector's view and the + allocator's cannot drift. ``start`` is not 0 whenever attention sinks are + configured: the first ``num_sink_tokens`` stay live however far below the + window they fall, and masking them would hide from the connector pages + attention still reads. + """ + stale = self._life_cycle_by_layer_group()[layer_group_id].get_stale_range( + history_length, self.tokens_per_block + ) + # Read `beg`/`end` rather than unpacking: only the pure-Python + # `HalfOpenRange` is a tuple subclass, the bound one is not. + return stale.beg, stale.end + + def get_page_indices_by_layer_group(self, request: LlmRequest) -> List[List[int]]: + """Page slot indices for ``request`` per layer group, by block ordinal. + + Indexed by layer group id, which is dense, and the outer length is always + the layer group count -- a request whose cache has already been released + reads back as one empty list per group. Callers route on that length, so + shortening it to zero would send a released request to a callback the + connector may not implement. + + Out-of-window blocks and blocks with no page both read back as + ``BAD_PAGE_INDEX`` in place rather than shortening the list, so ordinals + stay aligned to token ranges and an append-delta over successive calls + stays valid. + + The out-of-window entries are masked from ``history_length`` rather than + left to whatever page happens to still be attached, which the block + reuse policy would otherwise decide -- a connector must not be handed a + slot for a block the window has passed, and the sibling KV transfer path + drops the same range before sending. + """ + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return [[] for _ in range(len(self.impl.layer_grouping))] + history_length = kv_cache.history_length + by_group: List[List[int]] = [] + for layer_group_id in range(len(self.impl.layer_grouping)): + indices = list(kv_cache.get_aggregated_page_indices(layer_group_id, valid_only=False)) + stale_start, stale_end = self._stale_block_range(layer_group_id, history_length) + for ordinal in range(stale_start, min(stale_end, len(indices))): + indices[ordinal] = BAD_PAGE_INDEX + by_group.append(indices) + return by_group + @nvtx_range("prepare_resources_kv_cache_manager_v2") def prepare_resources(self, scheduled_batch: ScheduledRequests): if self.is_draft: @@ -3639,6 +3745,230 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self._prepare_draft_resources(scheduled_batch) return + # KV pages are allocated in `KVCacheV2Scheduler`, so by this point every + # scheduled request already has them and its page indices are readable. + # That is what makes this the place to drive the connector. + if self.kv_connector_manager is not None: + self._run_kv_connector_hooks(scheduled_batch) + + def _run_kv_connector_hooks(self, scheduled_batch: ScheduledRequests) -> None: + """Serve the connector prefix, then report the pages it may write into. + + Runs on the batch the forward pass will actually execute; see + ``_apply_connector_matched_prefix`` for why that placement is load-bearing. + """ + served_any = False + for request in scheduled_batch.context_requests: + # An allocation is asked about and reported exactly once, and it + # takes all three guards. The first two describe a fresh allocation + # -- a request whose asynchronous load has completed re-enters on + # its first chunk with the same pages and nothing left to load. + # `should_add_sequence` only goes false once that load *completes*, + # so a batch dropped between here and the forward pass brings every + # context request back with both predicates still true; the memo is + # what stops `update_state_after_alloc` reporting the same pages + # twice. + if not request.is_first_context_chunk: + continue + if not self.kv_connector_manager.should_add_sequence(request): + continue + if request.py_connector_allocation_reported: + continue + served_any |= self._apply_connector_matched_prefix(request) + request.py_connector_allocation_reported = True + # Both forms are offered; the manager picks the one this connector + # implements. With several layer groups the flat list is empty, + # because a page index is scoped to a group. + by_group = self.get_page_indices_by_layer_group(request) + flat = by_group[0] if len(by_group) == 1 else [] + self.kv_connector_manager.update_state_after_alloc(request, flat, by_group) + + if served_any: + # A served prefix can carry a request past its last chunk boundary, + # moving it between `context_requests_chunking` and + # `context_requests_last_chunk`. `build_scheduler_output` walks + # those two lists, so the split has to be rebuilt before it runs. + scheduled_batch.reset_context_requests() + + def report_batch_to_connector(self, scheduled_batch: ScheduledRequests) -> None: + """Report the batch to the KV connector. + + ``RequestData.num_scheduled_tokens`` describes the upcoming forward + pass, so this may only run once every resource manager has. That is why + ``ResourceManager.prepare_resources`` drives it rather than + ``prepare_resources`` here, which also gives the disagg-generation-init + mini-batch the same hook. + """ + if self.kv_connector_manager is not None and not self.is_draft: + self.kv_connector_manager.build_scheduler_output(scheduled_batch, self) + + # ---- KV connector prefix ---- + # + # The connector is asked from `prepare_resources`, downstream of every stage + # that can still drop a request (`_can_queue`, batch waiting, attention-DP + # balancing, the mamba-hybrid filter, the fp8 context-MLA cap). The + # connector ABC has no `cancel_load`, so that placement is required: an + # asked request must always reach `request_finished`. Moving the ask into + # the scheduling pass would buy the scheduler a budget that accounts for the + # served prefix and break the guarantee. + + def _connector_may_serve(self, req: LlmRequest) -> bool: + """Whether the connector is allowed to serve a prefix for ``req``.""" + if self.kv_connector_manager is None or self.is_draft: + return False + if req.is_dummy: + # A dummy request has no prompt to serve. + return False + if req.is_generation_only_request: + # The connector API rejects these outright; a disagg generation + # server gets its prompt KV from the context worker. + return False + return not req.is_disagg_generation_init_state + + def _apply_connector_matched_prefix(self, req: LlmRequest) -> bool: + """Ask the connector for a prefix and skip the runtime past what it takes. + + Returns whether the request's chunk was moved. + + The offer is clamped twice: to ``prompt_len - 1``, because the first + generation step consumes the last prompt token's activations, and to + what this iteration's pages can cover. Only the honoured amount is + committed, so the connector is never told to write into a page that + does not exist. + """ + if not self._connector_may_serve(req) or req.py_connector_allocation_reported: + return False + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + return False + + local_end = kv_cache.num_committed_tokens + num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( + req, local_end + ) + offered_end = min(local_end + num_tokens, req.prompt_len - 1) + + position = local_end + if offered_end > local_end: + fitted = self._reserve_connector_prefix(req, kv_cache, local_end, offered_end) + if fitted is not None: + position, chunk_size = fitted + self._mark_connector_prefix_populated(req, position, chunk_size) + + # The honoured amount, not `num_tokens`: `build_scheduler_output` + # reports `context_current_position - recorded` as the range computed + # locally, so over-reporting points the connector at an offset it has + # no page for. + if position == local_end and load_async: + logger.debug( + "req %s: connector offered %d tokens asynchronously and none " + "were honoured; the load is awaited and then discarded", + req.py_request_id, + num_tokens, + ) + self.kv_connector_manager.commit_new_matched_tokens(req, position - local_end, load_async) + return position > local_end + + def _reserve_connector_prefix( + self, req: LlmRequest, kv_cache, local_end: int, offered_end: int + ) -> Optional[Tuple[int, int]]: + """Fit ``[local_end, offered_end)`` into this iteration's allocation. + + Returns ``(position, chunk_size)``, or None when nothing can be honoured + while leaving the forward pass a token to compute. An offer inside the + chunk the scheduler already sized only moves the chunk's start, so it + allocates nothing and preserves the block alignment of a non-last + chunk's end; an offer past it shifts the whole window and grows the + allocation, which is why that arm needs a fallback. + """ + tokens_per_block = self.tokens_per_block + prompt_len = req.prompt_len + scheduler_chunk = req.context_chunk_size + chunk_end = min(local_end + scheduler_chunk, prompt_len) + + candidates = [] + if offered_end < chunk_end: + candidates.append((offered_end, chunk_end)) + else: + # `==` belongs here, not above: an offer that reaches the chunk end + # exactly leaves the forward pass nothing to compute unless the + # window moves. + shifted_end = min(offered_end + scheduler_chunk, prompt_len) + if shifted_end < prompt_len: + # A non-final chunk must end on a block boundary or the next + # chunk fragments the cache. + shifted_end = (shifted_end // tokens_per_block) * tokens_per_block + candidates.append((offered_end, shifted_end)) + # Fall back to the largest whole-block prefix the pages that already + # exist can hold while still leaving a chunk to compute. Whole + # blocks because that is the granularity a connector transfers in. + capped = ((chunk_end - 1) // tokens_per_block) * tokens_per_block + candidates.append((capped, chunk_end)) + + for position, end in candidates: + if position <= local_end or end <= position: + continue + if self._resize_for_connector_prefix(req, kv_cache, position, end): + return position, end - position + return None + + def _resize_for_connector_prefix( + self, req: LlmRequest, kv_cache, position: int, end: int + ) -> bool: + """Cover ``[0, end)`` and mark ``[0, position)`` as already valid. + + ``position`` becomes the cache's ``history_length``, which is the sole + input to the sliding-window stale-range computation -- passing the + served end is what stops the prefix costing a page per block in a + sliding-window layer group. When the capacity is unchanged the call + cannot fail, since no new slot is requested. + """ + target = end + self.num_extra_kv_tokens + if end >= req.prompt_len: + target += get_draft_token_length(req) + pre_cap = kv_cache.capacity + capacity = max(pre_cap, target) + if not kv_cache.resize(capacity, position): + logger.debug( + "req %s: could not cover a connector prefix through %d (capacity %d -> %d)", + req.py_request_id, + position, + pre_cap, + capacity, + ) + return False + if capacity > pre_cap and req.py_ctx_pre_resize_cap is None: + # `resize_context` records the capacity to roll back to; when it + # grew nothing there is no record, and this growth would then + # survive a `revert_allocate_context`. + req.py_ctx_pre_resize_cap = pre_cap + return True + + def _mark_connector_prefix_populated( + self, req: LlmRequest, position: int, chunk_size: int + ) -> None: + """Skip the request past ``position`` and give it ``chunk_size`` to compute. + + ``set_prepopulated_prompt_len`` has to run: ``is_first_context_chunk`` + is ``context_current_position == prepopulated_prompt_len`` and both must + move together. Its own chunk arithmetic assumes the whole prompt is + allocated and would push the chunk end past the reserved pages, so it is + handed a chunk spanning to the end of the prompt -- where its flooring + and block-alignment assertion are both skipped -- and the real chunk is + assigned afterwards. + """ + # `set_prepopulated_prompt_len` only advances when the length is + # non-zero, and a served prefix always is: it is strictly above the + # local match, which is never negative. + assert position > 0, f"req {req.py_request_id}: served prefix ends at {position}" + req.context_chunk_size = req.prompt_len - req.context_current_position + req.set_prepopulated_prompt_len(position, self.tokens_per_block) + req.context_chunk_size = chunk_size + # The request can leave the batch from here -- parked for an asynchronous + # load, or dropped with the whole batch -- and be settled again on the way + # back. Carry the served end across that gap. + req.py_connector_served_position = position + def _mirror_draft_kv_cache(self, req: LlmRequest): """The draft manager's entry for ``req``, created on first sight. @@ -4602,6 +4932,19 @@ def release_index_slot(self, request_id: int) -> None: self._early_freed_index_requests.add(request_id) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + # The promise to the connector is one ask per allocation, not one per + # request, so a replay of this request after a destructive pause may be + # asked and reported again. + request.py_connector_allocation_reported = False + # Same allocation, same lifetime. The pages the served position vouches + # for are gone, so a replay recomputes from its own reuse match. + request.py_connector_served_position = 0 + if self.kv_connector_manager is not None and not self.is_draft: + # The scheduler-output deltas are keyed to the allocation too. Left + # behind, the replay lands in `cached_requests` carrying a block-id + # delta against pages that no longer exist, and a connector that + # walks only `new_requests` issues no load at all. + self.kv_connector_manager.reset_request_state(request) if self.conversation_manager is not None: self.conversation_manager.finish_request(request) self._allocated_draft_lens.pop(request.py_request_id, None) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index aeba29dbe2b5..2dbcf732dcce 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1050,6 +1050,18 @@ def __init__( self.py_num_connector_matched_tokens = 0 + # Whether the KV connector has been asked about, and told about, this + # request's current KV allocation. The promise is at most once per + # allocation, not once per request, so this is cleared in + # `free_resources` -- the one place an allocation dies. + self.py_connector_allocation_reported = False + + # End of the prefix a KV connector populated for the current allocation, + # or 0. The cache holds those tokens but never commits them, so a context + # request that re-enters cannot recover the end from the cache's own + # reuse depth. + self.py_connector_served_position = 0 + self.py_result = PyResult( prompt_len=self.py_prompt_len, max_new_tokens=self.py_max_new_tokens, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index dc2f60fceb6d..aac79992db47 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -42,7 +42,8 @@ from tensorrt_llm.executor.request import TruncateKVCacheRequest from tensorrt_llm.inputs.multimodal import strip_mm_data_for_generation from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata -from tensorrt_llm.llmapi.llm_args import (ExecutorMemoryType, PeftCacheConfig, +from tensorrt_llm.llmapi.llm_args import (CapacitySchedulerPolicy, + ExecutorMemoryType, PeftCacheConfig, WaitingQueuePolicy) from tensorrt_llm.llmapi.utils import \ _reapply_current_thread_affinity_to_all_threads @@ -74,6 +75,7 @@ from ..speculative.speculation_gate import SpeculationGate from .adp_iter_stats import ADPIterStatsBuffer from .connectors.kv_cache_connector import KvCacheConnectorManager +from .connectors.kv_cache_layout import build_kv_cache_layout_v2 from .disagg_adapter import PyExecutorEffects, PyExecutorRequestRegistry from .dwdp import DwdpManager from .error_classification import ErrorBudget @@ -1067,8 +1069,43 @@ def _maybe_init_kv_connector_manager(self): "connector scheduler / worker hooks and are not " "distinguished from real requests.") + if self.kv_cache_manager is None: + raise ValueError( + "KV Cache Connector requires a KV Cache Manager.") + + is_kv_cache_manager_v2 = isinstance(self.kv_cache_manager, + KVCacheManagerV2) + + if is_kv_cache_manager_v2 and self.max_draft_len > 0: + raise NotImplementedError( + "KV Cache Connector is not supported with speculative " + "decoding. Rejected draft tokens " + "shrink a request's page list, and the freed slot goes to " + "whichever request allocates next. The connector is only " + "told about pages appended since the last report, so it " + "would keep addressing a slot another request now owns. " + "Disable speculative decoding to run a connector.") + + scheduler_config = getattr(self.llm_args, 'scheduler_config', None) + capacity_scheduler_policy = getattr(scheduler_config, + 'capacity_scheduler_policy', + None) + if (not is_kv_cache_manager_v2 + and capacity_scheduler_policy is not None + and capacity_scheduler_policy + != CapacitySchedulerPolicy.GUARANTEED_NO_EVICT): + raise NotImplementedError( + "KV Cache Connector in this configuration is only " + "supported with the GUARANTEED_NO_EVICT capacity " + "scheduler policy. A policy that destroys and replays a " + "live request leaves the connector's per-request block " + "delta measured against pages that were freed with it. " + "Set kv_cache_config.use_kv_cache_manager_v2=True, which " + "drops the delta when an allocation dies.") + kv_cache_config = getattr(self.llm_args, 'kv_cache_config', None) - if kv_cache_config is not None and kv_cache_config.host_cache_size: + if not is_kv_cache_manager_v2 and (kv_cache_config is not None and + kv_cache_config.host_cache_size): raise NotImplementedError( "KV Cache Connector is not supported with KV cache host " "offloading (KvCacheConfig.host_cache_size). The connector " @@ -1078,11 +1115,11 @@ def _maybe_init_kv_connector_manager(self): "streams are not synchronized with the internal " "onboard/offload streams.") - if self.kv_cache_manager is None: - raise ValueError( - "KV Cache Connector requires a KV Cache Manager.") - - if getattr(self.kv_cache_manager, 'is_vswa', False): + # VSWA allocates one pool per window size, which a single-tensor + # registration cannot describe. `register_kv_cache_layout` can, so + # the rejection only applies to the flat-tensor path. + if not is_kv_cache_manager_v2 and getattr(self.kv_cache_manager, + 'is_vswa', False): raise NotImplementedError( "KV Cache Connector is not supported with variable " "sliding-window attention (per-layer max_attention_window " @@ -1098,8 +1135,24 @@ def _maybe_init_kv_connector_manager(self): "per-layer load/save hooks have nothing meaningful to " "transfer for those layers.") - kv_tensor = self.kv_cache_manager.get_unique_primary_pool() - self.kv_connector_manager.worker.register_kv_caches(kv_tensor) + if is_kv_cache_manager_v2: + # Registered regions are device addresses, so every page has + # to stay pinned to GPU for as long as the connector holds + # them. + self._reject_non_gpu_cache_tiers(self.kv_cache_manager) + self._reject_connector_prefix_without_block_reuse( + self.kv_cache_manager) + layout = build_kv_cache_layout_v2(self.kv_cache_manager) + self.kv_connector_manager.reject_flat_only_scheduler( + len(layout.groups)) + window = (layout.groups[0].window_size + if len(layout.groups) == 1 else None) + self.kv_connector_manager.warn_flat_scheduler_under_swa(window) + self.kv_connector_manager.worker.register_kv_cache_layout( + layout) + else: + kv_tensor = self.kv_cache_manager.get_unique_primary_pool() + self.kv_connector_manager.worker.register_kv_caches(kv_tensor) # For each of our layers, we need to register the pre/post hooks. # These are used for methods like `wait_for_layer_load` and `save_kv_layer`. @@ -1112,6 +1165,66 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() + @staticmethod + def _reject_connector_prefix_without_block_reuse(kv_cache_manager) -> None: + """Reject a connector on KVCacheManagerV2 with block reuse disabled. + + V2 honours a connector-served prefix whatever ``enable_block_reuse`` + says -- the flag governs the local radix tree, and the connector is a + separate source. With reuse off it honours the offer and the restored + KV is wrong: prefill is skipped for the served range and generation + drifts within a few tokens. Measured on the reference connector with a + plain full-attention model, so it is neither VSWA- nor + connector-specific. + + Read off the manager rather than ``KvCacheConfig.enable_block_reuse``, + which ``create_py_executor`` coerces to False for several unrelated + reasons (SM version, KV cache quantization algorithm, hybrid linear + models). A user who never touched the flag can still land here, and a + config-level check would not see it. + + V1 is deliberately not guarded. It never honours the offer in this + configuration -- it asks the connector, then schedules the whole prompt + anyway -- so it miscomputes nothing. That it silently discards the + connector's work, and reports a negative ``computed_position`` while + doing so, is a separate defect tracked in the backlog. + """ + if getattr(kv_cache_manager, "enable_block_reuse", True): + return + raise NotImplementedError( + "KV Cache Connector is not supported with KV cache block reuse " + "disabled. The connector's prefix is honoured " + "regardless of that setting, and restores incorrect KV without it, " + "so the request would silently produce wrong output. Set " + "kv_cache_config.enable_block_reuse=True. Note that block reuse is " + "also disabled automatically for some quantization and model " + "combinations, so this can trigger without the flag being set " + "explicitly.") + + @staticmethod + def _reject_non_gpu_cache_tiers(kv_cache_manager) -> None: + """Reject cache tiers below GPU while a connector runs. + + A registered region is only a valid device address while its page is + pinned to GPU, and eviction to another tier reassigns that page's slot. + The resolved tier list is read from the manager rather than from + ``KvCacheConfig.host_cache_size``, whose default of ``None`` is falsy + but still yields a host tier. + """ + from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheTier + + cache_tiers = kv_cache_manager.impl.init_config.cache_tiers + extra = [tier for tier in cache_tiers if tier.tier != CacheTier.GPU_MEM] + if extra: + names = ", ".join(str(tier.tier) for tier in extra) + raise NotImplementedError( + "KV Cache Connector is not supported with " + f"cache tiers below GPU (found: {names}). Pages evicted to " + "another tier have their GPU slot reassigned, which would " + "invalidate the addresses registered with the connector. Set " + "KvCacheConfig.host_cache_size=0 and " + "KvCacheConfig.disk_cache_size=0 to run GPU-only.") + def _release_transfer(self, request: LlmRequest) -> None: """Release one transfer claim held by the KV connector or transceiver. @@ -4621,9 +4734,9 @@ def _sync_and_process_resource_governor_queue(self): def _can_pause_for_rebalance(self) -> bool: """Gate KV pool rebalance to the cases the hook supports. - Scope: no in-flight disagg transfer, no beam search, no drafter, not - during warmup or shutdown. Honors the ``enable_kv_pool_rebalance`` - opt-in flag (default off). + Scope: no in-flight disagg transfer, no KV connector, no beam search, + no drafter, not during warmup or shutdown. Honors the + ``enable_kv_pool_rebalance`` opt-in flag (default off). Pipeline parallelism *is* supported, but not in the same shape as the other two loops. ``_executor_loop`` and ``_executor_loop_overlap`` @@ -4649,6 +4762,14 @@ def _can_pause_for_rebalance(self) -> bool: return False if self.kv_cache_transceiver is not None: return False + if self.kv_connector_manager is not None: + # Rebalance suspends every active request and runs a defragmenting + # migration, which reassigns slot_id -- the same field a tier + # eviction reassigns, and the reason tiers below GPU are rejected + # at bring-up. The addresses a connector holds come from + # update_state_after_alloc and outlive the iteration, so every + # connector run is exposed, synchronous included. + return False if self.is_warmup: return False if self.is_shutdown: @@ -7171,8 +7292,8 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): # so the connector's per-request state advances exactly as before. kv_cache_manager = self.resource_manager.resource_managers.get( ResourceManagerType.KV_CACHE_MANAGER) - if hasattr(kv_cache_manager, "publish_connector_scheduler_output"): - kv_cache_manager.publish_connector_scheduler_output( + if hasattr(kv_cache_manager, "report_batch_to_connector"): + kv_cache_manager.report_batch_to_connector( disagg_gen_init_to_prepare) # Trigger KV cache exchange for new disagg_gen_init_requests @@ -7442,15 +7563,27 @@ def _save_kv_to_connector_async( return def kv_connector_request_finished(req: LlmRequest): + by_layer_group = None try: - cache_block_ids = self.kv_cache_manager.get_cache_indices(req) + # KVCacheManagerV2 has no primary-pool block list, so + # `get_cache_indices` raises there. The warning path below + # swallows that, leaving `request_finished` uncalled and the + # connector never told to save anything. + if isinstance(self.kv_cache_manager, KVCacheManagerV2): + by_layer_group = self.kv_cache_manager.get_page_indices_by_layer_group( + req) + cache_block_ids = by_layer_group[0] if len( + by_layer_group) == 1 else [] + else: + cache_block_ids = self.kv_cache_manager.get_cache_indices( + req) except Exception as e: logger.warning( f"Unable to get cache blocks for request {req.py_request_id}. Skipping asynchronous saving: {e}" ) else: if self.kv_connector_manager.request_finished( - req, cache_block_ids): + req, cache_block_ids, by_layer_group): self.async_transfer_manager.start_transfer(req) if not self.disable_overlap_scheduler: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 45211ed7d641..c7f6e0e861f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -779,15 +779,44 @@ def allocation_scope(current_stage: ExecutorMemoryType): logger.info( f"Initializing kv connector with config: {kv_connector_config}") - if scheduler_config.capacity_scheduler_policy != CapacitySchedulerPolicy.GUARANTEED_NO_EVICT: + # `use_kv_cache_manager_v2` is tri-state and under "auto" the manager is + # not chosen until model loading, so the three manager-dependent + # rejections below fire here only when the config names the manager + # outright, sparing an explicit config a model load it cannot use. + # `_maybe_init_kv_connector_manager` repeats all three against the + # manager that was actually built. + v2_selection = kv_cache_config.use_kv_cache_manager_v2 + + # A policy that destroys and replays a live request leaves the + # connector's per-request block delta measured against pages that were + # freed with it. Only KVCacheManagerV2 drops that delta on replay. + if (scheduler_config.capacity_scheduler_policy + != CapacitySchedulerPolicy.GUARANTEED_NO_EVICT + and v2_selection is False): raise NotImplementedError( - "KV connector is only supported with guaranteed no evict scheduler policy." + "KV connector in this configuration is only supported with the " + "GUARANTEED_NO_EVICT capacity scheduler policy. Set " + "kv_cache_config.use_kv_cache_manager_v2=True to use another policy." ) + # Rejected draft tokens shrink a request's page list, and the freed slot + # goes to whichever request allocates next. The connector is only told + # about pages appended since the last report, so it would keep + # addressing a slot another request now owns. + if (spec_config is not None and spec_config.max_draft_len > 0 + and v2_selection is True): + raise NotImplementedError( + "KV connector is not supported with speculative decoding. " + "Disable speculative decoding to run a connector.") + + # VSWA allocates one pool per window size, which only + # `register_kv_cache_layout` can describe. max_attention_window = kv_cache_config.max_attention_window - if uses_vswa_kv_cache_layout(max_attention_window): + if (uses_vswa_kv_cache_layout(max_attention_window) + and v2_selection is False): raise NotImplementedError( - "KV connector is not supported with VSWA (Variable Sliding Window Attention)." + "KV connector is not supported with VSWA (Variable Sliding Window Attention) " + "in this configuration. Set kv_cache_config.use_kv_cache_manager_v2=True." ) if mapping.enable_attention_dp: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 427cd7d19dbd..aff58aae784b 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1198,8 +1198,8 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # reuse, so we rebuild the context request lists here. scheduled_batch.reset_context_requests() - def publish_connector_scheduler_output( - self, scheduled_batch: ScheduledRequests) -> None: + def report_batch_to_connector(self, + scheduled_batch: ScheduledRequests) -> None: """Report the batch to the KV connector. Driven by ``ResourceManager.prepare_resources`` *after* the token-budget @@ -3043,11 +3043,11 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self.maybe_fit_token_budget(scheduled_batch) # Strictly after the trim: the connector is told how many tokens the # forward pass will compute, which is only settled once the trim has - # run. See KVCacheManager.publish_connector_scheduler_output. + # run. See `report_batch_to_connector`. kv_cache_manager = self.resource_managers.get( ResourceManagerType.KV_CACHE_MANAGER) - if hasattr(kv_cache_manager, "publish_connector_scheduler_output"): - kv_cache_manager.publish_connector_scheduler_output(scheduled_batch) + if hasattr(kv_cache_manager, "report_batch_to_connector"): + kv_cache_manager.report_batch_to_connector(scheduled_batch) @nvtx_range("update_resources") def update_resources( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 1996d219425a..a5b13e1a5177 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -527,14 +527,28 @@ def _schedule_loop(self, active_requests, inflight_request_ids): and not recompute_paused and not inflight_request_ids ): + # A connector rejects every tier below GPU at bring-up + # (`PyExecutor._reject_non_gpu_cache_tiers`), so offering those + # two settings there is advice the user cannot act on. + if getattr(self.kv_cache_manager, "kv_connector_manager", None) is not None: + remedy = ( + "A KV connector is attached, which requires a GPU-only cache, " + "so no secondary tier can be configured. Increase " + "kv_cache_config.max_tokens or kv_cache_config." + "free_gpu_memory_fraction, or lower max_num_tokens to hand " + "memory back to the KV pool." + ) + else: + remedy = ( + "Configure kv_cache_config.host_cache_size or " + "kv_cache_config.disk_cache_size, or increase " + "kv_cache_config.max_tokens." + ) raise RuntimeError( f"V2 scheduler deadlock: {num_gen_candidates} generation " f"request(s) active but none could be scheduled or " f"evicted or recompute-paused. KV cache pool is likely exhausted with no " - f"secondary cache tier for suspend/resume offload. " - f"Configure kv_cache_config.host_cache_size or " - f"kv_cache_config.disk_cache_size, or increase " - f"kv_cache_config.max_tokens." + f"secondary cache tier for suspend/resume offload. {remedy}" ) return ( @@ -712,6 +726,18 @@ def _try_schedule_context( Returns ``(action, tokens, chunking_flag)``. *tokens* and *chunking_flag* are meaningful only when *action* is ``SCHEDULED``. + + Deliberately connector-blind: the connector is asked later, in + the cache manager's ``prepare_resources``, so the budget here assumes it + serves nothing. That is safe because honouring an offer only ever + removes tokens from the forward pass, and it costs only that a served + prefix frees no budget for another request in the same iteration. + + ``should_add_sequence`` must not gate scheduling either: it stays false + from the moment an asynchronous load completes until + ``request_finished``, so a request gated on it would be skipped forever + and never run the prefill the load was for. A loading request is kept + out of the batch by its ``DISAGG_GENERATION_TRANS_IN_PROGRESS`` state. """ if self.chunking_enabled: return self._try_schedule_context_chunked(req, budget) @@ -891,6 +917,8 @@ def _reuse_claim_limit(self, req: LlmRequest) -> int | None: def _prepare_context_pair(self, req: LlmRequest) -> bool: """Prepare target/draft caches with one verified logical reuse depth.""" + from ..kv_cache.kv_cache_manager_v2 import _settle_context_cursor + draft_manager = self._joint_draft_manager if draft_manager is None: return self.kv_cache_manager.prepare_context(req) @@ -935,11 +963,7 @@ def _prepare_context_pair(self, req: LlmRequest) -> bool: # No enable_block_reuse guard: reaching here means a paired draft pool # exists, and KvCacheCreator only pairs when block reuse is on. - if common_reuse < req.context_current_position: - # setPrepopulatedPromptLen only moves forward, so a match shallower - # than a previous attempt's has to be rewound by hand. - req.context_current_position = common_reuse - req.set_prepopulated_prompt_len(common_reuse, self.kv_cache_manager.tokens_per_block) + _settle_context_cursor(req, common_reuse, self.kv_cache_manager.tokens_per_block) return True def _try_allocate_context(self, req: LlmRequest, num_tokens: int) -> bool: diff --git a/tests/integration/defs/llmapi/data/kv_connector_vswa_prompt.txt b/tests/integration/defs/llmapi/data/kv_connector_vswa_prompt.txt new file mode 100644 index 000000000000..281619f508cf --- /dev/null +++ b/tests/integration/defs/llmapi/data/kv_connector_vswa_prompt.txt @@ -0,0 +1,43 @@ +You are reading a technical brief. Answer the question at the end in one short paragraph. + +GPU inference systems keep a key-value cache so that attention does not recompute the +representations of tokens it has already seen. Each request occupies a set of cache pages, and a +page holds the keys and values for a fixed number of tokens across the layers that share a cache +life cycle. When a model applies the same attention pattern to every layer, one pool of pages is +enough, and a page index means the same thing everywhere in the model. + +Modern models increasingly break that assumption. Interleaved architectures alternate between +layers that attend to a short sliding window and layers that attend to the whole sequence. The +sliding layers need far fewer pages per request, because once the window has moved past a block of +tokens those keys and values are never read again and the page can be reclaimed. The full-attention +layers need a page for every block for as long as the request lives. Sizing one pool for both +wastes memory on the sliding layers or starves the full-attention ones, so the cache allocates a +separate pool group per distinct window size and gives each its own slot address space. + +That decision has a consequence that matters to anything reading or writing the cache from outside. +A page index is no longer meaningful on its own. Slot seven of the sliding pool and slot seven of +the full-attention pool are different memory holding the keys and values of different layers. Any +component that transfers cache contents must therefore carry the layer group alongside the slot, +and must key whatever external store it maintains on the pair rather than on the token range alone. +Two layer groups covering the same tokens hold genuinely different data, and a store keyed only by +tokens will overwrite one with the other and later restore it into the wrong layers. + +The sliding window introduces a second consequence. Because pages are reclaimed once the window has +passed them, a component that saves cache contents can persist at most one window of tokens for a +sliding layer group, not the whole prompt. Blocks the window has passed are reported with a sentinel +index rather than being dropped from the list, so that positions stay aligned with token ranges +across every layer group, and the component is expected to skip them. Skipping is not merely +tolerated but required: the sentinel is a valid subscript in most array libraries, so treating it as +an index silently reads or writes the last page of the pool, corrupting an unrelated request. + +Restoring a previously computed prefix interacts with all of this. When an external source supplies +the keys and values for a leading range of tokens, the runtime advances its notion of how much +history the request already has, and skips computing that range. For a full-attention layer group +every block of the restored prefix must be supplied. For a sliding layer group only the blocks that +remain inside the window need to be supplied, because the earlier ones will never be read. A correct +implementation therefore transfers different amounts of data to different layer groups for the very +same restored prefix, and a test that does not observe that difference has not observed the feature +working at all. + +Question: why is a page index meaningless without its layer group, and what does that imply for a +component that caches attention state outside the runtime? diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 15d2d1cf9021..aac0ade227b4 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -13,8 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib +import logging import math import os +import shutil +import sys import tempfile import time from unittest.mock import MagicMock, patch @@ -22,15 +26,87 @@ import pytest from tensorrt_llm import LLM, DisaggregatedParams, SamplingParams -from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, KvCacheConfig, - KvCacheConnectorConfig) +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + V2_RETENTION_IGNORED_LOG_KEY, KvCacheConnectorWorker) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import \ + KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, + CapacitySchedulerPolicy, + KvCacheConfig, KvCacheConnectorConfig, + NGramDecodingConfig, SchedulerConfig) from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig +from tensorrt_llm.logger import logger as trtllm_logger_singleton +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX + +from ..conftest import get_sm_version, llm_models_root + +# Name of the TensorRT-LLM logger. It sets `propagate = False` +# (tensorrt_llm/logger.py:186-187), so pytest's `caplog` only sees its records +# once `caplog.handler` is attached to it directly. +TRTLLM_LOGGER_NAME = "TRT-LLM" + +# Emitted by `_fallback_if_unsupported_kv_cache_manager_v2` +# (tensorrt_llm/_torch/pyexecutor/_util.py:629-631) when a connector run is +# downgraded from V2 to V1. +FALLBACK_WARNING_FRAGMENT = "Falling back to KVCacheManager" + +# V1 `KVCacheManager` methods reached on the KV connector path proper. Each is +# defined in tensorrt_llm/_torch/pyexecutor/resource_manager.py and wraps a +# nanobind method on the C++ manager. `KVCacheManagerV2` implements none of +# them and is not meant to: each assumes a single flat block-id space over one +# primary pool, which cannot describe memory whose page indices are scoped to a +# layer group. V2's connector contract is `register_kv_cache_layout` plus +# `get_page_indices_by_layer_group` instead. +CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS = ( + # Connector bring-up: hands the worker the single primary pool tensor. + # `PyExecutor._maybe_init_kv_connector_manager`. + "get_unique_primary_pool", + # `KvCacheConnectorSchedulerOutputRequest.update_and_build_data` and + # `PyExecutor.kv_connector_request_finished`. + "get_cache_indices", + # `update_and_build_data`, for `RequestData.block_hashes`. + "commit_and_get_block_hashes", + # `update_and_build_data`, for `RequestData.priorities`. + "get_priority_by_block_id", +) -from ..conftest import llm_models_root +# Reached only when the connector coexists with disaggregated serving +# (`test_connector_disagg_prefill`), via AsyncTransferManager and the V1 cache +# reuse adapter - not from `KvCacheConnectorManager` itself. None of them is a +# gap for V2, because V2 does not traverse those paths: +# `enable_partial_reuse_for_disagg` excludes V2, so AsyncTransferManager never +# reaches the pin/unpin pair, and the Python transceiver (the only one V2 can +# be driven by) resolves block ids through `_CacheReuseAdapterV2`. +DISAGG_PATH_KV_CACHE_MANAGER_METHODS = ( + "store_blocks_for_reuse", + "unpin_blocks_by_id", + "get_memory_pool_block_indices", + "pin_blocks", # no Python caller today +) @pytest.fixture(scope="function") -def model_with_connector(): +def use_kv_cache_manager_v2(request): + """Run each connector test under both KV cache managers. + + Parametrized by an explicit `@pytest.mark.parametrize(..., indirect=True)` + on each test, applied as the innermost decorator so the manager lands first + in the generated test id. It is spelled out per test rather than set as a + fixture `params=` because the test-list validator + (scripts/check_test_list.py) resolves ids from parametrize decorators via + AST and cannot see fixture-level parametrization. + + Selecting V2 must actually reach V2: `_fallback_if_unsupported_kv_cache_manager_v2` + silently substitutes the V1 manager for combinations it cannot serve, and a + connector test that ran on V1 while claiming to test V2 would pass while + exercising nothing. `test_connector_runs_on_kv_cache_manager_v2` guards that. + """ + return request.param + + +@pytest.fixture(scope="function") +def model_with_connector(use_kv_cache_manager_v2): with patch("tensorrt_llm._torch.pyexecutor.py_executor_creator.importlib" ) as importlib_mock: mock_scheduler = MagicMock() @@ -39,6 +115,28 @@ def model_with_connector(): importlib_mock.import_module.return_value.KvConnectorScheduler.return_value = mock_scheduler importlib_mock.import_module.return_value.KvConnectorWorker.return_value = mock_worker + # A cache that reports per layer group always calls the per-group form, + # so mirror what a real connector's default does with it: fold a single + # group back onto the flat callback, and leave several groups to a test + # that asserts the per-group form directly. Without this every mock + # would answer the per-group call with a truthy `Mock`, which reads as + # "saving asynchronously" and parks every request. + def _alloc_by_group(request, block_ids_by_layer_group): + if len(block_ids_by_layer_group) == 1: + mock_scheduler.update_state_after_alloc( + request, block_ids_by_layer_group[0]) + + def _finished_by_group(request, cache_block_ids_by_layer_group): + if len(cache_block_ids_by_layer_group) == 1: + return mock_scheduler.request_finished( + request, cache_block_ids_by_layer_group[0]) + return False + + mock_scheduler.update_state_after_alloc_by_layer_group.side_effect = ( + _alloc_by_group) + mock_scheduler.request_finished_by_layer_group.side_effect = ( + _finished_by_group) + kv_connector_config = KvCacheConnectorConfig( connector_module="", connector_scheduler_class="KvConnectorScheduler", @@ -58,7 +156,16 @@ def model_fn(*args, **kwargs): "max_seq_len": 1024, } - return LLM(*args, **{**default_kwargs, **kwargs}) + merged_kwargs = {**default_kwargs, **kwargs} + + # Tests that supply their own `KvCacheConfig` must still honour the + # manager under test, otherwise the V2 parametrization silently + # degrades into a second V1 run. + kv_cache_config = merged_kwargs.get("kv_cache_config") + if kv_cache_config is not None: + kv_cache_config.use_kv_cache_manager_v2 = use_kv_cache_manager_v2 + + return LLM(*args, **merged_kwargs) yield model_fn, mock_scheduler, mock_worker @@ -70,27 +177,226 @@ def enforce_single_worker(monkeypatch): yield -def generate_and_sleep(model, *args, **kwargs): - # Some KV connector API calls are made after a full response is returned. We want to be able to track these calls. - # However, we don't have any indication of when all the calls are complete. - # To compensate for this, we sleep between the generate call and the return of the outputs. - # TODO(jthomson04): Surely there's a better way to do this? +# Some KV connector API calls are made after a full response is returned +# (`request_finished`, the trailing `get_finished` polls and asynchronous +# saves), and there is no public signal for when they are complete. Instead of +# sleeping a fixed amount, wait until the connector mocks stop recording new +# calls. That returns as soon as the connector goes quiet, and - unlike a fixed +# sleep - stretches automatically when a slower path lengthens the tail. +CONNECTOR_QUIESCE_TIMEOUT_S = 60.0 +CONNECTOR_QUIET_PERIOD_S = 0.5 +CONNECTOR_POLL_INTERVAL_S = 0.01 + +# Fraction of generated tokens a connector-warmed run must reproduce exactly. +# See test_connector_e2e_persistent_cache for why this is not 1.0. +E2E_MIN_TOKEN_AGREEMENT = 0.75 + + +def assert_kv_caches_registered(worker, use_kv_cache_manager_v2): + """The two managers hand the worker its pools through different entry points. + + V1 passes a single pool tensor to `register_kv_caches`; V2 has no such + tensor and passes a `KvCacheLayout` to `register_kv_cache_layout` instead. + Asserting the V1 method unconditionally would silently pass on V2 only if + the connector were never registered at all. + """ + if use_kv_cache_manager_v2: + assert worker.register_kv_cache_layout.call_count == 1 + assert worker.register_kv_caches.call_count == 0 + else: + assert worker.register_kv_caches.call_count == 1 + assert worker.register_kv_cache_layout.call_count == 0 + + +def wait_for_connector_quiescence(scheduler, + worker, + timeout=CONNECTOR_QUIESCE_TIMEOUT_S, + quiet_period=CONNECTOR_QUIET_PERIOD_S): + """Block until no new connector callback lands for `quiet_period` seconds. + + `MagicMock.mock_calls` records every call made on the mock and its children, + so its length is a monotonic progress counter for connector activity. + """ + deadline = time.monotonic() + timeout + + def total_calls(): + return len(scheduler.mock_calls) + len(worker.mock_calls) + + last_seen = total_calls() + quiet_since = time.monotonic() + + while time.monotonic() < deadline: + time.sleep(CONNECTOR_POLL_INTERVAL_S) + + current = total_calls() + if current != last_seen: + last_seen = current + quiet_since = time.monotonic() + elif time.monotonic() - quiet_since >= quiet_period: + return + + raise AssertionError( + f"KV connector callbacks did not go quiet within {timeout}s " + f"({last_seen} calls recorded). The connector is still active or a " + "callback is blocked.") + + +def generate_and_wait(model, scheduler, worker, *args, **kwargs): + """`model.generate`, then block until the connector callbacks settle.""" outputs = model.generate(*args, **kwargs) - time.sleep(1) + wait_for_connector_quiescence(scheduler, worker) return outputs +def test_v2_connector_contract_does_not_reuse_the_v1_methods(): + """The V2 connector path implements none of the V1 accessors, by design. + + Something depends on that, and it does not ask: `update_and_build_data` + reports `block_hashes` and `priorities` empty on V2 by branching on + `isinstance(manager, KVCacheManagerV2)`, not on `hasattr`. Those + short-circuits are only correct while V2 genuinely has no such accessor - + the day one is added (retention priorities are a known gap; see + `test_connector_priorities`) the branch keeps reporting nothing while the + data exists, and this is what says so. + + A static check rather than an end-to-end run: under V2 the connector would + die at the *first* method it reached, so no run can report more than one at + a time. + + Needs no GPU. + """ + stale = [ + name for name in CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS + + DISAGG_PATH_KV_CACHE_MANAGER_METHODS + if not hasattr(KVCacheManager, name) + ] + assert stale == [], ( + f"{stale} are not defined on the V1 KVCacheManager either, so this " + "test is measuring a stale method list rather than a real difference.") + + implemented = [ + name for name in CONNECTOR_V1_ONLY_KV_CACHE_MANAGER_METHODS + if hasattr(KVCacheManagerV2, name) + ] + assert implemented == [], ( + f"KVCacheManagerV2 now implements {implemented}. A V1-shaped accessor " + "on V2 is not automatically the right answer - a flat block-id list " + "cannot describe more than one layer group - but if it is, revisit the " + "`is_v2` short-circuits in " + "`KvCacheConnectorSchedulerOutputRequest.update_and_build_data`, which " + "report nothing on the strength of these methods being absent.") + + # The other half of the contract: what V2 offers instead. Every connector + # path on V2 goes through this one accessor - `update_and_build_data`, + # `kv_connector_request_finished` and `_run_kv_connector_hooks` each call it + # and derive the flat list from `[0]` when there is a single layer group. + assert hasattr(KVCacheManagerV2, "get_page_indices_by_layer_group"), ( + "KVCacheManagerV2.get_page_indices_by_layer_group is the V2 " + "replacement for the V1 block-id accessors and every connector path on " + "V2 goes through it.") + assert hasattr(KvCacheConnectorWorker, "register_kv_cache_layout"), ( + "The worker ABC must keep a default `register_kv_cache_layout`, or " + "every existing connector becomes abstract and fails to instantiate.") + + +@pytest.mark.threadleak(enabled=False) +def test_connector_runs_on_kv_cache_manager_v2(enforce_single_worker, + monkeypatch, caplog): + """Anti-vacuity guard for the `kv_cache_manager_v2` parametrization. + + Without this, every V2-parametrized test below could pass green while + `_fallback_if_unsupported_kv_cache_manager_v2` silently swapped in the V1 + manager. It asserts positively that V2 is constructed, and that the + downgrade warning is absent. + + Any construction failure is recorded rather than asserted on: what must + hold is that the run reached V2 rather than being papered over by a + fallback, which stays true regardless of how far bring-up gets. + """ + constructed = [] + + def record_construction(cls): + original_init = cls.__init__ + + def recording_init(self, *args, **kwargs): + constructed.append(cls.__name__) + return original_init(self, *args, **kwargs) + + monkeypatch.setattr(cls, "__init__", recording_init) + + record_construction(KVCacheManagerV2) + record_construction(KVCacheManager) + + # The TensorRT-LLM logger sets `propagate = False`, so caplog only sees its + # records once its handler is attached to that logger directly. + trtllm_logger = logging.getLogger(TRTLLM_LOGGER_NAME) + trtllm_logger.addHandler(caplog.handler) + + construction_error = None + llm = None + try: + with patch( + "tensorrt_llm._torch.pyexecutor.py_executor_creator.importlib" + ) as importlib_mock: + connector_module = importlib_mock.import_module.return_value + connector_module.KvConnectorScheduler.return_value = MagicMock() + connector_module.KvConnectorWorker.return_value = MagicMock() + + try: + llm = LLM( + model=f"{llm_models_root()}/Qwen2-0.5B", + backend="pytorch", + kv_connector_config=KvCacheConnectorConfig( + connector_module="", + connector_scheduler_class="KvConnectorScheduler", + connector_worker_class="KvConnectorWorker", + ), + cuda_graph_config=None, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1, + use_kv_cache_manager_v2=True), + ) + # The V2 connector path is knowingly incomplete, so any construction + # failure is an acceptable outcome. What must hold is that the run + # reached V2 instead of being papered over by a V1 fallback. + except Exception as exc: # noqa: BLE001 + construction_error = exc + finally: + trtllm_logger.removeHandler(caplog.handler) + if llm is not None: + llm.shutdown() + + # Report the current frontier so the failure mode is visible in CI output + # instead of being silently swallowed by the except above. + print(f"\n[connector+V2 frontier] construction_error=" + f"{type(construction_error).__name__ if construction_error else None}" + f": {construction_error}") + + assert "KVCacheManagerV2" in constructed, ( + "KVCacheManagerV2 was never constructed, so the connector run silently " + f"fell back to V1 (managers constructed: {constructed}; construction " + f"error: {construction_error!r}). Every kv_cache_manager_v2-" + "parametrized connector test in this file is vacuous until this passes." + ) + + assert FALLBACK_WARNING_FRAGMENT not in caplog.text, ( + f"{FALLBACK_WARNING_FRAGMENT!r} was logged, so the connector was " + "downgraded to the V1 KV cache manager.") + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_simple(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -98,7 +404,8 @@ def test_connector_simple(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, ["Hello, world"], sampling_params) + generate_and_wait(model, scheduler, worker, ["Hello, world"], + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 @@ -156,22 +463,26 @@ def test_connector_simple(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_async_onboard(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 16, True worker.get_finished.side_effect = lambda finished_gen, load_async: ( finished_gen, load_async) - generate_and_sleep(model, [ + generate_and_wait(model, scheduler, worker, [ "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua." ], SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True)) @@ -185,15 +496,18 @@ def test_connector_async_onboard(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_async_save(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, use_kv_cache_manager_v2): NUM_TOKENS = 8 model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -204,7 +518,8 @@ def test_connector_async_save(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, ["Hello, world"], sampling_params) + generate_and_wait(model, scheduler, worker, ["Hello, world"], + sampling_params) assert scheduler.request_finished.call_count == 1 @@ -226,8 +541,12 @@ def test_connector_async_save(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_scheduler_output(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): NUM_INPUT_TOKENS = 48 NUM_TOKENS = 32 BLOCK_SIZE = 32 @@ -236,7 +555,7 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, model = model_fn(disable_overlap_scheduler=not use_overlap_scheduler, ) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -244,7 +563,8 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) - generate_and_sleep(model, [0] * NUM_INPUT_TOKENS, sampling_params) + generate_and_wait(model, scheduler, worker, [0] * NUM_INPUT_TOKENS, + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 assert len( @@ -294,7 +614,8 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, assert len(scheduler.request_finished.call_args.args[1]) == math.ceil( (NUM_INPUT_TOKENS + NUM_TOKENS) / BLOCK_SIZE) - generate_and_sleep(model, [1] * NUM_INPUT_TOKENS, sampling_params) + generate_and_wait(model, scheduler, worker, [1] * NUM_INPUT_TOKENS, + sampling_params) # The initial computed position should be 0, since we haven't yet onboarded any blocks. assert scheduler.build_connector_meta.call_args_list[0].args[ @@ -303,9 +624,13 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_scheduler_output_chunked_context(enforce_single_worker, model_with_connector, - use_overlap_scheduler): + use_overlap_scheduler, + use_kv_cache_manager_v2): model_fn, scheduler, worker = model_with_connector CHUNK_SIZE = 128 @@ -315,7 +640,7 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, enable_chunked_prefill=True, max_num_tokens=CHUNK_SIZE) - assert worker.register_kv_caches.call_count == 1 + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False @@ -323,13 +648,22 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, sampling_params = SamplingParams(max_tokens=BLOCK_SIZE, ignore_eos=True) - generate_and_sleep(model, [0] * (CHUNK_SIZE * 2), sampling_params) + generate_and_wait(model, scheduler, worker, [0] * (CHUNK_SIZE * 2), + sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 - assert len( - scheduler.update_state_after_alloc.call_args.args[1]) == math.ceil( - CHUNK_SIZE * 2 / BLOCK_SIZE) + # V1 allocates for the whole prompt when the sequence is added, so every + # block exists on the first chunk. V2 allocates per chunk, which is the + # lower-peak-memory behaviour and the one to keep; the remaining blocks + # arrive as append-deltas in `new_block_ids` on the next chunk, so no + # information is lost. The expectation is split rather than V2 changed. + total_blocks = math.ceil(CHUNK_SIZE * 2 / BLOCK_SIZE) + first_chunk_blocks = (math.ceil(CHUNK_SIZE / BLOCK_SIZE) + if use_kv_cache_manager_v2 else total_blocks) + + assert len(scheduler.update_state_after_alloc.call_args.args[1] + ) == first_chunk_blocks for i, call in enumerate(scheduler.build_connector_meta.call_args_list): sched_output = call.args[0] @@ -344,18 +678,17 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, req = sched_output.cached_requests[0] if i == 0: - # The first prefill chunk. - # All of the prefill tokens and all the blocks should be provided upfront. + # The first prefill chunk. All of the prefill tokens are provided + # upfront; the blocks are whatever has been allocated so far. assert req.computed_position == 0 assert len(req.new_tokens) == CHUNK_SIZE * 2 - assert len(req.new_block_ids) == math.ceil(CHUNK_SIZE * 2 / - BLOCK_SIZE) + assert len(req.new_block_ids) == first_chunk_blocks assert req.num_scheduled_tokens == CHUNK_SIZE elif i == 1: # The second prefill chunk. assert req.computed_position == CHUNK_SIZE assert len(req.new_tokens) == 0 - assert len(req.new_block_ids) == 0 + assert len(req.new_block_ids) == total_blocks - first_chunk_blocks assert req.num_scheduled_tokens == CHUNK_SIZE elif i == 2 and use_overlap_scheduler: assert len(req.new_tokens) == 0 @@ -367,19 +700,454 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, (CHUNK_SIZE * 2 + BLOCK_SIZE) / BLOCK_SIZE) +# Sliding-window coverage. `KvCacheConfig.max_attention_window` is repeated +# cyclically across layers (llm_args.py:3761-3765), so a one-element list gives +# every layer the same window -- one layer group with a live window -- while +# a two-element list alternates and produces two. A window equal to +# `max_seq_len` is normalised to "no window" (kv_cache_manager_v2.py:856-858), +# which is how the full-attention half of the VSWA pair is spelled. +SWA_WINDOW = 64 +SWA_MAX_SEQ_LEN = 512 +SWA_NUM_INPUT_TOKENS = 256 + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], + ids=["kv_cache_manager_v2"], + indirect=True) +def test_connector_uniform_sliding_window(enforce_single_worker, + model_with_connector, + use_kv_cache_manager_v2): + """Connector against a KV cache in which every layer slides. + + There was no sliding-window connector coverage before this test. Uniform + SWA is a single layer group, which is the configuration where the + connector's flat `new_block_ids` list is still well defined -- so this pins + the block-reporting contract, and the VSWA test below pins what happens once + that assumption breaks. + + V2 only. The V1 guard rejects *variable* windows + (py_executor_creator.py:845-850), so uniform SWA reaches V1's connector path + and then dies inside it: `commit_and_get_block_hashes` (kv_cache_connector.py:386) + raises "commitAndGetBlockHashesForRequest does not support sliding-window + attention with detached front blocks" (kvCacheManager.cpp:4645) as soon as + the window drops a front block. That is a pre-existing V1 limitation, not + something this work introduces, so it is recorded rather than asserted here. + """ + model_fn, scheduler, worker = model_with_connector + + model = model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW])) + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + scheduler.get_num_new_matched_tokens.return_value = 0, False + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + sched_output = scheduler.build_connector_meta.call_args_list[0].args[0] + assert len(sched_output.new_requests) == 1 + req = sched_output.new_requests[0] + + assert req.computed_position == 0 + assert req.num_scheduled_tokens == SWA_NUM_INPUT_TOKENS + # A single layer group keeps the flat list meaningful. + assert req.new_block_ids + + if use_kv_cache_manager_v2: + # Anti-vacuity: prove the window really did collapse to one layer + # group, otherwise the assertion above would hold for the wrong reason. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 1 + assert layout.groups[0].window_size == SWA_WINDOW + assert len(req.new_block_ids_by_layer_group) == 1 + assert req.new_block_ids_by_layer_group[0] == req.new_block_ids + else: + assert req.new_block_ids_by_layer_group == [] + + +# Half the prompt, and well past the window, so the pages the offer does *not* +# need are a large enough fraction to assert on. +SWA_OFFER_TOKENS = 128 + + +# The mock scheduler answers every query with the same offer; this records what +# it was asked and when, so a test can assert the connector was consulted +# exactly once per request rather than assuming it. +def record_connector_queries(scheduler, num_matched, load_async=False): + """Answer every query with `num_matched`, recording when each one arrived. + + The third element of each record is `build_connector_meta.call_count` at + query time -- the number of iterations whose connector hooks had already + run. The connector is asked from `prepare_resources`, in the same iteration + the request runs, so a request that runs in iteration `n` records `n`. A + speculative ask would record a smaller number. + """ + queries = [] + + def side_effect(request, num_computed_tokens): + queries.append((request.request_id, num_computed_tokens, + scheduler.build_connector_meta.call_count)) + return num_matched, load_async + + scheduler.get_num_new_matched_tokens.side_effect = side_effect + return queries + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_prefix_is_asked_once_and_shrinks_the_forward_pass( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """The offer is asked for once and removes work from the forward pass. + + Both managers ask on the batch that runs, so both record the same query + count and the same iteration index, and both report a prefill shortened by + the offer. This is the parity assertion for A0: a connector cannot tell the + two managers apart from what it is asked and what it is told. + """ + model_fn, scheduler, worker = model_with_connector + + OFFER_TOKENS = 64 + NUM_INPUT_TOKENS = 256 + + model = model_fn(disable_overlap_scheduler=True) + + queries = record_connector_queries(scheduler, OFFER_TOKENS) + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + # Once per request, anchored at the local match, in the first iteration -- + # which is also the iteration the request ran in. + assert len(queries) == 1 + _, num_computed_tokens, iterations_before = queries[0] + assert num_computed_tokens == 0 + assert iterations_before == 0 + + req = scheduler.build_connector_meta.call_args_list[0].args[0].new_requests[ + 0] + # The runtime rolls the reported position back to the local match, so the + # connector sees where its load begins; the offer shows up as work removed. + assert req.computed_position == 0 + assert req.num_scheduled_tokens == NUM_INPUT_TOKENS - OFFER_TOKENS + + +# Chunked-prefill prefix sizes, chosen either side of the chunk boundary so the +# two arms of the fit are both exercised: 64 lands inside the first chunk and +# only moves its start, 192 lands past it and has to shift the chunk window and +# grow the allocation. +CHUNKED_PREFIX_CHUNK_SIZE = 128 +CHUNKED_PREFIX_NUM_INPUT_TOKENS = 256 + + +def _context_reports(scheduler): + """`(computed_position, num_scheduled_tokens)` per prefill iteration. + + A generation iteration schedules exactly one token, which is what separates + the two here. Each `build_connector_meta` call is one iteration and carries + one request, in `new_requests` on its first report and `cached_requests` + after. + """ + reports = [] + for call in scheduler.build_connector_meta.call_args_list: + sched_output = call.args[0] + requests = sched_output.new_requests + sched_output.cached_requests + assert len(requests) == 1 + req = requests[0] + if req.num_scheduled_tokens > 1: + reports.append((req.computed_position, req.num_scheduled_tokens)) + return reports + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("offer_tokens", [64, 192], + ids=["offer_inside_chunk", "offer_past_chunk"]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_prefix_under_chunked_prefill(enforce_single_worker, + model_with_connector, + use_kv_cache_manager_v2, + offer_tokens): + """A served prefix removes exactly its own tokens without breaking chunking. + + Chunked prefill is where the two managers allocate differently -- one covers + the whole prompt when the sequence is added, the other only the chunk it is + about to compute -- so an offer reaching past the chunk is the case where a + prefix can outrun the pages that exist. What a connector observes has to be + the same either way: asked once at the local match, a prefill shortened by + the whole offer, and no chunk larger than `max_num_tokens`. + """ + model_fn, scheduler, worker = model_with_connector + + model = model_fn(disable_overlap_scheduler=True, + enable_chunked_prefill=True, + max_num_tokens=CHUNKED_PREFIX_CHUNK_SIZE) + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + queries = record_connector_queries(scheduler, offer_tokens) + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, + [0] * CHUNKED_PREFIX_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + # Once per request, anchored at the local match, in the iteration the + # request ran in -- chunking must not turn one allocation into one ask per + # chunk. + assert len(queries) == 1 + _, num_computed_tokens, iterations_before = queries[0] + assert num_computed_tokens == 0 + assert iterations_before == 0 + + reports = _context_reports(scheduler) + assert reports, "no prefill iteration was reported to the connector" + + # The reported position starts at the local match, so the connector knows + # where its load begins rather than where the runtime resumed. + assert reports[0][0] == 0 + + # Every prefill chunk still fits the token budget. Honouring an offer past + # the chunk moves the window, it does not widen it. + assert all(scheduled <= CHUNKED_PREFIX_CHUNK_SIZE + for _, scheduled in reports), reports + + # The offer removed exactly its own tokens from the forward pass, whether it + # was honoured in one chunk or spread across several. + total_scheduled = sum(scheduled for _, scheduled in reports) + assert total_scheduled == CHUNKED_PREFIX_NUM_INPUT_TOKENS - offer_tokens + + # Positions are contiguous over what was actually computed. Only the first + # report is rolled back to the local match, so from the second onwards each + # one resumes where the previous chunk stopped, offer included. How the + # first chunk is sized is left to the manager: covering the whole prompt up + # front leaves the chunk at its full width, allocating per chunk keeps the + # end where the scheduler put it and only moves the start. + expected = offer_tokens + reports[0][1] + for position, scheduled in reports[1:]: + assert position == expected, reports + expected += scheduled + + # The save at the end covers the whole prompt, offer included -- the served + # range is not silently dropped from what the connector is told to persist. + assert len(scheduler.request_finished.call_args.args[1]) >= math.ceil( + CHUNKED_PREFIX_NUM_INPUT_TOKENS / 32) + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], + ids=["kv_cache_manager_v2"], + indirect=True) +def test_connector_sliding_window_prefix_is_backed_by_real_pages( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """A served prefix is reported as pages for the window and `-1` before it. + + The offer reaches back past the window, so the served range splits: blocks + inside the live window carry real distinct pages, and blocks the window has + passed carry `BAD_PAGE_INDEX`. Skipping the compute for the whole offer + while moving bytes only for the window is the point of the design, and the + split is what a connector has to see to do it. + + The block reuse policy is deliberately left at its default. It decides + whether an out-of-window page survives the `history_length` bump as a + holder, which is exactly what the reported list must not depend on. + """ + model_fn, scheduler, worker = model_with_connector + + model = model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW])) + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + record_connector_queries(scheduler, SWA_OFFER_TOKENS) + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + # Anti-vacuity, as in the test above: one layer group with a live window, + # otherwise the page count below is being read off full attention. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 1 + assert layout.groups[0].window_size == SWA_WINDOW + + req = scheduler.build_connector_meta.call_args_list[0].args[0].new_requests[ + 0] + + # The offer was materialized: the position advanced over it, and the + # runtime rolled the reported position back to the local match. + assert req.computed_position == 0 + assert req.num_scheduled_tokens == SWA_NUM_INPUT_TOKENS - SWA_OFFER_TOKENS + + all_blocks = math.ceil(SWA_NUM_INPUT_TOKENS / 32) + page_indices = scheduler.update_state_after_alloc.call_args.args[1] + assert page_indices == req.new_block_ids + # One entry per block ordinal either way -- an out-of-window block reads back + # as BAD_PAGE_INDEX in place rather than shortening the list, so that block + # ordinals stay aligned to token ranges. The count is therefore not the + # signal; *which* entries are bad is. + assert len(page_indices) == all_blocks + + # The stale end the manager masks to, recomputed here from the sizes rather + # than read back, so a change to either one fails instead of adapting. + stale_end = max(0, (SWA_OFFER_TOKENS + 1 - SWA_WINDOW) // 32) + assert 0 < stale_end < SWA_OFFER_TOKENS // 32, ( + "test sizes no longer split the served range across the window edge") + + assert all(index == BAD_PAGE_INDEX for index in page_indices[:stale_end]), ( + f"a block the window has passed was reported as a page: {page_indices}") + live = page_indices[stale_end:] + assert all(index != BAD_PAGE_INDEX for index in live), ( + f"an in-window block was reported with no page: {page_indices}") + assert len(set(live)) == len(live), ( + f"page slots reported to the connector are not distinct: {page_indices}" + ) + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_vswa_reports_page_indices_per_layer_group( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """VSWA is where the connector's flat block list stops working. + + V1 cannot describe VSWA to a connector at all: it registers a single primary + pool, but VSWA allocates one pool per window size. V2's layout describes one + region set per layer group, so the combination runs there -- but a page index + is scoped to a layer group, and several groups cannot be flattened into one + list. Every block-id callback therefore switches to its per-layer-group + form, on the way in and on the way out. + + Both halves are asserted here because the rejection is deliberately + conditional on the manager. Pinning only the V2 half would let the V1 guard + silently disappear. + """ + model_fn, scheduler, worker = model_with_connector + + def build(): + return model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW, SWA_MAX_SEQ_LEN])) + + if not use_kv_cache_manager_v2: + with pytest.raises(NotImplementedError, match="VSWA"): + build() + return + + model = build() + + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + + # Alternating windows must actually produce two groups, one sliding and one + # full-attention, or the rest of this test is vacuous. + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 2 + assert {group.window_size for group in layout.groups} == {SWA_WINDOW, None} + + scheduler.get_num_new_matched_tokens.return_value = 0, False + worker.get_finished.return_value = [], [] + + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, + SamplingParams(max_tokens=4, ignore_eos=True)) + + # The flat list is empty with several groups, and the per-group list is what + # carries the pages. A connector reading only `new_block_ids` sees nothing, + # which is why implementing the per-group callbacks is required at bring-up + # for this model -- `MagicMock` satisfies that check. + assert scheduler.update_state_after_alloc_by_layer_group.call_count == 1 + by_group = scheduler.update_state_after_alloc_by_layer_group.call_args.args[ + 1] + assert len(by_group) == 2 + + sched_output = scheduler.build_connector_meta.call_args_list[0].args[0] + assert len(sched_output.new_requests) == 1 + req = sched_output.new_requests[0] + + assert req.new_block_ids == [] + assert len(req.new_block_ids_by_layer_group) == 2 + # Ordinals stay positionally aligned across groups: a block with no page in + # the sliding group, or one the window has passed, reads back as + # BAD_PAGE_INDEX in place rather than shortening the list. + lengths = {len(indices) for indices in req.new_block_ids_by_layer_group} + assert len(lengths) == 1 + + # The save direction switches with it. Without the per-group form the + # connector is handed an empty flat list at the end of the request and can + # persist nothing at all, which no other assertion here would catch. + assert scheduler.request_finished.call_count == 0 + assert scheduler.request_finished_by_layer_group.call_count == 1 + saved = scheduler.request_finished_by_layer_group.call_args.args[1] + assert len(saved) == 2 + + # Whatever the sliding group offers to save sits inside its window. The + # full-attention group keeps everything, which is what makes the comparison + # non-vacuous. + windows = [group.window_size for group in layout.groups] + sliding = windows.index(SWA_WINDOW) + full = windows.index(None) + live = [ + len([index for index in saved[group] if index != BAD_PAGE_INDEX]) + for group in (sliding, full) + ] + assert live[0] <= math.ceil(SWA_WINDOW / 32) + 1, ( + f"the sliding group offered more than its window to save: {saved[sliding]}" + ) + assert live[1] > live[0], ( + f"the full-attention group should keep more than the sliding one: {saved}" + ) + + +def _disagg_transceiver_config(use_kv_cache_manager_v2): + """The transceiver each KV cache manager can actually be driven by. + + `CacheTransceiverCpp` is bound to the V1 `BaseKVCacheManager`, while + `KVCacheManagerV2.impl` is the Python V2 core's manager, so V2 can only use + the Python transceiver -- which in turn only supports NIXL + (kv_cache_transceiver.py, `create_kv_cache_transceiver`). This is spelled + out per manager rather than left at the default because + `transceiver_runtime` defaults to "auto", and "auto" is resolved from the + *model's* preference (llm_utils._resolve_transceiver_runtime_auto), which + knows nothing about which cache manager will be built. Qwen2 declares no + preference, so the default resolves to the C++ transceiver, which V2 cannot + use. + """ + if use_kv_cache_manager_v2: + return CacheTransceiverConfig(backend="NIXL", + transceiver_runtime="PYTHON") + return CacheTransceiverConfig(backend="DEFAULT") + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("save_async", [False, True]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, - save_async): + save_async, use_kv_cache_manager_v2): model_fn, scheduler, worker = model_with_connector - prefill_worker = model_fn( - disable_overlap_scheduler=True, - cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT")) + transceiver_config = _disagg_transceiver_config(use_kv_cache_manager_v2) + + prefill_worker = model_fn(disable_overlap_scheduler=True, + cache_transceiver_config=transceiver_config) - decode_worker = model_fn( - cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), - kv_connector_config=None) + decode_worker = model_fn(cache_transceiver_config=transceiver_config, + kv_connector_config=None) sampling_params = SamplingParams(ignore_eos=True, max_tokens=16) @@ -396,16 +1164,20 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, scheduler.request_finished.return_value = False worker.get_finished.return_value = [], [] - result = generate_and_sleep(prefill_worker, [0] * 48, - sampling_params=sampling_params, - disaggregated_params=disaggregated_params) + result = generate_and_wait(prefill_worker, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=disaggregated_params) gen_disagg_params = result.disaggregated_params gen_disagg_params.request_type = "generation_only" - generate_and_sleep(decode_worker, [0] * 48, - sampling_params=sampling_params, - disaggregated_params=gen_disagg_params) + generate_and_wait(decode_worker, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=gen_disagg_params) assert scheduler.build_connector_meta.call_count == 1 @@ -424,6 +1196,9 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_multi_request(enforce_single_worker, model_with_connector): model_fn, scheduler, worker = model_with_connector @@ -448,12 +1223,38 @@ def test_connector_multi_request(enforce_single_worker, model_with_connector): @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [ + pytest.param(False, id="kv_cache_manager_v1"), + pytest.param( + True, + id="kv_cache_manager_v2", + marks=pytest.mark.xfail( + strict=True, + reason= + "KvCacheRetentionConfig does not reach KVCacheManagerV2 at all " + "(per-page priority comes from custom_priority_callback, which " + "V2 never overrides), so the connector reports priorities=None."), + ), +], + indirect=True) def test_connector_priorities(enforce_single_worker, model_with_connector): """Test that retention priorities flow through the connector correctly. This test verifies that when KvCacheRetentionConfig is provided, the RequestData.priorities field is populated with the correct per-block priorities based on the token ranges. + + KNOWN GAP -- `xfail(strict=True)` on `kv_cache_manager_v2`. + `KvCacheRetentionConfig` does not reach KVCacheManagerV2 at all: V2's + per-page priority comes from `custom_priority_callback` + (kv_cache_manager_v2/_core/_kv_cache_manager.py), which KVCacheManagerV2 + never overrides, so every page carries the default priority and the + connector reports `priorities=None`. A user who sets a retention config on + V2 silently gets none of it -- not only through the connector. The + assertions below stay the correct expectation for both managers rather than + being relaxed per manager, so wiring retention into V2 turns this green + instead of needing the test rewritten; `strict=True` is what makes it fail + loudly on that day rather than passing silently. """ BLOCK_SIZE = 32 NUM_INPUT_TOKENS = 64 # 2 blocks @@ -489,9 +1290,11 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - generate_and_sleep(model, [0] * NUM_INPUT_TOKENS, - sampling_params=sampling_params, - kv_cache_retention_config=retention_config) + generate_and_wait(model, + scheduler, + worker, [0] * NUM_INPUT_TOKENS, + sampling_params=sampling_params, + kv_cache_retention_config=retention_config) # Verify that build_connector_meta was called assert scheduler.build_connector_meta.call_count >= 1 @@ -519,6 +1322,86 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): @pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], + ids=["kv_cache_manager_v2"], + indirect=True) +def test_connector_warns_that_retention_is_ignored_on_v2( + enforce_single_worker, model_with_connector, caplog): + """A retention config that has no effect must say so. + + `KvCacheRetentionConfig` does not reach KVCacheManagerV2 at all, so the + connector reports `priorities=None` there. `test_connector_priorities` + pins that gap as `xfail(strict=True)`; this pins the diagnostic, which is + the only thing standing between a user's retention config and it being + dropped in silence. + + V2 only: on V1 the config is honoured and there is nothing to warn about. + """ + model_fn, scheduler, worker = model_with_connector + model = model_fn(disable_overlap_scheduler=True) + + scheduler.get_num_new_matched_tokens.return_value = 0, False + worker.get_finished.return_value = [], [] + + retention_config = KvCacheRetentionConfig(token_range_retention_configs=[ + KvCacheRetentionConfig.TokenRangeRetentionConfig(token_start=0, + token_end=32, + priority=80) + ], + decode_retention_priority=10) + + # Two gates sit between `logger.warning_once` and `caplog`, and both are + # process-global rather than per-test: + # + # 1. The TensorRT-LLM logger defaults to `error` (tensorrt_llm/logger.py:163), + # and `Logger.log` drops anything below that before it reaches Python + # logging at all. The LLM constructor raises the level to `info` only + # while it parses arguments and restores it afterwards, so a warning + # emitted during `generate` is discarded unless the level is raised here. + # 2. `log_once` records its key before consulting the level, so the key is + # consumed by whichever test drove this path first in the process -- + # `test_connector_priorities[kv_cache_manager_v2]` above sets a + # retention config on V2 too, and it runs first. + # + # Raise the level and clear the key so the assertion below observes this + # test's own emission instead of the leftovers of test ordering. + previous_level = trtllm_logger_singleton.level + trtllm_logger_singleton.set_level("warning") + trtllm_logger_singleton._appeared_keys.discard(V2_RETENTION_IGNORED_LOG_KEY) + + # The TensorRT-LLM logger sets `propagate = False`, so caplog only sees its + # records once its handler is attached to that logger by name. + trtllm_logger = logging.getLogger(TRTLLM_LOGGER_NAME) + trtllm_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=TRTLLM_LOGGER_NAME): + generate_and_wait(model, + scheduler, + worker, [0] * 64, + sampling_params=SamplingParams(max_tokens=4, + ignore_eos=True), + kv_cache_retention_config=retention_config) + finally: + trtllm_logger.removeHandler(caplog.handler) + trtllm_logger_singleton.set_level(previous_level) + + assert "KvCacheRetentionConfig has no effect" in caplog.text, ( + "A retention config was set on KVCacheManagerV2 and nothing said it " + "would be ignored. The user's configuration is silently dropped:\n" + f"{caplog.text}") + + # The other half: the warning describes what actually happened. + request = scheduler.build_connector_meta.call_args_list[0].args[ + 0].new_requests[0] + assert request.priorities is None, ( + "priorities are populated on V2 after all, so the warning is now " + "wrong -- revisit it together with `test_connector_priorities`.") + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_priorities_default(enforce_single_worker, model_with_connector): """Test that priorities are None when no retention config is provided.""" @@ -532,7 +1415,10 @@ def test_connector_priorities_default(enforce_single_worker, sampling_params = SamplingParams(max_tokens=4, ignore_eos=True) # Generate without retention config - generate_and_sleep(model, [0] * 48, sampling_params=sampling_params) + generate_and_wait(model, + scheduler, + worker, [0] * 48, + sampling_params=sampling_params) first_call = scheduler.build_connector_meta.call_args_list[0] sched_output = first_call.args[0] @@ -544,47 +1430,183 @@ def test_connector_priorities_default(enforce_single_worker, assert request.priorities is None +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_block_reuse_off_is_rejected_on_v2_only( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """`enable_block_reuse=False` with a connector: V2 refuses, V1 does not. + + The asymmetry is deliberate and both halves are pinned here, because each + is wrong on its own. + + **V2 must refuse.** It honours a connector-served prefix whatever + `enable_block_reuse` says -- that flag governs the local radix tree, and the + connector is a separate source -- and with reuse off the restored KV is + wrong. Measured with the reference connector on a plain full-attention + model: the offer is honoured (18 of 82 tokens scheduled rather than 82) and + generation drifts within a few tokens. Refusing at startup is what turns + silent wrong output into a message. + + **V1 must not refuse.** It never honours the offer in this configuration -- + it asks the connector, then schedules all 82 tokens anyway -- so nothing + miscomputes and there is no correctness reason to reject a setup that + works. It does waste the connector's lookup and hand it a negative + `computed_position`; that is a separate defect, tracked in the backlog + rather than papered over with a guard here. + + Asserting V1 still constructs is the load-bearing half: a guard written + against the config instead of the manager would reject both, and this is + what says so. + """ + model_fn, _, _ = model_with_connector + llm_kwargs = dict(kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, enable_block_reuse=False)) + + if use_kv_cache_manager_v2: + with pytest.raises(NotImplementedError, match="block reuse disabled"): + model_fn(**llm_kwargs) + else: + model = model_fn(**llm_kwargs) + model.shutdown() + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_spec_dec_is_rejected_on_v2_only(enforce_single_worker, + model_with_connector, + use_kv_cache_manager_v2): + """Speculative decoding with a connector: V2 refuses, V1 still constructs. + + A rejected draft token shrinks the request's page list, and the slot the + tail block held is recycled to whichever request allocates next. + `RequestData` carries only the pages appended since the last report, so the + connector is never told the block moved and keeps addressing a slot another + request now owns. + + Asserting that V1 still constructs is the load-bearing half. The same + shrink reaches V1, so the combination is no safer there -- but a guard + written against the config rather than the manager would reject both, and + this is what says which one is gated. + """ + model_fn, _, _ = model_with_connector + llm_kwargs = dict(speculative_config=NGramDecodingConfig( + max_draft_len=4, + max_matching_ngram_size=2, + is_keep_all=True, + is_use_oldest=True, + is_public_pool=True, + )) + + if use_kv_cache_manager_v2: + with pytest.raises(NotImplementedError, match="speculative decoding"): + model_fn(**llm_kwargs) + else: + model = model_fn(**llm_kwargs) + model.shutdown() + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_max_utilization_is_rejected_on_v1_only( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + """A policy other than `GUARANTEED_NO_EVICT`: V1 refuses, V2 accepts. + + A policy that destroys and replays a live request brings it back on + different pages. `KVCacheManagerV2` drops the connector's per-request block + delta when the allocation dies, so the replay reports from scratch; V1 + keeps the delta and would measure it against pages that are gone. + + V2 accepting is the load-bearing half. `KVCacheV2Scheduler` runs + `MAX_UTILIZATION` whatever policy is configured, so a guard that still + refused here would reject the policy V2 actually uses while admitting the + one it discards. + """ + model_fn, _, _ = model_with_connector + llm_kwargs = dict(scheduler_config=SchedulerConfig( + capacity_scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION)) + + if use_kv_cache_manager_v2: + model = model_fn(**llm_kwargs) + model.shutdown() + else: + with pytest.raises(NotImplementedError, match="GUARANTEED_NO_EVICT"): + model_fn(**llm_kwargs) + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize( - "llm_kwargs,match", + "llm_kwargs,match_v1,match_v2", [ + # The two managers refuse offloading for the same reason -- a page that + # leaves GPU has its slot reassigned, invalidating what the connector + # registered -- but say so differently, and V2 names the resolved tier + # rather than the config field so it also catches the disk tier. Match + # the manager-specific wording: both messages contain the bare word + # "host", so matching that would still pass if a silent fallback to V1 + # ever crept back in, which is the one thing this parametrization + # exists to rule out. pytest.param( dict(kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1, host_cache_size=1024**3)), - "host", + "host offloading", + "cache tiers below GPU", id="host_offloading", ), pytest.param( dict(max_beam_width=2), "beam", + "beam", id="beam_search", ), pytest.param( dict(enable_attention_dp=True), "attention data parallelism", + "attention data parallelism", id="attention_dp", ), ], ) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) def test_connector_rejects_unsupported_config(enforce_single_worker, - model_with_connector, llm_kwargs, - match): + model_with_connector, + use_kv_cache_manager_v2, + llm_kwargs, match_v1, match_v2): # Configurations the connector cannot handle today must fail loudly at # construction time rather than silently miscompute. This pins the set of # constructor-time exclusions in `_maybe_init_kv_connector_manager`. model_fn, _, _ = model_with_connector + match = match_v2 if use_kv_cache_manager_v2 else match_v1 with pytest.raises(NotImplementedError, match=match): model_fn(**llm_kwargs) @pytest.mark.threadleak(enabled=False) -def test_connector_e2e_persistent_cache(enforce_single_worker, monkeypatch): - """Test e2e KV cache connector using PersistentKvCacheConnector from examples. - - Runs generation twice with separate LLM instances sharing a disk-based - connector cache, verifying that outputs are identical (proving cache - save/load works end-to-end). +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], + indirect=True) +def test_connector_e2e_persistent_cache(enforce_single_worker, + use_kv_cache_manager_v2, monkeypatch): + """End-to-end KV connector test using PersistentKvCacheConnector from examples. + + Runs the same prompt through two separate LLM instances sharing a + disk-backed connector cache and asserts that: + + 1. the first (cold) run matches nothing and writes cache files, + 2. the second (warm) run actually reads blocks back from disk, and + 3. both runs produce identical text and token ids. + + (3) on its own proves nothing - two deterministic runs of the same prompt + agree whether or not the cache is ever consulted - so (2) is what makes + this a real correctness test rather than a tautology. """ # sys.path, not __extra_import_path__: the connector module is imported by # name from inside tensorrt_llm, not by this file. monkeypatch restores the @@ -595,9 +1617,29 @@ def test_connector_e2e_persistent_cache(enforce_single_worker, monkeypatch): monkeypatch.syspath_prepend(examples_dir) cache_dir = tempfile.mkdtemp() - os.environ["CONNECTOR_CACHE_FOLDER"] = cache_dir + monkeypatch.setenv("CONNECTOR_CACHE_FOLDER", cache_dir) try: + import llm_kv_cache_connector + + # Record how many tokens the connector served from disk on each run. + # The leader logs this, but the TensorRT-LLM logger does not propagate + # to the root logger, so read it from the return value instead. + matched_tokens = [] + leader_cls = llm_kv_cache_connector.PersistentKvCacheConnectorLeader + original_get_num_new_matched_tokens = ( + leader_cls.get_num_new_matched_tokens) + + def recording_get_num_new_matched_tokens(self, request, + num_computed_tokens): + result = original_get_num_new_matched_tokens( + self, request, num_computed_tokens) + matched_tokens.append(result[0]) + return result + + monkeypatch.setattr(leader_cls, "get_num_new_matched_tokens", + recording_get_num_new_matched_tokens) + kv_connector_config = KvCacheConnectorConfig( connector_module="llm_kv_cache_connector", connector_scheduler_class="PersistentKvCacheConnectorLeader", @@ -610,7 +1652,9 @@ def test_connector_e2e_persistent_cache(enforce_single_worker, monkeypatch): kv_connector_config=kv_connector_config, cuda_graph_config=None, disable_overlap_scheduler=True, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1), + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + use_kv_cache_manager_v2=use_kv_cache_manager_v2), ) prompt = ( @@ -625,18 +1669,960 @@ def test_connector_e2e_persistent_cache(enforce_single_worker, monkeypatch): sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) llm1 = LLM(**llm_kwargs) - output1 = llm1.generate([prompt], sampling_params) - output1[0].outputs[0].text - del llm1 + try: + output1 = llm1.generate([prompt], sampling_params) + cold_text = output1[0].outputs[0].text + cold_token_ids = list(output1[0].outputs[0].token_ids) + finally: + llm1.shutdown() + + assert matched_tokens and all(count == 0 for count in matched_tokens), ( + "The first run should be a cold miss, but the connector reported " + f"matched token counts {matched_tokens}. The cache directory was " + "not clean, so the comparison below is meaningless.") cache_files = [f for f in os.listdir(cache_dir) if f.endswith(".pt")] assert len(cache_files) > 0, "No cache files written by connector" + matched_tokens.clear() + llm2 = LLM(**llm_kwargs) - llm2.generate([prompt], sampling_params) - del llm2 + try: + output2 = llm2.generate([prompt], sampling_params) + warm_text = output2[0].outputs[0].text + warm_token_ids = list(output2[0].outputs[0].token_ids) + finally: + llm2.shutdown() + + assert matched_tokens and max(matched_tokens) > 0, ( + "The second run read nothing back from the connector cache " + f"(matched token counts {matched_tokens}), so the comparisons " + "below would pass just as well with the connector disabled.") + + assert len(warm_token_ids) == len(cold_token_ids), ( + f"Generation length changed: cold {len(cold_token_ids)} tokens, " + f"warm {len(warm_token_ids)} tokens.") + + # Exact equality is NOT asserted. Reusing cached KV skips prefill for + # the matched blocks, which changes the attention reduction order, so + # the logits differ in the last bits even though the restored K/V are + # bit-identical (the connector round-trips them through torch.save / + # torch.load). Greedy decoding turns a near-tie into a different token. + # Observed: the two runs agreed on 31 of 32 tokens and split on + # the final one ("The company's" vs "The company is"). + # + # A corrupted or misaddressed cache does not look like that - it + # diverges early and degenerates - so requiring a long common prefix + # keeps the test meaningful without making it a coin flip. + common_prefix = 0 + for cold_id, warm_id in zip(cold_token_ids, warm_token_ids): + if cold_id != warm_id: + break + common_prefix += 1 + + min_common_prefix = math.floor( + len(cold_token_ids) * E2E_MIN_TOKEN_AGREEMENT) + assert common_prefix >= min_common_prefix, ( + f"Connector cache reuse diverged at token {common_prefix} of " + f"{len(cold_token_ids)}, below the {min_common_prefix}-token " + "floor. Early divergence indicates the restored KV is wrong, not " + "just numerically different.\n" + f" cold run: {cold_text!r}\n" + f" warm run: {warm_text!r}\n" + f" cold ids: {cold_token_ids}\n" + f" warm ids: {warm_token_ids}") finally: - os.environ.pop("CONNECTOR_CACHE_FOLDER", None) + if examples_dir in sys.path: + sys.path.remove(examples_dir) + + shutil.rmtree(cache_dir, ignore_errors=True) + + +# The VSWA end-to-end sizes. The window is deliberately larger than the whole +# run (prompt + generation), so nothing goes out of window and the save/load +# round trip is the only thing under test. `test_connector_vswa_reports_page_ +# indices_per_layer_group` covers the routing when blocks *do* go stale, and +# `test_connector_sliding_window_prefix_is_backed_by_real_pages` covers the +# masking arithmetic; mixing all three into one test would leave a failure +# ambiguous. +VSWA_E2E_WINDOW = 256 +VSWA_E2E_MAX_SEQ_LEN = 512 + + +@pytest.mark.threadleak(enabled=False) +def test_connector_multi_pool_e2e_persistent_cache(enforce_single_worker, + monkeypatch): + """The multi-pool data path, end to end, with a connector that moves real bytes. + + Qwen3-0.6B has a single attention type; the two windows here are imposed by + `max_attention_window`, so this covers a VSWA *cache* over a uniform model + -- the cache plumbing, on a small fast model. The interleaved-attention + case, where the model itself has two attention types and the layer groups + come out different sizes, is `test_connector_vswa_e2e_gemma3`. + + Every other multi-group test in this file drives a `MagicMock` connector, so + they assert routing and shape and never touch memory. This one runs + `examples/llm-api/llm_kv_cache_connector_vswa.py` -- which addresses pages + through `KvCacheLayout` regions and nothing else -- across two LLM + instances sharing a disk cache, and asserts: + + 1. the cold run writes one file per (block, layer group), not per block, + 2. no two of those files hold the same bytes, and + 3. the warm run reads blocks back and reproduces the cold run's tokens. + + (1) is what fails if `layer_group_id` is dropped from the cache key: the + groups collide and half the files disappear. (2) is what fails if the + layout reports the same pool base for both groups -- the connector would + then read one group's pages twice and write identical bytes under two + names, while every shape assertion still passed. (3) is what fails if the + page slots address the wrong pool on the way back in. + + V2 only: VSWA has no connector path on the V1 manager, which registers a + single primary pool. + """ + examples_dir = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", + "examples", "llm-api")) + sys.path.insert(0, examples_dir) - import shutil + cache_dir = tempfile.mkdtemp() + monkeypatch.setenv("CONNECTOR_CACHE_FOLDER", cache_dir) + + try: + import llm_kv_cache_connector_vswa as vswa + + leader_cls = vswa.VswaKvCacheConnectorLeader + + matched_tokens = [] + original_query = leader_cls.get_num_new_matched_tokens + + def recording_query(self, request, num_computed_tokens): + result = original_query(self, request, num_computed_tokens) + matched_tokens.append(result[0]) + return result + + saves = [] + original_build = leader_cls.build_connector_meta + + def recording_build(self, scheduler_output): + metadata = original_build(self, scheduler_output) + saves.extend(metadata.save) + return metadata + + monkeypatch.setattr(leader_cls, "get_num_new_matched_tokens", + recording_query) + monkeypatch.setattr(leader_cls, "build_connector_meta", recording_build) + + def build(): + return vswa.build_llm( + model=f"{llm_models_root()}/Qwen3/Qwen3-0.6B", + # Two distinct windows: one sliding layer group and one + # full-attention layer group. An entry equal to max_seq_len + # normalizes to None. + max_attention_window=[VSWA_E2E_WINDOW, VSWA_E2E_MAX_SEQ_LEN], + max_seq_len=VSWA_E2E_MAX_SEQ_LEN, + free_gpu_memory_fraction=0.1, + ) + + prompt = ( + "Nvidia Corporation is an American technology company " + "headquartered in Santa Clara, California. Founded in 1993 by " + "Jensen Huang, Chris Malachowsky, and Curtis Priem, it develops " + "graphics processing units (GPUs), system on a chips (SoCs), and " + "application programming interfaces (APIs) for data science, " + "high-performance computing, and mobile and automotive " + "applications. Tell me about the company.") + sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) + + llm1 = build() + try: + output1 = llm1.generate([prompt], sampling_params) + cold_text = output1[0].outputs[0].text + cold_token_ids = list(output1[0].outputs[0].token_ids) + finally: + llm1.shutdown() + + assert matched_tokens and all(count == 0 for count in matched_tokens), ( + "The first run should be a cold miss, but the connector reported " + f"matched token counts {matched_tokens}. The cache directory was " + "not clean, so the comparison below is meaningless.") + + # (1) One save per (block ordinal, layer group). Two groups, so the + # ordinals must repeat exactly twice and every key must be distinct. + assert saves, "the connector saved nothing on the cold run" + groups = {group_id for _, group_id, _ in saves} + assert groups == { + 0, 1 + }, (f"expected saves in both layer groups, got groups {sorted(groups)}. " + "The run is not VSWA, so nothing below tests per-group addressing.") + paths = [path for path, _, _ in saves] + assert len(set(paths)) == len(paths), ( + "two saves shared a cache key. The layer group is missing from the " + "key, so one group's KV overwrites the other's:\n" + f" {sorted(paths)}") + per_group = {group_id: 0 for group_id in groups} + for _, group_id, _ in saves: + per_group[group_id] += 1 + assert len(set(per_group.values())) == 1, ( + f"layer groups saved different numbers of blocks: {per_group}. " + "Block ordinals must stay aligned across groups.") + + cache_files = sorted(f for f in os.listdir(cache_dir) + if f.endswith(".pt")) + assert len(cache_files) == len(set(paths)), ( + f"{len(set(paths))} distinct save keys produced {len(cache_files)} " + "files on disk.") + + # (2) Distinct bytes per file. Identical content across two groups + # would mean the connector read the same pool twice. + digests = {} + for name in cache_files: + with open(os.path.join(cache_dir, name), "rb") as handle: + digests.setdefault( + hashlib.sha256(handle.read()).hexdigest(), []).append(name) + collisions = { + digest: names + for digest, names in digests.items() if len(names) > 1 + } + assert not collisions, ( + "two cache files hold identical bytes, so the same pool was read " + "for more than one (block, layer group). Layer group g's page " + f"slots are not addressing layer group g's pool: {collisions}") + + matched_tokens.clear() + + llm2 = build() + try: + output2 = llm2.generate([prompt], sampling_params) + warm_text = output2[0].outputs[0].text + warm_token_ids = list(output2[0].outputs[0].token_ids) + finally: + llm2.shutdown() + + # (3) The load direction. Without this the test proves only that bytes + # were written somewhere. + assert matched_tokens and max(matched_tokens) > 0, ( + "The second run read nothing back from the connector cache " + f"(matched token counts {matched_tokens}), so the comparison " + "below would pass just as well with the connector disabled.") + + assert len(warm_token_ids) == len(cold_token_ids), ( + f"Generation length changed: cold {len(cold_token_ids)} tokens, " + f"warm {len(warm_token_ids)} tokens.") + + # Exact equality is not asserted, for the reason spelled out in + # test_connector_e2e_persistent_cache: skipping prefill for the matched + # blocks changes the attention reduction order, so greedy decoding can + # split on a near-tie. Misaddressed KV diverges early instead. + common_prefix = 0 + for cold_id, warm_id in zip(cold_token_ids, warm_token_ids): + if cold_id != warm_id: + break + common_prefix += 1 + + min_common_prefix = math.floor( + len(cold_token_ids) * E2E_MIN_TOKEN_AGREEMENT) + assert common_prefix >= min_common_prefix, ( + f"VSWA connector cache reuse diverged at token {common_prefix} of " + f"{len(cold_token_ids)}, below the {min_common_prefix}-token " + "floor. Early divergence means the restored KV went into the " + "wrong layer group's pages, not that it is numerically " + "different.\n" + f" cold run: {cold_text!r}\n" + f" warm run: {warm_text!r}\n" + f" cold ids: {cold_token_ids}\n" + f" warm ids: {warm_token_ids}") + finally: + if examples_dir in sys.path: + sys.path.remove(examples_dir) + shutil.rmtree(cache_dir, ignore_errors=True) + + +# Gemma-3-1B interleaves sliding and full attention on a 6-layer cycle +# (`sliding_window_pattern: 6`, `sliding_window: 512`), so 26 layers split +# 22 sliding / 4 full. That uneven split is the thing a forced +# `max_attention_window` on a uniform model cannot produce, and it is what +# makes this an interleaved-attention run rather than a VSWA cache config. +GEMMA3_SLIDING_WINDOW = 512 +GEMMA3_GLOBAL_WINDOW = 32768 +GEMMA3_NUM_LAYERS = 26 +GEMMA3_CYCLE = 6 +GEMMA3_MAX_SEQ_LEN = 2048 + + +@pytest.mark.threadleak(enabled=False) +def test_connector_vswa_e2e_gemma3(enforce_single_worker, monkeypatch): + """A KV connector on a model whose architecture is variable-window. + + `test_connector_multi_pool_e2e_persistent_cache` imposes two windows on a + uniform-attention model, which exercises the cache plumbing on something + small and fast. This runs Gemma-3-1B, which interleaves sliding and full + attention natively, and which asks for `KVCacheManagerV2` itself + (`Gemma3ForCausalLM.get_preferred_kv_cache_manager_version` returns "V2" + for exactly this layout). Before this series that combination was rejected + at bring-up, so it is the case the change exists to enable. + + `use_kv_cache_manager_v2` is left at "auto" on purpose: the model's own + preference has to be what lands the run on V2, or a user gets the rejection + without knowing to override anything. The assertions check that it did. + + Two things are asserted that the uniform-model test cannot show: + + 1. **The layer groups are different sizes.** 22 layers slide and 4 do not. + A forced window list on a uniform model splits evenly, so a builder that + partitioned by layer index rather than by window would pass there and + fail here. + 2. **Every layer is covered exactly once**, against the model's real layer + count rather than a number the test chose. + + Then the same save/load round trip: one cache file per (block, layer + group), no two files holding the same bytes, and a warm run that reads + blocks back and reproduces the cold run's tokens. + + The warm run is a separate process, so its radix tree starts empty and a + warm hit can only have come from the connector. + + Scope: the sliding window never engages here. The whole run fits inside the + 512-token window, so this covers per-layer-group addressing and the save / + load round trip on a model with two real attention types, not the + out-of-window `-1` path. See the comment on `enable_block_reuse` below. + """ + examples_dir = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", + "examples", "llm-api")) + sys.path.insert(0, examples_dir) + + cache_dir = tempfile.mkdtemp() + monkeypatch.setenv("CONNECTOR_CACHE_FOLDER", cache_dir) + + try: + import llm_kv_cache_connector_vswa as vswa + + leader_cls = vswa.VswaKvCacheConnectorLeader + worker_cls = vswa.VswaKvCacheConnectorWorker + + matched_tokens = [] + original_query = leader_cls.get_num_new_matched_tokens + + def recording_query(self, request, num_computed_tokens): + result = original_query(self, request, num_computed_tokens) + matched_tokens.append(result[0]) + return result + + saves = [] + # `num_scheduled_tokens` is how many tokens the upcoming forward pass + # will compute. It is the only thing separating "the connector was + # asked" from "the connector's answer was used": if the offer were + # ignored the prompt would just be recomputed, and recomputation + # produces the same tokens as the cold run. Matching output therefore + # cannot show the prefix was honoured, and this can. + scheduled = [] + original_build = leader_cls.build_connector_meta + + def recording_build(self, scheduler_output): + metadata = original_build(self, scheduler_output) + saves.extend(metadata.save) + scheduled.extend((rd.num_scheduled_tokens, len(rd.new_tokens)) + for rd in scheduler_output.new_requests) + return metadata + + layouts = [] + original_register = worker_cls.register_kv_cache_layout + + def recording_register(self, layout): + layouts.append(layout) + return original_register(self, layout) + + monkeypatch.setattr(leader_cls, "get_num_new_matched_tokens", + recording_query) + monkeypatch.setattr(leader_cls, "build_connector_meta", recording_build) + monkeypatch.setattr(worker_cls, "register_kv_cache_layout", + recording_register) + + # The 5-sliding : 1-global cycle Gemma-3 actually uses. The final entry + # is clamped to max_seq_len and normalizes to None, i.e. full attention. + windows = [GEMMA3_SLIDING_WINDOW] * (GEMMA3_CYCLE - 1) + [ + GEMMA3_GLOBAL_WINDOW + ] + + def build(): + return vswa.build_llm( + model=f"{llm_models_root()}/gemma/gemma-3-1b-it", + max_attention_window=windows, + max_seq_len=GEMMA3_MAX_SEQ_LEN, + free_gpu_memory_fraction=0.3, + # The model asks for V2 itself; "auto" is what a user gets. + use_kv_cache_manager_v2="auto", + # Block reuse stays on, which needs two justifications. + # + # It does not weaken the test: the two runs are separate + # processes, so the second starts with an empty radix tree and a + # warm hit can only have come from the connector. + # + # It also does not walk into the WAR the Gemma-3 accuracy suite + # carries ("gaps in kernel support for Gemma3's non-inclusive + # sliding window size", test_llm_api_pytorch.py). That boundary + # is only reached once a sequence outgrows the window, and this + # one does not come close: window 512 against a prompt under + # 96 tokens plus 32 generated. The stale range stays empty + # throughout, so the non-inclusive boundary is never evaluated. + # A longer prompt here would exercise the `-1` masking end to + # end, and would also collide with that WAR -- worth doing, but + # as its own test rather than by stretching this one. + # + # Turning reuse off would make the connector prefix restore + # wrong KV; that is a separate defect, tracked in the backlog, + # not something for this test to carry. + enable_block_reuse=True, + ) + + prompt = ( + "Nvidia Corporation is an American technology company " + "headquartered in Santa Clara, California. Founded in 1993 by " + "Jensen Huang, Chris Malachowsky, and Curtis Priem, it develops " + "graphics processing units (GPUs), system on a chips (SoCs), and " + "application programming interfaces (APIs) for data science, " + "high-performance computing, and mobile and automotive " + "applications. Tell me about the company.") + sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) + + llm1 = build() + try: + output1 = llm1.generate([prompt], sampling_params) + cold_text = output1[0].outputs[0].text + cold_token_ids = list(output1[0].outputs[0].token_ids) + prompt_len = len(output1[0].prompt_token_ids) + finally: + llm1.shutdown() + + # Control for the honoured-prefix check below: with an empty cache the + # cold run must schedule the whole prompt. + cold_sched = [n for n, tok in scheduled if tok == prompt_len] + assert cold_sched == [ + prompt_len + ], (f"cold run scheduled {cold_sched} for a {prompt_len}-token prompt; " + "something served a prefix on a cold cache, so the warm comparison " + "is not measuring the connector.") + + # (0) "auto" landed on V2 and the layout path ran. Without this the + # rest could be describing a run that never reached KVCacheManagerV2. + assert layouts, ( + "register_kv_cache_layout was never called, so this run did not go " + "through KVCacheManagerV2 and nothing below tests the V2 path.") + layout = layouts[0] + + # (1) The interleave, read off the layout rather than assumed. + assert len(layout.groups) == 2, ( + f"Gemma-3 must produce one sliding and one full-attention layer " + f"group, got {len(layout.groups)}: " + f"{[(g.layer_group_id, g.window_size) for g in layout.groups]}") + by_window = {group.window_size: group for group in layout.groups} + assert set(by_window) == { + GEMMA3_SLIDING_WINDOW, None + }, (f"unexpected window set {set(by_window)}; the final entry should " + "clamp to max_seq_len and normalize to None") + + sliding = by_window[GEMMA3_SLIDING_WINDOW] + full = by_window[None] + expected_full = GEMMA3_NUM_LAYERS // GEMMA3_CYCLE + assert len(full.layer_ids) == expected_full, ( + f"{len(full.layer_ids)} full-attention layers, expected " + f"{expected_full} for a {GEMMA3_CYCLE}-layer cycle over " + f"{GEMMA3_NUM_LAYERS} layers: {full.layer_ids}") + assert len(sliding.layer_ids) == GEMMA3_NUM_LAYERS - expected_full + assert len(sliding.layer_ids) != len(full.layer_ids), ( + "the two groups came out the same size, which a real interleaved " + "model does not do -- this is the uniform-model case, not Gemma-3") + + # (2) Every layer covered exactly once, against the model's own count. + covered = sorted(list(sliding.layer_ids) + list(full.layer_ids)) + assert covered == list(range(GEMMA3_NUM_LAYERS)), ( + f"layer coverage is not the full model: {covered}") + + # (3) Cold run: a miss, and one save per (block, layer group). + assert matched_tokens and all(count == 0 for count in matched_tokens), ( + "The first run should be a cold miss, but the connector reported " + f"matched token counts {matched_tokens}. The cache directory was " + "not clean, so the comparison below is meaningless.") + assert saves, "the connector saved nothing on the cold run" + groups = {group_id for _, group_id, _ in saves} + assert groups == {group.layer_group_id + for group in layout.groups + }, (f"saves reached layer groups {sorted(groups)}, " + "not every group in the layout") + paths = [path for path, _, _ in saves] + assert len(set(paths)) == len(paths), ( + "two saves shared a cache key, so the layer group is missing from " + f"the key and one group's KV overwrites the other's:\n{sorted(paths)}" + ) + + cache_files = sorted(f for f in os.listdir(cache_dir) + if f.endswith(".pt")) + assert len(cache_files) == len(set(paths)) + + # (4) Distinct bytes per file. The sliding group and the full group hold + # different layers, so identical content would mean the connector read + # one group's pool for both. + digests = {} + for name in cache_files: + with open(os.path.join(cache_dir, name), "rb") as handle: + digests.setdefault( + hashlib.sha256(handle.read()).hexdigest(), []).append(name) + collisions = { + digest: names + for digest, names in digests.items() if len(names) > 1 + } + assert not collisions, ( + "two cache files hold identical bytes, so the same pool was read " + "for more than one (block, layer group). Layer group g's page " + f"slots are not addressing layer group g's pool: {collisions}") + + # The two groups hold different numbers of layers, so their pages differ + # in size. That is only visible on an interleaved model. + assert sliding.bytes_per_page != full.bytes_per_page, ( + "the two layer groups report the same bytes per page, which " + f"{len(sliding.layer_ids)} and {len(full.layer_ids)} layers cannot " + "both produce") + + matched_tokens.clear() + scheduled.clear() + + llm2 = build() + try: + output2 = llm2.generate([prompt], sampling_params) + warm_text = output2[0].outputs[0].text + warm_token_ids = list(output2[0].outputs[0].token_ids) + finally: + llm2.shutdown() + + # (5) The load direction. Block reuse is off, so this can only be the + # connector. + assert matched_tokens and max(matched_tokens) > 0, ( + "The second run read nothing back from the connector cache " + f"(matched token counts {matched_tokens}), so the comparison " + "below would pass just as well with the connector disabled.") + + # The offer above is what the connector proposed. This is what the + # runtime did with it: prefill must be skipped for the served range, or + # the token comparison below is satisfied by plain recomputation and + # says nothing about the connector. + warm_sched = [n for n, tok in scheduled if tok == prompt_len] + assert warm_sched, ( + f"no context request with a {prompt_len}-token prompt reached the " + f"connector on the warm run: {scheduled}") + assert min(warm_sched) < prompt_len, ( + f"the warm run scheduled {warm_sched} tokens for a {prompt_len}-" + f"token prompt after the connector offered {matched_tokens}. The " + "offer was made but not honoured, so the whole prompt was " + "recomputed and the token agreement below is vacuous.") + assert min(warm_sched) == prompt_len - max(matched_tokens), ( + f"warm run scheduled {min(warm_sched)}; expected {prompt_len} - " + f"{max(matched_tokens)} for a fully honoured offer.") + + assert len(warm_token_ids) == len(cold_token_ids), ( + f"Generation length changed: cold {len(cold_token_ids)} tokens, " + f"warm {len(warm_token_ids)} tokens.") + + common_prefix = 0 + for cold_id, warm_id in zip(cold_token_ids, warm_token_ids): + if cold_id != warm_id: + break + common_prefix += 1 + + min_common_prefix = math.floor( + len(cold_token_ids) * E2E_MIN_TOKEN_AGREEMENT) + assert common_prefix >= min_common_prefix, ( + f"Gemma-3 connector cache reuse diverged at token {common_prefix} " + f"of {len(cold_token_ids)}, below the {min_common_prefix}-token " + "floor. Early divergence means the restored KV went into the wrong " + "layer group's pages, not that it is numerically different.\n" + f" cold run: {cold_text!r}\n" + f" warm run: {warm_text!r}\n" + f" cold ids: {cold_token_ids}\n" + f" warm ids: {warm_token_ids}") + finally: + if examples_dir in sys.path: + sys.path.remove(examples_dir) + shutil.rmtree(cache_dir, ignore_errors=True) + + +# The sliding window has to actually pass blocks for this test to mean +# anything, so the window is set well below the prompt length rather than above +# it as in `test_connector_vswa_e2e_gemma3`. 82 prompt + 64 generated with a +# 64-token window and 32-token blocks puts the boundary at +# max(0, (146 + 1 - 64) // 32) = 2 blocks out of window. +GEMMA3_ENGAGED_WINDOW = 64 +GEMMA3_ENGAGED_GEN_TOKENS = 64 + + +@pytest.mark.threadleak(enabled=False) +def test_connector_vswa_out_of_window_blocks_reach_the_connector( + enforce_single_worker, monkeypatch): + """The sliding window passing blocks, delivered per layer group. + + `test_connector_vswa_e2e_gemma3` keeps the whole run inside the window on + purpose, so it covers addressing and the save/load round trip and never the + out-of-window path. This is the other half: the window is set below the + prompt length so blocks genuinely go stale, and the assertions are about + what the connector is handed for them. + + **Why this cannot also be a round trip.** Staleness is a prefix property -- + the earliest blocks go first -- and a prefix cache must serve from ordinal + 0. Once the window has passed block 0 the sequence can never be served back + as a prefix, whatever the full-attention group still holds. An engaged + window and a warm hit are mutually exclusive for a sliding group, so this + test asserts the masking and deliberately makes no round-trip claim. + + **The control.** Gemma-3's accuracy suite disables block reuse citing "gaps + in kernel support for Gemma3's non-inclusive sliding window size", and that + boundary is reached only when the window engages -- which is what this test + does on purpose. So the same configuration is run without a connector and + the outputs are compared. If they disagree, the model is not self-consistent + at this window and nothing here can be attributed to the connector. + """ + examples_dir = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", + "examples", "llm-api")) + sys.path.insert(0, examples_dir) + cache_dir = tempfile.mkdtemp() + monkeypatch.setenv("CONNECTOR_CACHE_FOLDER", cache_dir) + + try: + import llm_kv_cache_connector_vswa as vswa + + leader_cls = vswa.VswaKvCacheConnectorLeader + worker_cls = vswa.VswaKvCacheConnectorWorker + + finished = [] + original_finished = leader_cls.request_finished_by_layer_group + + def recording_finished(self, request, cache_block_ids_by_layer_group): + finished.append([list(g) for g in cache_block_ids_by_layer_group]) + return original_finished(self, request, + cache_block_ids_by_layer_group) + + layouts = [] + original_register = worker_cls.register_kv_cache_layout + + def recording_register(self, layout): + layouts.append(layout) + return original_register(self, layout) + + monkeypatch.setattr(leader_cls, "request_finished_by_layer_group", + recording_finished) + monkeypatch.setattr(worker_cls, "register_kv_cache_layout", + recording_register) + + windows = [GEMMA3_ENGAGED_WINDOW] * (GEMMA3_CYCLE - 1) + [ + GEMMA3_MAX_SEQ_LEN + ] + model_path = f"{llm_models_root()}/gemma/gemma-3-1b-it" + prompt = ( + "Nvidia Corporation is an American technology company " + "headquartered in Santa Clara, California. Founded in 1993 by " + "Jensen Huang, Chris Malachowsky, and Curtis Priem, it develops " + "graphics processing units (GPUs), system on a chips (SoCs), and " + "application programming interfaces (APIs) for data science, " + "high-performance computing, and mobile and automotive " + "applications. Tell me about the company.") + sampling_params = SamplingParams(max_tokens=GEMMA3_ENGAGED_GEN_TOKENS, + ignore_eos=True) + + llm = vswa.build_llm(model=model_path, + max_attention_window=windows, + max_seq_len=GEMMA3_MAX_SEQ_LEN, + free_gpu_memory_fraction=0.3, + use_kv_cache_manager_v2="auto", + enable_block_reuse=True) + try: + out = llm.generate([prompt], sampling_params)[0] + with_connector = list(out.outputs[0].token_ids) + prompt_len = len(out.prompt_token_ids) + finally: + llm.shutdown() + + # --- The control, first: is the model self-consistent at this window? - + # Run identically with no connector at all. If this disagrees, the + # non-inclusive-window kernel gap is in play and nothing below can be + # read as a statement about the connector. + baseline_llm = LLM(model=model_path, + backend="pytorch", + cuda_graph_config=None, + disable_overlap_scheduler=True, + max_seq_len=GEMMA3_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.3, + enable_block_reuse=True, + max_attention_window=windows, + use_kv_cache_manager_v2=True)) + try: + baseline = list( + baseline_llm.generate([prompt], + sampling_params)[0].outputs[0].token_ids) + finally: + baseline_llm.shutdown() + + assert with_connector == baseline, ( + "the connector run and the connector-free run disagree at an " + "engaged sliding window, so the model is not self-consistent here " + "and the masking assertions below cannot be attributed to the " + f"connector.\n with connector: {with_connector}\n baseline: " + f"{baseline}") + + # --- The layout is genuinely two groups ------------------------------ + assert layouts, "register_kv_cache_layout never ran; this is not a V2 run" + layout = layouts[0] + by_window = {g.window_size: g for g in layout.groups} + assert set(by_window) == { + GEMMA3_ENGAGED_WINDOW, None + }, (f"expected a sliding and a full-attention group, got " + f"{[(g.layer_group_id, g.window_size) for g in layout.groups]}") + sliding = by_window[GEMMA3_ENGAGED_WINDOW].layer_group_id + full = by_window[None].layer_group_id + + # --- What request_finished offered, per group ------------------------ + assert finished, "request_finished_by_layer_group was never called" + offered = finished[-1] + assert len(offered) == 2 + + stale = [ + i for i, slot in enumerate(offered[sliding]) + if slot == BAD_PAGE_INDEX + ] + live = [slot for slot in offered[sliding] if slot != BAD_PAGE_INDEX] + + # Anti-vacuity: if the window never engaged, nothing here is being + # tested. This is the assertion that fails if the masking is removed. + assert stale, ( + "no block went out of window, so the sliding group offered its " + f"whole history and this test proves nothing: {offered[sliding]}") + + # The stale entries are a prefix, in place, not dropped from the list. + assert stale == list(range(len(stale))), ( + f"out-of-window entries are not a leading run: {offered[sliding]}") + assert len(offered[sliding]) == len(offered[full]), ( + "block ordinals are not aligned across groups, so an append-delta " + f"over these lists would be invalid: {offered}") + + # The full-attention group keeps everything. Without this the test + # would pass if masking were applied to every group indiscriminately. + assert BAD_PAGE_INDEX not in offered[full], ( + f"the full-attention group lost blocks to a window it does not " + f"have: {offered[full]}") + assert len(live) < len(offered[full]), ( + "the sliding group offered as much as the full-attention group, so " + "the window did not bound it") + + # The live range is bounded by the window, not merely smaller. + max_live = math.ceil(GEMMA3_ENGAGED_WINDOW / 32) + 1 + assert len(live) <= max_live, ( + f"the sliding group offered {len(live)} live blocks for a " + f"{GEMMA3_ENGAGED_WINDOW}-token window; at most {max_live} can hold " + f"readable KV: {offered[sliding]}") + assert len(set(live)) == len(live), ( + f"live page slots are not distinct: {live}") + + # And the prompt really did outrun the window, which is what put the + # boundary inside the sequence rather than at either end. + assert prompt_len > GEMMA3_ENGAGED_WINDOW, ( + f"prompt is {prompt_len} tokens against a " + f"{GEMMA3_ENGAGED_WINDOW}-token window; the window cannot engage") + finally: + if examples_dir in sys.path: + sys.path.remove(examples_dir) + shutil.rmtree(cache_dir, ignore_errors=True) + + +# Long enough that a 64-token window (2 blocks of 32) leaves most of the prompt +# behind once a served prefix moves history to the end of it. +SELECTIVE_PARA = ( + "Nvidia Corporation is an American technology company headquartered in " + "Santa Clara, California. Founded in 1993 by Jensen Huang, Chris " + "Malachowsky, and Curtis Priem, it develops graphics processing units " + "(GPUs), system on a chips (SoCs), and application programming interfaces " + "(APIs) for data science, high-performance computing, and mobile and " + "automotive applications. ") +SELECTIVE_PROMPT = SELECTIVE_PARA * 3 + "Tell me about the company." + +# A longer, self-contained prompt for models whose window is wide enough that +# the paragraph above would not outrun it. Kept in a file so the same text can +# be reused by hand when eyeballing output quality. +SELECTIVE_PROMPT_FILE = os.path.join(os.path.dirname(__file__), "data", + "kv_connector_vswa_prompt.txt") +SELECTIVE_WINDOW = 64 + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize( + "model_rel,assert_round_trip", + [ + # Real interleaved attention. Output correctness is NOT asserted: at an + # engaged window Gemma-3 does not survive prefix reuse at all, with or + # without a connector. Measured, same window and prompt, local reuse and + # no connector: 0/16 token agreement, against 16/16 with the window + # never engaging. That is the gap the accuracy suite's WAR names ("gaps + # in kernel support for Gemma3's non-inclusive sliding window size"), so + # a round-trip assertion here would fail for a reason that has nothing + # to do with the connector. + pytest.param("gemma/gemma-3-1b-it", False, id="gemma3_interleaved"), + # Uniform attention with the windows imposed, which is where the + # restoration can actually be checked: same window and prompt under + # local reuse gives 16/16. Weaker as a VSWA claim, stronger as a + # correctness claim, so the pair covers what neither does alone. + pytest.param("Qwen3/Qwen3-0.6B", True, id="qwen3_forced_windows"), + # The case this test wants most: a natively interleaved model whose + # prefix-reuse path is sound, so the round trip IS assertable on real + # VSWA. GPT-OSS applies a 128-token window to every even layer + # (modeling_gpt_oss.py). + # + # NEVER EXECUTED. It requires SM100: the model is MXFP4 (no Ampere MoE + # kernel) and forces attn_backend="TRTLLM" because of its attention + # sinks. Gated below rather than deleted so the case is queued rather + # than forgotten -- but nothing here has run, and it should not be + # counted as coverage until it has. + pytest.param("gpt_oss/gpt-oss-20b", + True, + id="gpt_oss_interleaved", + marks=pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason="gpt-oss-20b is MXFP4 and needs SM100/SM103")), + ]) +def test_connector_transfers_only_in_window_blocks_to_the_sliding_group( + enforce_single_worker, monkeypatch, model_rel, assert_round_trip): + """The connector transfers each layer group only what that group can read. + + Target: under a sliding window that has genuinely passed part of the + prompt, the connector loads **only the in-window blocks** into the sliding + layer group and **every block** into the full-attention group. + + The window engages because the prefix is served, not despite it: honouring + the offer moves `history_length` to the end of the served range, which is + what puts the earlier blocks out of the sliding group's window. They are + then reported as `-1`, and the connector must skip them -- there is no + readable KV there and attention will never look. + """ + examples_dir = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", + "examples", "llm-api")) + sys.path.insert(0, examples_dir) + cache_dir = tempfile.mkdtemp() + monkeypatch.setenv("CONNECTOR_CACHE_FOLDER", cache_dir) + + try: + import llm_kv_cache_connector_vswa as vswa + + L, Wk = vswa.VswaKvCacheConnectorLeader, vswa.VswaKvCacheConnectorWorker + layouts, loads, offers, scheduled, by_group = [], [], [], [], [] + original_build = L.build_connector_meta + original_query = L.get_num_new_matched_tokens + original_register = Wk.register_kv_cache_layout + + def recording_build(self, scheduler_output): + metadata = original_build(self, scheduler_output) + loads.extend(metadata.load) + for rd in scheduler_output.new_requests: + scheduled.append((rd.num_scheduled_tokens, len(rd.new_tokens))) + by_group.append( + [list(g) for g in rd.new_block_ids_by_layer_group]) + return metadata + + def recording_query(self, request, num_computed_tokens): + result = original_query(self, request, num_computed_tokens) + offers.append(result[0]) + return result + + def recording_register(self, layout): + layouts.append(layout) + return original_register(self, layout) + + monkeypatch.setattr(L, "build_connector_meta", recording_build) + monkeypatch.setattr(L, "get_num_new_matched_tokens", recording_query) + monkeypatch.setattr(Wk, "register_kv_cache_layout", recording_register) + + sampling_params = SamplingParams(max_tokens=16, ignore_eos=True) + + def gen(): + llm = vswa.build_llm(model=f"{llm_models_root()}/{model_rel}", + max_attention_window=[SELECTIVE_WINDOW] * + (GEMMA3_CYCLE - 1) + [GEMMA3_MAX_SEQ_LEN], + max_seq_len=GEMMA3_MAX_SEQ_LEN, + free_gpu_memory_fraction=0.3, + use_kv_cache_manager_v2=True, + enable_block_reuse=True) + try: + out = llm.generate([SELECTIVE_PROMPT], sampling_params)[0] + return (list(out.outputs[0].token_ids), + len(out.prompt_token_ids)) + finally: + llm.shutdown() + + cold, prompt_len = gen() + assert prompt_len > 3 * SELECTIVE_WINDOW, ( + f"prompt is {prompt_len} tokens against a {SELECTIVE_WINDOW}-token " + "window; too short for the window to leave most of it behind") + for recorder in (loads, offers, scheduled, by_group): + recorder.clear() + + warm, _ = gen() + + layout = layouts[0] + by_window = {g.window_size: g.layer_group_id for g in layout.groups} + assert set(by_window) == { + SELECTIVE_WINDOW, None + }, (f"expected a sliding and a full group, got {by_window}") + sliding, full = by_window[SELECTIVE_WINDOW], by_window[None] + + # The prefix was honoured -- otherwise the window never moves and the + # rest of this measures nothing. + warm_sched = [n for n, tok in scheduled if tok == prompt_len] + assert warm_sched and min(warm_sched) < prompt_len, ( + f"the warm run scheduled {warm_sched} of {prompt_len} tokens after " + f"the connector offered {offers}; the offer was not honoured, so " + "history never advanced and no block went out of window") + + # The masking: a leading run of -1 in the sliding group, none in the + # full group. The boundary is recomputed from the offer rather than read + # back from the implementation. + slots = by_group[0] + stale = [i for i, s in enumerate(slots[sliding]) if s == BAD_PAGE_INDEX] + expected_stale = max(0, (max(offers) + 1 - SELECTIVE_WINDOW) // + layout.tokens_per_block) + assert stale == list(range(len(stale))), ( + f"out-of-window entries are not a leading run: {slots[sliding]}") + assert len(stale) == expected_stale, ( + f"{len(stale)} blocks out of window; expected {expected_stale} for " + f"a {SELECTIVE_WINDOW}-token window at history {max(offers)}") + # Not just "engaged" -- the window must leave a substantial part of + # the prompt behind, or the two groups barely differ and the selective + # transfer below is not really being exercised. + assert expected_stale >= 2, ( + f"only {expected_stale} block(s) went out of window at history " + f"{max(offers)}; size the prompt so the window leaves more behind") + assert BAD_PAGE_INDEX not in slots[full], ( + f"the full-attention group lost blocks it should keep: {slots[full]}" + ) + + # The point of the test: what was actually transferred, per group. + sliding_loads = [x for x in loads if x[1] == sliding] + full_loads = [x for x in loads if x[1] == full] + assert sliding_loads, ( + "nothing was loaded into the sliding group; it still holds an " + "in-window range and skipping it entirely is not correct either") + assert len(sliding_loads) < len(full_loads), ( + f"the sliding group was loaded {len(sliding_loads)} blocks and the " + f"full group {len(full_loads)}; the window did not bound the " + "transfer") + max_live = math.ceil(SELECTIVE_WINDOW / layout.tokens_per_block) + 1 + assert len(sliding_loads) <= max_live, ( + f"{len(sliding_loads)} blocks loaded into a " + f"{SELECTIVE_WINDOW}-token window; at most {max_live} can be read") + assert len({x[2] for x in sliding_loads}) == len(sliding_loads) + + if assert_round_trip: + assert warm == cold, ( + "the selectively restored prefix did not reproduce the cold " + f"run.\n cold: {cold}\n warm: {warm}") + finally: + if examples_dir in sys.path: + sys.path.remove(examples_dir) shutil.rmtree(cache_dir, ignore_errors=True) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 5d9a61d5e8af..843424b386ac 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -37,6 +37,8 @@ l0_a10: - unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py - unittest/_torch/executor/kv_cache/test_kv_cache_compression_manager.py - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py + - unittest/_torch/executor/test_kv_cache_layout.py + - unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/executor/test_resource_manager.py - unittest/_torch/executor/test_profile_endpoints.py @@ -149,25 +151,86 @@ l0_a10: # usage telemetry - unittest/llmapi/test_llm_telemetry.py::TestTelemetryPyTorchBackend - unittest/llmapi/test_llm_telemetry.py::TestTelemetryArchitectureExtraction - - llmapi/test_llm_api_connector.py::test_connector_simple[True] - - llmapi/test_llm_api_connector.py::test_connector_simple[False] - - llmapi/test_llm_api_connector.py::test_connector_async_onboard[True] - - llmapi/test_llm_api_connector.py::test_connector_async_onboard[False] - - llmapi/test_llm_api_connector.py::test_connector_async_save[True] - - llmapi/test_llm_api_connector.py::test_connector_async_save[False] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[True] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[False] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[True] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[False] - - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[False] - - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[True] - - llmapi/test_llm_api_connector.py::test_connector_multi_request - - llmapi/test_llm_api_connector.py::test_connector_priorities - - llmapi/test_llm_api_connector.py::test_connector_priorities_default - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[host_offloading] - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[beam_search] - - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[attention_dp] - - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache + # KV connector. The suite is parametrized over the KV cache manager and both + # halves are gated: the connector is supported on KVCacheManagerV2, so a V2 + # regression has to be a merge-gate failure rather than something found later + # by hand. test_connector_runs_on_kv_cache_manager_v2 is what makes the V2 + # entries meaningful -- the creator silently falls back to the V1 manager for + # combinations it cannot serve, so without it every V2 id below could pass + # while running V1. test_connector_priorities[kv_cache_manager_v2] is + # xfail(strict=True) in the suite, not omitted here, so the retention gap + # stays counted. The two e2e cases below are the only multi-layer-group ones + # that move bytes -- the rest drive a MagicMock and assert routing. + # test_connector_vswa_e2e_gemma3 is the interleaved-attention run (Gemma-3 + # asks for KVCacheManagerV2 itself, and its layer groups come out different + # sizes); test_connector_multi_pool_e2e_persistent_cache is the same data + # path on a uniform model with the windows imposed, which is faster and + # isolates the cache plumbing from the model. Those two keep the run inside + # the window on purpose; test_connector_vswa_out_of_window_blocks_reach_the_connector + # is the complement, with the window set below the prompt length so blocks + # genuinely go stale. + - llmapi/test_llm_api_connector.py::test_v2_connector_contract_does_not_reuse_the_v1_methods + - llmapi/test_llm_api_connector.py::test_connector_runs_on_kv_cache_manager_v2 + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v1-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v1-True] + - llmapi/test_llm_api_connector.py::test_connector_multi_request[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-host_offloading] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-beam_search] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v1-attention_dp] + - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_prefix_is_asked_once_and_shrinks_the_forward_pass[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_prefix_under_chunked_prefill[kv_cache_manager_v1-offer_inside_chunk] + - llmapi/test_llm_api_connector.py::test_connector_prefix_under_chunked_prefill[kv_cache_manager_v1-offer_past_chunk] + - llmapi/test_llm_api_connector.py::test_connector_vswa_reports_page_indices_per_layer_group[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_simple[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_async_save[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v2-False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v2-True] + - llmapi/test_llm_api_connector.py::test_connector_multi_request[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_warns_that_retention_is_ignored_on_v2[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-host_offloading] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-beam_search] + - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-attention_dp] + - llmapi/test_llm_api_connector.py::test_connector_block_reuse_off_is_rejected_on_v2_only[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_block_reuse_off_is_rejected_on_v2_only[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_spec_dec_is_rejected_on_v2_only[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_spec_dec_is_rejected_on_v2_only[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_max_utilization_is_rejected_on_v1_only[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_max_utilization_is_rejected_on_v1_only[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_e2e_persistent_cache[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_prefix_is_asked_once_and_shrinks_the_forward_pass[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_prefix_under_chunked_prefill[kv_cache_manager_v2-offer_inside_chunk] + - llmapi/test_llm_api_connector.py::test_connector_prefix_under_chunked_prefill[kv_cache_manager_v2-offer_past_chunk] + - llmapi/test_llm_api_connector.py::test_connector_uniform_sliding_window[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_sliding_window_prefix_is_backed_by_real_pages[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_vswa_reports_page_indices_per_layer_group[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_multi_pool_e2e_persistent_cache + - llmapi/test_llm_api_connector.py::test_connector_vswa_e2e_gemma3 + - llmapi/test_llm_api_connector.py::test_connector_vswa_out_of_window_blocks_reach_the_connector + - llmapi/test_llm_api_connector.py::test_connector_transfers_only_in_window_blocks_to_the_sliding_group[gemma3_interleaved] + - llmapi/test_llm_api_connector.py::test_connector_transfers_only_in_window_blocks_to_the_sliding_group[qwen3_forced_windows] # third-party policy checks CPU-only - thirdparty/test_cmake_third_party.py::test_cmake_listfiles - thirdparty/test_git_modules.py::test_gitmodules diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index a437e30ed4c3..fea26c840a48 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py @@ -714,6 +714,7 @@ def test_prepare_context_cache_records_lookup_without_mutating_cursor( kv_cache = Mock(num_committed_tokens=2) manager = object.__new__(KVCacheManagerV2) manager.conversation_manager = None + manager.kv_connector_manager = None manager.enable_block_reuse = True manager._has_cp_helix = False manager.kv_cache_map = {} if fresh_cache else {request.py_request_id: kv_cache} @@ -1225,6 +1226,7 @@ def _make_publishing_manager(policy: BlockReusePolicy) -> KVCacheManagerV2: manager._can_publish_block_reuse = True manager.block_reuse_policy = policy manager.conversation_manager = None + manager.kv_connector_manager = None manager.kv_cache_map = {} return manager @@ -1336,6 +1338,7 @@ class _ContextRequest: is_dummy_request: bool = False return_perf_metrics: bool = False context_current_position: int = 0 + py_connector_served_position: int = 0 prepopulated_prompt: tuple[int, int] | None = None multimodal_hashes: None = None multimodal_positions: None = None diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py index fca8d538e0b8..b6e03e781a1d 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py @@ -44,11 +44,11 @@ def make_stub_manager( - tokens_per_block=TOKENS_PER_BLOCK, - enable_block_reuse=True, - num_reusable=0, - reuse_match_backoff=0, -): + tokens_per_block: int = TOKENS_PER_BLOCK, + enable_block_reuse: bool = True, + num_reusable: int = 0, + reuse_match_backoff: int = 0, +) -> KVCacheManagerV2: """A KVCacheManagerV2 reduced to what the two token paths need.""" mgr = object.__new__(KVCacheManagerV2) mgr.tokens_per_block = tokens_per_block @@ -58,6 +58,7 @@ def make_stub_manager( mgr.reuse_match_backoff = reuse_match_backoff mgr.vocab_size = 32000 mgr.conversation_manager = None + mgr.kv_connector_manager = None mgr.kv_cache_map = {} mgr.index_mapper = Mock() mgr.index_mapper.num_free_slots.return_value = 1 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 72cf11c31fcd..76eae5b5a3ed 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -64,15 +64,15 @@ def make_gen_request( def make_ctx_request( - request_id, - context_remaining_length, - prompt_len=None, - num_draft_tokens=0, - is_first_context_chunk=True, - is_last_context_chunk=True, - lora_task_id=None, - encoder_output_len=None, -): + request_id: int, + context_remaining_length: int, + prompt_len: int | None = None, + num_draft_tokens: int = 0, + is_first_context_chunk: bool = True, + is_last_context_chunk: bool = True, + lora_task_id: int | None = None, + encoder_output_len: int | None = None, +) -> Mock: req = Mock() req.request_id = request_id req.py_request_id = request_id @@ -80,6 +80,7 @@ def make_ctx_request( req.context_remaining_length = context_remaining_length req.prompt_len = prompt_len or context_remaining_length req.context_current_position = 0 + req.py_connector_served_position = 0 req.expect_snapshot_points = [] req.num_draft_tokens = num_draft_tokens req.has_draft_tokens = num_draft_tokens > 0 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py index b17a86ee744a..828ffef21e68 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_pool_rebalance.py @@ -57,6 +57,7 @@ def _make_executor( cp_size: int = 1, enable_attention_dp: bool = False, kv_cache_transceiver=None, + kv_connector_manager=None, is_warmup: bool = False, is_shutdown: bool = False, max_beam_width: int = 1, @@ -90,6 +91,7 @@ def _make_executor( exe.dist = MagicMock(pp_size=pp_size, tp_size=tp_size, cp_size=cp_size) exe.enable_attention_dp = enable_attention_dp exe.kv_cache_transceiver = kv_cache_transceiver + exe.kv_connector_manager = kv_connector_manager exe.is_warmup = is_warmup exe.is_shutdown = is_shutdown exe.drafter = drafter @@ -203,6 +205,13 @@ def test_transceiver_present_returns_false(self): exe = _make_executor(kv_cache_transceiver=MagicMock()) assert PyExecutor._can_pause_for_rebalance(exe) is False + def test_connector_present_returns_false(self): + """A rebalance reassigns ``slot_id``, and a connector holds addresses + derived from it across iterations (``update_state_after_alloc``), so + every connector run is exposed -- synchronous ones included.""" + exe = _make_executor(kv_connector_manager=MagicMock()) + assert PyExecutor._can_pause_for_rebalance(exe) is False + def test_warmup_returns_false(self): exe = _make_executor(is_warmup=True) assert PyExecutor._can_pause_for_rebalance(exe) is False diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 20fb07a41a5f..249fb590db63 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -981,8 +981,13 @@ def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): "max_beam_width, has_connector, expected", [ (2, False, "max_beam_width > 1"), - (1, True, "kv_connector_manager"), - (2, True, "kv_connector_manager, max_beam_width > 1"), + # A KV connector alone no longer forces a fallback: it is supported + # through the pool-layout registration path, so the manager is returned + # unchanged and nothing is raised. + (1, True, None), + # With beam search still incompatible, the connector must not appear in + # the reason list -- it is not what makes this configuration unsupported. + (2, True, "max_beam_width > 1"), ], ) def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( @@ -1001,6 +1006,15 @@ def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( creator._kv_connector_manager = object() if has_connector else None creator._max_beam_width = max_beam_width + if expected is None: + assert ( + creator._validate_or_fallback_kv_cache_manager_v2( + MambaHybridCacheManagerV2, model_config, KvCacheConfig() + ) + is MambaHybridCacheManagerV2 + ) + return + with pytest.raises(NotImplementedError, match=expected): creator._validate_or_fallback_kv_cache_manager_v2( MambaHybridCacheManagerV2, model_config, KvCacheConfig() diff --git a/tests/unittest/_torch/executor/test_kv_cache_layout.py b/tests/unittest/_torch/executor/test_kv_cache_layout.py new file mode 100644 index 000000000000..b4dd2bd40730 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_layout.py @@ -0,0 +1,908 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Tests for the KV cache layout description handed to a KV connector +# (``connectors/kv_cache_layout.py``). +# +# The layout replaces the single-pool-tensor registration for a cache the tensor +# cannot express: one slot address space per pool, and one page-index space per +# layer group. +# +# Region-arithmetic and lifecycle-configuration tests need no GPU. The manager +# layout tests construct a real cache manager and allocate device memory pools. + +import gc +import unittest +import unittest.mock + +import pytest +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import ( + KvCacheBufferRef, + KvCacheLayerGroupLayout, + KvCacheLayout, + KvCacheRegion, + build_kv_cache_layout_v2, + valid_page_slots, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, AttnLifeCycle + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + + +def _make_kwargs( + *, + num_layers: int = 4, + num_kv_heads=4, + head_dim=128, + tokens_per_block: int = 8, + max_seq_len: int = 256, + max_batch_size: int = 4, + max_tokens: int = 2048, + dtype=DataType.HALF, + kv_cache_type=CacheType.SELF, + vocab_size: int = 32000, + kv_cache_config=None, +): + return dict( + kv_cache_config=kv_cache_config + or KvCacheConfigV2(max_tokens=max_tokens, enable_block_reuse=False), + kv_cache_type=kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=dtype, + vocab_size=vocab_size, + ) + + +def _make_request(request_id: int = 1, prompt_len: int = 120): + return LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=list(range(prompt_len)), + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("window_size", "num_sink_tokens", "expected_sink_blocks", "expected_stale_range"), + [ + (None, None, 0, (0, 0)), + (None, 0, 0, (0, 0)), + (16, None, 0, (0, 6)), + (16, 0, 0, (0, 6)), + (16, 9, 2, (2, 6)), + ], +) +def test_attention_life_cycle_optional_configuration( + window_size: int | None, + num_sink_tokens: int | None, + expected_sink_blocks: int, + expected_stale_range: tuple[int, int], +) -> None: + """Both runtime backends preserve absent windows and optional sink tokens.""" + life_cycle = AttnLifeCycle.make(window_size, num_sink_tokens, 8) + assert life_cycle.window_size == window_size + assert life_cycle.num_sink_blocks == expected_sink_blocks + stale = life_cycle.get_stale_range(64, 8) + assert (stale.beg, stale.end) == expected_stale_range + assert life_cycle == AttnLifeCycle(window_size, expected_sink_blocks) + + +class TestKvCacheRegionArithmetic(unittest.TestCase): + """Address arithmetic and lookup helpers. No GPU required.""" + + def _region(self, base=4096, size=256, stride=1024, num_slots=8): + return KvCacheRegion( + base=base, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=0, role="key"),), + ) + + def test_address_of_follows_stride(self): + region = self._region() + self.assertEqual(region.address_of(0), 4096) + self.assertEqual(region.address_of(1), 4096 + 1024) + self.assertEqual(region.address_of(7), 4096 + 7 * 1024) + + def test_address_of_rejects_out_of_range_slot(self): + region = self._region(num_slots=8) + for bad in (-1, 8, 99): + with self.assertRaises(IndexError): + region.address_of(bad) + + def test_slot_tensor_rejects_out_of_range_slot(self): + # The check happens before any address is formed, so an invalid page + # index cannot reach device memory. -1 is the one that matters: + # BAD_PAGE_INDEX marks a block with no page, and as a subscript it + # names the pool's last slot. + region = self._region(num_slots=8) + for bad in (BAD_PAGE_INDEX, -1, -2, 8, 99): + with self.assertRaises(IndexError): + region.slot_tensor(bad) + + def test_slot_tensor_rejects_dtype_not_dividing_extent(self): + region = self._region(size=6, stride=6) + with self.assertRaises(ValueError): + region.slot_tensor(0, dtype=torch.float32) + + def test_as_tensor_rejects_dtype_not_dividing_extent(self): + # size/stride are byte counts; a dtype whose itemsize does not divide + # them cannot produce a correct view, so it must fail loudly rather + # than silently truncate. + region = self._region(size=6, stride=6) + with self.assertRaises(ValueError): + region.as_tensor(dtype=torch.float32) + + def test_bytes_per_page_sums_regions(self): + group = KvCacheLayerGroupLayout( + layer_group_id=0, + layer_ids=(0, 1), + window_size=None, + regions=(self._region(size=256), self._region(base=8192, size=128)), + ) + self.assertEqual(group.bytes_per_page, 384) + + def _pool_layout( + self, + num_layers=2, + kv_factor=2, + block_bytes=64, + num_slots=8, + expansion=1, + group_layer_ids=None, + ): + """A layout shaped the way a single-pool cache reports itself. + + Buffers are laid out layer-major and ascending, the way the storage + config builds a coalesced buffer. ``group_layer_ids`` sets the group's + membership list, which the manager reports in ``impl.layer_grouping`` + order rather than memory order. + """ + layer_ids = tuple(range(num_layers)) + order = [lid for lid in layer_ids for _ in range(kv_factor)] + buffers = tuple( + KvCacheBufferRef( + layer_id=lid, role="key" if i % kv_factor == 0 else "value", expansion=expansion + ) + for i, lid in enumerate(order) + ) + region = KvCacheRegion( + base=4096, + size=block_bytes * len(buffers), + stride=4096, + num_slots=num_slots, + buffers=buffers, + ) + group = KvCacheLayerGroupLayout( + 0, layer_ids if group_layer_ids is None else tuple(group_layer_ids), None, (region,) + ) + return KvCacheLayout(tokens_per_block=8, groups=(group,), dtype=torch.float16) + + def test_single_pool_view_has_the_v1_pool_shape(self): + # The shape `register_kv_caches` has always received, so a connector + # written against it needs no changes when the same cache reports + # itself as a layout instead. + layout = self._pool_layout(num_layers=3, kv_factor=2, block_bytes=64, num_slots=5) + backing = torch.zeros(5 * 4096 // 2, dtype=torch.float16) + region = layout.groups[0].regions[0] + flat = backing.as_strided((5, 3 * 2 * 32), (4096 // 2, 1)) + with unittest.mock.patch.object(KvCacheRegion, "as_tensor", return_value=flat): + view = layout.as_single_pool_tensor() + self.assertIsNotNone(view) + # block_bytes=64 at float16 is 32 elements per (layer, role). + self.assertEqual(tuple(view.shape), (5, 3, 2, 32)) + self.assertEqual(region.size, 6 * 64) + + def test_single_pool_view_ignores_the_layer_grouping_order(self): + # `layer_ids` is a membership list carrying `impl.layer_grouping`'s + # order, which comes off an unordered_map: Qwen3-0.6B reports layers + # 27..13 then 0..12 for a cache whose buffers run 0..27. The view is + # defined by the buffers alone, so consulting `layer_ids` for order must + # not creep back in -- it once did, and sent every connector run + # without a `register_kv_cache_layout` override to NotImplementedError. + layout = self._pool_layout( + num_layers=3, kv_factor=2, block_bytes=64, num_slots=5, group_layer_ids=(2, 0, 1) + ) + backing = torch.zeros(5 * 4096 // 2, dtype=torch.float16) + flat = backing.as_strided((5, 3 * 2 * 32), (4096 // 2, 1)) + with unittest.mock.patch.object(KvCacheRegion, "as_tensor", return_value=flat): + view = layout.as_single_pool_tensor() + self.assertIsNotNone(view) + self.assertEqual(tuple(view.shape), (5, 3, 2, 32)) + + def test_single_pool_view_declines_what_it_cannot_describe(self): + # Each of these is a real cache shape; the point is that the default + # registration path refuses rather than mislabelling the bytes. + two_groups = KvCacheLayout( + tokens_per_block=8, + groups=( + KvCacheLayerGroupLayout(0, (0,), None, ()), + KvCacheLayerGroupLayout(1, (1,), 128, ()), + ), + ) + self.assertIsNone(two_groups.as_single_pool_tensor()) + + region = self._pool_layout().groups[0].regions[0] + two_regions = KvCacheLayout( + tokens_per_block=8, groups=(KvCacheLayerGroupLayout(0, (0, 1), None, (region, region)),) + ) + self.assertIsNone(two_regions.as_single_pool_tensor()) + + # Buffer order is deliberately not a decline case. A coalesced buffer + # is laid out layer-major and ascending when the storage config walks + # `config.layers`, so dimension 1 is layer-indexed and there is nothing + # here to re-derive. + + # A page expansion factor breaks the uniform grid. + self.assertIsNone(self._pool_layout(expansion=2).as_single_pool_tensor()) + + def test_layout_lookup_by_group_and_layer(self): + group_a = KvCacheLayerGroupLayout(0, (0, 2), None, ()) + group_b = KvCacheLayerGroupLayout(1, (1, 3), 128, ()) + layout = KvCacheLayout(tokens_per_block=8, groups=(group_a, group_b)) + + self.assertIs(layout.group(1), group_b) + self.assertIs(layout.group_of_layer(2), group_a) + self.assertIs(layout.group_of_layer(3), group_b) + with self.assertRaises(KeyError): + layout.group(7) + with self.assertRaises(KeyError): + layout.group_of_layer(99) + + +class TestValidPageSlots(unittest.TestCase): + """The filter a connector builds transfer targets through. No GPU required.""" + + def test_preserves_the_original_ordinal(self): + # The ordinal is what maps a page back to its token range, so the + # filter must report the position in the input list, not a position + # in its own output. Compacting here would silently re-point every + # surviving block at an earlier token range. + page_indices = [7, BAD_PAGE_INDEX, 9, BAD_PAGE_INDEX, BAD_PAGE_INDEX, 3] + + self.assertEqual(list(valid_page_slots(page_indices)), [(0, 7), (2, 9), (5, 3)]) + + def test_drops_every_entry_without_a_page(self): + page_indices = [BAD_PAGE_INDEX, 4, BAD_PAGE_INDEX, 0, BAD_PAGE_INDEX] + + yielded = list(valid_page_slots(page_indices)) + + self.assertEqual([slot for _, slot in yielded], [4, 0]) + self.assertTrue( + all(slot >= 0 for _, slot in yielded), + "an entry that addresses no page survived the filter", + ) + + def test_slot_zero_is_a_page_and_survives(self): + # `if slot:` instead of `if slot >= 0:` would drop page slot 0, which is + # a real page and usually the first one a fresh pool hands out. + self.assertEqual(list(valid_page_slots([0])), [(0, 0)]) + + def test_a_list_with_no_pages_yields_nothing(self): + self.assertEqual(list(valid_page_slots([BAD_PAGE_INDEX] * 4)), []) + self.assertEqual(list(valid_page_slots([])), []) + + +# The classes below call `torch.cuda.init()`. `unittest/_torch/executor` is +# listed in the GPU-less l0_cpu stage, so the requirement is declared rather +# than left to fail there. It is per class, not module-level: the arithmetic +# and filter tests above are pure Python and must keep running. +@unittest.skipUnless(torch.cuda.is_available(), "allocates real device memory") +class TestKvCacheRegionAliasing(unittest.TestCase): + """as_tensor must alias the exact bytes address_of names.""" + + def setUp(self): + torch.cuda.init() + + def test_as_tensor_aliases_strided_slots(self): + # Lay out 4 "slots" of 32 bytes each, and describe the middle 8 bytes + # of every slot as a region. Writing through the view must land at + # base + stride * i, and must not disturb neighbouring bytes. + num_slots, stride, offset, size = 4, 32, 8, 8 + backing = torch.zeros(num_slots * stride, dtype=torch.uint8, device="cuda") + + region = KvCacheRegion( + base=backing.data_ptr() + offset, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=0, role="key"),), + ) + view = region.as_tensor() + self.assertEqual(tuple(view.shape), (num_slots, size)) + + for slot in range(num_slots): + view[slot] = slot + 1 + + flat = backing.cpu() + for slot in range(num_slots): + start = slot * stride + self.assertTrue( + bool((flat[start + offset : start + offset + size] == slot + 1).all()), + f"slot {slot} payload not written at the address address_of() names", + ) + # Bytes outside the described range must be untouched. + self.assertTrue(bool((flat[start : start + offset] == 0).all())) + self.assertTrue(bool((flat[start + offset + size : start + stride] == 0).all())) + + def test_slot_tensor_aliases_the_row_as_tensor_names(self): + # Same bytes, not a copy: a write through one form is visible through + # the other, and the pointer is the one address_of names. + num_slots, stride, offset, size = 4, 32, 8, 8 + backing = torch.zeros(num_slots * stride, dtype=torch.uint8, device="cuda") + region = KvCacheRegion( + base=backing.data_ptr() + offset, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=0, role="key"),), + ) + + for slot in range(num_slots): + region.slot_tensor(slot).fill_(slot + 1) + + view = region.as_tensor() + for slot in range(num_slots): + self.assertEqual(region.slot_tensor(slot).data_ptr(), region.address_of(slot)) + self.assertTrue( + bool((view[slot] == slot + 1).all()), + f"slot_tensor({slot}) did not write the bytes as_tensor()[{slot}] reads", + ) + + def test_guarded_transfer_never_reaches_the_page_bad_page_index_names(self): + """The corruption the guarded path prevents, asserted from both sides. + + A page-index list reports ``BAD_PAGE_INDEX`` for a block with no page. + Handed to a strided view that entry is a legal subscript naming the + pool's *last* slot -- a live page holding another request's KV. This + test first shows that hazard is live for the unguarded form, so the + guarded assertions that follow are not vacuous. + """ + num_slots, stride, size = 6, 16, 16 + backing = torch.zeros(num_slots * stride, dtype=torch.uint8, device="cuda") + region = KvCacheRegion( + base=backing.data_ptr(), + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=0, role="key"),), + ) + + # A distinct payload per page slot; the last one is what -1 names. + view = region.as_tensor() + for slot in range(num_slots): + view[slot] = 10 + slot + poison = 10 + num_slots - 1 + + # What a connector is handed for a sequence whose window has moved on: + # ordinals 0..2 hold no page, 3..5 hold pages 0..2. + page_indices = [BAD_PAGE_INDEX, BAD_PAGE_INDEX, BAD_PAGE_INDEX, 0, 1, 2] + + # The hazard, live. Without this the guarded assertion proves nothing. + unguarded = [int(view[slot][0].item()) for slot in page_indices] + self.assertEqual( + unguarded[:3], + [poison] * 3, + "indexing the raw list no longer reads the pool's last page, so the " + "guarded assertions below would pass for the wrong reason", + ) + + # Addressing the same entry through the guard raises instead. + with self.assertRaises(IndexError): + region.slot_tensor(page_indices[0]) + + # And the recommended shape never forms that address at all. + saved = [ + (ordinal, int(region.slot_tensor(slot)[0].item())) + for ordinal, slot in valid_page_slots(page_indices) + ] + self.assertEqual( + saved, + [(3, 10), (4, 11), (5, 12)], + "the guarded loop must transfer exactly the blocks that hold a page, " + "under their original ordinals", + ) + self.assertTrue( + all(payload != poison for _, payload in saved), + f"a transfer read the page BAD_PAGE_INDEX names: {saved}", + ) + + +@unittest.skipUnless(torch.cuda.is_available(), "allocates real device memory") +class TestBuildKvCacheLayoutV2(unittest.TestCase): + """The builder against a real cache manager.""" + + def setUp(self): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + + def tearDown(self): + gc.collect() + torch.cuda.empty_cache() + + def test_layout_covers_every_layer_exactly_once(self): + num_layers = 4 + mgr = KVCacheManagerV2(**_make_kwargs(num_layers=num_layers)) + try: + layout = build_kv_cache_layout_v2(mgr) + + self.assertEqual(layout.tokens_per_block, mgr.tokens_per_block) + self.assertTrue(layout.groups, "layout must describe at least one layer group") + + covered = [lid for group in layout.groups for lid in group.layer_ids] + self.assertCountEqual( + covered, + list(mgr.pp_layers), + "every local layer must appear in exactly one layer group", + ) + + # Every layer that owns storage must be reachable through a region. + in_regions = [ + ref.layer_id + for group in layout.groups + for region in group.regions + for ref in region.buffers + ] + self.assertCountEqual(set(in_regions), set(covered)) + finally: + mgr.shutdown() + del mgr + + def test_regions_are_disjoint_and_inside_the_slot(self): + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + pool_groups = list(mgr.impl.pool_group_descs) + self.assertEqual(len(pool_groups), 1, "test config should yield one pool group") + pool_group = pool_groups[0] + self.assertEqual(len(pool_group.pools), 1, "test config should yield one pool") + pool = pool_group.pools[0] + + for group in layout.groups: + spans = [] + for region in group.regions: + self.assertEqual(region.num_slots, int(pool_group.num_slots)) + self.assertEqual(region.stride, int(pool.slot_bytes)) + + offset = region.base - int(pool.base_address) + self.assertGreaterEqual(offset, 0) + self.assertLessEqual( + offset + region.size, + int(pool.slot_bytes), + "a region must lie inside one slot", + ) + spans.append((offset, offset + region.size)) + + spans.sort() + for (_, prev_end), (next_start, _) in zip(spans, spans[1:]): + self.assertLessEqual(prev_end, next_start, "regions must not overlap") + finally: + mgr.shutdown() + del mgr + + def test_region_addresses_agree_with_pool_descriptor(self): + # Cross-check: region base/stride come from get_aggregated_pages, while + # pool base_address/slot_bytes come from pool_group_descs. These are + # independent public APIs and must agree on where slot i lives. + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + pool = list(mgr.impl.pool_group_descs)[0].pools[0] + pool_base, slot_bytes = int(pool.base_address), int(pool.slot_bytes) + + for group in layout.groups: + for region in group.regions: + offset = region.base - pool_base + for slot in (0, 1, region.num_slots - 1): + self.assertEqual( + region.address_of(slot), + pool_base + slot_bytes * slot + offset, + ) + finally: + mgr.shutdown() + del mgr + + def test_uniform_model_yields_one_full_slot_region(self): + # With uniform K/V sizes every buffer in a layer group is adjacent, so + # coalescing should collapse them into a single region spanning the + # whole slot -- the efficient whole-page transfer, derived rather than + # assumed. + mgr = KVCacheManagerV2(**_make_kwargs(num_layers=4)) + try: + layout = build_kv_cache_layout_v2(mgr) + pool = list(mgr.impl.pool_group_descs)[0].pools[0] + + self.assertEqual(len(layout.groups), 1) + group = layout.groups[0] + self.assertEqual(len(group.regions), 1) + region = group.regions[0] + self.assertEqual(region.size, int(pool.slot_bytes)) + self.assertEqual(region.base, int(pool.base_address)) + # 4 layers * (K, V) + self.assertEqual(len(region.buffers), 8) + self.assertEqual( + [ref.role for ref in region.buffers], + ["key", "value"] * 4, + ) + finally: + mgr.shutdown() + del mgr + + def test_full_attention_reports_no_window(self): + mgr = KVCacheManagerV2(**_make_kwargs()) + try: + layout = build_kv_cache_layout_v2(mgr) + for group in layout.groups: + self.assertIsNone(group.window_size) + finally: + mgr.shutdown() + del mgr + + def test_mla_layout_has_no_value_buffers(self): + # SELFKONLY carries a single compressed latent per token. Nothing in the + # layout counts K against V, so this must simply come out as a layout + # with no "value" role rather than needing a kv-factor special case. + mgr = KVCacheManagerV2(**_make_kwargs(kv_cache_type=CacheType.SELFKONLY)) + try: + layout = build_kv_cache_layout_v2(mgr) + roles = { + ref.role + for group in layout.groups + for region in group.regions + for ref in region.buffers + } + self.assertIn("key", roles) + self.assertNotIn("value", roles) + finally: + mgr.shutdown() + del mgr + + +@unittest.skipUnless(torch.cuda.is_available(), "allocates real device memory") +class TestBuildKvCacheLayoutV2Vswa(unittest.TestCase): + """The builder against a real cache manager with more than one window. + + This is the case the single-tensor registration cannot describe, and the + case every per-layer-group callback exists for. The single-group tests + above cannot catch a builder that mixes up which pool backs which group, + because there is only one pool to get right. + """ + + #: Two distinct windows over 4 layers: the cyclic assignment gives + #: layers 0, 2 the sliding window and layers 1, 3 full attention (``max_attention_window`` entries equal to + #: ``max_seq_len`` normalize to ``None``). + WINDOW = 64 + MAX_SEQ_LEN = 256 + NUM_LAYERS = 4 + + def setUp(self): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + + def tearDown(self): + gc.collect() + torch.cuda.empty_cache() + + def _vswa_manager(self): + return KVCacheManagerV2( + **_make_kwargs( + num_layers=self.NUM_LAYERS, + max_seq_len=self.MAX_SEQ_LEN, + kv_cache_config=KvCacheConfigV2( + max_tokens=2048, + enable_block_reuse=False, + max_attention_window=[self.WINDOW, self.MAX_SEQ_LEN], + ), + ) + ) + + def _expected_window(self, layer_id: int): + return None if layer_id % 2 else self.WINDOW + + def test_vswa_yields_one_group_per_window(self): + mgr = self._vswa_manager() + try: + layout = build_kv_cache_layout_v2(mgr) + + self.assertEqual( + len(layout.groups), + 2, + "two distinct windows must produce two layer groups, or every " + "assertion below is a single-group test in disguise", + ) + self.assertEqual( + {group.window_size for group in layout.groups}, + {self.WINDOW, None}, + "one group must be the sliding window and one full attention", + ) + + for group in layout.groups: + for layer_id in group.layer_ids: + self.assertEqual( + group.window_size, + self._expected_window(layer_id), + f"layer {layer_id} landed in a group whose window is {group.window_size}", + ) + + covered = [lid for group in layout.groups for lid in group.layer_ids] + self.assertCountEqual(covered, list(mgr.pp_layers)) + finally: + mgr.shutdown() + del mgr + + def test_group_of_layer_resolves_to_that_layer_s_window(self): + # The per-layer connector hooks receive a model layer index and have + # nothing else to route on. If this mapping is wrong, a sliding layer + # waits on the full-attention group's pages. + mgr = self._vswa_manager() + try: + layout = build_kv_cache_layout_v2(mgr) + for layer_id in mgr.pp_layers: + group = layout.group_of_layer(int(layer_id)) + self.assertEqual(group.window_size, self._expected_window(int(layer_id))) + self.assertIn(int(layer_id), group.layer_ids) + finally: + mgr.shutdown() + del mgr + + def test_a_page_index_means_nothing_without_its_layer_group(self): + """Why the flat block-id list cannot exist under VSWA. + + When two layer groups hold buffers of different sizes they are backed by + different pool groups, each with its own slot address space starting at + zero. Both groups then report page index 0, and the two index 0s are + different pools. A connector handed one flat list could not tell them + apart, and would write one group's KV over the other's. + """ + # Heterogeneous KV head counts make the per-layer buffers differ in + # size, which is what splits the two windows across two pool groups. + mgr = KVCacheManagerV2( + **_make_kwargs( + num_layers=self.NUM_LAYERS, + num_kv_heads=[4, 8, 4, 8], + max_seq_len=self.MAX_SEQ_LEN, + kv_cache_config=KvCacheConfigV2( + max_tokens=2048, + enable_block_reuse=False, + max_attention_window=[self.WINDOW, self.MAX_SEQ_LEN], + ), + ) + ) + try: + layout = build_kv_cache_layout_v2(mgr) + self.assertEqual(len(layout.groups), 2) + + pool_groups = list(mgr.impl.pool_group_descs) + self.assertEqual( + len(pool_groups), + 2, + "layers of differing size must land in separate pool groups, or " + "this test is the shared-slot-space case instead", + ) + + bases = { + group.layer_group_id: {region.base for region in group.regions} + for group in layout.groups + } + self.assertTrue( + bases[0].isdisjoint(bases[1]), + f"the two layer groups report the same pool base: {bases}", + ) + + # The same page index in both groups, addressing different memory, + # is the whole reason the callbacks are per layer group. + for slot in (0, 1): + addresses = { + layout.group(group_id).regions[0].address_of(slot) for group_id in (0, 1) + } + self.assertEqual( + len(addresses), + 2, + f"page index {slot} resolves to one address across two " + f"layer groups: {addresses}", + ) + finally: + mgr.shutdown() + del mgr + + def test_filter_agrees_with_the_manager_valid_only_filter(self): + """One definition of "holds a page", applied on both sides. + + The cache drops entries at the source with ``valid_only=True``. A + connector cannot, because its list stays aligned to block ordinals, so + it filters downstream with ``valid_page_slots`` instead. The two must + select the same pages, or a connector is filtering against a different + rule than the cache allocates by. + """ + mgr = self._vswa_manager() + try: + prompt_len, tokens_per_block = 120, 8 + request = _make_request(prompt_len=prompt_len) + self.assertTrue(mgr.prepare_context(request)) + mgr.resize_context(request, request.context_remaining_length) + batch = ScheduledRequests() + batch.append_context_request(request) + mgr.prepare_resources(batch) + + # Advance the sequence past the sliding window, which is the state a + # connector is reported at `request_finished`: the window has moved + # on, so the early ordinals hold no readable KV. Without this the + # lists below have nothing for the filter to remove. + kv_cache = mgr.kv_cache_map[request.py_request_id] + kv_cache.history_length = prompt_len + stale_end = max(0, (prompt_len + 1 - self.WINDOW) // tokens_per_block) + self.assertGreater( + stale_end, + 0, + "the sizes no longer carry the window past a whole block, so " + "nothing here exercises a block without a page", + ) + + for group_id in range(len(mgr.impl.layer_grouping)): + every = list(kv_cache.get_aggregated_page_indices(group_id, valid_only=False)) + only_valid = list(kv_cache.get_aggregated_page_indices(group_id, valid_only=True)) + self.assertEqual( + [slot for _, slot in valid_page_slots(every)], + only_valid, + f"group {group_id}: valid_page_slots and valid_only=True " + f"disagree on which blocks hold a page", + ) + + # The same filter over the list a connector is actually handed, which + # carries the out-of-window mask on top of the missing-page entries. + by_group = mgr.get_page_indices_by_layer_group(request) + gaps = sum(1 for indices in by_group for slot in indices if slot < 0) + self.assertGreaterEqual( + gaps, + stale_end, + f"the sliding group must report at least {stale_end} ordinals " + f"without a page, so the filter has something to remove: {by_group}", + ) + for group_id, indices in enumerate(by_group): + allocated = set(kv_cache.get_aggregated_page_indices(group_id, valid_only=True)) + for ordinal, slot in valid_page_slots(indices): + self.assertEqual( + indices[ordinal], + slot, + f"group {group_id}: ordinal {ordinal} does not index " + f"back to the slot reported for it", + ) + self.assertIn( + slot, + allocated, + f"group {group_id}: page slot {slot} survived the filter " + f"but the cache does not hold it for this request", + ) + finally: + mgr.shutdown() + del mgr + + def test_live_pages_never_overlap_in_memory(self): + """The invariant a connector transfers against, in the shared case. + + When both layer groups draw from one pool group they share a slot + address space, so the same region base describes both and the page + indices handed out are disjoint instead. Either way, no two live + ``(layer group, page slot)`` pairs name the same bytes -- which is what + makes a save keyed on that pair safe. + """ + mgr = self._vswa_manager() + try: + layout = build_kv_cache_layout_v2(mgr) + request = _make_request(prompt_len=120) + self.assertTrue(mgr.prepare_context(request)) + mgr.resize_context(request, request.context_remaining_length) + batch = ScheduledRequests() + batch.append_context_request(request) + mgr.prepare_resources(batch) + + by_group = mgr.get_page_indices_by_layer_group(request) + self.assertEqual(len(by_group), 2) + + spans = [] + for layer_group_id, indices in enumerate(by_group): + group = layout.group(layer_group_id) + for slot in indices: + if slot == BAD_PAGE_INDEX: + continue + for region in group.regions: + start = region.address_of(slot) + spans.append((start, start + region.size, layer_group_id, slot)) + self.assertTrue(spans, "the request holds no live pages") + + spans.sort() + for previous, current in zip(spans, spans[1:]): + self.assertLessEqual( + previous[1], + current[0], + f"live page (group {previous[2]}, slot {previous[3]}) at " + f"[{previous[0]}, {previous[1]}) overlaps (group " + f"{current[2]}, slot {current[3]}) at [{current[0]}, " + f"{current[1]})", + ) + finally: + mgr.shutdown() + del mgr + + def test_vswa_layout_declines_the_single_pool_view(self): + # The back-compat shim must refuse here rather than reconstruct a + # tensor that silently covers one group. `register_kv_cache_layout`'s + # default turns this None into the startup error that names the method + # a VSWA connector has to implement. + mgr = self._vswa_manager() + try: + layout = build_kv_cache_layout_v2(mgr) + self.assertIsNone(layout.as_single_pool_tensor()) + finally: + mgr.shutdown() + del mgr + + def test_each_group_addresses_the_pool_that_backs_it(self): + # Cross-check against pool_group_descs, which names the layer groups a + # pool group serves through slot_desc.variants. region base/stride come + # from get_aggregated_pages; the two APIs must agree per group, not + # just in aggregate. + mgr = self._vswa_manager() + try: + layout = build_kv_cache_layout_v2(mgr) + + pools_by_group = {} + slots_by_group = {} + for pool_group in mgr.impl.pool_group_descs: + for variant in pool_group.slot_desc.variants: + group_id = int(variant.layer_group_id) + pools_by_group.setdefault(group_id, []).extend(pool_group.pools) + slots_by_group[group_id] = int(pool_group.num_slots) + + self.assertEqual( + set(pools_by_group), + {group.layer_group_id for group in layout.groups}, + "every layer group in the layout must be backed by a pool group", + ) + + for group in layout.groups: + bases = {int(pool.base_address) for pool in pools_by_group[group.layer_group_id]} + slot_bytes = {int(pool.slot_bytes) for pool in pools_by_group[group.layer_group_id]} + for region in group.regions: + self.assertEqual(region.num_slots, slots_by_group[group.layer_group_id]) + self.assertIn( + region.stride, + slot_bytes, + f"layer group {group.layer_group_id} region stride " + f"{region.stride} matches no pool backing it", + ) + offsets = [region.base - base for base in bases] + self.assertTrue( + any(0 <= offset < region.stride for offset in offsets), + f"layer group {group.layer_group_id} region base " + f"{region.base} lies in no pool backing it (pool bases " + f"{sorted(bases)})", + ) + finally: + mgr.shutdown() + del mgr + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py new file mode 100644 index 000000000000..f35e33dd6e6a --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py @@ -0,0 +1,796 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the KV connector prefix. + +The connector is asked in ``prepare_resources``, on the batch the forward pass +will run, which is downstream of every stage that can drop a request. The +*asked => scheduled => eventually request_finished* invariant therefore holds, +an offer is never abandoned, and no ``cancel_load`` is needed. + +Two properties of the allocation are what these tests pin. + +* **Pages are allocated per context chunk**, deliberately -- that is what + chunked prefill is for -- so an offer reaching beyond the chunk needs a + bounded grow, and the grow can fail. +* **The local match is token-granular**: ``num_committed_tokens`` is not + floored to whole shared blocks, so the arithmetic can go negative. + +``FakeRequest`` reproduces ``LlmRequest``'s chunk arithmetic including +``setContextChunkSize``'s non-negative check and ``setPrepopulatedPromptLen``'s +block-alignment assertion, so a version of this code that violates either fails +here rather than only on hardware. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + +TOKENS_PER_BLOCK = 32 +PROMPT_LEN = 256 + + +class FakeKvCache: + """The parts of ``_KVCache`` the connector path touches. + + ``resize`` reproduces the real invariants -- history never decreases, + history never exceeds capacity -- and ``grow_ok`` models + ``OutOfPagesError``, which the real ``resize`` reports by returning False + after restoring the state it unlocked. + """ + + def __init__(self, committed=0, capacity=None, grow_ok=True): + self.num_committed_tokens = committed + self.capacity = committed if capacity is None else capacity + self.history_length = committed + self.enable_swa_scratch_reuse = True + self.is_active = True + self.grow_ok = grow_ok + self.resize_calls = [] + + def resume(self, cuda_stream): + self.is_active = True + return True + + def suspend(self): + self.is_active = False + + def resize(self, capacity, history_length=None): + assert self.is_active, "resize on a suspended cache" + self.resize_calls.append((capacity, history_length)) + growing = capacity is not None and capacity > self.capacity + if growing and not self.grow_ok: + return False + if history_length is not None: + if history_length < self.history_length: + raise ValueError("History length cannot be decreased") + if capacity is not None and capacity < history_length: + raise ValueError("History length cannot be greater than capacity") + self.history_length = history_length + if capacity is not None: + self.capacity = capacity + return True + + +class FakeRequest: + """``LlmRequest``'s context-chunk arithmetic, reproduced faithfully. + + The properties mirror ``llmRequest.h``: ``mContextChunkSize`` defaults to + ``mPromptLen``, ``setContextChunkSize`` rejects a negative size and clamps + to the remaining length, ``isFirstContextChunk`` is + ``contextCurrentPosition == prepopulatedPromptLen``, and + ``setPrepopulatedPromptLen`` floors a non-final chunk's end to a block + boundary and then asserts it. + """ + + def __init__(self, request_id=0, prompt_len=PROMPT_LEN): + self.py_request_id = request_id + self.request_id = request_id + self.prompt_len = prompt_len + self.context_current_position = 0 + self.prepopulated_prompt_len = 0 + self.is_dummy = False + self.is_generation_only_request = False + self.is_disagg_generation_init_state = False + self.py_num_connector_matched_tokens = 0 + self.py_connector_allocation_reported = False + self.py_connector_served_position = 0 + self.py_ctx_pre_resize_cap = None + self.py_draft_tokens = [] + self._context_chunk_size = prompt_len + + @property + def context_remaining_length(self): + return self.prompt_len - self.context_current_position + + @property + def context_chunk_size(self): + return self._context_chunk_size + + @context_chunk_size.setter + def context_chunk_size(self, size): + assert size >= 0, f"The chunk size of context ({size}) can't be negative." + self._context_chunk_size = min(size, self.context_remaining_length) + + @property + def is_first_context_chunk(self): + return self.context_current_position == self.prepopulated_prompt_len + + @property + def is_last_context_chunk(self): + return self.context_current_position + self.context_chunk_size == self.prompt_len + + def set_prepopulated_prompt_len(self, prepopulated_prompt_len, tokens_per_block): + assert prepopulated_prompt_len < self.prompt_len, ( + f"prepopulatedPromptLen ({prepopulated_prompt_len}) >= promptLen ({self.prompt_len})" + ) + self.prepopulated_prompt_len = prepopulated_prompt_len + if prepopulated_prompt_len > 0: + chunk_size = self.context_chunk_size + if prepopulated_prompt_len + chunk_size < self.prompt_len: + floored = ( + (prepopulated_prompt_len + chunk_size) // tokens_per_block * tokens_per_block + ) + chunk_size = floored - prepopulated_prompt_len + self.context_current_position = prepopulated_prompt_len + self.context_chunk_size = chunk_size + if not self.is_last_context_chunk: + assert ( + self.context_current_position + self.context_chunk_size + ) % tokens_per_block == 0, ( + "the context position after the current chunk must be block-aligned" + ) + + +class FakeConnectorManager: + """Records the calls the prefix path makes, in order.""" + + def __init__(self, num_matched=0, load_async=False, add_sequence=True): + self.num_matched = num_matched + self.load_async = load_async + self.add_sequence = add_sequence + self.queries = [] + self.commits = [] + self.allocs = [] + self.alloc_by_group = [] + self.forgotten = [] + + def query_num_new_matched_tokens(self, request, num_computed_tokens): + self.queries.append((request.request_id, num_computed_tokens)) + return self.num_matched, self.load_async + + def commit_new_matched_tokens(self, request, num_tokens, load_kv_async): + self.commits.append((request.request_id, num_tokens, load_kv_async)) + request.py_num_connector_matched_tokens = num_tokens + + def should_add_sequence(self, request): + return self.add_sequence + + def reset_request_state(self, request): + self.forgotten.append(request.request_id) + + def update_state_after_alloc(self, request, page_indices, by_layer_group=None): + self.allocs.append((request.request_id, tuple(page_indices))) + self.alloc_by_group.append(by_layer_group) + + def build_scheduler_output(self, scheduled_batch, kv_cache_manager): + pass + + +def make_manager(connector, num_extra_kv_tokens=0, is_draft=False): + """A cache manager carrying only the fields the prefix path reads. + + Constructing a real one needs a GPU and a pool allocation; what is under + test is request/cache bookkeeping, so bypass ``__init__`` rather than + turning this into an integration test. ``test_kv_connector_v2_prefix_real_manager`` + covers the same path against real pools. + """ + manager = object.__new__(KVCacheManagerV2) + manager.kv_connector_manager = connector + manager.is_draft = is_draft + manager.tokens_per_block = TOKENS_PER_BLOCK + manager.num_extra_kv_tokens = num_extra_kv_tokens + manager.kv_cache_map = {} + manager.enable_block_reuse = True + manager.conversation_manager = None + manager._stream = SimpleNamespace(cuda_stream=0) + # One layer group, the shape every non-VSWA, non-hybrid model has. The real + # accessor reads `impl.layer_grouping`, which only a pool allocation fills + # in, so it is stubbed here rather than bypassed -- `_run_kv_connector_hooks` + # derives the flat list from this and hands the connector both forms. + manager.get_page_indices_by_layer_group = lambda request: [[]] + return manager + + +def scheduled(*requests): + batch = SimpleNamespace(context_requests=list(requests), reset_calls=0) + + def reset_context_requests(): + batch.reset_calls += 1 + + batch.reset_context_requests = reset_context_requests + return batch + + +def serve(manager, req, kv_cache): + """Drive the real ``_apply_connector_matched_prefix`` for one scheduled request.""" + manager.kv_cache_map[req.py_request_id] = kv_cache + req.context_current_position = kv_cache.num_committed_tokens + req.set_prepopulated_prompt_len(kv_cache.num_committed_tokens, TOKENS_PER_BLOCK) + return manager._apply_connector_matched_prefix(req) + + +class TestAskTiming: + """The ask happens on the final batch, and only there. + + This is the whole of A0: the query is downstream of every drop, which is + what guarantees *asked => scheduled*. + """ + + def test_the_scheduling_pass_does_not_ask(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=64, capacity=PROMPT_LEN) + manager.kv_cache_map[req.py_request_id] = kv_cache + + assert manager.prepare_context(req) + + assert connector.queries == [] + assert req.context_current_position == 64 + + def test_the_connector_is_asked_once_per_allocation(self): + """Serving leaves ``is_first_context_chunk`` true, by design -- it is + ``context_current_position == prepopulated_prompt_len`` and + ``set_prepopulated_prompt_len`` moves both. So the first-chunk test + cannot be the ask-once guard, and a re-entry before the forward pass + would otherwise take remote ownership twice for one allocation.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=PROMPT_LEN) + manager.kv_cache_map[req.py_request_id] = kv_cache + + manager._run_kv_connector_hooks(scheduled(req)) + assert req.is_first_context_chunk + manager._run_kv_connector_hooks(scheduled(req)) + + assert len(connector.queries) == 1 + + def test_a_second_pass_reports_one_allocation(self): + """``should_add_sequence`` is not enough on its own. + + It only goes false once an asynchronous load has *completed*. A batch + dropped between ``prepare_resources`` and the forward pass -- the second + ``_can_queue``, which a parked request can flip -- brings every context + request back on its first chunk with the predicate still true. The ask + is idempotent by memo; a second ``update_state_after_alloc`` would + report the same pages twice. + """ + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + + manager._run_kv_connector_hooks(scheduled(req)) + manager._run_kv_connector_hooks(scheduled(req)) + + assert len(connector.queries) == 1 + assert len(connector.allocs) == 1 + + def test_an_excluded_request_is_still_reported_once(self): + """The allocation is reported for every first-chunk sequence added, + including the kinds the connector is never asked about.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + req.is_disagg_generation_init_state = True + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + + manager._run_kv_connector_hooks(scheduled(req)) + manager._run_kv_connector_hooks(scheduled(req)) + + assert connector.queries == [] + assert len(connector.allocs) == 1 + + def test_a_destroyed_allocation_is_asked_again(self): + """A destructive pause replays the sequence add and with it the query, + because the pages the first answer described are gone.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + + manager._run_kv_connector_hooks(scheduled(req)) + req.py_connector_allocation_reported = False # what free_resources does + req.context_current_position = 0 + req.prepopulated_prompt_len = 0 + manager._run_kv_connector_hooks(scheduled(req)) + + assert len(connector.queries) == 2 + + def test_a_completed_async_load_is_not_asked_again(self): + """``should_add_sequence`` is the re-entry gate. + + A request whose asynchronous load finished re-enters the batch still on + its first context chunk, with the same pages and nothing left to load. + Asking again would take ownership twice for one allocation. + """ + connector = FakeConnectorManager(num_matched=64, add_sequence=False) + manager = make_manager(connector) + req = FakeRequest() + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + + manager._run_kv_connector_hooks(scheduled(req)) + + assert connector.queries == [] + assert connector.allocs == [] + + +class TestOfferIsClamped: + def test_the_last_prompt_token_stays_local(self): + """The first generation step consumes its activations, so it must be computed.""" + connector = FakeConnectorManager(num_matched=PROMPT_LEN) + manager = make_manager(connector) + req = FakeRequest() + + assert serve(manager, req, FakeKvCache(committed=0, capacity=PROMPT_LEN)) + + assert req.context_current_position == PROMPT_LEN - 1 + assert req.context_chunk_size == 1 + assert connector.commits == [(0, PROMPT_LEN - 1, False)] + + def test_the_query_is_anchored_at_the_local_match(self): + """The connector is asked what it can serve *past* what the radix tree + already holds, and the offer is an extent measured from there.""" + connector = FakeConnectorManager(num_matched=16) + manager = make_manager(connector) + req = FakeRequest() + + assert serve(manager, req, FakeKvCache(committed=128, capacity=PROMPT_LEN)) + + assert connector.queries == [(0, 128)] + assert req.context_current_position == 144 + assert connector.commits == [(0, 16, False)] + + def test_an_empty_offer_touches_nothing(self): + connector = FakeConnectorManager(num_matched=0) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=64, capacity=PROMPT_LEN) + + assert not serve(manager, req, kv_cache) + + assert req.context_current_position == 64 + assert kv_cache.resize_calls == [] + assert connector.commits == [(0, 0, False)] + + +class TestUnchunkedIsPureShrink: + """With chunking off, ``resize_context`` already covered the whole prompt. + + The offer is inside that by construction, so the chunk's end stays where + the scheduler put it, its start moves up, and no page is allocated. + """ + + def test_the_chunk_end_does_not_move(self): + connector = FakeConnectorManager(num_matched=96) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=32, capacity=PROMPT_LEN) + + assert serve(manager, req, kv_cache) + + assert req.context_current_position == 128 + assert req.context_current_position + req.context_chunk_size == PROMPT_LEN + assert req.context_chunk_size == PROMPT_LEN - 128 + + def test_no_capacity_is_allocated(self): + connector = FakeConnectorManager(num_matched=96) + manager = make_manager(connector) + kv_cache = FakeKvCache(committed=32, capacity=PROMPT_LEN) + + serve(manager, FakeRequest(), kv_cache) + + assert kv_cache.capacity == PROMPT_LEN + assert kv_cache.resize_calls == [(PROMPT_LEN, 128)] + + def test_history_is_raised_to_the_served_position(self): + """The sole input to the stale-range computation, so a sliding-window + layer group does not keep a page per served block.""" + connector = FakeConnectorManager(num_matched=96) + manager = make_manager(connector) + kv_cache = FakeKvCache(committed=32, capacity=PROMPT_LEN) + + serve(manager, FakeRequest(), kv_cache) + + assert kv_cache.history_length == 128 + + def test_the_request_stays_on_its_first_chunk(self): + connector = FakeConnectorManager(num_matched=96) + manager = make_manager(connector) + req = FakeRequest() + + serve(manager, req, FakeKvCache(committed=32, capacity=PROMPT_LEN)) + + assert req.is_first_context_chunk + + +class TestChunkedInsideTheChunk: + def test_the_scheduler_s_chunk_end_is_preserved(self): + connector = FakeConnectorManager(num_matched=32) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=64) + req.context_chunk_size = 64 # the scheduler's choice this iteration + + assert serve(manager, req, kv_cache) + + assert req.context_current_position == 32 + assert req.context_current_position + req.context_chunk_size == 64 + assert kv_cache.capacity == 64 + + +class TestChunkedBeyondTheChunk: + """Only reachable with chunked prefill, and the only place a grow happens. + + The offer costs pages whatever the query timing -- the connector writes real + KV into them -- so the chunk window shifts forward and the allocation grows + by the offer plus the compute the scheduler budgeted, bounded by one + chunk. + """ + + def test_the_window_shifts_and_the_allocation_grows(self): + connector = FakeConnectorManager(num_matched=128) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=64) + req.context_chunk_size = 64 + + assert serve(manager, req, kv_cache) + + assert req.context_current_position == 128 + assert req.context_chunk_size == 64 + assert kv_cache.capacity == 192 + + def test_a_non_final_chunk_end_stays_block_aligned(self): + """Otherwise the next chunk fragments the cache.""" + connector = FakeConnectorManager(num_matched=100) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=48) + req.context_chunk_size = 48 + + assert serve(manager, req, kv_cache) + + assert req.context_current_position == 100 + end = req.context_current_position + req.context_chunk_size + assert end % TOKENS_PER_BLOCK == 0 + assert end == 128 + + def test_a_failed_grow_falls_back_to_the_pages_that_exist(self): + connector = FakeConnectorManager(num_matched=128) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=64, grow_ok=False) + req.context_chunk_size = 64 + + assert serve(manager, req, kv_cache) + + # The largest whole block inside the 64 tokens the pages already hold + # that still leaves the forward pass a chunk to compute. + assert req.context_current_position == 32 + assert req.context_chunk_size == 32 + assert kv_cache.capacity == 64 + + def test_a_failed_grow_commits_only_what_was_honoured(self): + """Over-reporting would point the connector at an offset it has no page for. + + The unconsumed tail of the offer is recomputed locally and released at + ``request_finished``. + """ + connector = FakeConnectorManager(num_matched=128) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=64, grow_ok=False) + req.context_chunk_size = 64 + + serve(manager, req, kv_cache) + + assert connector.commits == [(0, 32, False)] + + def test_an_offer_that_reaches_the_chunk_end_exactly_shifts(self): + """Honouring it in place would leave a zero-token context chunk.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=64) + req.context_chunk_size = 64 + + assert serve(manager, req, kv_cache) + + assert req.context_current_position == 64 + assert req.context_chunk_size == 64 + assert kv_cache.capacity == 128 + + def test_a_fallback_that_cannot_advance_honours_nothing(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + # One block of pages: the offer needs a grow, the grow fails, and the + # capped fallback lands back on the local match. + kv_cache = FakeKvCache(committed=0, capacity=TOKENS_PER_BLOCK, grow_ok=False) + req.context_chunk_size = TOKENS_PER_BLOCK + + assert not serve(manager, req, kv_cache) + + assert req.context_current_position == 0 + assert connector.commits == [(0, 0, False)] + + def test_a_grow_that_cannot_leave_a_token_to_compute_is_refused(self): + connector = FakeConnectorManager(num_matched=PROMPT_LEN) + manager = make_manager(connector) + req = FakeRequest() + # The whole prompt is already committed but for its last token, so the + # clamped offer lands exactly on the chunk end and nothing is left. + kv_cache = FakeKvCache(committed=PROMPT_LEN - 1, capacity=PROMPT_LEN, grow_ok=False) + req.context_chunk_size = 1 + + assert not serve(manager, req, kv_cache) + + assert req.context_current_position == PROMPT_LEN - 1 + assert connector.commits == [(0, 0, False)] + + +class TestExtraKvTokens: + def test_the_grow_target_reserves_the_extra_tokens(self): + connector = FakeConnectorManager(num_matched=128) + manager = make_manager(connector, num_extra_kv_tokens=4) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=64) + req.context_chunk_size = 64 + + serve(manager, req, kv_cache) + + assert kv_cache.capacity == 192 + 4 + + +class TestExclusions: + @pytest.mark.parametrize( + "attribute, excluded", + [ + ("is_dummy", True), + ("is_generation_only_request", True), + ("is_disagg_generation_init_state", True), + ], + ) + def test_request_kinds_that_are_never_asked(self, attribute, excluded): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + setattr(req, attribute, excluded) + + assert not serve(manager, req, FakeKvCache(committed=0, capacity=PROMPT_LEN)) + + assert connector.queries == [] + assert connector.commits == [] + + def test_no_connector_is_a_no_op(self): + manager = make_manager(None) + req = FakeRequest() + + assert not serve(manager, req, FakeKvCache(committed=0, capacity=PROMPT_LEN)) + + assert req.context_current_position == 0 + + def test_the_draft_manager_never_asks(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector, is_draft=True) + req = FakeRequest() + + assert not serve(manager, req, FakeKvCache(committed=0, capacity=PROMPT_LEN)) + + assert connector.queries == [] + + +class TestAsyncLoad: + def test_the_async_flag_reaches_the_commit(self): + connector = FakeConnectorManager(num_matched=64, load_async=True) + manager = make_manager(connector) + + serve(manager, FakeRequest(), FakeKvCache(committed=0, capacity=PROMPT_LEN)) + + assert connector.commits == [(0, 64, True)] + + def test_the_async_hold_survives_an_under_honoured_offer(self): + """The connector has already started the transfer, so the request must + still be parked even though the runtime took less than it offered.""" + connector = FakeConnectorManager(num_matched=128, load_async=True) + manager = make_manager(connector) + req = FakeRequest() + req.context_chunk_size = 64 + + serve(manager, req, FakeKvCache(committed=0, capacity=64, grow_ok=False)) + + assert connector.commits == [(0, 32, True)] + + def test_the_async_hold_survives_an_offer_honoured_for_nothing(self): + """Zero honoured still parks the request, and the pass runs nothing. + + This is the one pair ``query_num_new_matched_tokens`` rejects on the way + in, and it is deliberate on the way out. The connector started the + transfer as soon as it answered asynchronously -- a parked request never + reaches ``start_load_kv``, so that is the only place it can start -- and + nothing tells it the offer was declined. Committing ``False`` here would + leave the request in the batch and let prefill write the pages the + transfer is landing in. + """ + connector = FakeConnectorManager(num_matched=PROMPT_LEN, load_async=True) + manager = make_manager(connector) + req = FakeRequest() + + served = serve( + manager, req, FakeKvCache(committed=0, capacity=TOKENS_PER_BLOCK, grow_ok=False) + ) + + assert not served + assert connector.commits == [(0, 0, True)] + assert req.context_current_position == 0, ( + "nothing was honoured, so the forward pass must still compute the " + "whole prompt once the load completes" + ) + + +class TestBatchBookkeeping: + def test_a_served_prefix_rebuilds_the_chunking_split(self): + """A shift can carry a request onto its last chunk, which moves it + between ``context_requests_chunking`` and ``context_requests_last_chunk`` + -- the lists ``build_scheduler_output`` walks.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + batch = scheduled(req) + + manager._run_kv_connector_hooks(batch) + + assert batch.reset_calls == 1 + + def test_an_unserved_batch_leaves_the_split_alone(self): + connector = FakeConnectorManager(num_matched=0) + manager = make_manager(connector) + req = FakeRequest() + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + batch = scheduled(req) + + manager._run_kv_connector_hooks(batch) + + assert batch.reset_calls == 0 + + def test_allocation_is_reported_after_the_prefix_is_served(self): + """The pages the connector may write into include the ones the grow + added, so the report has to follow the serve.""" + seen = [] + connector = FakeConnectorManager(num_matched=128) + manager = make_manager(connector) + req = FakeRequest() + req.context_chunk_size = 64 + kv_cache = FakeKvCache(committed=0, capacity=64) + manager.kv_cache_map[req.py_request_id] = kv_cache + manager.get_page_indices_by_layer_group = lambda request: seen.append( + kv_cache.capacity + ) or [[]] + + manager._run_kv_connector_hooks(scheduled(req)) + + assert seen == [192] + # The per-layer-group form is what reaches the connector. Asserting it + # arrives keeps the hook wired to the accessor the manager actually + # implements: overriding only the flat one leaves the real accessor to + # run against a manager with no pools, which raises. + assert connector.alloc_by_group == [[[]]] + + +class TestReEntryAfterAServe: + """A served request that comes back before it runs. + + Two paths reach it. An asynchronous load parks the request out of the batch, + and the scheduler runs ``prepare_context`` on it again when it returns; a + batch dropped after ``prepare_resources`` brings a synchronous serve back + the same way. + """ + + def _serve_then_re_enter(self, load_async=True): + connector = FakeConnectorManager(num_matched=64, load_async=load_async) + manager = make_manager(connector) + req = FakeRequest() + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + assert manager.prepare_context(req) + manager._run_kv_connector_hooks(scheduled(req)) + served_position = req.context_current_position + assert manager.prepare_context(req) + return req, served_position + + def test_the_chunk_still_spans_to_the_end_of_the_prompt(self): + """The position and the chunk have to move together. + + ``prepare_context`` re-derives both. Leaving the chunk narrowed by the + offer makes the pair describe two different ranges: ``position + chunk`` + stops reaching ``prompt_len``, so a non-chunked request silently becomes + a chunked one and the forward pass is handed inconsistent metadata. + """ + req, served_position = self._serve_then_re_enter() + + assert served_position == 64, "the serve must have moved the position" + assert req.context_current_position + req.context_chunk_size == req.prompt_len + assert req.is_last_context_chunk + + def test_the_same_holds_for_a_synchronous_serve(self): + req, _ = self._serve_then_re_enter(load_async=False) + + assert req.context_current_position + req.context_chunk_size == req.prompt_len + + def test_the_served_position_survives_re_entry(self): + """The cache commits the local match alone, so settling on it would put + the request back below the prefix the connector already wrote.""" + req, served_position = self._serve_then_re_enter() + + assert req.context_current_position == served_position + assert req.is_first_context_chunk + + def test_a_dead_allocation_drops_the_floor(self): + """A replay recomputes from its own match: the served position vouches + for pages that died with the allocation.""" + connector = FakeConnectorManager(num_matched=64, load_async=True) + manager = make_manager(connector) + req = FakeRequest() + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + assert manager.prepare_context(req) + manager._run_kv_connector_hooks(scheduled(req)) + assert req.py_connector_served_position == 64 + + req.py_connector_allocation_reported = False # what free_resources does + req.py_connector_served_position = 0 # ... and this + + assert manager.prepare_context(req) + + assert req.context_current_position == 0 + + +class TestSwaScratchReuse: + def test_scratch_reuse_is_disabled_before_the_scheduler_can_take_slots(self): + """The flag has to be cleared in ``prepare_context``: by the time the + connector is asked, ``resize_context`` has already run and may have + taken scratch slots for blocks the connector is about to write.""" + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=PROMPT_LEN) + manager.kv_cache_map[req.py_request_id] = kv_cache + + assert manager.prepare_context(req) + + assert kv_cache.enable_swa_scratch_reuse is False + + def test_scratch_reuse_survives_without_a_connector(self): + manager = make_manager(None) + req = FakeRequest() + kv_cache = FakeKvCache(committed=0, capacity=PROMPT_LEN) + manager.kv_cache_map[req.py_request_id] = kv_cache + + assert manager.prepare_context(req) + + assert kv_cache.enable_swa_scratch_reuse is True diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py new file mode 100644 index 000000000000..db2f83697e23 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py @@ -0,0 +1,636 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KV connector prefix against a *real* cache manager. + +``test_kv_connector_v2_prefix.py`` drives the same code against a stub cache. +That is what makes it fast and exhaustive, and it is the right place for the +arithmetic and the ask-once rules -- but a stub cannot show that a real +``_KVCache`` survives the sequence: that ``resize`` finds real pages for the +offered prefix, that ``history_length`` really moves, that the grow the chunked +path needs succeeds against real pools, and that the page slots handed to the +connector are distinct and real. + +The engine-level suite cannot show the scheduling-order claims either, because +whether a request is dropped after being prepared depends on which pass it +reaches the scheduler in -- a race. Preparation and delivery are therefore +driven directly here: ``prepare_context`` plus ``resize_context`` is one +scheduling pass, and ``prepare_resources`` is the batch actually running. + +These tests allocate device memory pools. +""" + +import gc + +import pytest +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import valid_page_slots +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, AttnLifeCycle + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + +# These build a real manager, which allocates device pools. The directory is +# listed in the GPU-less l0_cpu stage, so the requirement is declared rather +# than left to fail at `torch.cuda.init()`. +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="allocates real KV cache pools" +) + +TOKENS_PER_BLOCK = 32 +PROMPT_LEN = 96 +OFFER_TOKENS = 32 + + +class FakeConnectorManager: + """Records what the prefix path tells the connector, in order.""" + + def __init__(self, num_matched=OFFER_TOKENS, load_async=False): + self.num_matched = num_matched + self.load_async = load_async + self.queries = [] + self.commits = [] + self.allocs = [] + self.allocs_by_group = [] + self.forgotten = [] + + def query_num_new_matched_tokens(self, request, num_computed_tokens): + self.queries.append((request.py_request_id, num_computed_tokens)) + return self.num_matched, self.load_async + + def commit_new_matched_tokens(self, request, num_tokens, load_kv_async): + self.commits.append((request.py_request_id, num_tokens, load_kv_async)) + request.py_num_connector_matched_tokens = num_tokens + + def should_add_sequence(self, request): + return True + + def reset_request_state(self, request): + self.forgotten.append(request.py_request_id) + + def update_state_after_alloc(self, request, block_ids, by_layer_group=None): + self.allocs.append((request.py_request_id, list(block_ids))) + self.allocs_by_group.append( + ( + request.py_request_id, + None if by_layer_group is None else [list(g) for g in by_layer_group], + ) + ) + + def build_scheduler_output(self, scheduled_batch, kv_cache_manager): + pass + + +def make_manager(connector, **overrides): + kwargs = dict( + kv_cache_config=KvCacheConfig(max_tokens=2048, enable_block_reuse=True), + kv_cache_type=CacheType.SELF, + num_layers=2, + num_kv_heads=4, + head_dim=64, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=256, + max_batch_size=4, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.HALF, + vocab_size=32000, + kv_connector_manager=connector, + ) + kwargs.update(overrides) + return KVCacheManagerV2(**kwargs) + + +def make_request(request_id=1, prompt_len=PROMPT_LEN): + return LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=list(range(prompt_len)), + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + + +def schedule(manager, request, num_tokens=None): + """One scheduling pass: prepare the cache and size it for the chunk.""" + assert manager.prepare_context(request) + if num_tokens is None: + num_tokens = request.context_remaining_length + return manager.resize_context(request, num_tokens) + + +def run(manager, *requests): + """One ``prepare_resources``, i.e. the requests reached the final batch.""" + batch = ScheduledRequests() + for request in requests: + batch.append_context_request(request) + manager.prepare_resources(batch) + return batch + + +@pytest.fixture +def connector(): + return FakeConnectorManager() + + +@pytest.fixture +def manager(connector): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + mgr = make_manager(connector) + yield mgr + mgr.shutdown() + del mgr + gc.collect() + torch.cuda.empty_cache() + + +def test_a_request_dropped_before_the_batch_is_never_asked(manager, connector): + """The whole point of asking in ``prepare_resources``. + + A request can be prepared and sized and then lose the token budget, fail + multimodal alignment, or be dropped when the batch cannot be queued. None + of that can strand an offer, because the query comes after all of it. + """ + request = make_request() + + assert schedule(manager, request) + + assert connector.queries == [] + assert connector.commits == [] + assert request.context_current_position == 0 + + +def test_offer_is_backed_by_real_pages(manager, connector): + """Read capacity, history and page slots back off the cache the forward + pass would use -- the difference between "resize was called correctly" and + "the offered prefix is resident".""" + request = make_request() + assert schedule(manager, request) + + run(manager, request) + + kv_cache = manager.kv_cache_map[request.py_request_id] + assert request.context_current_position == OFFER_TOKENS + assert kv_cache.history_length == OFFER_TOKENS + assert kv_cache.capacity >= PROMPT_LEN + assert kv_cache.is_active + + assert connector.commits == [(request.py_request_id, OFFER_TOKENS, False)] + assert len(connector.allocs) == 1 + + _, page_indices = connector.allocs[0] + assert len(page_indices) >= PROMPT_LEN // TOKENS_PER_BLOCK + assert all(index != BAD_PAGE_INDEX for index in page_indices) + assert len(set(page_indices)) == len(page_indices) + + +def test_the_unchunked_path_allocates_nothing_for_the_prefix(manager, connector): + """``resize_context`` already covered the whole prompt, so honouring the + offer only moves the request's start.""" + request = make_request() + assert schedule(manager, request) + kv_cache = manager.kv_cache_map[request.py_request_id] + before = kv_cache.capacity + + run(manager, request) + + assert kv_cache.capacity == before + assert request.context_current_position + request.context_chunk_size == PROMPT_LEN + + +def test_a_chunked_offer_beyond_the_chunk_grows_and_shifts(manager, connector): + """Chunked prefill keeps its per-chunk allocation, so an offer past the + chunk has to grow the cache before it can be honoured.""" + connector.num_matched = 64 + request = make_request() + request.context_chunk_size = TOKENS_PER_BLOCK + assert schedule(manager, request, num_tokens=TOKENS_PER_BLOCK) + kv_cache = manager.kv_cache_map[request.py_request_id] + assert kv_cache.capacity < 64 + TOKENS_PER_BLOCK + + run(manager, request) + + assert request.context_current_position == 64 + assert request.context_chunk_size == TOKENS_PER_BLOCK + assert kv_cache.capacity >= 64 + TOKENS_PER_BLOCK + assert kv_cache.history_length == 64 + assert connector.commits == [(request.py_request_id, 64, False)] + + +def test_a_second_pass_reports_one_allocation(manager, connector): + """An asynchronously loaded request re-enters on its first context chunk, + with the same pages and nothing left to load.""" + request = make_request() + assert schedule(manager, request) + + run(manager, request) + run(manager, request) + + assert len(connector.queries) == 1 + assert len(connector.commits) == 1 + assert len(connector.allocs) == 1 + + +def test_a_served_prefix_survives_the_scheduling_pass_that_brings_it_back(manager, connector): + """Two ``run`` calls are not the re-entry. + + An asynchronous load parks the request, and the scheduler calls + ``prepare_context`` again when it returns -- that is where the cursor is + settled, from a commit depth that describes the local match alone. + """ + connector.load_async = True + request = make_request() + assert schedule(manager, request) + run(manager, request) + served = request.context_current_position + assert served > 0, "the serve must have moved the position" + + assert schedule(manager, request) + + assert request.context_current_position == served + assert request.context_current_position + request.context_chunk_size == request.prompt_len + + kv_cache = manager.kv_cache_map[request.py_request_id] + assert kv_cache.history_length == served + assert kv_cache.num_committed_tokens < served, ( + "held but not committed -- which is why the cursor needs a floor" + ) + + +def test_freeing_the_allocation_makes_the_request_askable_again(manager, connector): + """A destructive pause replays the sequence add and with it the query, + because the pages the first answer described are gone.""" + request = make_request() + assert schedule(manager, request) + run(manager, request) + manager.free_resources(request) + + # Everything keyed to the dead allocation goes with it: the ask memo, and + # the scheduler-output deltas whose block ids describe pages that no longer + # exist. Leaving the latter is D1, which reports the replay as a cached + # request a `new_requests`-only connector never loads for. + assert connector.forgotten == [request.py_request_id] + + request.reset_for_recompute(PROMPT_LEN) + assert schedule(manager, request) + run(manager, request) + + assert len(connector.queries) == 2 + + +# --------------------------------------------------------------------------- +# Variable sliding-window attention. +# +# Two distinct windows over two layers give two layer groups, which is the only +# shape where the per-layer-group callbacks are reachable and the only shape +# where `_stale_block_range` has to pick a window rather than being handed the +# one there is. VSWA_PROMPT_LEN and VSWA_OFFER are sized so the offered prefix +# straddles the sliding window's edge: some block ordinals fall out of window +# and some stay live, in the same request. +# --------------------------------------------------------------------------- + +VSWA_WINDOW = 64 +VSWA_MAX_SEQ_LEN = 256 +VSWA_PROMPT_LEN = 160 +VSWA_OFFER = 128 + + +def make_vswa_manager(connector, **overrides): + return make_manager( + connector, + kv_cache_config=KvCacheConfig( + max_tokens=2048, + enable_block_reuse=True, + # Layer 0 slides, layer 1 is full attention: an entry equal to + # max_seq_len normalizes to None. + max_attention_window=[VSWA_WINDOW, VSWA_MAX_SEQ_LEN], + ), + max_seq_len=VSWA_MAX_SEQ_LEN, + **overrides, + ) + + +@pytest.fixture +def vswa_connector(): + return FakeConnectorManager(num_matched=VSWA_OFFER) + + +@pytest.fixture +def vswa_manager(vswa_connector): + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + mgr = make_vswa_manager(vswa_connector) + yield mgr + mgr.shutdown() + del mgr + gc.collect() + torch.cuda.empty_cache() + + +def _sliding_and_full(manager): + """Layer group ids of the sliding and the full-attention group.""" + windows = [lc.window_size for lc in manager._life_cycle_by_layer_group()] + assert len(windows) == 2, f"expected two layer groups, got {windows}" + return windows.index(VSWA_WINDOW), windows.index(None) + + +def test_window_size_is_read_per_layer_group(vswa_manager): + """Each group carries its own life cycle, and the pair is not interchangeable. + + Everything downstream -- the masking boundary, what a connector is offered + to save -- is derived from this list by index, so an off-by-one here is + silent and total. + """ + life_cycles = vswa_manager._life_cycle_by_layer_group() + windows = [lc.window_size for lc in life_cycles] + + assert len(windows) == 2 + assert sorted(windows, key=lambda w: (w is None, w)) == [VSWA_WINDOW, None] + + layers = vswa_manager.kv_cache_manager_py_config.layers + for layer_group_id, local_layer_ids in enumerate(vswa_manager.impl.layer_grouping): + for local_layer_id in local_layer_ids: + assert windows[layer_group_id] == layers[int(local_layer_id)].window_size + + +def test_stale_block_range_uses_each_group_s_own_window(vswa_manager): + """The full-attention group must never be masked, whatever the history. + + Computed here from the window rather than copied from the implementation: + a test that reuses `_stale_block_range` to predict `_stale_block_range` + cannot catch a wrong window being selected for the group. + """ + sliding, full = _sliding_and_full(vswa_manager) + tokens_per_block = vswa_manager.tokens_per_block + + for history_length in (0, 32, 64, 96, 128, 200): + expected_end = max(0, (history_length + 1 - VSWA_WINDOW) // tokens_per_block) + assert vswa_manager._stale_block_range(sliding, history_length) == (0, expected_end) + assert vswa_manager._stale_block_range(full, history_length) == (0, 0), ( + "a full-attention group has no stale range; masking it would hide " + "pages the connector is entitled to save" + ) + + +def test_attention_sinks_are_never_masked(vswa_manager: KVCacheManagerV2) -> None: + """Sink blocks stay live below the window, so they must not be masked. + + The stale range starts at `num_sink_blocks`, not at 0. Attention reads the + sink positions on every step however far they fall behind the window, so a + connector handed `BAD_PAGE_INDEX` for them would neither save nor restore + KV the model goes on to read. + + Both `AttentionLayerConfig` construction sites pass `num_sink_tokens=None` + today, so the life cycle is substituted here rather than configured. This + pins the masking loop's use of the range's lower bound, which is the part + that would silently do the wrong thing once sinks are wired up. + """ + sliding, full = _sliding_and_full(vswa_manager) + tokens_per_block = vswa_manager.tokens_per_block + # Keep one stale non-sink block so the masking assertion is exercised. + num_sink_tokens = tokens_per_block + + life_cycles = list(vswa_manager._life_cycle_by_layer_group()) + life_cycles[sliding] = AttnLifeCycle.make(VSWA_WINDOW, num_sink_tokens, tokens_per_block) + vswa_manager._connector_life_cycle_by_group = life_cycles + + request = make_request(prompt_len=VSWA_PROMPT_LEN) + assert schedule(vswa_manager, request) + run(vswa_manager, request) + + by_group = vswa_manager.get_page_indices_by_layer_group(request) + num_sink_blocks = num_sink_tokens // tokens_per_block + stale_end = max(0, (VSWA_OFFER + 1 - VSWA_WINDOW) // tokens_per_block) + assert num_sink_blocks < stale_end, ( + f"test sizes put the sinks outside the stale range " + f"(sinks={num_sink_blocks}, stale_end={stale_end}), so nothing is proven" + ) + + assert all(index != BAD_PAGE_INDEX for index in by_group[sliding][:num_sink_blocks]), ( + f"a sink block was masked: {by_group[sliding]}" + ) + assert all(index == BAD_PAGE_INDEX for index in by_group[sliding][num_sink_blocks:stale_end]), ( + f"a block the window has passed was reported as a page: {by_group[sliding]}" + ) + assert all(index != BAD_PAGE_INDEX for index in by_group[full]), ( + f"the full-attention group must keep every block: {by_group[full]}" + ) + + +def test_page_indices_mask_only_the_group_whose_window_passed(vswa_manager, vswa_connector): + """The mask is per group, in place, and against real pages. + + The offered prefix straddles the sliding window's edge, so the same request + has out-of-window ordinals in one group and live pages at those same + ordinals in the other. That is the case a single flat block-id list cannot + describe, and the case a `-1`-blind connector corrupts. + """ + request = make_request(prompt_len=VSWA_PROMPT_LEN) + assert schedule(vswa_manager, request) + run(vswa_manager, request) + + kv_cache = vswa_manager.kv_cache_map[request.py_request_id] + assert kv_cache.history_length == VSWA_OFFER, ( + "the prefix was not honoured in full, so the masking boundary below " + "is not the one this test was sized for" + ) + + by_group = vswa_manager.get_page_indices_by_layer_group(request) + assert len(by_group) == 2 + + sliding, full = _sliding_and_full(vswa_manager) + tokens_per_block = vswa_manager.tokens_per_block + stale_end = max(0, (VSWA_OFFER + 1 - VSWA_WINDOW) // tokens_per_block) + assert 0 < stale_end < len(by_group[sliding]), ( + f"test sizes no longer split the request across the window edge " + f"(stale_end={stale_end}, blocks={len(by_group[sliding])})" + ) + + # Ordinals stay positionally aligned across groups, which is what makes an + # append-delta over successive calls valid. + assert len(by_group[sliding]) == len(by_group[full]) + + assert all(index == BAD_PAGE_INDEX for index in by_group[sliding][:stale_end]), ( + f"a block the sliding window has passed was reported as a page: {by_group[sliding]}" + ) + live = by_group[sliding][stale_end:] + assert all(index != BAD_PAGE_INDEX for index in live), ( + f"an in-window block was reported with no page: {by_group[sliding]}" + ) + assert len(set(live)) == len(live), f"page slots are not distinct: {live}" + + assert all(index != BAD_PAGE_INDEX for index in by_group[full]), ( + f"the full-attention group must keep every block: {by_group[full]}" + ) + assert len(set(by_group[full])) == len(by_group[full]) + + # The pair is not the same list read twice: at the masked ordinals one + # group has pages and the other does not. + assert by_group[sliding] != by_group[full] + + +def test_alloc_is_reported_per_layer_group_and_the_flat_list_is_empty(vswa_manager, vswa_connector): + """A page index is scoped to a group, so the flat list must be withheld. + + Reporting group 0's indices as `block_ids` would look right to a connector + that never checks, and address the wrong pool for every layer outside that + group. + """ + request = make_request(prompt_len=VSWA_PROMPT_LEN) + assert schedule(vswa_manager, request) + run(vswa_manager, request) + + assert len(vswa_connector.allocs) == 1 + _, flat = vswa_connector.allocs[0] + assert flat == [], "the flat block-id list must be empty with several layer groups" + + assert len(vswa_connector.allocs_by_group) == 1 + _, by_group = vswa_connector.allocs_by_group[0] + assert len(by_group) == 2 + assert by_group == vswa_manager.get_page_indices_by_layer_group(request) + + +def test_a_released_request_still_reports_one_list_per_layer_group(vswa_manager): + """A released allocation reads back as empty lists, not as no groups at all. + + The connector callbacks are routed by this outer length, so collapsing a + released request to `[]` sends it to the flat form, which a connector + written for VSWA does not define. + """ + request = make_request(prompt_len=VSWA_PROMPT_LEN) + assert schedule(vswa_manager, request) + run(vswa_manager, request) + + num_groups = len(vswa_manager.impl.layer_grouping) + assert num_groups == 2, "the fixture must keep two windows for this to mean anything" + assert len(vswa_manager.get_page_indices_by_layer_group(request)) == num_groups + + vswa_manager.free_resources(request) + + assert vswa_manager.get_page_indices_by_layer_group(request) == [[], []] + + +def test_a_single_window_still_reports_the_flat_list(connector): + """Sibling check: the same code path with one group keeps the flat shape. + + Withholding the flat list is conditional on the group count, so the + single-group arm has to be pinned here or a change to that condition + breaks every existing connector without failing a VSWA test. + """ + torch.cuda.init() + mgr = make_manager(connector) + try: + request = make_request() + assert schedule(mgr, request) + run(mgr, request) + + assert len(connector.allocs) == 1 + _, flat = connector.allocs[0] + assert flat, "a single-group cache must still report the flat block ids" + + _, by_group = connector.allocs_by_group[0] + assert len(by_group) == 1 + assert by_group[0] == flat + finally: + mgr.shutdown() + del mgr + gc.collect() + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# One sliding window across every layer: a single layer group that still +# reports the flat list, which is the arm `reject_flat_only_scheduler` lets +# through and the only one where that list can carry a sentinel. +# --------------------------------------------------------------------------- + + +def make_swa_manager(connector, **overrides): + return make_manager( + connector, + kv_cache_config=KvCacheConfig( + max_tokens=2048, + enable_block_reuse=True, + # One window for every layer: one life cycle, so one layer group. + max_attention_window=[VSWA_WINDOW], + ), + max_seq_len=VSWA_MAX_SEQ_LEN, + **overrides, + ) + + +def test_a_sliding_single_group_reports_sentinels_in_the_flat_list(): + """The flat list is still reported, and it carries sentinels. + + Blocks the window has passed are held in place as ``BAD_PAGE_INDEX`` so the + ordinals stay aligned, which means a connector has to filter -- that is + what `valid_page_slots` is for. + """ + torch.cuda.init() + connector = FakeConnectorManager(num_matched=VSWA_OFFER) + mgr = make_swa_manager(connector) + try: + request = make_request(prompt_len=VSWA_PROMPT_LEN) + assert schedule(mgr, request) + run(mgr, request) + + assert len(mgr.impl.layer_grouping) == 1, "one window is one layer group" + _, flat = connector.allocs[0] + assert flat, "a single-group cache must still report the flat block ids" + + stale_beg, stale_end = mgr._stale_block_range(0, request.context_current_position) + assert stale_end > stale_beg, "the window must have passed a block for this to bite" + assert BAD_PAGE_INDEX in flat + live = [ordinal for ordinal, _ in valid_page_slots(flat)] + assert all(ordinal < stale_beg or ordinal >= stale_end for ordinal in live) + finally: + mgr.shutdown() + del mgr + gc.collect() + torch.cuda.empty_cache() + + +def test_a_served_prefix_survives_re_entry_under_a_sliding_window(): + """Here the rewind is not only wasted work. + + ``_resize_for_connector_prefix`` sets the cache's ``history_length`` to the + served end, and raising history unlocks the blocks the window has passed. A + cursor settled below that end points the forward pass at ordinals whose + pages are gone. + """ + torch.cuda.init() + connector = FakeConnectorManager(num_matched=VSWA_OFFER, load_async=True) + mgr = make_swa_manager(connector) + try: + request = make_request(prompt_len=VSWA_PROMPT_LEN) + assert schedule(mgr, request) + run(mgr, request) + served = request.context_current_position + stale_beg, stale_end = mgr._stale_block_range(0, served) + assert stale_end > stale_beg, "the window must have passed a block for this to bite" + + assert schedule(mgr, request) + + assert request.context_current_position == served + indices = mgr.get_page_indices_by_layer_group(request)[0] + first = request.context_current_position // TOKENS_PER_BLOCK + last = ( + request.context_current_position + request.context_chunk_size - 1 + ) // TOKENS_PER_BLOCK + assert all(indices[ordinal] != BAD_PAGE_INDEX for ordinal in range(first, last + 1)), ( + "the forward pass must not be pointed at a block the window has released" + ) + finally: + mgr.shutdown() + del mgr + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/unittest/_torch/executor/test_send_kv_async_split.py b/tests/unittest/_torch/executor/test_send_kv_async_split.py index 060e93b5ae4a..53de7a70e876 100644 --- a/tests/unittest/_torch/executor/test_send_kv_async_split.py +++ b/tests/unittest/_torch/executor/test_send_kv_async_split.py @@ -105,7 +105,7 @@ def test_connector_save_uses_previous_batch_with_overlap_scheduler() -> None: PyExecutor._save_kv_to_connector_async(executor, [current_req]) - executor.kv_connector_manager.request_finished.assert_called_once_with(prev_req, [7]) + executor.kv_connector_manager.request_finished.assert_called_once_with(prev_req, [7], None) executor.async_transfer_manager.start_transfer.assert_called_once_with(prev_req) @@ -117,7 +117,7 @@ def test_connector_save_uses_scheduled_batch_without_overlap_scheduler() -> None PyExecutor._save_kv_to_connector_async(executor, [finished, running]) - executor.kv_connector_manager.request_finished.assert_called_once_with(finished, [7]) + executor.kv_connector_manager.request_finished.assert_called_once_with(finished, [7], None) executor.async_transfer_manager.start_transfer.assert_called_once_with(finished) diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py index e45827e238da..a3a65ee3dc42 100644 --- a/tests/unittest/_torch/executor/test_token_budget_fallback.py +++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py @@ -88,7 +88,7 @@ def _make_manager(max_num_tokens, tokens_per_block, enable_chunked_prefill=True) # covered explicitly below. mgr.enable_chunked_prefill = enable_chunked_prefill mgr.is_draft = False - # Read by publish_connector_scheduler_output; most tests run without a + # Read by report_batch_to_connector; most tests run without a # connector attached. mgr.kv_connector_manager = None return mgr @@ -535,7 +535,7 @@ def maybe_fit_token_budget(self, scheduled_batch): req.context_chunk_size = self._shrink_to self._log.append(("trim", self._chunks())) - def publish_connector_scheduler_output(self, scheduled_batch): + def report_batch_to_connector(self, scheduled_batch): self._log.append(("publish", self._chunks())) def _chunks(self): @@ -571,7 +571,7 @@ def test_managers_without_a_connector_hook_are_skipped(self): def test_publishing_is_a_no_op_without_a_connector(self): mgr = _make_manager(max_num_tokens=128, tokens_per_block=16) mgr.kv_connector_manager = None - mgr.publish_connector_scheduler_output(_make_batch()) # must not raise + mgr.report_batch_to_connector(_make_batch()) # must not raise def test_publishing_forwards_the_batch_to_the_connector(self): class _FakeConnector: @@ -585,7 +585,7 @@ def build_scheduler_output(self, scheduled_batch, kv_cache_manager): mgr.kv_connector_manager = _FakeConnector() batch = _make_batch([_FakeRequest(context_chunk_size=16, prompt_len=16)]) - mgr.publish_connector_scheduler_output(batch) + mgr.report_batch_to_connector(batch) self.assertEqual(mgr.kv_connector_manager.calls, [(batch, mgr)]) diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py index 31ce6b825d39..91ec477e17fb 100644 --- a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -100,6 +100,7 @@ def _make_executor( exe.enable_attention_dp = enable_attention_dp exe.enable_kv_pool_rebalance = True exe.kv_cache_transceiver = None + exe.kv_connector_manager = None exe.is_warmup = False exe.is_shutdown = False exe.drafter = None diff --git a/tests/unittest/_torch/test_connector.py b/tests/unittest/_torch/test_connector.py index 01675416c5ea..8d6b7d611d3e 100644 --- a/tests/unittest/_torch/test_connector.py +++ b/tests/unittest/_torch/test_connector.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,10 +23,11 @@ from tensorrt_llm import mpi_rank from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( - AsyncRequests, KvCacheConnectorManager, - KvCacheConnectorSchedulerOutputManager) + AsyncRequests, KvCacheConnectorManager, KvCacheConnectorScheduler, + KvCacheConnectorSchedulerOutputManager, KvCacheConnectorWorker) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.logger import logger pytestmark = pytest.mark.cpu_only @@ -177,6 +178,161 @@ def test(): run_across_mpi(mpi_pool_executor, test, 2) +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_connector_manager_query_is_side_effect_free(mpi_pool_executor): + """The query and the commit are separable, and the query records nothing. + + A cache that allocates per context chunk can be offered more than the + pages the scheduler reserved. It therefore commits the amount it + honours rather than the amount it was offered, which is only possible if + asking is inert: `external_loads` is what tells the connector where its + load begins, and a request registered as loading is dropped from the batch. + """ + + def test(): + worker = MagicMock() + + if mpi_rank() == 0: + scheduler = MagicMock() + scheduler.get_num_new_matched_tokens.return_value = (16, True) + else: + scheduler = None + + manager = KvCacheConnectorManager(worker, scheduler=scheduler) + + req = MagicMock() + req.request_id = 42 + req.is_generation_only_request = False + req.py_num_connector_matched_tokens = 0 + + assert manager.query_num_new_matched_tokens(req, 32) == (16, True) + + assert manager.new_async_requests.loading_ids == set() + assert manager.scheduler_output_manager.external_loads == {} + assert req.py_num_connector_matched_tokens == 0 + + manager.commit_new_matched_tokens(req, 16, True) + + assert manager.new_async_requests.loading_ids == {42} + assert manager.scheduler_output_manager.external_loads == {42: 16} + assert req.py_num_connector_matched_tokens == 16 + + if mpi_rank() == 0: + assert scheduler.get_num_new_matched_tokens.call_count == 1 + + run_across_mpi(mpi_pool_executor, test, 2) + + +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_connector_manager_commits_only_what_is_honoured(mpi_pool_executor): + """Committing less than the offer is legal and is what gets reported. + + The unconsumed tail is recomputed locally; the connector releases its + ownership of the whole request at `request_finished`. + """ + + def test(): + worker = MagicMock() + + if mpi_rank() == 0: + scheduler = MagicMock() + scheduler.get_num_new_matched_tokens.return_value = (128, False) + else: + scheduler = None + + manager = KvCacheConnectorManager(worker, scheduler=scheduler) + + req = MagicMock() + req.request_id = 7 + req.is_generation_only_request = False + req.py_num_connector_matched_tokens = 0 + + num_tokens, load_async = manager.query_num_new_matched_tokens(req, 0) + assert num_tokens == 128 + manager.commit_new_matched_tokens(req, 32, load_async) + + assert manager.scheduler_output_manager.external_loads == {7: 32} + assert req.py_num_connector_matched_tokens == 32 + + run_across_mpi(mpi_pool_executor, test, 2) + + +def test_scheduler_output_resets_a_destroyed_allocation(): + """A replayed request must be reported as new, with its whole block list. + + `block_ids` and `tokens` are cumulative deltas, which is right while an + allocation lives and wrong the moment one is destroyed. A destructive + recompute pause does exactly that and the request replays -- reachable only + where the scheduler can pause a live request. Left stale, the replay lands in + `cached_requests` with a delta against pages that no longer exist, and a + connector that walks only `new_requests` issues no load. + """ + manager = KvCacheConnectorSchedulerOutputManager() + kv_cache_manager = MagicMock() + kv_cache_manager.get_cache_indices.return_value = [0, 1, 2, 3] + kv_cache_manager.commit_and_get_block_hashes.return_value = [] + kv_cache_manager.get_priority_by_block_id.return_value = 0 + + req = MagicMock() + req.request_id = 7 + req.state = LlmRequestState.CONTEXT_INIT + req.get_tokens.return_value = list(range(64)) + req.context_current_position = 0 + req.context_remaining_length = 64 + req.context_chunk_size = 64 + req.kv_cache_retention_config = None + req.cache_salt = None + + batch = ScheduledRequests() + batch.context_requests_last_chunk = [req] + + first = manager.build_scheduler_output(batch, AsyncRequests(dict(), dict()), + kv_cache_manager) + assert len(first.new_requests) == 1 + assert first.new_requests[0].new_block_ids == [0, 1, 2, 3] + + # Without the drop this is a `cached_request` carrying an empty delta. + manager.reset_request(req.request_id) + + replay = manager.build_scheduler_output(batch, + AsyncRequests(dict(), dict()), + kv_cache_manager) + assert len(replay.cached_requests) == 0 + assert len(replay.new_requests) == 1 + assert replay.new_requests[0].new_block_ids == [0, 1, 2, 3] + + +def test_scheduler_output_keeps_deltas_while_the_allocation_lives(): + """The drop is scoped to a destroyed allocation, not to every re-report.""" + manager = KvCacheConnectorSchedulerOutputManager() + kv_cache_manager = MagicMock() + kv_cache_manager.get_cache_indices.return_value = [0, 1, 2, 3] + kv_cache_manager.commit_and_get_block_hashes.return_value = [] + kv_cache_manager.get_priority_by_block_id.return_value = 0 + + req = MagicMock() + req.request_id = 8 + req.state = LlmRequestState.CONTEXT_INIT + req.get_tokens.return_value = list(range(64)) + req.context_current_position = 0 + req.context_remaining_length = 64 + req.context_chunk_size = 64 + req.kv_cache_retention_config = None + req.cache_salt = None + + batch = ScheduledRequests() + batch.context_requests_last_chunk = [req] + + manager.build_scheduler_output(batch, AsyncRequests(dict(), dict()), + kv_cache_manager) + again = manager.build_scheduler_output(batch, AsyncRequests(dict(), dict()), + kv_cache_manager) + + assert len(again.new_requests) == 0 + assert len(again.cached_requests) == 1 + assert again.cached_requests[0].new_block_ids == [] + + def test_scheduler_output_num_scheduled_tokens_with_mtp(): """Test that num_scheduled_tokens is correctly set for MTP (multi-token prediction).""" NUM_DRAFT_TOKENS = 3 @@ -252,3 +408,273 @@ def test_scheduler_output_block_hashes_read_through(): assert kv_cache_manager.commit_and_get_block_hashes.call_count == 2 for call in kv_cache_manager.commit_and_get_block_hashes.call_args_list: assert call.args == (req, ) + + +class _FlatOnlyScheduler(KvCacheConnectorScheduler): + """A connector written against the flat API.""" + + def __init__(self): + super().__init__(llm_args=None) + self.finished = [] + self.allocs = [] + + def build_connector_meta(self, scheduler_output): + return None + + def get_num_new_matched_tokens(self, request, num_computed_tokens): + return 0, False + + def request_finished(self, request, cache_block_ids): + self.finished.append(list(cache_block_ids)) + return False + + def update_state_after_alloc(self, request, block_ids): + self.allocs.append(list(block_ids)) + + +class _GroupedOnlyScheduler(KvCacheConnectorScheduler): + """A connector written against the per-layer-group API, as VSWA requires.""" + + def __init__(self): + super().__init__(llm_args=None) + self.finished = [] + self.allocs = [] + + def build_connector_meta(self, scheduler_output): + return None + + def get_num_new_matched_tokens(self, request, num_computed_tokens): + return 0, False + + def request_finished_by_layer_group(self, request, + cache_block_ids_by_layer_group): + self.finished.append( + [list(group) for group in cache_block_ids_by_layer_group]) + return False + + def update_state_after_alloc_by_layer_group(self, request, + block_ids_by_layer_group): + self.allocs.append([list(group) for group in block_ids_by_layer_group]) + + +def test_a_released_request_reaches_the_per_layer_group_callbacks(): + """Empty per-group lists route by length, so a released request still lands. + + A grouped-only connector has no flat method of its own, only the stand-in + that raises. Routing a released request there kills the executor loop on + the leader and strands every other rank in the broadcast that was waiting + on its answer. + """ + scheduler = _GroupedOnlyScheduler() + manager = KvCacheConnectorManager(MagicMock(), scheduler) + + req = MagicMock() + req.request_id = 4 + + manager.update_state_after_alloc(req, [], [[], []]) + assert scheduler.allocs == [[[], []]] + + assert manager.request_finished(req, [], [[], []]) is False + assert scheduler.finished == [[[], []]] + + +def test_an_empty_group_list_reaches_the_flat_callbacks(): + """No per-group view of the cache means the flat list is the description. + + A caller that cannot report per layer group leaves the grouped argument + unset. Dispatching that to the per-layer-group form would reach the ABC + default's "0 layer groups" refusal instead of the flat callback the + connector implements. + """ + scheduler = _FlatOnlyScheduler() + manager = KvCacheConnectorManager(MagicMock(), scheduler) + + req = MagicMock() + req.request_id = 3 + + manager.update_state_after_alloc(req, [], []) + assert scheduler.allocs == [[]] + + assert manager.request_finished(req, [], []) is False + assert scheduler.finished == [[]] + + +def test_a_per_layer_group_connector_needs_no_flat_stubs(): + """Overriding the grouped form is enough to instantiate. + + The flat methods are the abstract ones and abstractness is tracked per + method name, so overriding `request_finished_by_layer_group` does not clear + the flag on `request_finished`. Without the stand-in, a VSWA-only connector + would carry three dead methods purely to construct -- which is a trap, + because the failure is a `TypeError` at bring-up naming methods the + connector has no reason to implement. + """ + + class VswaWorker(KvCacheConnectorWorker): + + def register_kv_cache_layout(self, layout): + pass + + def start_load_kv(self, stream): + pass + + def wait_for_layer_load(self, layer_idx, stream): + pass + + def save_kv_layer(self, layer_idx, stream): + pass + + def wait_for_save(self, stream): + pass + + def get_finished(self, finished_gen_req_ids, started_loading_req_ids): + return [], [] + + class VswaScheduler(KvCacheConnectorScheduler): + + def build_connector_meta(self, scheduler_output): + return None + + def get_num_new_matched_tokens(self, request, num_computed_tokens): + return 0, False + + def request_finished_by_layer_group(self, request, + cache_block_ids_by_layer_group): + return False + + def update_state_after_alloc_by_layer_group(self, request, + block_ids_by_layer_group): + pass + + worker = VswaWorker(llm_args=None) + scheduler = VswaScheduler(llm_args=None) + + # The flat forms are still reachable, and say which form this connector + # implements rather than reporting nothing. + with pytest.raises(NotImplementedError, match="register_kv_cache_layout"): + worker.register_kv_caches(None) + with pytest.raises(NotImplementedError, + match="request_finished_by_layer_group"): + scheduler.request_finished(None, []) + with pytest.raises(NotImplementedError, + match="update_state_after_alloc_by_layer_group"): + scheduler.update_state_after_alloc(None, []) + + +def test_a_connector_implementing_neither_form_still_fails_at_construction(): + """The stand-in only fires for a connector that implements the grouped form. + + A connector implementing neither must keep failing where it always has, at + construction, naming the method to implement. + """ + + class NeitherScheduler(KvCacheConnectorScheduler): + + def build_connector_meta(self, scheduler_output): + return None + + def get_num_new_matched_tokens(self, request, num_computed_tokens): + return 0, False + + with pytest.raises(TypeError, match="request_finished"): + NeitherScheduler(llm_args=None) + + +def _grouped_scheduler(*implemented): + """A scheduler class overriding exactly the named per-layer-group forms.""" + + body = { + "build_connector_meta": lambda self, scheduler_output: None, + "get_num_new_matched_tokens": lambda self, request, num_computed: + (0, False), + "request_finished": lambda self, request, cache_block_ids: False, + "update_state_after_alloc": lambda self, request, block_ids: None, + } + body.update({ + name: lambda self, request, by_group: None + for name in implemented + }) + return type("PartialScheduler", (KvCacheConnectorScheduler, ), body) + + +@pytest.mark.parametrize("implemented", [ + (), + ("request_finished_by_layer_group", ), + ("update_state_after_alloc_by_layer_group", ), +]) +def test_a_partial_grouped_scheduler_is_rejected_at_bring_up(implemented): + """Mixing the two forms must fail before the model is loaded. + + The two pairs clear their abstract flags independently, so a connector + implementing one grouped form and leaving the other flat constructs + cleanly. Left to the base defaults it would then raise on the first + scheduled request -- after model load, which is what this check buys back. + """ + scheduler = _grouped_scheduler(*implemented)(llm_args=None) + manager = KvCacheConnectorManager(MagicMock(), scheduler=scheduler) + + with pytest.raises(NotImplementedError, match="by_layer_group"): + manager.reject_flat_only_scheduler(2) + + +def test_a_grouped_scheduler_and_a_single_group_cache_are_accepted(): + """The two ways to pass: both forms implemented, or only one layer group. + + The single-group arm is the compatibility promise -- every existing + connector defines only the flat forms, and must keep starting up. + """ + both = _grouped_scheduler("request_finished_by_layer_group", + "update_state_after_alloc_by_layer_group") + KvCacheConnectorManager( + MagicMock(), + scheduler=both(llm_args=None)).reject_flat_only_scheduler(2) + + flat_only = _grouped_scheduler() + KvCacheConnectorManager( + MagicMock(), + scheduler=flat_only(llm_args=None)).reject_flat_only_scheduler(1) + + +def _record_warnings(monkeypatch): + warnings = [] + monkeypatch.setattr(logger, "warning", + lambda msg, *args, **kwargs: warnings.append(msg)) + return warnings + + +def test_a_flat_scheduler_under_a_sliding_window_warns_but_runs(monkeypatch): + """The compatibility promise and the hazard, together. + + One window across every layer is a single layer group, so a flat-only + connector is not refused -- refusing would stop connectors that work + today. What such a connector does not otherwise get is a `-1` where the + window has passed a block, so the difference is logged instead. + """ + warnings = _record_warnings(monkeypatch) + flat_only = _grouped_scheduler() + manager = KvCacheConnectorManager(MagicMock(), + scheduler=flat_only(llm_args=None)) + + manager.reject_flat_only_scheduler(1) + manager.warn_flat_scheduler_under_swa(4096) + + assert len(warnings) == 1 + assert "valid_page_slots" in warnings[0] + + +def test_nothing_is_said_where_nothing_changed(monkeypatch): + """A window is not enough on its own, and neither is a flat scheduler.""" + warnings = _record_warnings(monkeypatch) + + flat_only = _grouped_scheduler() + KvCacheConnectorManager( + MagicMock(), + scheduler=flat_only(llm_args=None)).warn_flat_scheduler_under_swa(None) + + both = _grouped_scheduler("request_finished_by_layer_group", + "update_state_after_alloc_by_layer_group") + KvCacheConnectorManager( + MagicMock(), + scheduler=both(llm_args=None)).warn_flat_scheduler_under_swa(4096) + + assert warnings == [] diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index e83ec29939a2..c45f825d08d6 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py @@ -62,6 +62,7 @@ class _StatsRequest: draft_tokens: list[int] = field(default_factory=list) state: LlmRequestState = LlmRequestState.GENERATION_IN_PROGRESS context_current_position: int = 0 + py_connector_served_position: int = 0 context_chunk_size: int = 0 expect_snapshot_points: list[int] = field(default_factory=list) prepopulated_prompt: tuple[int, int] | None = None