[https://nvbugs/6550275][fix] Bound V2 residency for non-droppable state pools - #17211
Open
trtllm-agent wants to merge 2 commits into
Open
[https://nvbugs/6550275][fix] Bound V2 residency for non-droppable state pools#17211trtllm-agent wants to merge 2 commits into
trtllm-agent wants to merge 2 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>
trtllm-agent
requested review from
allisonlim-nv,
liji-nv,
nvpohanh and
tongyuantongyu
August 3, 2026 14:47
Contributor
Walkthrough
ChangesResident-sequence admission
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
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}
Contributor
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>
Collaborator
|
/bot run |
Contributor
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
Comment on lines
+43
to
+61
| 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 |
Contributor
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/.
Collaborator
|
PR_Github #63725 [ run ] triggered by Bot. Commit: |
VALLIS-NERIA
changed the title
Collaborator
|
PR_Github #63725 [ run ] completed with state
|
2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Noneand keep existing behavior.Test plan
Links
Dev Engineer Review
KVCacheManagerV2.max_resident_sequences().Nonefor attention-only models.MAX_UTILIZATIONexceeds resident KV capacity.QA Engineer Review
tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py.tests/unittest/_torch/executor/test_kv_cache_manager_v2.py.tests/integration/test_lists/were identified.