[None][feat] Batched physical KV-cache compaction for KV cache compression - #16836
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: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds BF16 sparse KV-cache compaction through a CUDA kernel, Torch operator, Python/Triton orchestration, and validation 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 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>
|
PR_Github #61550 [ run ] completed with state
|
…drop stale comment Signed-off-by: tianruih <tianruih@nvidia.com>
The kernel entry point is a function template defined only in the template header; without an explicit instantiation no CUDA TU emits invokeSparseKvCacheCompactLayers<__nv_bfloat16>, so a fresh build fails to resolve the symbol used by sparseKvCacheCompactOp.cpp. Signed-off-by: tianruih <tianruih@nvidia.com>
… geometry Signed-off-by: tianruih <tianruih@nvidia.com>
…undary Reject kept-ordinal row counts that do not divide the request count evenly and decision layouts other than broadcast, per-head, or per-layer-per-head, so a mismatched cache geometry fails loudly instead of packing uninitialized move indices. Signed-off-by: tianruih <tianruih@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #61814 [ run ] triggered by Bot. Commit: |
|
PR_Github #61814 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61889 [ run ] triggered by Bot. Commit: |
|
PR_Github #61889 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61924 [ run ] triggered by Bot. Commit: |
|
PR_Github #61924 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62051 [ run ] triggered by Bot. Commit: |
|
PR_Github #62051 [ run ] completed with state |
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Regression axis checks out: C++ is purely additive (only - lines are copyright bumps), and the one shared-path change — the kv_compression_manages_history guard in model_engine.py — defaults False, so no existing manager changes behavior.
Not assessing kernel numerics — leaving that to the attention/runtime owners.
pengbowang-nv
left a comment
pengbowang-nv
left a comment
There was a problem hiding this comment.
Attention part LGTM
10 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 pools, focused on SM100 BF16 “native” pools ([pages, K/V, heads, tokens, dim]).invokeSparseKvCacheCompactLayerskernel entry-point declaration/wiring incpp/tensorrt_llm/kernels/unfusedAttentionKernels.handunfusedAttentionKernels_2_template.h.cp.async(double-buffered pipeline) for the layered V2 “sparse KV compaction” fast path, including explicit rejection of unsupported launch geometry combinations.sparse_kv_cache_compact_layers(cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp) with strict runtime validation of tensor dtype/device/shape/contiguity/layout invariants and BF16-only dispatch to the new compact layers kernel path.cpp/tensorrt_llm/thop/CMakeLists.txt) and added the BF16 explicit instantiation needed for symbol availability (unfusedAttentionKernels_2_bf16_bf16.cu).tensorrt_llm/_torch/kv_cache_compression/compaction.pywithbuild_compaction_params(...)andcompact(...)to pre-bind compaction tensors/geometry, run a Triton packing kernel, and invoke the native op per prepared group bundle.kv_cache_managerreportskv_compression_manages_history=True(tensorrt_llm/_torch/pyexecutor/model_engine.py).BaseKVCacheCompressionManager→KVCacheCompressionManagerand updated related typing/util references and unit tests (tensorrt_llm/_torch/pyexecutor/resource_manager.py,tensorrt_llm/_torch/pyexecutor/_util.py, corresponding unittest updates).QA Engineer Review
Test files changed/added (unit tests):
tests/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.pyNew/updated test coverage:
test_kv_cache_compression_manager.py.Coverage in
tests/integration/test_lists/:tests/integration/test_lists/entries were modified/detected in the workspace check; therefore, CI/manual coverage mapping for these new unit tests needs follow-up.Verdict: needs follow-up.