[None][feat] Add TensorRT-LLM runtime integration for KV cache compression#15697
[None][feat] Add TensorRT-LLM runtime integration for KV cache compression#15697Hudayday wants to merge 6 commits into
Conversation
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthrough
ChangesKV-cache compression reclaim
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 2497-2502: The no-free fallback in
kv_cache_manager_v2._KVCacheManagerV2 should not ignore the boolean result from
kv_cache.resize(None, req.max_beam_num_tokens - evicted). After the warning in
the fallback branch, check the return value just like the normal resize path and
treat a failed resize as fatal or otherwise handle it consistently so the
request state stays aligned with the live _KVCache state.
In `@tests/unittest/_torch/pyexecutor/test_kv_cache_v2_compression_reclaim.py`:
- Around line 77-83: The current test coverage only exercises completion with
evicted=0, so the combined evicted-and-completing path in the reclaim logic is
still untested. Add a test in test_kv_cache_v2_compression_reclaim.py using
_fake_manager, _run, and _req with evicted > 0 and a completing state such as
LlmRequestState.GENERATION_COMPLETE or CONTEXT_INIT, and assert that the request
still calls resize(None, max_beam - 1) while fork() is not called. This should
verify the guard in the kv cache reclaim behavior when both conditions are
present.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f47017d9-645a-4086-ab61-1067b72e308b
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytests/integration/test_lists/test-db/l0_a10.ymltests/unittest/_torch/pyexecutor/test_kv_cache_v2_compression_reclaim.py
|
PR_Github #56237 [ run ] triggered by Bot. Commit: |
|
PR_Github #56237 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56247 [ run ] triggered by Bot. Commit: |
|
PR_Github #56247 [ run ] completed with state |
5cdbbf3 to
095f9ff
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #56292 [ run ] triggered by Bot. Commit: |
2f10293 to
c11ca19
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #56307 [ run ] triggered by Bot. Commit: |
|
PR_Github #56292 [ run ] completed with state |
|
PR_Github #56307 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56337 [ run ] triggered by Bot. Commit: |
|
PR_Github #56337 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56371 [ run ] triggered by Bot. Commit: |
|
PR_Github #56371 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #59608 [ run ] triggered by Bot. Commit: |
Signed-off-by: tianruih <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #59620 [ run ] triggered by Bot. Commit: |
|
PR_Github #59608 [ run ] completed with state |
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
LGTM. Inert runtime plumbing — ships no compression algorithm and dense runs are byte-identical (py_num_compressed_tokens defaults to 0). Independently double-checked the one API-surface item: the new kv_cache_compression_mode is a @property (not a Pydantic field), so it isn't snapshotted by the api_stability harness and needs no golden-manifest regen; the only gate is the api-compatible label, which its CI check already passed. The dispatch change (isinstance→resource_type==) is analyzed safe (draft/cross managers don't consume the withheld args) and even fixes a latent V2 arg-passing case.
|
[by Codex] @lowsfer Friendly reminder to review this PR as the primary KV-cache-manager reviewer when you have a chance. Thanks! |
|
PR_Github #59620 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59700 [ run ] triggered by Bot. Commit: |
|
PR_Github #59700 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59742 [ run ] triggered by Bot. Commit: |
|
PR_Github #59742 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59822 [ run ] triggered by Bot. Commit: |
|
PR_Github #59822 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59928 [ run ] triggered by Bot. Commit: |
|
PR_Github #59928 [ run ] completed with state |
8 participants
Description
Eviction-based KV cache compression makes the target cache's physical KV
length shorter than the request's logical generation length. Before this
change the runtime could not represent that:
KVCacheManagerV2pushed bothcapacity and history back to the logical token count on every generation
step (undoing any compaction), and attention inputs always assumed logical
lengths.
This PR adds the minimal runtime support for generation-length-changing KV
cache compression, on top of the compression-manager framework merged in
#15106. The whole mechanism is one integer carried on the request:
Request-carried compressed length
LlmRequest.py_num_compressed_tokens(default 0), written only by acompression manager: how many of the request's tokens have been physically
evicted.
num_cached_tokens_per_seq,so attention kernels read the compacted KV length and append new tokens at
the right slot. Position ids and the
cached_tokensstat keep the logicalcount.
KVCacheManagerV2 length policy
kv_compression_manages_historyflag: when set by a compressionmanager, generation-step resize updates capacity only and preserves the
committed history, instead of forcing it back to the logical length.
Compression manager base and speculative gate
BaseKVCacheCompressionManageraccepts an optional draftKVCacheManagerV2and applies the same length policy to it, so analgorithm can evict the draft cache together with the target. Block reuse
is rejected at attach (compression rewrites stored K/V).
on_context_step_endreceives the requests whose prefill finished thisiteration, so prefill-end methods process the cohort in one launch.
KvCacheCompressionModecarries algorithm traits the same waySpeculativeDecodingModedoes; configs map theiralgorithmstring to itand callers read
is_*predicates.manager is created: evicting methods require one-model MTP/EAGLE3 draft
KV; otherwise the run stays uncompressed with a warning.
ModelConfig.kv_cache_compression_configplumb so attention modules canread method traits.
Explicitly not in this PR: no scheduler, attention, or
attention-metadata changes; no concrete eviction algorithm and no draft
compaction (follow-up algorithm PR).
Activation and scope
Inert unless
kv_cache_compression_configselects an algorithm; the factorythen registers the manager as a regular resource manager.
KVCacheManagerV2only.Test Coverage
test_kv_cache_compression_manager.py: base-class contract, resource-manager-to-hook translation, target/draft length-policy attach, call-site
speculative gate, factory registration.
test_kv_cache_v2_capacity_only.py: default V2 history behaviorunchanged; capacity-only generation updates preserve history. Both
registered in
l0_a10.