[https://nvbugs/6438658][fix] Fix MLA KV cache estimation sizing - #16311
Conversation
Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
jiaganc
changed the title
📝 WalkthroughWalkthrough
ChangesKV cache runtime updates
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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
879-888: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winRestore
pool_ratio/avg_seq_lenbefore calling_cal_max_memory, not after.
_cal_max_memory(L505) callsself._get_kv_size_per_token()whilekv_cache_config.pool_ratio/avg_seq_lenstill hold the estimation-time override values (set intry_prepare_estimation), since the restore at L886-887 happens after the_cal_max_memorycall at L880. The returnedavailable_kv_memdoesn't itself depend onkv_size_per_token, so the finalmax_gpu_total_bytes/max_tokensaren't corrupted — but thekv_size_per_tokenvalue baked into the log line at L509-514 ("kv size per token is {kv_size_per_token}") is computed from the wrong config, which is exactly the kind of diagnostic this PR is trying to make trustworthy for KV-cache OOM debugging.Moving the restore to before the
_cal_max_memorycall (or to the very top of the function) also incidentally hardens against the case where an exception is raised somewhere between function entry and L884, which currently leaveskv_cache_configin the temporary estimation state.Note the new test
test_estimation_temporarily_uses_inferred_pool_sizingmocks_cal_max_memoryentirely, so it wouldn't have caught this.🩹 Proposed fix: restore before computing max memory
- # calculate max memory from peak memory and free gpu memory fraction - kv_cache_max_memory = self._cal_max_memory(peak_memory, - total_gpu_memory, fraction, - allocated_bytes) - - # Estimation uses inferred pool sizing; the final manager uses the - # user-provided configuration. - self._kv_cache_config.pool_ratio = self._pool_ratio_in - self._kv_cache_config.avg_seq_len = self._avg_seq_len_in + # Estimation uses inferred pool sizing; the final manager uses the + # user-provided configuration. Restore before any capacity/logging + # computation below reads kv_cache_config, so log output and sizing + # reflect the user's real settings. + self._kv_cache_config.pool_ratio = self._pool_ratio_in + self._kv_cache_config.avg_seq_len = self._avg_seq_len_in + + # calculate max memory from peak memory and free gpu memory fraction + kv_cache_max_memory = self._cal_max_memory(peak_memory, + total_gpu_memory, fraction, + allocated_bytes)🤖 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 `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 879 - 888, Move the assignments restoring self._kv_cache_config.pool_ratio and self._kv_cache_config.avg_seq_len to before the _cal_max_memory call in the surrounding method. Ensure _cal_max_memory and its kv-size diagnostic use the user-provided configuration, and preserve the existing restoration values from _pool_ratio_in and _avg_seq_len_in.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
1532-1561: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLGTM!
The added warning correctly includes
label/bs/draft_lencontext for diagnosing skipped CUDA-graph capture shapes.Optional/out-of-scope nitpick:
_run_attention_warmup(L1240-1244) and_capture_piecewise_cuda_graphs(L1622, L1652) still silentlycontinueonbatch is Nonewithout a similar warning. Not a defect (PR scope is explicitly "general and CUDA-graph warmup shapes"), but for consistent observability across all warmup skip-paths, the same warning pattern could be applied there in a 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 1532 - 1561, Optional follow-up: add equivalent warnings before the batch-is-None continue paths in _run_attention_warmup and _capture_piecewise_cuda_graphs, including the relevant warmup label and shape context such as batch size and draft length. Keep the existing skip behavior unchanged.tests/unittest/_torch/executor/test_kv_cache_estimation.py (1)
548-610: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winGood coverage for the V2 override/restore contract; consider extending for the V1 path and the
_cal_max_memoryinteraction.This test correctly verifies
try_prepare_estimation/configure_kv_cache_capacityoverride-then-restore semantics, but only for theKVCacheManagerV2path (_get_model_kv_cache_manager_clsis patched to always returnKVCacheManagerV2). The override intry_prepare_estimation(L758-759 in_util.py) is unconditional regardless of V1/V2, so a second test using the plainKVCacheManagerclass would confirm the same restore contract holds there too.Also,
_cal_max_memoryis fully mocked here (patch.object(creator, "_cal_max_memory", return_value=512)), so this test can't catch ordering bugs where_cal_max_memory's internals readkv_cache_configbefore the restore runs (see the corresponding comment intensorrt_llm/_torch/pyexecutor/_util.py). A follow-up test that lets_cal_max_memory/_get_kv_size_per_tokenrun for real (mocking onlytorch.cuda/the model config) would give stronger coverage of the restore-ordering contract.As per path instructions, coverage here is sufficient for the primary regression scenario this PR targets, but could use a follow-up test for the V1 manager path and the
_cal_max_memoryordering interaction, ideally in this same file.🤖 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_estimation.py` around lines 548 - 610, Add follow-up coverage in test_estimation_temporarily_uses_inferred_pool_sizing for the plain KVCacheManager path by patching _get_model_kv_cache_manager_cls accordingly and verifying the same override, configure, and restoration contract. Also add coverage that runs _cal_max_memory and _get_kv_size_per_token without mocking them, while mocking only CUDA and model-config dependencies, to verify they observe the temporary estimation configuration before restoration.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 `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 879-888: Move the assignments restoring
self._kv_cache_config.pool_ratio and self._kv_cache_config.avg_seq_len to before
the _cal_max_memory call in the surrounding method. Ensure _cal_max_memory and
its kv-size diagnostic use the user-provided configuration, and preserve the
existing restoration values from _pool_ratio_in and _avg_seq_len_in.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 1532-1561: Optional follow-up: add equivalent warnings before the
batch-is-None continue paths in _run_attention_warmup and
_capture_piecewise_cuda_graphs, including the relevant warmup label and shape
context such as batch size and draft length. Keep the existing skip behavior
unchanged.
In `@tests/unittest/_torch/executor/test_kv_cache_estimation.py`:
- Around line 548-610: Add follow-up coverage in
test_estimation_temporarily_uses_inferred_pool_sizing for the plain
KVCacheManager path by patching _get_model_kv_cache_manager_cls accordingly and
verifying the same override, configure, and restoration contract. Also add
coverage that runs _cal_max_memory and _get_kv_size_per_token without mocking
them, while mocking only CUDA and model-config dependencies, to verify they
observe the temporary estimation configuration before restoration.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: adb8d732-2df1-4de2-a7fa-0ce2391042b8
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_kv_cache_estimation.py
|
/bot run --disable-fail-fast |
|
PR_Github #58969 [ run ] triggered by Bot. Commit: |
|
PR_Github #58969 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59002 [ run ] triggered by Bot. Commit: |
|
PR_Github #59002 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59010 [ run ] triggered by Bot. Commit: |
|
PR_Github #59010 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59017 [ run ] triggered by Bot. Commit: |
|
PR_Github #59017 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59049 [ run ] triggered by Bot. Commit: |
|
PR_Github #59049 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59103 [ run ] triggered by Bot. Commit: |
|
PR_Github #59103 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59132 [ run ] triggered by Bot. Commit: |
|
PR_Github #59132 [ run ] completed with state |
5 participants
Summary by CodeRabbit
Bug Fixes
Tests
Description
The MLA branch of
_create_kv_cache_managerdid not forwardmax_num_tokens, so MLA managers could size the temporary estimation cache from the full token estimate instead of the runtime chunk size.This change forwards
max_num_tokensand makes estimation use inferred pool sizing. User-providedpool_ratioandavg_seq_lencan underprovision a pool in the tightly sized temporary manager, causing warmup to hang or fail, so estimation temporarily usespool_ratio=Noneandavg_seq_len=max_seq_len. The user values are restored before constructing the final manager. Warnings also make skipped general and CUDA-graph warmup shapes visible.The change intentionally does not modify warmup batch sizing or KV-cache-manager-v2 internals.
Test Coverage
tests/unittest/_torch/executor/test_kv_cache_estimation.py::test_mla_branch_forwards_max_num_tokens_to_managertests/unittest/_torch/executor/test_kv_cache_estimation.py::test_estimation_temporarily_uses_inferred_pool_sizingPR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
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. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
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.