[https://nvbugs/6550275][fix] Add KVCacheManagerV2.max_resident_sequences() from per-pool-group page counts… - #17211
Open
trtllm-agent wants to merge 1 commit into
Open
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) | ~20 minutes Possibly related PRs
Suggested labels: 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
Comment on lines
+2177
to
+2188
| # Charging every attention pool the full max_seq_len is deliberately | ||
| # worst-case: a genuinely sliding-window pool needs fewer pages per | ||
| # sequence, so the bound errs low rather than over-admitting. | ||
| pages_per_seq = math.ceil(self.max_seq_len / self.tokens_per_block) | ||
| # A state pool holds one fixed slot per sequence; an attention pool | ||
| # must hold every page of each resident sequence. | ||
| return max( | ||
| 1, | ||
| min( | ||
| stat.total // (1 if pool_group_id in state_pool_groups else pages_per_seq) | ||
| for pool_group_id, stat in enumerate(self._get_storage_statistics(GPU_LEVEL)) | ||
| ), |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Charge attention pools for the full maximum sequence allocation.
Line 2180 excludes num_extra_kv_tokens, _kv_reserve_draft_tokens, and the base decode token. The manager reserves all of these in max_seq_capacity at Lines 1121-1127.
If an attention pool is the limiting pool, this calculation can return a limit that exceeds the pages required by fully allocated resident sequences. The scheduler can then over-admit and reintroduce the allocation deadlock that this limit is intended to prevent.
Use the same maximum capacity basis as max_blocks_per_seq. Add a unit test where the extra allocation crosses a block boundary.
Proposed fix
- pages_per_seq = math.ceil(self.max_seq_len / self.tokens_per_block)
+ max_sequence_capacity = (
+ self.max_seq_len
+ + self.num_extra_kv_tokens
+ + self._kv_reserve_draft_tokens
+ + 1
+ )
+ pages_per_seq = math.ceil(max_sequence_capacity / self.tokens_per_block)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # Charging every attention pool the full max_seq_len is deliberately | |
| # worst-case: a genuinely sliding-window pool needs fewer pages per | |
| # sequence, so the bound errs low rather than over-admitting. | |
| pages_per_seq = math.ceil(self.max_seq_len / self.tokens_per_block) | |
| # A state pool holds one fixed slot per sequence; an attention pool | |
| # must hold every page of each resident sequence. | |
| return max( | |
| 1, | |
| min( | |
| stat.total // (1 if pool_group_id in state_pool_groups else pages_per_seq) | |
| for pool_group_id, stat in enumerate(self._get_storage_statistics(GPU_LEVEL)) | |
| ), | |
| # Charging every attention pool the full max_seq_len is deliberately | |
| # worst-case: a genuinely sliding-window pool needs fewer pages per | |
| # sequence, so the bound errs low rather than over-admitting. | |
| max_sequence_capacity = ( | |
| self.max_seq_len | |
| self.num_extra_kv_tokens | |
| self._kv_reserve_draft_tokens | |
| 1 | |
| ) | |
| pages_per_seq = math.ceil(max_sequence_capacity / self.tokens_per_block) | |
| # A state pool holds one fixed slot per sequence; an attention pool | |
| # must hold every page of each resident sequence. | |
| return max( | |
| 1, | |
| min( | |
| stat.total // (1 if pool_group_id in state_pool_groups else pages_per_seq) | |
| for pool_group_id, stat in enumerate(self._get_storage_statistics(GPU_LEVEL)) | |
| ), |
🤖 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/kv_cache_manager_v2.py` around lines 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.
Comment on lines
+407
to
+425
| # | ||
| # Starting a new sequence is what grows the resident set, so the | ||
| # residency cap is enforced here. Requests already started keep | ||
| # their slot; the rest wait in the queue until one drains. | ||
| residency_cap = self.max_resident_sequences | ||
| num_started = ( | ||
| sum(1 for r in requests_list if self._is_started_request(r)) | ||
| if residency_cap is not None | ||
| else 0 | ||
| ) | ||
|
|
||
| for req in pending_ctx: | ||
| if budget.requests_full: | ||
| break | ||
| # Read before scheduling: _try_schedule_context advances the | ||
| # request past its first chunk. | ||
| is_new_sequence = residency_cap is not None and req.is_first_context_chunk | ||
| if is_new_sequence and num_started >= residency_cap: | ||
| break |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Apply the residency cap to disaggregated generation initialization.
DISAGG_GENERATION_INIT runs in Phase 1 at Lines 338-357. It calls prepare_disagg_gen_init(), which creates and resizes a primary KV cache. Lines 411-425 run only after every Phase 1 candidate has been processed.
A hybrid disaggregated generation worker can therefore admit new resident sequences without this cap. This path can exceed the non-droppable-state capacity and reproduce the suspend/resume deadlock.
Initialize resident accounting before Phase 1. Gate successful disaggregated generation initialization with the same limit. Include already initialized disaggregated requests in subsequent scheduler iterations. Add a regression test for capped DISAGG_GENERATION_INIT admission.
Also applies to: 441-442
🤖 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/scheduler/scheduler_v2.py` around lines 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.
Comment on lines
+2229
to
+2268
| class TestResidencyCap: | ||
| """A manager reporting max_resident_sequences bounds newly started | ||
| sequences, so a pool whose state cannot be evicted never over-subscribes. | ||
| """ | ||
|
|
||
| def test_none_leaves_admission_unbounded(self): | ||
| """Plain attention models report None and keep MAX_UTILIZATION behavior.""" | ||
| mgr = make_kv_cache_manager(max_resident_sequences=None) | ||
| sched = make_scheduler(mgr) | ||
| reqs = [make_ctx_request(i, context_remaining_length=10) for i in range(20)] | ||
| out = sched.schedule_request(reqs, set()) | ||
| assert len(out.context_requests) == 20 | ||
|
|
||
| def test_caps_newly_started_sequences(self): | ||
| mgr = make_kv_cache_manager(max_resident_sequences=3) | ||
| sched = make_scheduler(mgr) | ||
| reqs = [make_ctx_request(i, context_remaining_length=10) for i in range(20)] | ||
| out = sched.schedule_request(reqs, set()) | ||
| assert ids(out.context_requests) == [0, 1, 2] | ||
|
|
||
| def test_already_started_sequences_consume_the_cap(self): | ||
| """In-progress generation holds its slot, so no new context is admitted.""" | ||
| mgr = make_kv_cache_manager(max_resident_sequences=2) | ||
| sched = make_scheduler(mgr) | ||
| reqs = [ | ||
| make_gen_request(0), | ||
| make_gen_request(1), | ||
| make_ctx_request(2, context_remaining_length=10), | ||
| ] | ||
| out = sched.schedule_request(reqs, set()) | ||
| assert ids(out.generation_requests) == [0, 1] | ||
| assert ids(out.context_requests) == [] | ||
|
|
||
| def test_continuing_chunk_is_not_charged_again(self): | ||
| """Only a first chunk starts a sequence; later chunks keep their slot.""" | ||
| mgr = make_kv_cache_manager(max_resident_sequences=1) | ||
| sched = make_scheduler(mgr) | ||
| req = make_ctx_request(0, context_remaining_length=10, is_first_context_chunk=False) | ||
| out = sched.schedule_request([req], set()) | ||
| assert ids(out.context_requests) == [0] |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Test coverage is insufficient.
The added tests cover unbounded admission, capped new context admission, active generation occupancy, and continuing context chunks.
They do not test KVCacheManagerV2.max_resident_sequences() with storage statistics, block rounding, and extra KV allocation. They also do not test capped DISAGG_GENERATION_INIT admission.
Test coverage summary
- Added tests:
test_none_leaves_admission_unbounded,test_caps_newly_started_sequences,test_already_started_sequences_consume_the_cap, andtest_continuing_chunk_is_not_charged_again. - Test-list status: No relevant
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/file was included for review. - Verdict: insufficient.
Add direct manager calculation tests and a disaggregated generation initialization regression test.
🤖 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_v2_scheduler.py` around lines
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.
Source: Path instructions
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
2 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
Test plan
Links
Dev Engineer Review
KVCacheManagerV2.max_resident_sequences()for models with non-droppable state pools.Nonefor attention-only models and preserves existing scheduler behavior.MAX_UTILIZATIONexceeds resident KV capacity.QA Engineer Review
tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py.max_resident_sequences.tests/integration/test_lists/was identified.