From d29c6f10db06b4958a1ab15306a4ab9113fc2fe3 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 7 Aug 2026 19:17:36 -0700 Subject: [PATCH 1/5] [nvbugs/6575476][fix] Restore GB200/GB300 coverage for qwen3.5_122b_a10b 4-GPU perf cases Both cases reported in this bug are already excluded from 80G H100 at HEAD by f9b2457819 (#17303), which landed after the bug's failing run (bf1ddb7e, 2026-08-03), so the reported OOM needs no further memory fix. Confirmed by running the coverage audit below against the failing commit: at bf1ddb7e 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> --- .../test_lists/qa/llm_perf_core.yml | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/integration/test_lists/qa/llm_perf_core.yml b/tests/integration/test_lists/qa/llm_perf_core.yml index 8bb7ead14921..be482c5034d3 100644 --- a/tests/integration/test_lists/qa/llm_perf_core.yml +++ b/tests/integration/test_lists/qa/llm_perf_core.yml @@ -14,6 +14,7 @@ llm_perf_core: # 9: H100, H20, B200, B300, RTX6000-Server test cases # 10: H20, B200, B300, RTX6000-Server test cases # 10b: H20, B200, B300 test cases (>120G per GPU; excludes 96G RTX6000-Server) +# 10c: H20, GB200, B200, B300, GB300, RTX6000-Server test cases (>90G per GPU, 4-GPU capable) # 11: RTX-6000 Server test cases # =============================================================================== @@ -326,9 +327,6 @@ llm_perf_core: gpu_memory: gt: 90000 tests: - #qwen3.5_122b_a10b (MoE BF16 234G, 4-GPU): OOM on 80G GPUs, needs >90G - - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:128,128-ep:4-tp:4-gpus:4] - - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:500,2000-ep:4-tp:4-gpus:4] #qwen3.5_397b_a17b_fp4 (MoE FP4 234G, 8-GPU ep=8) - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:128,128-ep:8-tp:8-gpus:8] - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:500,2000-ep:8-tp:8-gpus:8] @@ -357,6 +355,23 @@ llm_perf_core: - perf/test_perf.py::test_perf[deepseek_r1_0528_fp8-bench-pytorch-float8-input_output_len:1000,2000-reqs:3000-ep:8-tp:8-gpus:8] TIMEOUT(100) +# 10c: H20, GB200, B200, B300, GB300, RTX6000-Server test cases - same memory gate +# as 10, but only 4 GPUs are required, so the 4-GPU GB200/GB300 nodes qualify too. +- condition: + ranges: + system_gpu_count: + gte: 4 + compute_capability: + gte: 9.0 + lte: 12.0 + gpu_memory: + gt: 90000 + tests: + #qwen3.5_122b_a10b (MoE BF16 234G, 4-GPU): OOM on 80G GPUs, needs >90G + - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:128,128-ep:4-tp:4-gpus:4] + - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:500,2000-ep:4-tp:4-gpus:4] + + # 11: RTX-6000 Server test cases - condition: ranges: From d153412fe2beb59f12775593b171d057cfdc746b Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:10:24 -0700 Subject: [PATCH 2/5] [nvbugs/6575476][fix] Restore H20/RTX6000 coverage for qwen3.5_397b_a17b_fp8 1000,2000 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. f9b2457819 (#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> --- tests/integration/test_lists/qa/llm_perf_core.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_lists/qa/llm_perf_core.yml b/tests/integration/test_lists/qa/llm_perf_core.yml index be482c5034d3..fc2dff6ff973 100644 --- a/tests/integration/test_lists/qa/llm_perf_core.yml +++ b/tests/integration/test_lists/qa/llm_perf_core.yml @@ -280,7 +280,6 @@ llm_perf_core: - perf/test_perf.py::test_perf[minimax_m3_fp4-bench-pytorch-float4-maxbs:1-input_output_len:1024,1024-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency - perf/test_perf.py::test_perf[minimax_m3_fp4-bench-pytorch-float4-maxbs:256-maxnt:8192-input_output_len:1024,1024-con:256-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-maxbs:384-maxnt:1536-input_output_len:1000,2000-reqs:10000-con:3072-ep:8-tp:8-gpus:8] TIMEOUT(120) #max throughput test - - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:1000,2000-ep:8-tp:8-gpus:8] # 8: B300 test cases @@ -327,6 +326,9 @@ llm_perf_core: gpu_memory: gt: 90000 tests: + #qwen3.5_397b_a17b_fp8 (MoE FP8 380G, 8-GPU ep=8): the 1000,2000 shape OOMs on + #80G GPUs, so it needs >90G; its siblings stay in condition 9. + - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:1000,2000-ep:8-tp:8-gpus:8] #qwen3.5_397b_a17b_fp4 (MoE FP4 234G, 8-GPU ep=8) - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:128,128-ep:8-tp:8-gpus:8] - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:500,2000-ep:8-tp:8-gpus:8] From 912d68eb72e39a6872e209b8d677caf0458cdc34 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Sat, 8 Aug 2026 02:58:24 -0700 Subject: [PATCH 3/5] [nvbugs/6575476][fix] Size hybrid KV residency on reachable capacity 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> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 10 ++ .../_torch/pyexecutor/mamba_cache_manager.py | 106 ++++++++++-- .../pyexecutor/scheduler/scheduler_v2.py | 25 ++- .../executor/test_kv_cache_v2_scheduler.py | 57 +++++++ .../executor/test_mamba_cache_manager.py | 160 +++++++++++++++++- 5 files changed, 343 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 2c567824a8cb..8a6298923b18 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2164,6 +2164,16 @@ def is_request_active(self, request_id: int) -> bool: kv_cache = self.kv_cache_map.get(request_id) return kv_cache is not None and kv_cache.is_active + def max_resident_sequences(self) -> Optional[int]: + """Cap on concurrently resident sequences, or ``None`` if unbounded. + + Attention pages are droppable, so pure-attention models let + suspend/resume absorb over-admission and need no cap. Managers that own + non-droppable per-sequence state (e.g. Mamba recurrent state) override + this so the scheduler stops admitting sequences the pool cannot hold. + """ + return None + def _effective_draft_len(self, req: LlmRequest) -> int: """Draft token length to use for next-step KV capacity calculation. diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 12df9681e6ab..3f4464f50ed7 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -28,6 +28,7 @@ from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig +from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( BlockReusePolicy, KVCacheManagerV2, Role) from tensorrt_llm._torch.pyexecutor.llm_request import ( @@ -2835,6 +2836,11 @@ class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): _supports_additional_snapshot_offsets = True + # Bound on concurrently resident sequences, set in _build_cache_config when + # the GPU cache quota cannot hold the requested max_batch_size. ``None`` + # means no bound applies and max_batch_size is used as requested. + _resident_sequence_cap: Optional[int] = None + def __init__( self, # mamba cache parameters @@ -3164,12 +3170,71 @@ def _get_pool_roles(self, return MambaRole.SSM_STATE, None return super()._get_pool_roles(pool_id) - def _max_resident_sequences(self) -> int: + def _requested_resident_sequences(self) -> int: + """Sequences the configured ``max_batch_size`` asks to keep resident.""" return self.max_batch_size * self.mapping.pp_size + def _max_resident_sequences(self) -> int: + if self._resident_sequence_cap is not None: + return self._resident_sequence_cap + return self._requested_resident_sequences() + + 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() + + def _resident_sequences_for_quota(self, gpu_quota: int) -> int: + """Return how many sequences the GPU quota can keep resident. + + A recurrent state is fixed-size per sequence and cannot be recomputed + from tokens, so every resident sequence permanently occupies one state + slot. ``max_batch_size`` is a ceiling, not an allocation contract, so a + request for more sequences than the quota can hold is bounded here + instead of becoming a hard allocation floor. + + The per-sequence cost counts live states *plus* the attention pages the + sequence occupies: sizing on the state cost alone would admit sequences + whose states consume the whole quota and leave nothing for attention. + """ + state_bytes = self._mamba_state_bytes_per_slot() + if state_bytes <= 0: + # Attention-only rank: nothing non-droppable to bound. Return the + # request itself so this rank stays neutral in the allreduce(MIN). + return self._requested_resident_sequences() + + attention_block_bytes = self._attention_block_bytes() + # Residency is a *guarantee*, so it must be sized on the capacity a + # sequence can actually reach, not the average one. ``typical_capacity`` + # is only a pool-ratio hint and defaults to max_seq_len / 2, which + # over-admits by ~2x on workloads where every sequence runs to + # max_seq_len: admission then fills the pool, and because recurrent + # state is non-droppable and resume() is refused past + # max_util_for_resume, the run has no way back out (nvbugs/6575476). + residency_capacity = max( + self._get_typical_request_capacity(self.kv_cache_config), + self.max_seq_len, + ) + num_states = self._num_ssm_states_per_typical_request( + residency_capacity, self.kv_cache_config) + attention_blocks = math.ceil(residency_capacity / self.tokens_per_block) + per_sequence = (num_states * state_bytes + + attention_blocks * attention_block_bytes) + + # Reserved dummy slots and one attention page are charged up front: + # they exist regardless of how many real sequences are resident. + reserved = (self._num_reserved_dummy_slots * state_bytes + + attention_block_bytes) + return (gpu_quota - reserved) // per_sequence + def _mamba_state_bytes_per_slot(self) -> int: return self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) + def _attention_block_bytes(self) -> int: + """Bytes of one attention page across all local attention layers.""" + return self._attention_cache_bytes_per_token() * self.tokens_per_block + def _num_ssm_snapshots_for_capacity( self, capacity: int, @@ -3251,8 +3316,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: # attention page per request lineage. This remains conservative when # the plan contains fewer than one non-live slot per lineage. extra_attention_quota = (num_request_lineages * - self._attention_cache_bytes_per_token() * - self.tokens_per_block + self._attention_block_bytes() if snapshot_slots > 0 else 0) return attention_quota + state_quota + extra_attention_quota @@ -3278,14 +3342,12 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: def _minimum_live_gpu_quota(self) -> int: """Return the minimum quota for live states and one attention page.""" - attention_block_quota = (self._attention_cache_bytes_per_token() * - self.tokens_per_block) num_state_slots = (self._max_resident_sequences() + self._num_reserved_dummy_slots) state_quota = num_state_slots * self._mamba_state_bytes_per_slot() return max( self._get_quota_from_max_tokens(0), - state_quota + attention_block_quota, + state_quota + self._attention_block_bytes(), ) def _build_cache_config( @@ -3293,6 +3355,24 @@ def _build_cache_config( kv_cache_config = self.kv_cache_config cache_tiers = config.cache_tiers gpu_quota = cache_tiers[0].quota + requested_resident = self._requested_resident_sequences() + affordable_resident = self._resident_sequences_for_quota(gpu_quota) + if self.mapping.world_size > 1: + # Quotas and per-rank state costs differ across ranks (uneven mamba + # layer splits, attention-only PP ranks). Every rank must apply the + # same bound or the schedulers admit different batches and desync. + affordable_resident = Distributed.get(self.mapping).allreduce( + affordable_resident, op=ReduceOp.MIN) + if affordable_resident < requested_resident: + self._resident_sequence_cap = max(1, affordable_resident) + logger.warning( + f"The V2 Mamba GPU cache quota ({gpu_quota} bytes) cannot keep " + f"{requested_resident} sequences resident: each one holds a " + f"fixed {self._mamba_state_bytes_per_slot()} bytes of recurrent " + "state that cannot be evicted. Limiting concurrently resident " + f"sequences to {self._resident_sequence_cap}. Reduce " + "max_batch_size, or raise free_gpu_memory_fraction / " + "max_gpu_total_bytes, to run the requested batch size.") minimum_live_quota = self._minimum_live_gpu_quota() if minimum_live_quota > gpu_quota: raise ValueError( @@ -3316,14 +3396,20 @@ def _build_cache_config( ], ) + max_resident = self._max_resident_sequences() dummy_requests = [ KVCacheDesc(capacity=0, history_length=0) for _ in range(self._num_reserved_dummy_slots) ] + # The base class sizes its warmup constraints from the requested + # max_batch_size. Every entry of a constraint costs one SSM slot (the + # planner never shares recurrent state between requests), so a + # constraint wider than the clamped residency would restore the very + # floor the clamp removed. Truncate to what can actually be resident. constraints = [ replace( batch, - kv_caches=[*batch.kv_caches, *dummy_requests], + kv_caches=[*batch.kv_caches[:max_resident], *dummy_requests], ) for batch in config.constraints ] @@ -3333,8 +3419,7 @@ def _build_cache_config( kv_cache_config) request_descs = self._typical_request_descs(typical_capacity, kv_cache_config) - typical_step = BatchDesc(request_descs * - self._max_resident_sequences() + + typical_step = BatchDesc(request_descs * max_resident + dummy_requests) # The recurrent (SSM) state pool must hold one slot per resident # sequence plus every reserved dummy slot. Unlike attention pages, a @@ -3346,8 +3431,7 @@ def _build_cache_config( # / __init__). Add a min-slots constraint of zero-capacity requests: # these cost no attention pages but reserve one SSM slot each. if any(isinstance(layer, SsmLayerConfig) for layer in layers): - ssm_floor_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) + ssm_floor_slots = max_resident + self._num_reserved_dummy_slots constraints = [ *constraints, BatchDesc([ diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 958afc59ac04..8be7456cb2e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -176,6 +176,11 @@ def __init__( scheduler_policy = CapacitySchedulerPolicy.MAX_UTILIZATION self.policy = scheduler_policy self.peft_cache_manager = peft_cache_manager + # Non-droppable per-sequence state (Mamba recurrent state) yields no + # evictable pages, so MAX_UTILIZATION's suspend/resume cannot recover + # from over-admission: once every resident sequence is suspended the + # pool can never drain. Bound admission instead. + self.max_resident_sequences = kv_cache_manager.max_resident_sequences() # Chunking config. self.chunking_enabled = False @@ -193,7 +198,8 @@ def __init__( f"KVCacheV2Scheduler: tokens_per_block={self.tokens_per_block}, " f"max_num_tokens={max_num_tokens}, max_batch_size={max_batch_size}, " f"draft_mgr={draft_mgr_name}, cross_mgr={cross_mgr_name}, " - f"enable_prefix_aware_scheduling={enable_prefix_aware_scheduling}" + f"enable_prefix_aware_scheduling={enable_prefix_aware_scheduling}, " + f"max_resident_sequences={self.max_resident_sequences}" ) if ctx_chunk_config is not None: self.chunking_enabled = True @@ -309,6 +315,16 @@ def _schedule_loop(self, active_requests, inflight_request_ids): if req.state_value == self._gen_to_complete_state_value: budget.pre_claim_peft(req) + # 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 + ) + # --- Phase 1: generation / disagg only --- while req_it < req_it_end: req = requests_list[req_it] @@ -403,6 +419,11 @@ def _schedule_loop(self, active_requests, inflight_request_ids): 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 peft_pages = budget.peft_pages_needed(req) if peft_pages is None: continue @@ -421,6 +442,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids): has_chunking = has_chunking or chunking_flag scheduled_ctx.append(req) budget.commit(req, tokens, peft_pages) + if starts_new_sequence: + num_resident += 1 # Deadlock detection: if generation requests exist but none were # scheduled and none were evicted, no forward pass will run and no diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 0a7c67e9e702..d5eb968a4103 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -160,6 +160,7 @@ def make_kv_cache_manager( resize_context_fn=None, prepare_disagg_gen_init_fn=None, try_allocate_generation_fn=None, + max_resident_sequences=None, ): mgr = Mock() mgr.tokens_per_block = tokens_per_block @@ -170,6 +171,9 @@ def make_kv_cache_manager( mgr.try_allocate_generation.side_effect = try_allocate_generation_fn or (lambda req: True) mgr.suspend_request.return_value = None mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active + # Pin explicitly: a bare Mock() would auto-vivify a truthy Mock here, which + # the residency gate compares against an int. + mgr.max_resident_sequences.return_value = max_resident_sequences return mgr @@ -2143,6 +2147,59 @@ def track_resize(req: Mock, num_tokens: int) -> bool: assert ids(out.context_requests) == [] +# =========================================================================== +# Residency cap (non-droppable per-sequence state, e.g. Mamba) +# =========================================================================== + + +class TestResidencyCap: + """A manager owning non-droppable per-sequence state bounds admission. + + Mamba recurrent state yields no evictable pages, so suspend/resume cannot + recover from over-admission (nvbugs/6550276). + """ + + def test_new_sequences_capped_at_max_resident(self): + mgr = make_kv_cache_manager(max_resident_sequences=2) + sched = make_scheduler(mgr, max_num_tokens=100000, max_batch_size=100) + reqs = [make_ctx_request(i, context_remaining_length=10) for i in range(5)] + + out = sched.schedule_request(reqs, set()) + + assert ids(out.context_requests) == [0, 1] + + def test_started_sequences_count_against_the_cap(self): + """A resident sequence keeps its state slot, so it consumes the cap.""" + mgr = make_kv_cache_manager(max_resident_sequences=2) + sched = make_scheduler(mgr, max_num_tokens=100000, max_batch_size=100) + gens = [make_gen_request(i) for i in range(2)] + ctxs = [make_ctx_request(10 + i, context_remaining_length=10) for i in range(2)] + + out = sched.schedule_request(gens + ctxs, set()) + + assert ids(out.generation_requests) == [0, 1] + assert ids(out.context_requests) == [] + + def test_non_first_chunk_is_not_charged_again(self): + """Only the first chunk starts a sequence; later chunks already hold a slot.""" + mgr = make_kv_cache_manager(max_resident_sequences=1) + sched = make_scheduler(mgr, max_num_tokens=100000, max_batch_size=100) + 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] + + def test_unbounded_manager_admits_every_request(self): + mgr = make_kv_cache_manager(max_resident_sequences=None) + sched = make_scheduler(mgr, max_num_tokens=100000, max_batch_size=100) + reqs = [make_ctx_request(i, context_remaining_length=10) for i in range(5)] + + out = sched.schedule_request(reqs, set()) + + assert ids(out.context_requests) == [0, 1, 2, 3, 4] + + # =========================================================================== # Edge Cases # =========================================================================== diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 81db2977398d..a413a0eccf6c 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1644,23 +1644,51 @@ def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): assert "workload's average total sequence length" in warnings_seen[0] -def test_v2_hybrid_rejects_quota_below_live_state_floor(): +def _residency_clamp_manager(max_batch_size): + """A hybrid V2 manager stub whose recurrent state dominates the quota.""" mgr = object.__new__(MambaHybridCacheManagerV2) - mgr.max_batch_size = 2 + mgr.max_batch_size = max_batch_size mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) mgr.local_num_mamba_layers = 1 mgr.ssm_bytes = 64 mgr.conv_bytes = 32 mgr._num_reserved_dummy_slots = 1 + mgr._resident_sequence_cap = None mgr.tokens_per_block = 32 mgr.num_local_layers = 2 mgr.pp_layers = [0, 1] - mgr.max_attention_window_vec = [128, 128] + mgr._mamba_layer_mask = [True, False] + # Full attention (no sliding window), matching the hybrid Qwen3.5 layout + # where the recurrent state is the only per-sequence fixed cost. + mgr.max_attention_window_vec = [None, None] + mgr.max_seq_len = 128 mgr.max_num_tokens = 128 mgr.enable_swa_scratch_reuse = False mgr.get_layer_bytes_per_token = lambda **kwargs: 8 mgr._attention_cache_bytes_per_token = lambda: 16 mgr.kv_cache_config = KvCacheConfig(enable_partial_reuse=False) + return mgr + + +def _quota_for_resident_sequences(num_sequences): + """Quota that lets ``_residency_clamp_manager`` hold exactly N sequences. + + A resident sequence costs one 96-byte state slot plus the attention pages + it can reach. Residency is a guarantee, so it is sized on max_seq_len (128 + tokens = 4 blocks of 16 B/token * 32 tokens), not on the avg_seq_len pool + hint. One dummy slot and one attention page are reserved regardless of + residency. + """ + state_bytes = 64 + 32 + attention_block_bytes = 16 * 32 + per_sequence = state_bytes + 4 * attention_block_bytes + reserved = state_bytes + attention_block_bytes + return reserved + num_sequences * per_sequence + + +def test_v2_hybrid_rejects_quota_below_single_live_state_floor(): + """A quota too small for even one sequence is still fatal.""" + mgr = _residency_clamp_manager(max_batch_size=1) minimum_quota = mgr._minimum_live_gpu_quota() base_config = KVCacheManagerConfig( @@ -1673,6 +1701,132 @@ def test_v2_hybrid_rejects_quota_below_live_state_floor(): mgr._build_cache_config(base_config) +def test_v2_hybrid_clamps_resident_sequences_to_quota(): + """max_batch_size is a ceiling, not an allocation floor: clamp, don't raise. + + The recurrent state of a hybrid sequence is fixed-size and non-droppable, so + a quota that cannot hold max_batch_size states must bound residency instead + of failing initialization (nvbugs/6550276). + """ + mgr = _residency_clamp_manager(max_batch_size=64) + # Size the quota for exactly 10 sequences, far below max_batch_size. + quota = _quota_for_resident_sequences(10) + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + layers=_base_attention_layer_configs(2), + ) + + config = mgr._build_cache_config(base_config) + + resident = mgr._max_resident_sequences() + assert resident == 10 + # The clamped floor now fits the quota, so initialization can proceed. + assert mgr._minimum_live_gpu_quota() <= quota + # Every sizing path agrees on the clamped residency. + assert mgr.max_resident_sequences() == resident + ssm_floor = [ + batch + for batch in config.constraints + if batch.kv_caches and all(kv.capacity == 0 for kv in batch.kv_caches) + ] + assert ssm_floor + assert all( + len(batch.kv_caches) == resident + mgr._num_reserved_dummy_slots for batch in ssm_floor + ) + + +def test_v2_hybrid_residency_sized_on_max_seq_len_not_avg_hint(): + """Residency must be sized on reachable capacity, not the avg_seq_len hint. + + ``_get_typical_request_capacity`` is a pool-*ratio* hint that falls back to + max_seq_len / 2 when avg_seq_len is unset. Sizing the residency guarantee on + it admits roughly twice the sequences the quota can actually hold once every + sequence grows to max_seq_len. Recurrent state is non-droppable and resume() + is refused past max_util_for_resume, so the over-admitted run cannot recover + (nvbugs/6575476). + """ + mgr = _residency_clamp_manager(max_batch_size=64) + assert mgr.kv_cache_config.avg_seq_len is None + # Sized on the avg hint (max_seq_len/2 = 2 blocks), this quota would look + # like it affords 10 sequences; at max_seq_len (4 blocks) it affords fewer. + hint_state_bytes = 64 + 32 + hint_block_bytes = 16 * 32 + quota_by_avg_hint = (hint_state_bytes + hint_block_bytes) + 10 * ( + hint_state_bytes + 2 * hint_block_bytes + ) + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota_by_avg_hint)], + layers=_base_attention_layer_configs(2), + ) + + mgr._build_cache_config(base_config) + + resident = mgr._max_resident_sequences() + assert resident < 10, "residency was sized on the avg_seq_len hint" + # Every resident sequence must still fit when it reaches max_seq_len. + blocks_at_max_seq_len = mgr.max_seq_len // 32 + worst_case = (resident + mgr._num_reserved_dummy_slots) * hint_state_bytes + ( + resident * blocks_at_max_seq_len * hint_block_bytes + ) + assert worst_case <= quota_by_avg_hint + + +def test_v2_hybrid_truncates_base_constraints_to_clamped_residency(): + """Inherited warmup constraints must not re-impose the unclamped floor. + + Every entry of a constraint costs one SSM slot, so a base-class constraint + built from the requested max_batch_size would restore the very state floor + the residency clamp removed (nvbugs/6550276). + """ + mgr = _residency_clamp_manager(max_batch_size=64) + quota = _quota_for_resident_sequences(10) + wide_constraint = BatchDesc( + [KVCacheDesc(capacity=32, history_length=0) for _ in range(mgr.max_batch_size)] + ) + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + layers=_base_attention_layer_configs(2), + constraints=[wide_constraint], + ) + + config = mgr._build_cache_config(base_config) + + resident = mgr._max_resident_sequences() + assert resident == 10 + truncated = config.constraints[0] + assert len(truncated.kv_caches) == resident + mgr._num_reserved_dummy_slots + # No constraint may demand more state slots than the pool floor allows. + assert all( + len(batch.kv_caches) <= resident + mgr._num_reserved_dummy_slots + for batch in config.constraints + ) + + +def test_v2_hybrid_keeps_requested_residency_when_quota_is_ample(): + mgr = _residency_clamp_manager(max_batch_size=4) + quota = mgr._minimum_live_gpu_quota() * 1000 + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + layers=_base_attention_layer_configs(2), + ) + + mgr._build_cache_config(base_config) + + assert mgr._resident_sequence_cap is None + assert mgr._max_resident_sequences() == 4 + + +def test_v2_attention_only_manager_reports_unbounded_residency(): + """Droppable attention pages need no residency cap.""" + mgr = object.__new__(KVCacheManagerV2) + + assert mgr.max_resident_sequences() is None + + def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): mgr = object.__new__(MambaHybridCacheManagerV2) mgr.max_batch_size = 2 From daf222b3db9683587b72e25a20fcf2958c477566 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Sat, 8 Aug 2026 03:10:59 -0700 Subject: [PATCH 4/5] [nvbugs/6575476][fix] Revert the perf-list re-gating for these two cases 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> --- .../_torch/pyexecutor/mamba_cache_manager.py | 18 +++-- .../test_lists/qa/llm_perf_core.yml | 25 ++----- .../executor/test_mamba_cache_manager.py | 71 +++++++------------ 3 files changed, 36 insertions(+), 78 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 3f4464f50ed7..9e5c47e57074 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -3206,16 +3206,14 @@ def _resident_sequences_for_quota(self, gpu_quota: int) -> int: attention_block_bytes = self._attention_block_bytes() # Residency is a *guarantee*, so it must be sized on the capacity a - # sequence can actually reach, not the average one. ``typical_capacity`` - # is only a pool-ratio hint and defaults to max_seq_len / 2, which - # over-admits by ~2x on workloads where every sequence runs to - # max_seq_len: admission then fills the pool, and because recurrent - # state is non-droppable and resume() is refused past - # max_util_for_resume, the run has no way back out (nvbugs/6575476). - residency_capacity = max( - self._get_typical_request_capacity(self.kv_cache_config), - self.max_seq_len, - ) + # sequence can actually reach -- max_seq_len -- not on the average one. + # ``_get_typical_request_capacity`` is only a pool-ratio hint and + # defaults to max_seq_len / 2, which over-admits by ~2x on workloads + # where every sequence runs to max_seq_len: admission then fills the + # pool, and because recurrent state is non-droppable and resume() is + # refused past max_util_for_resume, the run has no way back out + # (nvbugs/6575476). + residency_capacity = self.max_seq_len num_states = self._num_ssm_states_per_typical_request( residency_capacity, self.kv_cache_config) attention_blocks = math.ceil(residency_capacity / self.tokens_per_block) diff --git a/tests/integration/test_lists/qa/llm_perf_core.yml b/tests/integration/test_lists/qa/llm_perf_core.yml index fc2dff6ff973..8bb7ead14921 100644 --- a/tests/integration/test_lists/qa/llm_perf_core.yml +++ b/tests/integration/test_lists/qa/llm_perf_core.yml @@ -14,7 +14,6 @@ llm_perf_core: # 9: H100, H20, B200, B300, RTX6000-Server test cases # 10: H20, B200, B300, RTX6000-Server test cases # 10b: H20, B200, B300 test cases (>120G per GPU; excludes 96G RTX6000-Server) -# 10c: H20, GB200, B200, B300, GB300, RTX6000-Server test cases (>90G per GPU, 4-GPU capable) # 11: RTX-6000 Server test cases # =============================================================================== @@ -280,6 +279,7 @@ llm_perf_core: - perf/test_perf.py::test_perf[minimax_m3_fp4-bench-pytorch-float4-maxbs:1-input_output_len:1024,1024-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency - perf/test_perf.py::test_perf[minimax_m3_fp4-bench-pytorch-float4-maxbs:256-maxnt:8192-input_output_len:1024,1024-con:256-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-maxbs:384-maxnt:1536-input_output_len:1000,2000-reqs:10000-con:3072-ep:8-tp:8-gpus:8] TIMEOUT(120) #max throughput test + - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:1000,2000-ep:8-tp:8-gpus:8] # 8: B300 test cases @@ -326,9 +326,9 @@ llm_perf_core: gpu_memory: gt: 90000 tests: - #qwen3.5_397b_a17b_fp8 (MoE FP8 380G, 8-GPU ep=8): the 1000,2000 shape OOMs on - #80G GPUs, so it needs >90G; its siblings stay in condition 9. - - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:1000,2000-ep:8-tp:8-gpus:8] + #qwen3.5_122b_a10b (MoE BF16 234G, 4-GPU): OOM on 80G GPUs, needs >90G + - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:128,128-ep:4-tp:4-gpus:4] + - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:500,2000-ep:4-tp:4-gpus:4] #qwen3.5_397b_a17b_fp4 (MoE FP4 234G, 8-GPU ep=8) - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:128,128-ep:8-tp:8-gpus:8] - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:500,2000-ep:8-tp:8-gpus:8] @@ -357,23 +357,6 @@ llm_perf_core: - perf/test_perf.py::test_perf[deepseek_r1_0528_fp8-bench-pytorch-float8-input_output_len:1000,2000-reqs:3000-ep:8-tp:8-gpus:8] TIMEOUT(100) -# 10c: H20, GB200, B200, B300, GB300, RTX6000-Server test cases - same memory gate -# as 10, but only 4 GPUs are required, so the 4-GPU GB200/GB300 nodes qualify too. -- condition: - ranges: - system_gpu_count: - gte: 4 - compute_capability: - gte: 9.0 - lte: 12.0 - gpu_memory: - gt: 90000 - tests: - #qwen3.5_122b_a10b (MoE BF16 234G, 4-GPU): OOM on 80G GPUs, needs >90G - - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:128,128-ep:4-tp:4-gpus:4] - - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:500,2000-ep:4-tp:4-gpus:4] - - # 11: RTX-6000 Server test cases - condition: ranges: diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index a413a0eccf6c..b018b5a7f60f 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1670,32 +1670,38 @@ def _residency_clamp_manager(max_batch_size): return mgr -def _quota_for_resident_sequences(num_sequences): +def _quota_for_resident_sequences(num_sequences, blocks_per_sequence=4): """Quota that lets ``_residency_clamp_manager`` hold exactly N sequences. - A resident sequence costs one 96-byte state slot plus the attention pages - it can reach. Residency is a guarantee, so it is sized on max_seq_len (128 - tokens = 4 blocks of 16 B/token * 32 tokens), not on the avg_seq_len pool - hint. One dummy slot and one attention page are reserved regardless of - residency. + A resident sequence costs one 96-byte state slot (ssm 64 + conv 32) plus the + ``blocks_per_sequence`` attention pages it can reach, each 16 B/token * 32 + tokens. Residency is a guarantee, so the manager sizes it on max_seq_len + (128 tokens = 4 blocks), not on the avg_seq_len pool hint (2 blocks). One + dummy slot and one attention page are reserved regardless of residency. """ state_bytes = 64 + 32 attention_block_bytes = 16 * 32 - per_sequence = state_bytes + 4 * attention_block_bytes + per_sequence = state_bytes + blocks_per_sequence * attention_block_bytes reserved = state_bytes + attention_block_bytes return reserved + num_sequences * per_sequence +def _residency_clamp_config(quota, **kwargs): + """Base config for ``_residency_clamp_manager`` at the given GPU quota.""" + kwargs.setdefault("layers", _base_attention_layer_configs(2)) + return KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + **kwargs, + ) + + def test_v2_hybrid_rejects_quota_below_single_live_state_floor(): """A quota too small for even one sequence is still fatal.""" mgr = _residency_clamp_manager(max_batch_size=1) minimum_quota = mgr._minimum_live_gpu_quota() - base_config = KVCacheManagerConfig( - tokens_per_block=32, - cache_tiers=[GpuCacheTierConfig(quota=minimum_quota - 1)], - layers=[], - ) + base_config = _residency_clamp_config(minimum_quota - 1, layers=[]) with pytest.raises(ValueError, match="too small for live recurrent states"): mgr._build_cache_config(base_config) @@ -1711,13 +1717,8 @@ def test_v2_hybrid_clamps_resident_sequences_to_quota(): mgr = _residency_clamp_manager(max_batch_size=64) # Size the quota for exactly 10 sequences, far below max_batch_size. quota = _quota_for_resident_sequences(10) - base_config = KVCacheManagerConfig( - tokens_per_block=32, - cache_tiers=[GpuCacheTierConfig(quota=quota)], - layers=_base_attention_layer_configs(2), - ) - config = mgr._build_cache_config(base_config) + config = mgr._build_cache_config(_residency_clamp_config(quota)) resident = mgr._max_resident_sequences() assert resident == 10 @@ -1750,27 +1751,14 @@ def test_v2_hybrid_residency_sized_on_max_seq_len_not_avg_hint(): assert mgr.kv_cache_config.avg_seq_len is None # Sized on the avg hint (max_seq_len/2 = 2 blocks), this quota would look # like it affords 10 sequences; at max_seq_len (4 blocks) it affords fewer. - hint_state_bytes = 64 + 32 - hint_block_bytes = 16 * 32 - quota_by_avg_hint = (hint_state_bytes + hint_block_bytes) + 10 * ( - hint_state_bytes + 2 * hint_block_bytes - ) - base_config = KVCacheManagerConfig( - tokens_per_block=32, - cache_tiers=[GpuCacheTierConfig(quota=quota_by_avg_hint)], - layers=_base_attention_layer_configs(2), - ) + quota_by_avg_hint = _quota_for_resident_sequences(10, blocks_per_sequence=2) - mgr._build_cache_config(base_config) + mgr._build_cache_config(_residency_clamp_config(quota_by_avg_hint)) resident = mgr._max_resident_sequences() assert resident < 10, "residency was sized on the avg_seq_len hint" # Every resident sequence must still fit when it reaches max_seq_len. - blocks_at_max_seq_len = mgr.max_seq_len // 32 - worst_case = (resident + mgr._num_reserved_dummy_slots) * hint_state_bytes + ( - resident * blocks_at_max_seq_len * hint_block_bytes - ) - assert worst_case <= quota_by_avg_hint + assert _quota_for_resident_sequences(resident) <= quota_by_avg_hint def test_v2_hybrid_truncates_base_constraints_to_clamped_residency(): @@ -1785,14 +1773,8 @@ def test_v2_hybrid_truncates_base_constraints_to_clamped_residency(): wide_constraint = BatchDesc( [KVCacheDesc(capacity=32, history_length=0) for _ in range(mgr.max_batch_size)] ) - base_config = KVCacheManagerConfig( - tokens_per_block=32, - cache_tiers=[GpuCacheTierConfig(quota=quota)], - layers=_base_attention_layer_configs(2), - constraints=[wide_constraint], - ) - config = mgr._build_cache_config(base_config) + config = mgr._build_cache_config(_residency_clamp_config(quota, constraints=[wide_constraint])) resident = mgr._max_resident_sequences() assert resident == 10 @@ -1808,13 +1790,8 @@ def test_v2_hybrid_truncates_base_constraints_to_clamped_residency(): def test_v2_hybrid_keeps_requested_residency_when_quota_is_ample(): mgr = _residency_clamp_manager(max_batch_size=4) quota = mgr._minimum_live_gpu_quota() * 1000 - base_config = KVCacheManagerConfig( - tokens_per_block=32, - cache_tiers=[GpuCacheTierConfig(quota=quota)], - layers=_base_attention_layer_configs(2), - ) - mgr._build_cache_config(base_config) + mgr._build_cache_config(_residency_clamp_config(quota)) assert mgr._resident_sequence_cap is None assert mgr._max_resident_sequences() == 4 From 67a9327c54fd3700136e37fae405224e7b39330c Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Sat, 8 Aug 2026 07:34:05 -0700 Subject: [PATCH 5/5] [nvbugs/6575476][fix] Pin max_resident_sequences on the spec'd KV manager 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> --- .../_torch/pyexecutor/mamba_cache_manager.py | 41 ++++++++++--------- .../pyexecutor/scheduler/scheduler_v2.py | 20 ++++----- .../executor/test_dual_pool_kv_cache.py | 3 ++ 3 files changed, 34 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 9e5c47e57074..0b88502d6051 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -3197,6 +3197,14 @@ def _resident_sequences_for_quota(self, gpu_quota: int) -> int: The per-sequence cost counts live states *plus* the attention pages the sequence occupies: sizing on the state cost alone would admit sequences whose states consume the whole quota and leave nothing for attention. + + Residency is a *guarantee*, so it is sized on the capacity a sequence + can actually reach -- max_seq_len -- not on the average one. + ``_get_typical_request_capacity`` is only a pool-ratio hint and defaults + to max_seq_len / 2, which over-admits by ~2x on workloads where every + sequence runs to max_seq_len: admission then fills the pool, and because + recurrent state is non-droppable and resume() is refused past + max_util_for_resume, the run has no way back out (nvbugs/6575476). """ state_bytes = self._mamba_state_bytes_per_slot() if state_bytes <= 0: @@ -3204,27 +3212,15 @@ def _resident_sequences_for_quota(self, gpu_quota: int) -> int: # request itself so this rank stays neutral in the allreduce(MIN). return self._requested_resident_sequences() - attention_block_bytes = self._attention_block_bytes() - # Residency is a *guarantee*, so it must be sized on the capacity a - # sequence can actually reach -- max_seq_len -- not on the average one. - # ``_get_typical_request_capacity`` is only a pool-ratio hint and - # defaults to max_seq_len / 2, which over-admits by ~2x on workloads - # where every sequence runs to max_seq_len: admission then fills the - # pool, and because recurrent state is non-droppable and resume() is - # refused past max_util_for_resume, the run has no way back out - # (nvbugs/6575476). residency_capacity = self.max_seq_len num_states = self._num_ssm_states_per_typical_request( residency_capacity, self.kv_cache_config) attention_blocks = math.ceil(residency_capacity / self.tokens_per_block) per_sequence = (num_states * state_bytes + - attention_blocks * attention_block_bytes) - - # Reserved dummy slots and one attention page are charged up front: - # they exist regardless of how many real sequences are resident. - reserved = (self._num_reserved_dummy_slots * state_bytes + - attention_block_bytes) - return (gpu_quota - reserved) // per_sequence + attention_blocks * self._attention_block_bytes()) + # Subtract the same residency-independent floor _minimum_live_gpu_quota + # charges, so a residency this returns always clears that check. + return (gpu_quota - self._fixed_live_gpu_quota()) // per_sequence def _mamba_state_bytes_per_slot(self) -> int: return self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) @@ -3233,6 +3229,12 @@ def _attention_block_bytes(self) -> int: """Bytes of one attention page across all local attention layers.""" return self._attention_cache_bytes_per_token() * self.tokens_per_block + def _fixed_live_gpu_quota(self) -> int: + """Live-quota bytes that do not depend on how many sequences are resident.""" + return (self._num_reserved_dummy_slots * + self._mamba_state_bytes_per_slot() + + self._attention_block_bytes()) + def _num_ssm_snapshots_for_capacity( self, capacity: int, @@ -3340,12 +3342,11 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: def _minimum_live_gpu_quota(self) -> int: """Return the minimum quota for live states and one attention page.""" - num_state_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - state_quota = num_state_slots * self._mamba_state_bytes_per_slot() + resident_state_quota = (self._max_resident_sequences() * + self._mamba_state_bytes_per_slot()) return max( self._get_quota_from_max_tokens(0), - state_quota + self._attention_block_bytes(), + resident_state_quota + self._fixed_live_gpu_quota(), ) def _build_cache_config( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 8be7456cb2e0..945bd8ce2ddd 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -315,16 +315,6 @@ def _schedule_loop(self, active_requests, inflight_request_ids): if req.state_value == self._gen_to_complete_state_value: budget.pre_claim_peft(req) - # 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 - ) - # --- Phase 1: generation / disagg only --- while req_it < req_it_end: req = requests_list[req_it] @@ -416,6 +406,16 @@ def _schedule_loop(self, active_requests, inflight_request_ids): # --- Phase 2: schedule deferred context / encoder requests --- # Generation PEFT pages are now fully committed in the budget. + # + # 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 diff --git a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py index 06e0231a1a0b..8efed16c0716 100644 --- a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @@ -710,6 +710,9 @@ class TestKVCacheV2SchedulerCrossParam: def _make_mock_kv_mgr(self, tokens_per_block=64): mgr = Mock(spec=KVCacheManagerV2) mgr.tokens_per_block = tokens_per_block + # spec= auto-vivifies a truthy Mock, which the scheduler's residency + # gate would compare against an int; None is the real unbounded value. + mgr.max_resident_sequences.return_value = None return mgr def test_default_cross_is_none(self):