[https://nvbugs/6575476][fix] Keep the residency fix sizing on self.max_seq_len (scales per case: 2500 →… - #17440
Open
trtllm-agent wants to merge 5 commits into
Open
[https://nvbugs/6575476][fix] Keep the residency fix sizing on self.max_seq_len (scales per case: 2500 →…#17440trtllm-agent wants to merge 5 commits into
self.max_seq_len (scales per case: 2500 →…#17440trtllm-agent wants to merge 5 commits into
Conversation
…10b 4-GPU perf cases Both cases reported in this bug are already excluded from 80G H100 at HEAD by f9b2457 (NVIDIA#17303), which landed after the bug's failing run (bf1ddb7, 2026-08-03), so the reported OOM needs no further memory fix. Confirmed by running the coverage audit below against the failing commit: at bf1ddb7 both ids selected H100; at HEAD the 122b case loses H100 to a gpu_memory gate and the 397b case loses it to condition 7's compute_capability gte 10.0. That re-gate, however, parked a gpus:4 case in condition 10, which bundles system_gpu_count gte 8 alongside the gpu_memory gt 90000 gate it was moved for. GB200 and GB300 CI nodes are 4-GPU, so both qwen3.5_122b_a10b 4-GPU cases silently stopped running there despite having 186G+ per GPU. Add condition 10c (system_gpu_count gte 4 + gpu_memory gt 90000) and host the two cases there, keeping 80G H100 excluded while restoring GB200/GB300. Verified by parsing the yml per-platform before/after and diffing the whole file's coverage map: 164 tests tracked, exactly 2 changed, both purely additive (+GB200, +GB300), no platform lost anywhere. That the H100 exclusion is the correct remedy rather than a masked code bug is confirmed arithmetically from a run where the weights were already cached. On an 80G H100 at tp:4 the KV manager reports a 14.49 GiB device quota against a 9.33 GiB fixed cost, leaving 5.15 GiB = 900,702 tokens; the default maxbs:512 / max_seq_len:2500 shape needs 1,280,000 tokens, a 379,298-token shortfall. Under GUARANTEED_NO_EVICT with no host cache tier that surfaces as "V2 scheduler deadlock", matching the owning engineer's diagnosis on the parent bug 6550276 that the batch size, not the runtime, is at fault. Note this change is not observable through the reproduce command: --test-list= feeds node ids directly and bypasses yml condition gating entirely, so no EXIT_CODE can exercise it. The two cases also take 1897s + 1907s, exceeding the harness 2260s budget, so that command times out (rc=124) independently of any code change. The static coverage diff above is the applicable verification. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…17b_fp8 1000,2000 The companion commit restored GB200/GB300 for this bug's first reported case. Its second reported case, qwen3.5_397b_a17b_fp8 …1000,2000-ep:8-tp:8-gpus:8, has the same class of collateral loss, which that commit missed. f9b2457 (NVIDIA#17303) moved this case out of condition 9 (H100, H20, B200, B300, RTX6000-Server) into condition 7, gated compute_capability gte 10.0. That does exclude the 80G H100 where the bug reproduces, but it also drops H20 and RTX6000-Server, both of which are 8-GPU parts with ~95.6G per GPU. The bug's own regression table records this case PASSING on the smaller 79G H100 on 2026-07-31, so a 95.6G part is not the constraint; the reported failure is specific to running out of room on an 80G part. Express that as the memory gate it actually is: condition 10 is condition 9 plus gpu_memory gt 90000 (i.e. "same platforms, minus the 80G H100"), and already hosts the sibling qwen3.5_397b_a17b_fp4 8-GPU cases. Move the case there. Its own siblings at other ISL/OSL shapes stay in condition 9, matching how the first reported case was split from its siblings. Whole-file per-platform coverage audit, origin/main vs HEAD: 164 test ids tracked, exactly 3 changed, every one purely additive with nothing removed: 122b_a10b …128,128-ep:4 + [GB200, GB300] 122b_a10b …500,2000-ep:4 + [GB200, GB300] 397b_a17b_fp8 …1000,2000 + [H20, RTX6000-S] Both reported cases still exclude H100, which is the defect this bug reports. Condition 10's platform set is unchanged by the move, so its "# 10:" index comment stays accurate and needs no edit. That excluding the 80G H100 is the correct remedy rather than a masked runtime bug is confirmed by a forced run of the first case on H100, which reproduces the reported "CUDA out of memory" verbatim: the KV manager reports a 14.49 GiB device quota against a 10.02 GiB fixed cost and 6144 bytes/token, while the default maxbs:512 x max_seq_len:2500 shape needs 1,280,000 tokens. This matches the owning engineer's diagnosis on parent bug 6550276 that the batch size, not the runtime, is at fault. This change is not observable through the reproduce command: --test-list= feeds node ids directly and bypasses yml condition gating, so no EXIT_CODE exercises it. The two listed cases also take 1897s + 1907s against a 2260s budget, so that command times out regardless of any diff. The static coverage audit above is the applicable verification. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Qwen3.5 MoE runs on 80G H100 die mid-generation, matching the reported kv_cache_manager_v2.CuOOMError / "Error in event loop: CUDA out of memory" signature, or trip the V2 scheduler's deadlock guard. MambaHybridCacheManagerV2 bounds how many sequences may be concurrently resident, but it sized that bound with _get_typical_request_capacity, which is a pool-*ratio* hint and falls back to max_seq_len / 2 when avg_seq_len is unset. A residency bound is a guarantee, not an average, so on workloads where every sequence grows to max_seq_len the bound admits roughly twice the sequences the quota can hold: at tp=4 the quota affords 443 sequences at full length but 567 by the halved hint, so all 512 are admitted. A hybrid sequence's recurrent state is fixed-size and non-droppable and _KVCache.resume is refused past max_util_for_resume, so the over-admitted run cannot recover; it dies when physical KV allocation fails or when nothing is schedulable and nothing is evictable. Size the residency bound on max(typical_capacity, max_seq_len). avg_seq_len stays authoritative for the pool ratio; only the guarantee changes. The residency machinery this relies on was written for the same defect on a smaller model (nvbugs/6550276) but never merged, so it is included here: the max_resident_sequences() hook (None for plain-attention managers, which keep today's unbounded admission), the scheduler admission gate, and truncation of the inherited warmup constraints, without which a constraint built from the raw max_batch_size re-imposes the SSM floor the bound removes. The new sizing test fails on the halved hint and passes on reachable capacity. The shared quota helper of the two pre-existing residency tests moves from 2 blocks to 4 because it hard-coded the max_seq_len / 2 assumption; their assertions are unchanged. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Earlier commits on this branch moved the two reported cases into memory-gated conditions on the premise that 80G H100 cannot run them. With the residency fix that premise no longer holds: qwen3.5_122b_a10b input_output_len:500,2000-ep:4-tp:4-gpus:4 now passes on an 80G H100 at 1820.75 tok/s output / 2275.94 tok/s total, so gating it off that GPU would hide the coverage the fix restores. Restores llm_perf_core.yml to origin/main and leaves this branch as a single-concern runtime fix. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…ager mock Mock(spec=KVCacheManagerV2) auto-vivifies a truthy Mock for the new max_resident_sequences() hook. These two tests pass only because they never reach schedule_request; a future test that does would compare Mock >= int. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
trtllm-agent
requested review from
liji-nv,
mikeiovine,
yizhang-nv and
zhaoyangwang-nvidia
August 8, 2026 16:37
Contributor
WalkthroughThe change adds a resident-sequence capacity hook, derives quota-based limits for Mamba V2 caches, enforces those limits during scheduling, and adds regression tests for bounded and unbounded behavior. ChangesResident sequence capacity
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MambaCacheManagerV2
participant Distributed
participant KVCacheV2Scheduler
participant ContextRequests
MambaCacheManagerV2->>Distributed: Reduce minimum resident capacity
KVCacheV2Scheduler->>MambaCacheManagerV2: Read max_resident_sequences()
ContextRequests->>KVCacheV2Scheduler: Submit context requests
KVCacheV2Scheduler->>KVCacheV2Scheduler: Admit first chunks up to capacity
Possibly related PRs
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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/scheduler/scheduler_v2.py`:
- Around line 409-426: Update resident-slot accounting in the scheduling flow
around prepare_context() and DISAGG_GENERATION_INIT so every request that has
allocated KV/Mamba state is counted before admitting more context chunks. Do not
rely solely on _is_started_request(), which excludes first context chunks; use
actual resident-cache ownership or record the slot at first allocation,
including disaggregated initialization, while preserving the
max_resident_sequences cap.
🪄 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: 3e119a5e-b88b-434a-8c30-210738d8de45
📒 Files selected for processing (6)
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_dual_pool_kv_cache.pytests/unittest/_torch/executor/test_kv_cache_v2_scheduler.pytests/unittest/_torch/executor/test_mamba_cache_manager.py
Comment on lines
+409
to
+426
| # | ||
| # 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 | ||
| ) | ||
| for req in pending_ctx: | ||
| if budget.requests_full: | ||
| break | ||
| # 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 |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Count allocated resident slots before context admission.
_is_started_request() excludes a first context chunk. However, prepare_context() has already allocated its KV cache and Mamba state slot before that chunk completes. If that request remains in inflight_request_ids during an overlap iteration, num_resident omits it and the scheduler can admit up to another full cap of first chunks.
DISAGG_GENERATION_INIT also allocates before phase 2 and bypasses this gate. Count actual resident cache ownership, or claim a slot on every first allocation, including disaggregated initialization.
🤖 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 409 -
426, Update resident-slot accounting in the scheduling flow around
prepare_context() and DISAGG_GENERATION_INIT so every request that has allocated
KV/Mamba state is counted before admitting more context chunks. Do not rely
solely on _is_started_request(), which excludes first context chunks; use actual
resident-cache ownership or record the slot at first allocation, including
disaggregated initialization, while preserving the max_resident_sequences cap.
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
_get_typical_request_capacity, an average pool-ratio hint that falls back tomax_seq_len // 2, so admission allowed ~2x the sequences the quota can hold at full length; non-droppable SSM state plusresume()refused pastmax_util_for_resumemakes the over-admission unrecoverable.self.max_seq_len(scales per case: 2500 → 3000), and add the missing test hardening — pinmax_resident_sequences → NoneonMock(spec=KVCacheManagerV2)so the new hook cannot auto-vivify a truthy Mock into the scheduler's residency comparison.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()with an unlimited default.max_seq_len._attention_block_bytes()for consistent quota calculations.QA Engineer Review
Changed test files:
tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.pytests/unittest/_torch/executor/test_mamba_cache_manager.pymax_seq_lensizing, inherited constraints, ample quotas, and attention-only managers.tests/unittest/_torch/executor/test_dual_pool_kv_cache.pyNonefor unlimited residency.No corresponding
test-db/orqa/coverage entries were provided. Verdict: needs follow-up.