From 94d2555d655e5d4178af136dba2743d73150ed42 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 4 Aug 2026 09:37:53 -0700 Subject: [PATCH] [nvbugs/6550276][fix] Bound resident Mamba sequences to the KV cache 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> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 10 ++ .../_torch/pyexecutor/mamba_cache_manager.py | 108 +++++++++++++--- .../pyexecutor/scheduler/scheduler_v2.py | 25 +++- .../executor/test_kv_cache_v2_scheduler.py | 57 ++++++++ .../executor/test_mamba_cache_manager.py | 122 +++++++++++++++++- 5 files changed, 303 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 967c7d413f68..250203851c4e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2156,6 +2156,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 75eec8a8b5aa..8f53c3bcd8f9 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 ( @@ -2477,6 +2478,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 @@ -2740,12 +2746,62 @@ 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, + typical_descs: List[KVCacheDesc]) -> 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 reuses ``typical_step``'s request model + (``typical_descs``: one descriptor per live state, together spanning the + typical request capacity), so sizing here cannot drift from the pool + sizing it bounds. Charging 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() + typical_capacity = sum(desc.capacity for desc in typical_descs) + attention_blocks = math.ceil(typical_capacity / self.tokens_per_block) + per_sequence = (len(typical_descs) * 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, @@ -2827,8 +2883,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 @@ -2854,14 +2909,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( @@ -2869,6 +2922,31 @@ def _build_cache_config( kv_cache_config = self.kv_cache_config cache_tiers = config.cache_tiers gpu_quota = cache_tiers[0].quota + # One descriptor per live SSM state of a typical request. Shared with + # the typical_step sizing below so both agree on the per-sequence cost + # (and so the avg_seq_len fallback warning is emitted only once). + request_descs = self._typical_request_descs( + self._get_typical_request_capacity(kv_cache_config), + kv_cache_config) + requested_resident = self._requested_resident_sequences() + affordable_resident = self._resident_sequences_for_quota( + gpu_quota, request_descs) + 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( @@ -2892,25 +2970,26 @@ 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 ] typical_step = config.typical_step if config.initial_pool_ratio is None: - typical_capacity = self._get_typical_request_capacity( - 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 @@ -2922,8 +3001,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 aab465e0a475..c96f5fc15560 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1316,23 +1316,50 @@ 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 typical request costs one 96-byte state slot plus its attention pages: + avg_seq_len is unset, so typical capacity is max_seq_len / 2 = 64 tokens = + 2 blocks of 16 B/token * 32 tokens. 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 + 2 * 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( @@ -1345,6 +1372,95 @@ 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_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