[None][fix] Make auto host tier sizing rank-aware in KVCacheManagerV2 - #16106
Conversation
Signed-off-by: Eric Tsai <ertsai@nvidia.com>
📝 WalkthroughWalkthroughHost-tier quota sizing in KVCacheManagerV2 is updated to be rank-aware within a physical node by determining co-located ranks via hostname allgather with fallback logic, scaling the memory-based quota candidate accordingly, and guarding against non-positive results. ChangesHost Tier Quota Sizing
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 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
🤖 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/pyexecutor/kv_cache_manager_v2.py`:
- Around line 875-887: The fallback in KV cache manager v2 host sizing is
reintroducing the host OOM risk by resetting a nonpositive host_quota back to
quota. In kv_cache_manager_v2.py, update the auto-sizing logic around
candidates/host_quota so that a computed host_quota of 0 or less is left
unchanged and allowed to skip host tier creation via the downstream host tier
check, rather than falling back to the full quota. Keep the existing sizing path
and logger.info context intact, but remove the unsafe quota fallback in the host
quota calculation.
- Around line 844-864: The hostname allgather in kv_cache_manager_v2’s
local_ranks sizing logic is swallowing real distributed failures, which can
leave other ranks stuck in the collective. Update the logic around
Distributed.get(mapping).allgather(socket.gethostname()) so exceptions are not
broadly caught; either let the error propagate or restructure the fallback to
occur only outside the collective. Keep the rank-aware sizing behavior in place,
but avoid masking failures in this block.
🪄 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: 3dd1084e-1f1c-4bec-95bf-000868d376dd
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
…ric fallback Signed-off-by: Eric Tsai <ertsai@nvidia.com>
eopXD
left a comment
eopXD
left a comment
There was a problem hiding this comment.
Looks good to me with slight comments.
A unit test coverage would be good to validate your change.
…t tests Signed-off-by: Eric Tsai <ertsai@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #58669 [ run ] triggered by Bot. Commit: |
|
PR_Github #58669 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #58870 [ run ] triggered by Bot. Commit: |
|
PR_Github #58870 [ run ] completed with state |
… rank discovery Signed-off-by: Eric Tsai <ertsai@nvidia.com>
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #59894 [ run ] triggered by Bot. Commit: |
|
PR_Github #59894 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60216 [ run ] triggered by Bot. Commit: |
|
PR_Github #60216 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60289 [ run ] triggered by Bot. Commit: |
|
PR_Github #60289 [ run ] completed with state |
…NVIDIA#16106) Signed-off-by: Eric Tsai <ertsai@nvidia.com>
5 participants
Summary by CodeRabbit
Description
When
kv_cache_config.host_cache_sizeis not set,KVCacheManagerV2auto-provisions a host cache tier per rank, capped at 50% of the node's
available memory. The cap is computed by each rank independently from the
same node-level reading, so N co-located ranks can collectively reserve up
to N x 50% of node memory.
Observed on GB300 NVL72 (DSR1 FP4 disagg, decode TP16 = 4 nodes x 4 GPUs,
~173 GiB device quota, 975 GiB node RAM): the 4 ranks on each node each
provisioned a 161-173 GiB host tier (~650 GiB total vs ~440 GiB available),
and decode workers were reliably OOM-killed during engine init (6/6
reproductions). The same config with the v1 manager runs fine.
Fix: count ranks sharing the physical node (hostname allgather via the
existing
Distributedhelper, falling back tomapping.gpus_per_node) anddivide the available-memory cap by that count. Explicitly-set
host_cache_sizeis unaffected.Test Coverage
Manually validated on GB300 NVL72 with the previously-failing config:
local_ranks=4detected correctly; host tier reduced to ~55 GiB perrank; 4/4 runs completed 43010/43010 requests with zero oom_kill events
(previously 6/6 OOM).
PR 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.