[https://nvbugs/6550276][fix] Clamp residency to what the quota affords (allreduce-MIN across ranks) behind… - #17261
Conversation
…quota When max_batch_size asks for more fixed-size recurrent state than the GPU cache quota can hold, the V2 Mamba state pool is unbuildable at any free_gpu_memory_fraction (513 slots need 37.5 GiB of the 24.9 GiB an 80 GiB H100 has left after qwen3.5_27b weights). Bound the resident set to what the quota affords and have the scheduler queue the excess via a new max_resident_sequences() hook, which reports None (unbounded) for plain attention models and whenever the quota affords every requested sequence. The inherited warmup constraints are truncated to the clamped residency too: every constraint entry costs one SSM slot, so a constraint built from the raw max_batch_size would re-impose the very floor the clamp removes. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
WalkthroughThe change adds an optional resident-sequence cap to KV cache managers, computes quota-based residency for Mamba caches across ranks, enforces admission limits in the scheduler, and adds scheduler and cache-manager regression tests. ChangesKV cache residency control
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant KVCacheV2Scheduler
participant KVCacheManagerV2
participant MambaCacheManager
participant Distributed
KVCacheV2Scheduler->>KVCacheManagerV2: read max_resident_sequences()
MambaCacheManager->>Distributed: reduce quota-affordable capacity
Distributed-->>MambaCacheManager: return clamped capacity
KVCacheV2Scheduler->>KVCacheV2Scheduler: count resident started requests
KVCacheV2Scheduler->>KVCacheManagerV2: admit first context chunk when capacity remains
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 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: 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/mamba_cache_manager.py`:
- Around line 2758-2762: Update max_resident_sequences to return the allreduced
_resident_sequence_cap before checking local_num_mamba_layers, returning None
only when no distributed cap was inherited; preserve the local
_max_resident_sequences calculation for Mamba ranks without a cap. Add a mixed
pipeline regression covering one Mamba rank and one attention-only rank to
verify both ranks enforce the same admission cap.
In `@tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py`:
- Around line 318-326: Extend the residency accounting in the scheduler around
_is_started_request and the Phase 1 DISAGG_GENERATION_INIT path so an
in-progress disaggregated initialization consumes a resident slot before
prepare_disagg_gen_init() runs. Apply max_resident_sequences to admission at
that path, and retain the slot until the request releases its cache. Add a
regression test with cap=1 that submits two disaggregated generation
initializations and verifies the second is not admitted concurrently.
- Around line 422-426: The resident-limit guard in the context scheduling loop
must skip only the capped first-chunk request; replace the `break` under
`starts_new_sequence` with `continue` so later non-first chunks in `pending_ctx`
remain schedulable. Add a regression test covering enforced resident limits
where a capped first chunk precedes another request’s non-first chunk, verifying
the latter is scheduled and the scheduler does not deadlock.
🪄 Autofix
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: 01870fef-8a4a-49a7-a8d8-88cc5d99353f
📒 Files selected for processing (5)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pytests/unittest/_torch/executor/test_kv_cache_v2_scheduler.pytests/unittest/_torch/executor/test_mamba_cache_manager.py
| def max_resident_sequences(self) -> Optional[int]: | ||
| """Number of sequences whose recurrent state can be resident at once.""" | ||
| if self.local_num_mamba_layers == 0: | ||
| return None | ||
| return self._max_resident_sequences() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Keep the allreduced cap on attention-only pipeline ranks.
Lines 2934-2941 calculate and store the reduced cap on every rank. Line 2760 returns None before reading that cap when a hybrid pipeline stage has no local Mamba layers. The Mamba stage then limits admission while the attention-only stage admits an uncapped batch. This can desynchronize pipeline execution.
Return _resident_sequence_cap before the local-Mamba check. Return None only when the attention-only rank did not inherit a distributed cap. Add a mixed PP regression with one Mamba rank and one attention-only rank.
Proposed fix
def max_resident_sequences(self) -> Optional[int]:
"""Number of sequences whose recurrent state can be resident at once."""
+ if self._resident_sequence_cap is not None:
+ return self._resident_sequence_cap
if self.local_num_mamba_layers == 0:
return None
return self._max_resident_sequences()Also applies to: 2934-2941
🤖 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/mamba_cache_manager.py` around lines 2758 -
2762, Update max_resident_sequences to return the allreduced
_resident_sequence_cap before checking local_num_mamba_layers, returning None
only when no distributed cap was inherited; preserve the local
_max_resident_sequences calculation for Mamba ranks without a cap. Add a mixed
pipeline regression covering one Mamba rank and one attention-only rank to
verify both ranks enforce the same admission cap.
| # Sequences already holding a non-droppable state slot. Counted over all | ||
| # active requests (not just the ones scheduled this iteration) because a | ||
| # suspended sequence keeps its slot. | ||
| max_resident = self.max_resident_sequences | ||
| num_resident = ( | ||
| sum(1 for req in requests_list if self._is_started_request(req)) | ||
| if max_resident is not None | ||
| else 0 | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Apply the residency cap to DISAGG_GENERATION_INIT.
At Line 321, the scheduler counts only started context and generation requests. The only new admission gate starts at Line 425, after the Phase 1 DISAGG_GENERATION_INIT path. That path calls prepare_disagg_gen_init(), which creates and resizes the primary cache. For Mamba, this consumes non-droppable recurrent-state slots.
Gate disaggregated initialization before prepare_disagg_gen_init(). Count its slot until the request releases its cache. Add a cap=1 regression that submits two disaggregated generation initializations.
Also applies to: 422-426
🤖 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 318 -
326, Extend the residency accounting in the scheduler around _is_started_request
and the Phase 1 DISAGG_GENERATION_INIT path so an in-progress disaggregated
initialization consumes a resident slot before prepare_disagg_gen_init() runs.
Apply max_resident_sequences to admission at that path, and retain the slot
until the request releases its cache. Add a regression test with cap=1 that
submits two disaggregated generation initializations and verifies the second is
not admitted concurrently.
| # A first context chunk starts a new sequence and therefore claims a | ||
| # state slot for the rest of its lifetime. | ||
| starts_new_sequence = max_resident is not None and req.is_first_context_chunk | ||
| if starts_new_sequence and num_resident >= max_resident: | ||
| break |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 6 --glob '*.py' \
'active_requests|RequestList|schedule_request\(' \
tensorrt_llm/_torch/pyexecutorRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
# First, let's look at the specific file and lines mentioned
cd /tmp && find . -name "scheduler_v2.py" -type f 2>/dev/null | head -1 | xargs wc -lRepository: NVIDIA/TensorRT-LLM
Length of output: 159
🏁 Script executed:
# Read the scheduler_v2.py file around lines 422-426
sed -n '400,450p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2377
🏁 Script executed:
# Get a broader view of the function containing this code
sed -n '350,500p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 6700
🏁 Script executed:
# Search for context about how active_requests is populated and used in scheduler_v2
ast-grep outline tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py --view expandedRepository: NVIDIA/TensorRT-LLM
Length of output: 2072
🏁 Script executed:
# Look at the full _schedule_loop function signature and Phase 1 logic
sed -n '258,350p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 4545
🏁 Script executed:
# Check how pending_ctx is used and whether chunks are separate request objects
rg -n "pending_ctx|is_first_context_chunk|is_last_context_chunk" tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py | head -40Repository: NVIDIA/TensorRT-LLM
Length of output: 761
🏁 Script executed:
# Look for request progression through context chunks to understand the lifecycle
rg -n "context_current_position|chunk" tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py | head -30Repository: NVIDIA/TensorRT-LLM
Length of output: 1945
🏁 Script executed:
# Check the request list iteration to see if a single request can have multiple "entries"
sed -n '230,260p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1367
🏁 Script executed:
# Now let me understand how requests progress through chunks across iterations
# Check if is_first_context_chunk changes during a request's lifetime
rg -n "context_current_position|is_first_context_chunk" tensorrt_llm/_torch/pyexecutor -A 2 -B 2 | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 6145
🏁 Script executed:
# Let me look at how requests move between scheduling phases
# Understanding whether a request stays in active_requests or gets removed/re-added
sed -n '419,450p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1712
🏁 Script executed:
# Check what happens at the end of Phase 2 and how requests transition
sed -n '450,500p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2427
🏁 Script executed:
# Look for how context position advances to understand the chunk progression
rg -n "context_current_position\s*=" tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py -A 2 -B 2Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Let me understand if requests can appear multiple times in pending_ctx
# and how chunks advance across iterations
sed -n '530,650p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 5632
🏁 Script executed:
# Look at the chunked context scheduling to see how chunks progress
sed -n '591,700p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 5018
🏁 Script executed:
# Search for how context chunks are handled and if a request appears multiple times
rg -n "is_context_finished|context_remaining_length" tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 654
🏁 Script executed:
# Check if there are any comments or code about request ordering in pending_ctx
sed -n '300,320p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1357
🏁 Script executed:
# Look for any sorting or ordering logic applied to pending_ctx or active_requests
rg -n "sort|pending_ctx\|active_requests" tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py | grep -E "(sort|pending_ctx)"Repository: NVIDIA/TensorRT-LLM
Length of output: 788
🏁 Script executed:
# Verify: after a request is scheduled in Phase 2, does its state change to advance chunks?
# Look for what happens to context_current_position or chunk progression
sed -n '440,445p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 434
🏁 Script executed:
# Critical: verify whether the same request can appear multiple times in a single pending_ctx loop
# and whether is_first_context_chunk can change within the same iteration
rg -n "is_first_context_chunk\|context_current_position" tensorrt_llm/_torch/pyexecutor/llm_request.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Check how requests are ordered in active_requests and if first chunks are guaranteed to come first
# Look at the py_executor to see if active_requests maintains any ordering
rg -n "active_requests\s*=" tensorrt_llm/_torch/pyexecutor/py_executor.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1270
🏁 Script executed:
# Look for evidence of how requests progress through chunks
# Check if context_current_position is modified during scheduling
rg -n "context_current_position\s*=" tensorrt_llm/_torch/pyexecutor -B 2 -A 2 | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 3019
🏁 Script executed:
# Specifically check the update_requests function that advances context position
rg -n "update_requests|_update_request" tensorrt_llm/_torch/pyexecutor/py_executor.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1671
🏁 Script executed:
# Critical: trace whether the same request can appear in pending_ctx across iterations
# Check the flow: is the request removed from active_requests after scheduling a chunk?
rg -n "remove\|pop\|del.*active_requests" tensorrt_llm/_torch/pyexecutor/py_executor.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Verify the exact flow: if a request has chunks, does it appear in pending_ctx with
# is_first_context_chunk changing between iterations?
# Check how a request transitions from first chunk to non-first chunk
sed -n '1083,1100p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1002
🏁 Script executed:
# Look at _sort_requests to see how context_requests are ordered
sed -n '1083,1100p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1002
🏁 Script executed:
# Let me trace the actual problem: when is_first_context_chunk FALSE for a request?
rg -n "is_first_context_chunk.*=" tensorrt_llm/_torch/pyexecutor -B 3 -A 1 | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Confirm: does pending_ctx maintain FIFO order of active_requests?
# And can a chunked request have its first chunk in one iteration and subsequent in the next?
sed -n '385,395p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 567
🏁 Script executed:
# Understand: in a single active_requests list, can we have req1 (first chunk) then req2 (non-first chunk)?
# This would happen if req1 just arrived and req2 has already been partially scheduled
# Check if requests stay in active_requests during chunk progression
rg -n "is_context_finished\|context_remaining_length\|is_last_context_chunk" tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py -B 1 -A 1Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Final verification: trace through a multi-chunk scenario
# If request A schedules first chunk in iteration 1, does it stay in active_requests for iteration 2?
# and can its second chunk then encounter the resident cap?
sed -n '419,428p' tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 655
🏁 Script executed:
# Check if there's any guarantee about ordering: are first chunks always at the head?
# Or can they be interleaved?
rg -n "\.append\(req\)" tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 429
Skip capped first-chunk requests with continue, not break. A first context chunk that hits the resident-sequence limit should be skipped in that iteration, not exit the context phase entirely. If a non-first chunk of another request follows the capped first chunk in pending_ctx, the break statement prevents it from being scheduled. This can deadlock the scheduler: the first chunk's request holds a resident slot without being scheduled, while the non-first chunk cannot advance to free or reuse cache. Change line 426 from break to continue and add a regression test that schedules two requests with resident-sequence limits enforced, where an earlier request's non-first chunk is iterated after a later request's capped first chunk.
🤖 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 422 -
426, The resident-limit guard in the context scheduling loop must skip only the
capped first-chunk request; replace the `break` under `starts_new_sequence` with
`continue` so later non-first chunks in `pending_ctx` remain schedulable. Add a
regression test covering enforced resident limits where a capped first chunk
precedes another request’s non-first chunk, verifying the latter is scheduled
and the scheduler does not deadlock.
2 participants
Summary
_max_resident_sequences()returned rawmax_batch_size * pp_size(512), turning a batch-size ceiling into a hard 37.48 GiB floor of non-droppable recurrent-state slots that nofree_gpu_memory_fractioncan satisfy against the 24.86 GiB quota.max_resident_sequences()hook the V2 scheduler gates first-context-chunk admission on, and truncate the inherited warmup constraints to that clamp so all three sizing paths agree.pytest tests/integration/defs/perf/test_perf.py --perf --test-list=.repair-bot/perf_test_list.txt --output-dir=build/perf_output -vTest plan
Links
Dev Engineer Review
KVCacheManagerV2.max_resident_sequences()provides an optional scheduler cap.QA Engineer Review
max_batch_size.tests/integration/test_lists/coverage data is available.