[None][feat] Batched physical KV-cache compaction for KV cache compression#16836
[None][feat] Batched physical KV-cache compaction for KV cache compression#16836Hudayday wants to merge 5 commits into
Conversation
…ssion Signed-off-by: tianruih <tianruih@nvidia.com>
…at/compaction-for-kv-cache-compression
Signed-off-by: tianruih <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #61550 [ run ] triggered by Bot. Commit: |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughAdds BF16 sparse KV-cache compaction through a CUDA kernel, Torch operator, Python/Triton orchestration, and tests. Also renames the KV compression manager abstraction and guards the executor fast path when compression manages history. ChangesSparse KV cache compaction
Estimated code review effort: 5 (Critical) | ~120 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant PythonCompaction
participant TritonPackKernel
participant sparse_kv_cache_compact_layers
participant BF16CompactionKernel
PythonCompaction->>TritonPackKernel: pack keep decisions into move sources
PythonCompaction->>sparse_kv_cache_compact_layers: submit grouped pool compaction
sparse_kv_cache_compact_layers->>BF16CompactionKernel: launch validated BF16 geometry
BF16CompactionKernel->>BF16CompactionKernel: copy selected K/V tiles into destination pages
🚥 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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/kv_cache_compression/test_compaction.py (1)
1-325: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd the new compaction tests to CI/QA lists and guard the SM100-only cases.
- Added test functions:
test_eager_compaction_preserves_exact_selected_bytes_and_tail,test_eager_compaction_rebases_masked_swa_window_and_tail,test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle.- I don’t see entries for
tests/unittest/_torch/kv_cache_compression/test_compaction.pyintests/integration/test_lists/test-db/orqa/; add this file there.- Add
@requires_sm100totest_eager_compaction_preserves_exact_selected_bytes_and_tailandtest_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle; only the SWA test is currently guarded.Coverage verdict: needs follow-up.
🤖 Prompt for 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. In `@tests/unittest/_torch/kv_cache_compression/test_compaction.py` around lines 1 - 325, Add tests/unittest/_torch/kv_cache_compression/test_compaction.py to the relevant CI/QA test lists under tests/integration/test_lists/test-db/ and qa/. Apply the existing requires_sm100 marker to test_eager_compaction_preserves_exact_selected_bytes_and_tail and test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle, matching the guard already used by test_eager_compaction_rebases_masked_swa_window_and_tail.Source: Path instructions
🤖 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 `@tests/unittest/_torch/kv_cache_compression/test_compaction.py`:
- Around line 30-31: Add the existing `@requires_sm100` decorator to
test_eager_compaction_preserves_exact_selected_bytes_and_tail and
test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle. Keep their
current parametrization and test bodies unchanged so both GPU-only tests skip on
non-SM100 or unavailable CUDA environments.
In `@tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py`:
- Around line 189-442: Add the four new test
functions—test_sparse_kv_cache_compact_layers,
test_sparse_kv_cache_compact_layers_cuda_graph_replay,
test_sparse_kv_cache_compact_layers_fast_geometry, and
test_sparse_kv_cache_compact_layers_rejects_invalid_launch—to the appropriate
integration test-list entries under the test-db and QA lists. Preserve the
existing list format and ensure each test is registered for CI/manual QA
execution.
---
Outside diff comments:
In `@tests/unittest/_torch/kv_cache_compression/test_compaction.py`:
- Around line 1-325: Add
tests/unittest/_torch/kv_cache_compression/test_compaction.py to the relevant
CI/QA test lists under tests/integration/test_lists/test-db/ and qa/. Apply the
existing requires_sm100 marker to
test_eager_compaction_preserves_exact_selected_bytes_and_tail and
test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle, matching the
guard already used by test_eager_compaction_rebases_masked_swa_window_and_tail.
🪄 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: adf9b668-beac-434e-9638-815dd2f8713e
📒 Files selected for processing (12)
cpp/tensorrt_llm/kernels/unfusedAttentionKernels.hcpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpptensorrt_llm/_torch/kv_cache_compression/compaction.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytests/unittest/_torch/executor/test_kv_cache_compression_manager.pytests/unittest/_torch/kv_cache_compression/conftest.pytests/unittest/_torch/kv_cache_compression/test_compaction.pytests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py
Signed-off-by: tianruih <tianruih@nvidia.com>
There was a problem hiding this comment.
Looks like this compaction is used only for TriAttention, at least for now. It'd be better to create a TriAttention dir here and move this file to the dir.
There was a problem hiding this comment.
Good point. The current placement is intentional: compaction.py has no TriAttention concept in it. build_compaction_params() consumes only the V2 cache layout, block offset tables, and a generic "kept ordinals" decision, so any eviction based compression algorithm needs the same mover. That is why it sits at the kv_cache_compression framework level, so if you still prefer a TriAttention dir here, I am happy to move the file there.
| *cache_params.pack_args, **cache_params.pack_constexprs | ||
| ) | ||
| for args in cache_params.compact_args: | ||
| torch.ops.trtllm.sparse_kv_cache_compact_layers(*args) |
There was a problem hiding this comment.
I noticed a nested for-loop here. Are we actually launching the kernel this many times?
There was a problem hiding this comment.
The loop bounds are small constants independent of batch size, layer count, and tokens. The outer loop runs once per cache (target, plus draft under speculative decoding), and the inner one runs once per uniform pool group, which is 1 for homogeneous dense models and 2 for dense plus SWA mixes like GPT-OSS. Each sparse_kv_cache_compact_layers call compacts every layer, request, and head of its pool group in one batched launch, so the worst production case is 2 pack plus 2 compact launches per eviction round.
I have tested it that using this loop is faster than treating SWA as a fake dense layer and evicting them in one kernel.
|
|
||
|
|
||
| class BaseKVCacheCompressionManager(BaseResourceManager): | ||
| class KVCacheCompressionManager(BaseResourceManager): |
There was a problem hiding this comment.
Why this was renamed to KVCacheCompressionManager?
There was a problem hiding this comment.
This aligns the name with the other manager classes like KV Cache Manager: the Base prefix is reserved for abstract bases like BaseResourceManager, and this one is a concrete class with no subclass. Pure rename, no behavior change.
| and self._get_position_id_offset() == 0 | ||
| # KV compression shrinks the cache mid-generation; take the full prepare path. | ||
| and not getattr(kv_cache_manager, | ||
| "kv_compression_manages_history", False)): |
There was a problem hiding this comment.
Why do we need to skip this for the kv compression cases?
Nit: remove the comments above.
There was a problem hiding this comment.
The steady generation fast path memoizes prepare_inputs for pure decode batches and advances the cached sequence lengths by one per step, which assumes the KV history only grows. Compression physically shrinks the cache mid generation, so the memoized lengths go stale and the model silently runs with wrong history lengths (I hit IMA issue from exactly this). Managers that manage history themselves opt out via this flag; the dense path keeps the fast path. Comment removed as suggested.
|
PR_Github #61550 [ run ] completed with state
|
There was a problem hiding this comment.
Please rename this file to a more meaningful name, check others.
…drop stale comment Signed-off-by: tianruih <tianruih@nvidia.com>
5 participants
Description
First PR of the TriAttention KV-cache compression series (compaction primitive → eviction algorithm).
This adds an algorithm-neutral, batched physical compaction primitive for
KVCacheManagerV2paged KV pools: given per-request keep decisions (kept tokenordinals plus per-request move offsets), it packs the move sources and compacts
the pools in place, so a KV cache compression method can physically drop evicted
tokens instead of masking them.
CUDA / C++
unfusedAttentionKernels_2_template.h: batched in-place compaction kernels —one launch covers a uniform group of layer pools; per request and KV head the
moves are ascending and never overtake their sources, which makes the forward
tiled in-place copy safe.
thop/sparseKvCacheCompactOp.cpp:trtllm::sparse_kv_cache_compact_layers,taking device pointer arrays so one call compacts a whole pool group.
Python host API (
tensorrt_llm/_torch/kv_cache_compression/compaction.py)build_compaction_params(layout, ...) -> CompactionParams: called once percache geometry; pre-binds every launch argument (a Triton pack kernel plus the
native compact calls). Supports dense and sliding-window layer families, and
an optional draft cache co-compacted under the target's keep decision.
compact(params, request_count): the per-round entry — pure pre-boundlaunches, no host-side construction on the generation path.
Framework
BaseKVCacheCompressionManager→KVCacheCompressionManager.kv_compression_manages_historyis set (the cache shrinks mid-generation).Supported surface: BF16 paged pools in the native
[pages, K/V, heads, tokens, dim]layout; kernels target SM100 (Blackwell).Test Coverage
tests/unittest/_torch/kv_cache_compression/test_compaction.py— byte-exactoracles against torch references: dense keep set + protected tail, SWA window
rebase, and draft co-compaction broadcast over the draft's own heads/tables.
tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py— nativeop contract: supported geometries, argument validation, and CUDA-graph
capture safety.
tests/unittest/_torch/executor/test_kv_cache_compression_manager.py—manager rename covered by the existing framework tests.
Dev Engineer Review
KVCacheManagerV2paged KV pools (dense + masked sliding-window/SWA), with optional draft-cache co-compaction.sparse_kv_cache_compact_layerswith runtime validation and BF16-only dispatch.cp.async(and explicit rejection for unsupportedhead_dim/tokens_per_blockgeometry).build_compaction_params) and per-round launcher (compact) that packs keep decisions and prepares dense move offsets (plus optional SWA/draft move sources/bases).kv_compression_manages_history=True.BaseKVCacheCompressionManagertoKVCacheCompressionManager, updating typing imports and the corresponding unit tests.tests/integration/test_lists/in this PR’s diff.QA Engineer Review
Test files changed/added:
tests/unittest/_torch/executor/test_kv_cache_compression_manager.pytest_inherits_base_resource_managertest_four_hooks_default_nooptest_hooks_accept_extra_kwargstest_resource_counts_are_zerotest_length_adjustment_marks_target_and_draft_v2test_rejects_non_v2_ownershiptest_request_field_defaults_to_zerotest_target_update_receives_metadata_before_final_compressiontest_real_v2_target_receives_relocation_metadatatest_prepare_fires_init_on_first_chunk_onlytest_update_fires_context_end_on_last_chunktest_step_end_fires_once_per_iterationtest_free_fires_finishtest_returns_none_when_no_algorithm_registeredtest_warns_for_unregistered_algorithmtest_factory_accepts_independent_draft_managertest_eviction_method_predicate_defaults_falsetest_spec_gate_only_restricts_eviction_methodstest_names_importable_from_canonical_modulestest_names_not_in_sparse_moduletest_raises_when_reuse_ontest_ok_when_reuse_offtests/integration/test_lists/: not found/referenced.tests/unittest/_torch/kv_cache_compression/conftest.pytests/integration/test_lists/: not applicable (helpers only); no references found.tests/unittest/_torch/kv_cache_compression/test_compaction.pytest_eager_compaction_preserves_exact_selected_bytes_and_tailtest_eager_compaction_rebases_masked_swa_window_and_tailtest_draft_moves_and_pack_match_keep_broadcast_and_tail_oracletests/integration/test_lists/: not found/referenced.tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.pytest_sparse_kv_cache_compact_layerstest_sparse_kv_cache_compact_layers_cuda_graph_replaytest_sparse_kv_cache_compact_layers_fast_geometrytest_sparse_kv_cache_compact_layers_rejects_invalid_launchtests/integration/test_lists/: not found/referenced.Verdict: needs follow-up — confirm/add CI/manual test-list registration for these new/updated unittests (no entries detected under
tests/integration/test_lists/).