[https://nvbugs/6550275][fix] Bound V2 residency for non-droppable state pools - #17211
[https://nvbugs/6550275][fix] Bound V2 residency for non-droppable state pools#17211trtllm-agent wants to merge 3 commits into
Conversation
…state pools MAX_UTILIZATION admits new sequences up to max_batch_size and relies on suspend/resume to survive over-subscription. That recovery needs some resident sequence to still be evictable so its pages can be freed for a suspended one to resume. A hybrid Mamba recurrent state is fixed-size per sequence and cannot be recomputed from tokens, so such a sequence yields no evictable pages; once every sequence is suspended the pool can no longer drain and the scheduler stops making progress. On L40S, qwen3.5_9b at 500-in/2000-out admitted 485 sequences while the attention pool holds only ~118 at max_seq_len, then spun 8676 iterations scheduling nothing before raising 'V2 scheduler deadlock'. Add KVCacheManagerV2.max_resident_sequences(), derived from per-pool-group page counts, and gate new-sequence admission on it. Returns None (unbounded, unchanged behavior) for models without a non-droppable state pool. Pin the return value on the scheduler test's manager double: a bare Mock() is not None, so the new gate would otherwise compare an int against a child Mock and raise TypeError in every test that schedules a first context chunk. Add direct coverage for the cap, which had none. Two config levers suggested by triage were measured and do not fix this: avg_seq_len=2500 reproduces the failure identically, and max_util_for_resume=1.0 replaces the raise with an unbounded livelock. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Walkthrough
ChangesResident-sequence admission
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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: 3
🤖 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 2177-2188: Update the capacity calculation around
max_blocks_per_seq so attention pools are charged using the full maximum
per-sequence allocation, including num_extra_kv_tokens,
_kv_reserve_draft_tokens, and the base decode token, matching max_seq_capacity
and max_blocks_per_seq. Preserve the fixed one-slot divisor for state pools, and
add a unit test where the extra allocation crosses a block boundary to verify
the scheduler cannot over-admit.
In `@tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py`:
- Around line 407-425: Extend resident-sequence accounting to the Phase 1
DISAGG_GENERATION_INIT path: initialize the count before Phase 1, reject or
defer new disaggregated requests when the cap is reached, and increment/account
only successful prepare_disagg_gen_init() admissions. Ensure already initialized
disaggregated requests are included in num_started on later scheduler
iterations, and add a regression test covering capped DISAGG_GENERATION_INIT
admission.
In `@tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py`:
- Around line 2229-2268: Extend TestResidencyCap with direct
KVCacheManagerV2.max_resident_sequences() tests covering storage statistics,
block-size rounding, and extra KV allocation, asserting the calculated residency
cap. Add a scheduler regression test for capped DISAGG_GENERATION_INIT
admission, verifying that generation initialization requests respect
max_resident_sequences while preserving existing admission behavior.
🪄 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: bd455d53-e826-4bbf-8e40-e075cea0e56c
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pytests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_kv_cache_manager_v2.py (1)
64-103: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for the
None-return branch and themax(1, ...)floor.The PR objective states that
max_resident_sequences()"returns None, preserving existing behavior" for models without a non-droppable state pool. This test file only exercises the non-Nonebranch (models with an"ssm"life cycle). Add a test withlife_cycle_metadatacontaining only"attention"entries to confirm theNonereturn path. Also consider a case wherestorage_stats[...].total // slots_per_sequencewould compute to 0, to confirm themax(1, ...)floor is exercised.Do you want me to generate these two additional test cases?
🤖 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 64 - 103, Add tests for max_resident_sequences() covering both missing non-droppable pools and the minimum capacity floor: create an attention-only life_cycle_metadata setup and assert the method returns None, then create a setup where total available slots divided by slots_per_sequence is zero and assert the result is 1. Reuse _make_residency_manager and the existing pool configuration patterns.
🤖 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 `@tests/unittest/_torch/executor/test_kv_cache_manager_v2.py`:
- Around line 43-61: Extend max_resident_sequences() tests to cover the
all-"attention" lifecycle case returning None and the zero-capacity case
returning 1 via the max(1, ...) floor, using the existing
_make_residency_manager helper. Register
test_max_resident_sequences_uses_full_rounded_capacity and
test_max_resident_sequences_sums_coalesced_life_cycles in the appropriate
test-db and qa test list files under tests/integration/test_lists/.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_kv_cache_manager_v2.py`:
- Around line 64-103: Add tests for max_resident_sequences() covering both
missing non-droppable pools and the minimum capacity floor: create an
attention-only life_cycle_metadata setup and assert the method returns None,
then create a setup where total available slots divided by slots_per_sequence is
zero and assert the result is 1. Reuse _make_residency_manager and the existing
pool configuration patterns.
🪄 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: bb066dec-2d10-4b51-8566-cf3a3f6e37de
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pytests/unittest/_torch/executor/test_kv_cache_manager_v2.pytests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
- tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
- tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py
| def _make_residency_manager( | ||
| *, | ||
| max_seq_len: int, | ||
| tokens_per_block: int, | ||
| num_extra_kv_tokens: int, | ||
| reserve_draft_tokens: int, | ||
| life_cycle_metadata: dict[int, tuple[int, int | None, str]], | ||
| pool_group_totals: list[int], | ||
| ) -> KVCacheManagerV2: | ||
| manager = object.__new__(KVCacheManagerV2) | ||
| manager.max_seq_len = max_seq_len | ||
| manager.tokens_per_block = tokens_per_block | ||
| manager.num_extra_kv_tokens = num_extra_kv_tokens | ||
| manager._kv_reserve_draft_tokens = reserve_draft_tokens | ||
| manager._stats_life_cycle_metadata = lambda: life_cycle_metadata | ||
| manager._get_storage_statistics = lambda _level: [ | ||
| SimpleNamespace(total=total) for total in pool_group_totals | ||
| ] | ||
| return manager |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Locate and inspect the _max_sequence_capacity implementation.
ast-grep run --pattern 'def _max_sequence_capacity($$$) {
$$$
}' --lang python tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
rg -n -B2 -A15 'def _max_sequence_capacity' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1569
🏁 Script executed:
# Inspect full max_resident_sequences method to check for None return and max(1,...) floor
rg -n -A30 'def max_resident_sequences' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
# Check for test registration in test lists
find tests/integration/test_lists -type f -name '*.txt' -o -name '*.yaml' -o -name '*.json' | head -20
# Check if test_kv_cache_manager_v2.py is referenced
grep -r 'test_kv_cache_manager_v2' tests/integration/test_lists/ 2>/dev/null || echo "Not found in test lists"
# Verify copyright header year in the test file
head -40 tests/unittest/_torch/executor/test_kv_cache_manager_v2.py | grep -i copyrightRepository: NVIDIA/TensorRT-LLM
Length of output: 2584
🏁 Script executed:
# Read the full test file to verify test functions and coverage
wc -l tests/unittest/_torch/executor/test_kv_cache_manager_v2.py
# Show test functions and their parameters
ast-grep outline tests/unittest/_torch/executor/test_kv_cache_manager_v2.py --view expanded
# Show lines 43-110 to see the helper and parametrized tests
sed -n '43,110p' tests/unittest/_torch/executor/test_kv_cache_manager_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 5358
Add test coverage for None-return path and max(1, ...) floor in max_resident_sequences(); register new tests in test list files.
The arithmetic assumptions in both tests are correct: _max_sequence_capacity() returns max_seq_len + num_extra_kv_tokens + _kv_reserve_draft_tokens + 1, so test expectations (6 and 10) are valid. However, the test suite omits two branches:
- Attention-only path (line 2172-2173): When all lifecycle kinds are
"attention", the method returnsNone. This branch is explicitly documented as preserving unbounded MAX_UTILIZATION behavior and should be tested. - Floor at 1 (line 2183-2184): The
max(1, ...)guard returns at least 1 resident sequence even if computed capacity would be 0. This edge case is not covered.
Additionally, per test-code guidelines, the two new test functions (test_max_resident_sequences_uses_full_rounded_capacity and test_max_resident_sequences_sums_coalesced_life_cycles) must be registered in the appropriate test list files under tests/integration/test_lists/ (test-db/ for CI, qa/ for manual QA). They are currently absent.
Test coverage verdict: Insufficient. Missing coverage for attention-only models and zero-capacity floor, and test registration incomplete.
🤖 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 43 -
61, Extend max_resident_sequences() tests to cover the all-"attention" lifecycle
case returning None and the zero-capacity case returning 1 via the max(1, ...)
floor, using the existing _make_residency_manager helper. Register
test_max_resident_sequences_uses_full_rounded_capacity and
test_max_resident_sequences_sums_coalesced_life_cycles in the appropriate
test-db and qa test list files under tests/integration/test_lists/.
|
PR_Github #63725 [ run ] triggered by Bot. Commit: |
VALLIS-NERIA
changed the title
|
PR_Github #63725 [ run ] completed with state
|
|
similar to #17261 |
|
adjust test cases |
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
|
PR_Github #64773 [ run ] triggered by Bot. Commit: |
|
PR_Github #64773 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65048 [ run ] triggered by Bot. Commit: |
|
PR_Github #65048 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65066 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #65066 [ run ] completed with state
|
|
PR_Github #65102 [ run ] triggered by Bot. Commit: |
|
PR_Github #65102 [ run ] completed with state
|
6 participants
Summary
GUARANTEED_NO_EVICTpolicy and reserves the blocks needed to finish a request before admitting it. V2 currently replaces every non-MAX_UTILIZATIONpolicy withMAX_UTILIZATION, budgets only the next scheduling step, and relies on suspend/resume to recover from over-admission.Test plan
ce537324a3is in progress.Links