[https://nvbugs/6438658][fix] Fix KV cache estimation capacity - #16545
Conversation
fa320a4 to
5e8eea7
Compare
jiaganc
added
the
api-compatible
📝 WalkthroughWalkthroughKV cache configuration construction is split into base and subclass stages. Resumable quota synchronization and ChangesKV cache configuration changes
Sequence Diagram(s)sequenceDiagram
participant KVCacheManagerV2
participant BaseConfigBuilder
participant DeepseekV4CacheManager
participant CacheConfig
KVCacheManagerV2->>BaseConfigBuilder: build base cache configuration
BaseConfigBuilder->>CacheConfig: create layers and constraints
KVCacheManagerV2->>DeepseekV4CacheManager: apply subclass customization
DeepseekV4CacheManager->>CacheConfig: set layers and typical_step
DeepseekV4CacheManager-->>KVCacheManagerV2: return updated configuration
Possibly related PRs
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.
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/executor/test_kv_cache_manager_v2.py (1)
34-59: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCover nonzero extra-KV-token warmup capacity.
The helper fixes
num_extra_kv_tokensto zero, so the new constraints never verify either added extra-token term. Parameterize it and add a case asserting both decode and context capacity.Suggested coverage
def _make_cache_config_for_test( kv_cache_config: KvCacheConfig, *, ... max_draft_len: int = 0, + num_extra_kv_tokens: int = 0, ) -> KVCacheManagerConfig: ... - cache_manager.num_extra_kv_tokens = 0 + cache_manager.num_extra_kv_tokens = num_extra_kv_tokensAdd a test in
tests/unittest/_torch/executor/test_kv_cache_manager_v2.pywithmax_draft_len=2andnum_extra_kv_tokens=4; it should expect minimal decode capacities of7and a context constraint capacity of2052.As per path instructions,
tests/**feedback must identify concrete files and whether coverage is sufficient; coverage is insufficient here.🤖 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/executor/test_kv_cache_manager_v2.py` around lines 34 - 59, Update _make_cache_config_for_test to accept a num_extra_kv_tokens parameter and assign it to cache_manager instead of always using zero. Add the requested test case in test_kv_cache_manager_v2.py with max_draft_len=2 and num_extra_kv_tokens=4, asserting minimal decode capacity 7 and context constraint capacity 2052.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.
Outside diff comments:
In `@tests/unittest/_torch/executor/test_kv_cache_manager_v2.py`:
- Around line 34-59: Update _make_cache_config_for_test to accept a
num_extra_kv_tokens parameter and assign it to cache_manager instead of always
using zero. Add the requested test case in test_kv_cache_manager_v2.py with
max_draft_len=2 and num_extra_kv_tokens=4, asserting minimal decode capacity 7
and context constraint capacity 2052.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 167c3e6e-1c4d-4930-8017-818c68545fe2
📒 Files selected for processing (7)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/executor/test_kv_cache_manager_v2.pytests/unittest/llmapi/test_llm_args.py
5e8eea7 to
745369c
Compare
|
/bot run |
|
PR_Github #60299 [ run ] triggered by Bot. Commit: |
|
PR_Github #60299 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60315 [ run ] triggered by Bot. Commit: |
|
PR_Github #60315 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60396 [ run ] triggered by Bot. Commit: |
|
PR_Github #60396 [ run ] completed with state
|
|
/bot run |
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
5d4390f to
8604e52
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #60555 [ run ] triggered by Bot. Commit: |
|
PR_Github #60555 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60639 [ run ] triggered by Bot. Commit: |
|
PR_Github #60639 [ run ] completed with state
|
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
|
/bot run |
|
PR_Github #60664 [ run ] triggered by Bot. Commit: |
|
PR_Github #60664 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60679 [ run ] triggered by Bot. Commit: |
|
PR_Github #60679 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60698 [ run ] triggered by Bot. Commit: |
|
PR_Github #60698 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60864 [ run ] triggered by Bot. Commit: |
|
PR_Github #60864 [ run ] completed with state |
…A#16545) Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
6 participants
Summary by CodeRabbit
Improvements
Bug Fixes
max_util_for_resume=0values are now rejected.Documentation
Description
Fix KV cache capacity estimation so the resulting cache can satisfy runtime warmup and resume requirements.
Without these constraints, pool-ratio estimation can allocate a pool layout that cannot run CUDA graph warmup. Ignoring resume utilization during distributed quota synchronization can also remove the reserved headroom and overestimate usable token capacity.
max_util_for_resumeheadroom when synchronizing KV cache capacity across distributed ranks.KVCacheManagerV2. Let specialized cache managers update the base configuration.Test Coverage
max_util_for_resumevalidation.pre-commitpassed for all changed files.python3 -m py_compilepassed for the changed Python files.tensorrt_llm.bindings.PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, comment
/bot help.