[None][feat] Add paged KV cache support to Vanilla attention - #16714
Conversation
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
yihwang-nv
changed the title
|
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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughVanilla attention now supports paged KV-cache reads and writes using per-request block IDs, including MLA latent gathering and invalid-page handling. Backend configuration and capabilities are updated, with a sliding-window eviction test added. ChangesVanilla paged KV-cache support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant VanillaAttention
participant PagedKVCache
participant ScaledDotProductAttention
VanillaAttention->>PagedKVCache: write K/V tokens using block_ids
VanillaAttention->>PagedKVCache: gather logical cached K/V tokens
PagedKVCache-->>VanillaAttention: materialized K/V prefix
VanillaAttention->>ScaledDotProductAttention: run attention over gathered K/V
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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/attention_backend/vanilla.py`:
- Around line 185-219: Update the write loop in _single_request_update_kv_cache
to validate each block_ids-derived blk before indexing kv_cache_tensor. If blk
equals BAD_PAGE_INDEX, fail loudly with the established error mechanism instead
of writing through negative indexing; preserve the existing cache-write behavior
for valid block indices.
In
`@tests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py`:
- Around line 94-102: Replace the source-text regex assertion in
test_vanilla_preserves_configured_tokens_per_block with a behavioral test that
invokes create_py_executor, or the specific tokens_per_block resolution helper,
using attn_backend="VANILLA" and a configured page size. Assert that the
resulting tokens_per_block equals the configured value, and preserve the test’s
focus on preventing KV-cache linearization.
🪄 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: 9f4d69e1-6301-4757-b329-ba82c76343d1
📒 Files selected for processing (5)
tensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytests/unittest/_torch/attention/backend_capability.pytests/unittest/_torch/attention/test_vanilla_attention.pytests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
|
/bot run --disable-fail-fast |
|
PR_Github #60922 [ run ] triggered by Bot. Commit: |
|
PR_Github #60922 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61201 [ run ] triggered by Bot. Commit: |
…parsing test - _single_request_update_kv_cache: assert the destination block is not BAD_PAGE_INDEX before writing, so an inconsistent block_ids/metadata fails loudly instead of silently writing the last page via negative indexing. - _gather_paged_kv: the empty-slice (num_tokens<=0) path indexed block_ids[0], which may be BAD_PAGE_INDEX; index the always-valid block 0 instead. - Remove test_vanilla_preserves_configured_tokens_per_block: a regex-on-source assertion is brittle; the paged behavior is covered by the attention sweep. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
|
/bot run |
|
PR_Github #61257 [ run ] triggered by Bot. Commit: |
|
PR_Github #61201 [ run ] completed with state |
|
PR_Github #61257 [ run ] completed with state
|
|
/bot run |
|
/bot run --disable-fail-fast |
|
PR_Github #62756 [ run ] triggered by Bot. Commit: |
|
PR_Github #62756 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62804 [ run ] triggered by Bot. Commit: |
|
PR_Github #62804 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62912 [ run ] triggered by Bot. Commit: |
|
PR_Github #62912 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62988 [ run ] triggered by Bot. Commit: |
|
PR_Github #62988 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63279 [ run ] triggered by Bot. Commit: |
|
PR_Github #63279 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63333 [ run ] triggered by Bot. Commit: |
|
PR_Github #63333 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63366 [ run ] triggered by Bot. Commit: |
|
PR_Github #63366 [ run ] completed with state |
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
5 participants
Summary
tokens_per_blockin the production PyTorch executorBAD_PAGE_INDEXlogical holes for KVCacheManagerV2 sliding-window evictionTesting
Dev Engineer Review
BAD_PAGE_INDEXholes and sliding-window eviction.tokens_per_blockfor the VANILLA backend.QA Engineer Review
test_kv_cache_manager_v2_sliding_window_eviction.test-db/orqa/registration was identified; CI coverage data is unavailable.