[None][fix] Correct RocketKV KT cache byte accounting for FP8 KV cache#16225
[None][fix] Correct RocketKV KT cache byte accounting for FP8 KV cache#16225eopXD wants to merge 1 commit into
Conversation
The RocketKV KT (key-landmark) cache is a separate pool physically allocated at kt_cache_dtype, independent of the main KV cache dtype. Both get_cache_size_per_token (the pre-allocation estimate) and get_cache_bytes_per_token (used by calculate_max_num_blocks) sized the KT term by the main KV dtype byte width and folded the (min, max) factor-of-2 into a dtype-dependent kt_factor. This was correct only when the main KV cache was bf16; with an FP8 KV cache it under-counted the KT pool by 2x, over-estimating the block count (max_tokens = budget / bytes_per_token) and over-subscribing GPU memory. Size the KT contribution at its own dtype's byte width with an explicit factor of 2, matching the pool allocated in RocketKVCacheManager.__init__ regardless of the main KV dtype. Add a unit test that builds a RocketKVCacheManager with an FP8 main KV cache and a bfloat16 KT cache and asserts get_cache_bytes_per_token() matches the physically allocated KT pool (fails without the fix). Co-Authored-By: Yueh-Ting Chen <yueh.ting.chen@gmail.com> Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
|
/bot run |
📝 WalkthroughWalkthroughRocketKV cache byte estimation now separates main KV-cache dtype sizing from KT-cache dtype sizing. A new SM100-gated test validates mixed FP8 KV and bfloat16 KT cache accounting. ChangesRocketKV cache sizing
Estimated code review effort: 3 (Moderate) | ~20 minutes 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.
🧹 Nitpick comments (2)
tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py (2)
711-711: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new test return type.
Use
-> Nonefor this test function.Proposed fix
-def test_rocketkv_kt_cache_sized_by_kt_dtype_not_kv_dtype(): +def test_rocketkv_kt_cache_sized_by_kt_dtype_not_kv_dtype() -> None:🤖 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/attention/sparse/rocketkv/test_rocketkv.py` at line 711, Annotate the test function test_rocketkv_kt_cache_sized_by_kt_dtype with the return type -> None.Source: Coding guidelines
710-751: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the changed static estimator too.
This regression test covers
get_cache_bytes_per_token(), but not the separately modifiedRocketKVCacheManager.get_cache_size_per_token(). Coverage is insufficient unless another test already exercises that path; add a companion case intests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.pyor track it as 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/attention/sparse/rocketkv/test_rocketkv.py` around lines 710 - 751, Add a companion regression test for RocketKVCacheManager.get_cache_size_per_token() using differing main KV and KT dtypes, such as FP8 main KV with bfloat16 KT, and assert the KT contribution is sized from the KT pool dtype rather than the main KV dtype. Reuse the setup and cleanup pattern from test_rocketkv_kt_cache_sized_by_kt_dtype_not_kv_dtype, or explicitly track coverage as a follow-up if the static estimator cannot be exercised here.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.
Nitpick comments:
In `@tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py`:
- Line 711: Annotate the test function test_rocketkv_kt_cache_sized_by_kt_dtype
with the return type -> None.
- Around line 710-751: Add a companion regression test for
RocketKVCacheManager.get_cache_size_per_token() using differing main KV and KT
dtypes, such as FP8 main KV with bfloat16 KT, and assert the KT contribution is
sized from the KT pool dtype rather than the main KV dtype. Reuse the setup and
cleanup pattern from test_rocketkv_kt_cache_sized_by_kt_dtype_not_kv_dtype, or
explicitly track coverage as a follow-up if the static estimator cannot be
exercised here.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 371e25d8-915b-4a9a-a61b-89694ce5feeb
📒 Files selected for processing (2)
tensorrt_llm/_torch/attention_backend/sparse/rocket.pytests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py
|
PR_Github #58633 [ run ] triggered by Bot. Commit: |
|
PR_Github #58633 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58667 [ run ] triggered by Bot. Commit: |
|
PR_Github #58667 [ run ] completed with state |
2 participants
Description
The RocketKV KT (key-landmark) cache is a separate pool physically allocated at
kt_cache_dtype(seeRocketKVCacheManager.__init__:torch.empty((num_blocks, kt_tokens_per_block, num_kv_heads, head_dim * 2), dtype=kt_cache_dtype)),independent of the main KV cache dtype. Two accounting methods sized the KT term by the
main KV dtype byte width and folded the
(min, max)factor-of-2 into a dtype-dependentkt_factor:RocketKVCacheManager.get_cache_size_per_token— the pre-allocation memory estimate(KV-cache budget path in
_util.py).RocketKVCacheManager.get_cache_bytes_per_token— used bycalculate_max_num_blocks(
max_tokens = free_mem_fraction * free_mem / bytes_per_token).This was numerically correct only when the main KV cache is bf16. With an FP8 KV cache
the KT pool was under-counted by 2×, which over-estimates the block count and over-subscribes
GPU memory (OOM risk), since the KT pool is allocated with the same block count as the primary
pool.
This PR sizes the KT contribution at its own dtype's byte width with an explicit factor of 2,
matching the physically allocated pool regardless of the main KV dtype. The bf16 main-KV path
(the RocketKV default) is numerically unchanged.
Test Coverage
tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py::test_rocketkv_kt_cache_sized_by_kt_dtype_not_kv_dtypeConstructs a
RocketKVCacheManagerwith an FP8 main KV cache + bfloat16 KT cache (so the twodtypes differ) and asserts
get_cache_bytes_per_token()matches the physically allocatedkt_cache_pool_per_layerbytes. It fails on the old code (KT sized at the 1-byte main dtype) andpasses with the fix. The test runs in the L0 pre-merge B200 stage via the existing
unittest/_torch/attentiontest-db entry (SM100-gated;getSMVersion() < 100is skipped).PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why.
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. (N/A — internal cache-accounting only, no public API change.)Any new dependencies have been scanned for license and vulnerabilities (N/A — no new dependencies).
CODEOWNERS updated if ownership changes (N/A).
Documentation updated as needed (N/A — no behavior change to documented interfaces).
Update tava architecture diagram if there is a significant design change in PR (N/A).
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.Summary by CodeRabbit
Bug Fixes
Tests