[TRTLLM-12891][feat] Support the KV cache connector on KVCacheManagerV2 - #18700
Merged
brb-nv merged 10 commits intoSep 4, 2026
Merged
Conversation
…quest The C++ base exposes is_generation_only_request as a read-only property (def_prop_ro, nanobind/batch_manager/bindings.cpp), but the Python LlmRequest subclass redefined it as a plain method with no @Property. Reading it as an attribute therefore yields a bound method, which is always truthy. V1 never trips this because the connector path reaches the attribute from C++ with the C++ object, where the property is real. Every Python-side reader gets the method object instead, so `if request.is_generation_only_request` is unconditionally true and the corresponding guard is dead. Restore the property and drop the call parentheses at the six production call sites and in the test mocks. A local shim was rejected: the same attribute is reached from C++ with the C++ object and from Python with the Python one, so a shim would have to sniff callable(). Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
CacheTransceiverCpp is bound to the V1 BaseKVCacheManager, while KVCacheManagerV2.impl is the Python V2 core's manager. The combination reached BindKvCacheTransceiver and died on a raw nanobind signature mismatch that named neither the manager nor the way out. An equivalent check already existed for MambaHybridCacheManagerV2 but not for its base class, so plain KVCacheManagerV2 fell straight through. The new check sits after the hybrid one so the subclass keeps its more specific message; both messages share the phrase the existing tests match on, so that ordering is pinned by its own test. This makes the failure legible; it does not make V2 work with a default transceiver config. CacheTransceiverConfig.transceiver_runtime defaults to "auto", which is resolved from the model's preference (llm_utils._resolve_transceiver_runtime_auto) and knows nothing about which cache manager will be built, so most models land on the C++ transceiver. Working configuration is backend="NIXL" with transceiver_runtime="PYTHON". Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Every connector test now runs twice, once per KV cache manager, so the V2 integration that follows has somewhere to land. 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 spies on both managers' __init__ to prove positively which one was built, and asserts the fallback warning is absent. assert_kv_caches_registered is the per-test half of the same argument: the managers hand the worker their pools through different entry points, and asserting the V1 one unconditionally would pass on V2 exactly when the connector was never registered at all. Parametrization is spelled out per test as an explicit @pytest.mark.parametrize(..., indirect=True) rather than as a fixture params=, because scripts/check_test_list.py resolves ids from decorators via AST and cannot see fixture-level parametrization. Tests that supply their own KvCacheConfig have the manager forced onto it, otherwise the V2 parametrization degrades into a second V1 run. Also replace the fixed time.sleep(1) barrier -- which carried its author's TODO -- with a poll on the recorded mock call count until the connector goes quiet. That returns as soon as the callbacks settle and stretches automatically when a slower path lengthens the tail. The one end-to-end test evaluated the generated text and discarded it; it now asserts a cold miss, that cache files were written, a warm hit, and token agreement. Comparing two deterministic runs proves nothing on its own -- they agree whether or not the cache is consulted -- so the spy on the connector's matched-token count is what makes it non-tautological. Agreement is a prefix floor rather than equality: skipping prefill changes the attention reduction order, so reuse can legitimately move the last token. Every test id changed, so the CI list is retargeted onto the kv_cache_manager_v1 variants. The kv_cache_manager_v2 variants fail until the connector is implemented on V2, and are registered in no list until then. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Connector bring-up hands the worker a single tensor from get_unique_primary_pool. KVCacheManagerV2 has no such tensor and cannot grow one: it allocates pool groups of slots, a slot is the set of coalesced buffers belonging to one layer group, and there is one slot address space per pool and one page-index space per layer group. The first hop of bring-up was therefore a hard stop, and the fix is to describe the memory rather than point at it. KvCacheLayout carries, per layer group, the byte ranges its pages live in: slot i of a region is at base + stride * i for size bytes, or equivalently region.as_tensor()[i]. build_kv_cache_layout_v2 is assembly of V2's own public layout API -- layer_grouping, all_buffer_ids, get_aggregated_pages, pool_group_descs -- so coalescing is derived from the allocator rather than assumed by the consumer. For a uniform model a layer group's buffers merge into a single whole-slot region, which is why the example connector's load and save paths need no change. Regions are byte-oriented because one may span roles with different element types, and a connector moving bytes should not have to care. register_kv_cache_layout is non-abstract and raises by default, so existing connectors are untouched on V1 and get an actionable error on V2 rather than a crash. A registered address is only valid while its page is pinned to GPU, and eviction to another tier reassigns the page's slot, so reject any tier below GPU. 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 -- a truthiness check there is dead. That automatic tier is also skipped outright when a connector is attached: it exists only to give the MAX_UTILIZATION scheduler's suspend/resume somewhere to spill to, and a connector run is already restricted to GUARANTEED_NO_EVICT, so it is dead weight rather than a capability. The VSWA guard in bring-up is relaxed for V2 in passing. It exists because the single-tensor registration cannot describe one pool per window size, which is the ordinary case for a layout. The creator-level gate above it still rejects, and is dealt with separately. Removing kv_connector_manager from the V2 incompatible-feature list left the hybrid test asserting a reason that is no longer produced, so it now asserts V2 is returned unchanged; the operator-facing message no longer suggests disabling the connector when the connector is not what makes the configuration unsupported. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Registration alone leaves the connector inert: it receives its pools but no metadata, because its scheduler-side hooks are never invoked. The two managers place allocation differently. V1 allocates in KVCacheManager.prepare_resources and drives the connector from there. V2 allocates in KVCacheV2Scheduler, via prepare_context and resize_context, and its prepare_resources is a no-op for the non-draft path -- so the hooks had no home and build_connector_meta was never called. prepare_resources still runs after scheduling and before the forward pass, and by then every scheduled request has its pages, which is the same position in the iteration where V1 drives the connector. The hooks go there. Page indices come from the manager rather than from the connector reaching into kv_cache_map and impl.layer_grouping. They are reported per layer group and positionally aligned: valid_only=False yields one entry per block ordinal, so a block with no page in a group -- the sliding-window case -- reads back as BAD_PAGE_INDEX in place rather than shortening the list. That is what keeps the ordinal-to-token-range mapping intact and an append-delta over successive calls valid. RequestData carries them as new_block_ids_by_layer_group; with a single group -- every non-VSWA, non-hybrid model -- new_block_ids still carries that group's indices, so connectors that do not reason about layer groups keep working. request_finished is routed through the same accessor. Without it the V1 lookup raises, the surrounding warning path swallows the exception, and a connector on V2 is never told to save anything. Block hashes and retention priorities are reported empty rather than guessed at. V2 has no per-request accessor for the hash chain, and KvCacheRetentionConfig does not reach KVCacheManagerV2 at all -- per-page priority comes from custom_priority_callback, which V2 never overrides, so every page carries the default. Reporting those defaults would be worse than reporting nothing: a connector doing priority-based offload filtering would act on values unrelated to what the user configured. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
py_executor_creator refused any multi-valued max_attention_window before a cache manager was ever chosen, so the V2 carve-out in PyExecutor below it was unreachable and every multi-layer-group path was dead code. VSWA allocates one pool per window size, which the V1 single-tensor registration cannot describe. A KvCacheLayout can: it carries one region set per layer group. So reject here only when V2 is definitively off. kv_cache_config.use_kv_cache_manager_v2 is tri-state (True / False / "auto"), and under "auto" the manager is not chosen yet, so defer -- PyExecutor re-checks against the manager it actually built and rejects there if the selection landed on V1. Layered rather than duplicated. Two tests cover what this unblocks. A uniform sliding window, where every layer shares one window: the registered layout must still collapse to a single layer group and record the window. And VSWA, where V1 still raises and V2 reports two groups with equal-length, positionally aligned index lists -- the flat new_block_ids being empty is the correct answer there, since a page index is scoped to a layer group and there is no correct way to flatten several. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
get_num_new_matched_tokens is a protocol, not a query: five of its six clauses are obligations on the runtime. V1 discharges them in one step because its entire connector interaction happens in prepare_resources, on the final batch. V2 cannot. prepare_context runs during scheduling, and a prepared request can still be dropped at the token budget, at resize_context, at multimodal alignment or at cross attention. Rolling computed_position back on an iteration the request never runs in is silent corruption: the next iteration arrives with the position already advanced, the connector loads nothing, and the runtime skips computing tokens whose KV is garbage. V2 already separates match-and-take-ownership -- which needs only the token sequence and cannot fail for lack of memory -- from residency, which claims slots, migrates pages and can fail. The connector interaction is decomposed along the same seam, with transmission as a third phase sharing the forward pass's preconditions: phase 1 _connector_prefix_position ask, before residency exists phase 2 _reserve_connector_prefix cover the offer with pages, after resume phase 3 _deliver_connector_prefix record the external load, on the batch Phase 2 has to follow the resume because _KVCache.resize asserts the cache is ACTIVE and a freshly created or deferred-and-suspended one is not. Capacity and history move in one call: after a reuse match both equal the local match, so raising history alone trips "History length cannot be greater than capacity". Raising history is also what stops a served prefix allocating a page for every block in a sliding-window layer group, since it is the sole input to the stale-range computation. Phase 2 clears enable_swa_scratch_reuse for a served prefix: scratch slots are transient prefill storage and a connector writes real cache content into those blocks, the same reason the disagg generation path opts out. The memoised value is the absolute offer end, never the returned delta. A deferred request re-derives its local match from a tree another request's commit may have grown; adding the delta to the new match would place the position past the union of what is computed and what is loaded, leaving tokens that are neither. The offer end is also clamped below the prompt, because the first generation step consumes the last prompt position's activations -- a connector holding the whole prompt is the steady state of a repeat, not an error. KvCacheConnectorManager.get_num_new_matched_tokens is split into a side-effect-free query and a commit that registers the async hold and the external load. V1 keeps calling the fused entry point, which is correct there: it is invoked from C++ under the block manager's tree mutex, so the match and the query are atomic with respect to the tree and the answer can be committed immediately. cancel_load is additive with a no-op default. It hands an offer back when phase 2 cannot allocate, and when a request is asked and then cancelled, times out or fails before delivery -- otherwise the connector holds remote blocks for the life of the process with nothing left to release them. should_add_sequence stays out of the V2 scheduler. That predicate is false from the moment an asynchronous load completes until request_finished at the end of generation; in the V2 scheduler it means SKIP, so V2 would skip such a request forever and never run the prefill the load was for. What keeps a loading request out of the batch is its DISAGG_GENERATION_TRANS_IN_PROGRESS state, and what stops the connector being asked or told twice is the per-request state machine. Phase 3 asserts that what it records is covered by the position phase 2 advanced. Nothing downstream validates it: the runtime's subtraction is unguarded and connectors divide the result into block ordinals, so a mismatch would silently point a connector at the wrong offset inside its own code. The unit suite models a suspended cache and drives the real _prepare_context_impl rather than restating it, so the ordering is observed rather than declared. Verified by mutation: moving the reserve above the resume turns 17 of its tests red, where an earlier version of the same suite passed all 24 against that same defect. A second suite drives the phases against a real KVCacheManagerV2 with real pools and a deterministic deferral, which is the only place the ask-once path is exercised end to end -- an engine-level deferral test proved to be a race against the executor's request queue and was removed rather than weakened. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
BAD_PAGE_INDEX was undocumented. Under a sliding window a connector on V2 receives -1 entries in update_state_after_alloc, RequestData.new_block_ids and request_finished; one that treats them as page slots computes an address from -1. State it, along with the reason the entry is kept in place rather than dropped: ordinals stay aligned to token ranges and append-deltas stay valid. Document cancel_load and the speculative scheduling pass it exists for, and record what does not change -- get_num_new_matched_tokens is still called exactly once per request on both managers, including across a deferral. Correct its trigger list. The local match overtaking an offer cannot arise today: a request's local match is fixed when its cache is created, resume() does not re-match, and only the request's own completed forward passes extend it. The two cases that can occur -- allocation failure, and a request freed before delivery -- were not named at all. Note that update_state_after_alloc covers only the first chunk's blocks under chunked prefill on V2, since V2 allocates per chunk; the rest arrive as append-deltas through build_connector_meta. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The connector suite has run under both KV cache managers since it was parametrized, but only the kv_cache_manager_v1 half was registered, behind a comment saying the V2 variants were expected to fail. That is no longer true -- the connector is supported on V2 -- and coverage that runs only by hand decays the first time someone touches _util.py or the bring-up path. Register the 22 kv_cache_manager_v2 ids, plus the two that were never listed: test_connector_runs_on_kv_cache_manager_v2, which is what makes the rest meaningful (the creator silently falls back to the V1 manager for combinations it cannot serve, so without it every V2 id could pass while running V1), and the V1 half of the VSWA test. test_connector_priorities[kv_cache_manager_v2] is marked xfail(strict=True) rather than dropped from the list. KvCacheRetentionConfig does not reach KVCacheManagerV2 at all, so a retention config is silently ignored there -- not only through the connector. Its assertions stay the correct expectation for both managers, so wiring retention into V2 turns the test green rather than needing it rewritten, and strict=True makes that day loud instead of silent. The host_offloading rejection assertion matched the bare string "host", which both managers' messages contain, so it would have passed through exactly the silent V1 fallback the parametrization exists to rule out. Match per manager instead: V1 names the config field, V2 the resolved tier. Retarget the static contract test. It was written as a Phase 2 worklist -- "remove entries as they are implemented" -- and that is not what happened: KVCacheManagerV2 implements none of the V1 block-id methods and is not meant to, because a flat pool-wide block id cannot describe memory whose page indices are scoped to a layer group. What the check is actually worth keeping for is that update_and_build_data reports block_hashes and priorities empty on V2 by branching on the manager type, not on hasattr: those short-circuits are correct only while the accessors are genuinely absent. It now says that, asserts the V2-side contract it does have, and the disagg method list records why none of its entries is a V2 gap rather than describing them as unported. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The `prompt_len - 1` clamp in `_connector_prefix_position` shrank the offer without telling the connector. `_deliver_connector_prefix` and `_release_undelivered_connector_prefix` both read the already-clamped end, so the tail past it reached neither, and the connector kept ownership of those remote blocks for the life of the process -- the exact leak those two paths exist to prevent. Cancel the dropped range where it is dropped, and assert it in the two tests that exercise the clamp. `_reject_non_gpu_cache_tiers` rejects every tier below GPU, which includes a disk tier configured through `KvCacheConfig.disk_cache_size`, but told the user to set `host_cache_size=0`. Name both fields. Also add the missing NVIDIA header to `test_kv_connector_v2_prefix.py`. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
requested review from
JunyiXu-nv,
QiJune,
SimengLiu-nv,
Tabrizian,
chienchunhung,
fredricz-20070104,
liji-nv,
schetlur-nv and
zhaoyangwang-nvidia
September 4, 2026 04:05
brb-nv
changed the title
brb-nv
removed request for
a team,
EmmaQiaoCh,
JunyiXu-nv,
QiJune,
SimengLiu-nv,
StanleySun639,
Tabrizian,
chienchunhung,
fredricz-20070104,
liji-nv,
mlefeb01,
schetlur-nv and
zhaoyangwang-nvidia
September 4, 2026 04:06
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #71431 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #71431 [ run ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Cherry-pick of #17974.
Needed for mooncake integration: #18676
Test Coverage
N/A
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.