[None][fix] Seq-slot pool overlap headroom with consistent slot-indexed buffer sizing

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -1310,7 +1310,9 @@ def create_autodeploy_executor(
         )
         guided_decoder = GuidedDecoder(
             guided_decoding_config=guided_decoding_config,
-            max_num_sequences=ad_config.max_batch_size,
+            # Guided-decoder state is indexed by sequence slot; size it to the
+            # SeqSlotManager pool above (max_batch_size * pp_size).
+            max_num_sequences=max_num_sequences,
             vocab_size_padded=vocab_size_padded,
         )
 

tensorrt_llm/_torch/models/modeling_qwen2vl.py

@@ -180,6 +180,15 @@ def _prepare_qwen_vl_mrope_config(
 _MAX_PIXELS_TOKEN_PROBE = 1 << 31
 
 
+def _get_mrope_position_delta_cache_size(
+        model_config: ModelConfig[PretrainedConfig]) -> int:
+    """Return real sequence-slot capacity plus one reserved dummy slot."""
+    max_num_seq_slots = model_config.extra_attrs.get(
+        'max_num_seq_slots',
+        model_config.max_num_tokens * model_config.mapping.pp_size)
+    return max_num_seq_slots + 1
+
+
 class Qwen2VLInputProcessorBase(BaseMultimodalInputProcessor,
                                 BaseMultimodalDummyInputsBuilder):
 
@@ -1705,8 +1714,8 @@ def __init__(
         if not disable_fuse_rope:
             self.init_mrope_embedding(model_config)
             # Extra slot is reserved for CUDA graph / warmup dummy requests.
-            max_mrope_delta_slots = (
-                model_config.max_num_tokens * model_config.mapping.pp_size + 1)
+            max_mrope_delta_slots = _get_mrope_position_delta_cache_size(
+                model_config)
             self.register_buffer('mrope_position_deltas_cache',
                                  torch.zeros(max_mrope_delta_slots,
                                              dtype=torch.int32,

tensorrt_llm/_torch/models/modeling_qwen3vl.py

@@ -56,6 +56,7 @@
 from .modeling_qwen2vl import (
     Qwen2_5_VLVisionAttention,
     Qwen2VLInputProcessorBase,
+    _get_mrope_position_delta_cache_size,
     _prepare_qwen_vl_mrope_config,
     _prepare_qwen_vl_vision_attn_metadata,
 )
@@ -1206,7 +1207,7 @@ def __init__(
         if not disable_fuse_rope:
             self.init_mrope_embedding(model_config)
             # Extra slot is reserved for CUDA graph / warmup dummy requests.
-            max_mrope_delta_slots = model_config.max_num_tokens * model_config.mapping.pp_size + 1
+            max_mrope_delta_slots = _get_mrope_position_delta_cache_size(model_config)
             self.register_buffer(
                 "mrope_position_deltas_cache",
                 torch.zeros(

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -2046,6 +2046,24 @@ def create_kv_cache_compression_manager(
     return None
 
 
+def compute_max_num_sequences(mapping: Mapping, max_batch_size: int,
+                              disable_overlap_scheduler: bool) -> int:
+    """Size the sequence-slot pool (and the sampler state it indexes).
+
+    The slot pool needs headroom for two iterations under overlap (prev-iter
+    finished requests still hold slots when next-iter prepare_resources
+    runs). Without it, the V2 scheduler can backfill the budget seat vacated
+    by a GENERATION_TO_COMPLETE request whose slot is not freed yet and
+    exhaust the pool ("No free slots"). With pipeline parallelism the
+    micro-batch count is pp_size, which covers overlap already.
+    """
+    if mapping.has_pp():
+        num_micro_batches = mapping.pp_size
+    else:
+        num_micro_batches = 1 if disable_overlap_scheduler else 2
+    return max_batch_size * num_micro_batches
+
+
 def create_py_executor_instance(
     *,
     dist,
@@ -2080,7 +2098,8 @@ def create_py_executor_instance(
 
     spec_config = model_engine.spec_config
 
-    max_num_sequences = max_batch_size * mapping.pp_size
+    max_num_sequences = compute_max_num_sequences(
+        mapping, max_batch_size, llm_args.disable_overlap_scheduler)
 
     logger.info(
         f"max_seq_len={max_seq_len}, max_num_requests={max_num_sequences}, max_num_tokens={max_num_tokens}, max_batch_size={max_batch_size}"
@@ -2277,7 +2296,9 @@ def create_py_executor_instance(
 
     # When scheduler_capacity == 1, attention dp dummy request will prevent the scheduling of DISAGG_GENERATION_INIT.
     # Enlarge scheduler capacity to avoid DISAGG_GENERATION_INIT stuck in the scheduler.
-    scheduler_capacity = max_num_sequences
+    # V1 scheduler handles overlap via two_step_lookahead, so skip the
+    # slot-pool overlap factor here.
+    scheduler_capacity = max_batch_size * mapping.pp_size
     if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager:
         scheduler_capacity += 1
 
@@ -2434,7 +2455,10 @@ def create_torch_sampler_args(
     enable_async_worker: bool,
     enable_speculative_beam_history_d2h: bool,
 ):
-    max_num_sequences = max_batch_size * mapping.pp_size
+    # The sampler's per-slot state is indexed by sequence slots, so it must
+    # be sized identically to the executor's slot pool.
+    max_num_sequences = compute_max_num_sequences(mapping, max_batch_size,
+                                                  disable_overlap_scheduler)
     max_draft_len = (0 if speculative_config is None else
                      speculative_config.max_draft_len)
     max_total_draft_tokens = (0 if speculative_config is None else

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -1120,8 +1120,22 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False):
         KVCacheManager.free_resources(self, request, pin_on_release)
 
     def add_dummy_requests(self, request_ids: List[int], **kwargs):
-        MambaCacheManager.add_dummy_requests(self, request_ids)
-        return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs)
+        requests = KVCacheManager.add_dummy_requests(self, request_ids,
+                                                     **kwargs)
+        if not requests:
+            return requests
+        try:
+            MambaCacheManager.add_dummy_requests(self, request_ids)
+        except RuntimeError:
+            # KV allocation completed first; restore it if the independent
+            # Mamba pool cannot reserve the matching dummy slot.
+            draft_kv_cache_manager = kwargs.get('draft_kv_cache_manager')
+            for request in requests:
+                if draft_kv_cache_manager is not None:
+                    draft_kv_cache_manager.free_resources(request)
+                KVCacheManager.free_resources(self, request)
+            raise
+        return requests
 
     def shutdown(self):
         MambaCacheManager.shutdown(self)
@@ -1831,23 +1845,29 @@ def get_num_available_tokens(self,
         result = super().get_num_available_tokens(token_num_upper_bound,
                                                   max_num_draft_tokens,
                                                   **kwargs)
-        # Also bound by the recurrent-state pool capacity: each request needs
-        # roughly ceil(N / states_snapshot_interval) recurrent-state blocks
-        # (+1 for the corner case where N is a multiple of tokens_per_block).
-        # When block reuse is disabled, only one snapshot is needed per
-        # request, so no additional capping is required here.
-        interval = (self.linear_attention_metadata.states_snapshot_interval
-                    if self.linear_attention_metadata is not None else 0)
+        # Also bound by the recurrent-state pool capacity. With block reuse,
+        # each request needs roughly ceil(N / states_snapshot_interval) blocks
+        # (+1 for the final live state). Without reuse, it still needs one live
+        # recurrent-state block, so an exhausted pool must fail preflight.
+        if self.linear_attention_metadata is None:
+            # A PP rank with no local Mamba layers has no recurrent-state pool.
+            return max(result, 0)
+
+        interval = self.linear_attention_metadata.states_snapshot_interval
+        stats = self.impl.get_kv_cache_stats()
+        rs_free = stats.num_free_blocks_per_window_size.get(
+            LinearCacheType.RECURRENT_STATES.value, 0)
         if interval and interval > 0:
-            stats = self.impl.get_kv_cache_stats()
-            rs_free = stats.num_free_blocks_per_window_size.get(
-                LinearCacheType.RECURRENT_STATES.value, 0)
             # Reserve 1 block for the always-allocated last block (corner case
             # / final live state) so we don't promise more tokens than the
             # pool can actually back at allocation time.
             usable_rs_blocks = max(0, rs_free - 1)
-            rs_token_cap = usable_rs_blocks * interval
+            rs_token_cap = max(
+                0, usable_rs_blocks * interval - self.num_extra_kv_tokens -
+                max_num_draft_tokens)
             result = min(result, rs_token_cap)
+        elif rs_free == 0:
+            result = 0
         return max(result, 0)
 
     def get_ssm_states(self, layer_idx: int) -> torch.Tensor:
@@ -1968,8 +1988,7 @@ def _setup_state_indices(self, requests=None) -> None:
 
         self.cuda_state_indices.copy_(self._host_state_indices,
                                       non_blocking=True)
-        self._refresh_dummy_request_mask(
-            [req.is_dummy for req in self.requests])
+        self._refresh_dummy_request_mask([req.is_dummy for req in requests])
 
         # Build request_id → pool block offset mapping so that
         # get_state_indices can return indices in arbitrary request order.

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import bisect
 import contextlib
 import functools
@@ -296,6 +299,12 @@ def __init__(
         self.mapping = mapping
         if mapping.has_pp():
             init_pp_comm(mapping)
+        # The overlap scheduler can hold two iterations' requests at once.
+        # Model-side buffers indexed by py_seq_slot must span that full pool,
+        # not the token budget or scheduler admission batch size.
+        from ._util import compute_max_num_sequences
+        self.max_num_seq_slots = compute_max_num_sequences(
+            mapping, self.batch_size, llm_args.disable_overlap_scheduler)
         self.dist = dist
         if dist is not None:
             ExpertStatistic.create(self.dist.rank)
@@ -363,6 +372,7 @@ def __init__(
                 sparse_attention_config=self.sparse_attention_config,
                 max_num_tokens=self.max_num_tokens,
                 max_seq_len=self.max_seq_len,
+                max_num_seq_slots=self.max_num_seq_slots,
                 lora_config=lora_config,
                 model_weights_memory_tag=model_weights_memory_tag,
                 model_weights_restore_mode=model_weights_restore_mode,
@@ -373,6 +383,7 @@ def __init__(
                 setattr(self, "moe_load_balancer", moe_load_balancer)
         else:
             self.model = model
+        self._validate_mrope_position_delta_cache_capacity()
         if drafting_loop_wrapper is not None:
             self.model = drafting_loop_wrapper(self.model)
             self.model_is_wrapped = True
@@ -877,6 +888,34 @@ def no_cuda_graph(self):
         finally:
             self.cuda_graph_runner.enabled = _run_cuda_graphs
 
+    def _validate_mrope_position_delta_cache_capacity(self) -> None:
+        """Validate slot-indexed MRoPE state on preconstructed models.
+
+        Models created by ModelLoader receive ``max_num_seq_slots`` before
+        construction. A caller-supplied model bypasses that path, so fail
+        early instead of letting the overlap scheduler index past an old-sized
+        cache at runtime.
+        """
+        mrope_position_deltas_cache = getattr(self.model,
+                                              "mrope_position_deltas_cache",
+                                              None)
+        if mrope_position_deltas_cache is None:
+            mrope_position_deltas_cache = getattr(
+                getattr(self.model, "draft_model", None),
+                "mrope_position_deltas_cache", None)
+        if mrope_position_deltas_cache is None:
+            return
+
+        required_size = self.max_num_seq_slots + 1
+        actual_size = mrope_position_deltas_cache.shape[0]
+        if actual_size < required_size:
+            raise ValueError(
+                "The supplied model's MRoPE position-delta cache has "
+                f"{actual_size} slots, but this executor requires at least "
+                f"{required_size} ({self.max_num_seq_slots} runtime sequence "
+                "slots plus one reserved dummy slot). Rebuild the model with "
+                "the executor's sequence-slot capacity.")
+
     def _pad_batch_seed_mrope_delta_cache(
             self, padded_requests: ScheduledRequests) -> None:
         if not self.use_mrope or padded_requests.num_generation_requests == 0:
@@ -2155,6 +2194,9 @@ def _set_up_spec_metadata(
             spec_resource_manager: Optional[BaseResourceManager],
             no_cache=False):
         spec_config = self.spec_config if self.enable_spec_decode else None
+        # Slot-indexed metadata buffers must span the SeqSlotManager pool,
+        # which exceeds batch_size under the overlap scheduler.
+        num_seq_slots = self.max_num_seq_slots
         if no_cache:
             return get_spec_metadata(
                 spec_config,
@@ -2163,7 +2205,8 @@ def _set_up_spec_metadata(
                 max_num_tokens=self.max_num_tokens,
                 spec_resource_manager=spec_resource_manager,
                 is_draft_model=self.is_draft_model,
-                max_seq_len=self.max_seq_len)
+                max_seq_len=self.max_seq_len,
+                num_seq_slots=num_seq_slots)
 
         if self.spec_metadata is not None:
             return self.spec_metadata
@@ -2174,7 +2217,8 @@ def _set_up_spec_metadata(
             max_num_tokens=self.max_num_tokens,
             spec_resource_manager=spec_resource_manager,
             is_draft_model=self.is_draft_model,
-            max_seq_len=self.max_seq_len)
+            max_seq_len=self.max_seq_len,
+            num_seq_slots=num_seq_slots)
         return self.spec_metadata
 
     def cleanup(self) -> None:
@@ -3251,7 +3295,9 @@ def _prepare_tp_inputs(
         mrope_delta_read_seq_slots = []
         # Extra model-side cache slot reserved for CUDA graph / warmup dummy
         # requests, whose outputs are discarded.
-        mrope_dummy_seq_slot = self.max_num_tokens * self.mapping.pp_size
+        mrope_dummy_seq_slot = getattr(
+            self, "max_num_seq_slots",
+            self.max_num_tokens * self.mapping.pp_size)
         num_accepted_draft_tokens = []  # per request
         is_enc_dec = self._is_encoder_decoder_model()
         cross_encoder_hidden_states: List[torch.Tensor] = []

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -310,7 +310,8 @@ def __init__(self,
                  max_seq_len: Optional[int],
                  lora_config: Optional[LoraConfig] = None,
                  model_weights_memory_tag: Optional[ExecutorMemoryType] = None,
-                 model_weights_restore_mode: Optional[RestoreMode] = None):
+                 model_weights_restore_mode: Optional[RestoreMode] = None,
+                 max_num_seq_slots: Optional[int] = None):
         """
         Initializes the ModelLoader.
 
@@ -326,20 +327,29 @@ def __init__(self,
                 they can be released/materialized independently of buffers.
             model_weights_restore_mode: RestoreMode for the model weights
                 virtual-memory scope.
+            max_num_seq_slots: Capacity of model buffers indexed by sequence
+                slot. This can exceed the scheduler admission batch size when
+                overlap scheduling is enabled.
         """
         self.llm_args = llm_args
         self.mapping = mapping
         self.spec_config = spec_config
         self.sparse_attention_config = sparse_attention_config
         self.max_num_tokens = max_num_tokens
         self.max_seq_len = max_seq_len
+        self.max_num_seq_slots = max_num_seq_slots
         self.lora_config = lora_config
         self.model_weights_memory_tag = model_weights_memory_tag
         self.model_weights_restore_mode = model_weights_restore_mode
         self.weight_mapper = None
         self._weight_pool_proxy = None
         self._gms_backend = None
 
+    def _set_runtime_model_config_attrs(self, config: ModelConfig) -> None:
+        """Attach executor-only allocation sizes before model construction."""
+        if self.max_num_seq_slots is not None:
+            config.extra_attrs['max_num_seq_slots'] = self.max_num_seq_slots
+
     @staticmethod
     def load_config_and_apply_defaults(
             checkpoint_dir: str, llm_args: TorchLlmArgs,
@@ -1208,6 +1218,7 @@ def _load_and_validate_config(
             load_config_kwargs['model_kwargs'] = self.llm_args.model_kwargs
 
         config = checkpoint_loader.load_config(**load_config_kwargs)
+        self._set_runtime_model_config_attrs(config)
 
         # Store nvfp4 config in extra_attrs for Linear layer access
         config.extra_attrs[