[https://nvbugs/6487039][fix] Generalize ADP dummy lifecycle

tensorrt_llm/_torch/models/modeling_qwen2vl.py

@@ -187,6 +187,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):
 
@@ -1723,8 +1732,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

@@ -58,6 +58,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,
 )
@@ -1216,7 +1217,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

@@ -2333,41 +2333,24 @@ def create_kv_cache_compression_manager(
     return None
 
 
-def compute_max_num_sequences(mapping: Mapping,
-                              max_batch_size: int,
-                              disable_overlap_scheduler: bool,
-                              enable_overlap_headroom: bool = False) -> int:
+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).
 
-    ``enable_overlap_headroom`` is intentionally opt-in. DeepSeek-V4 needs a
-    second non-PP slot set because the V2 scheduler can backfill seats before
-    the overlap scheduler releases the previous iteration's terminal slots.
-    Other models retain their established sizing until that behavior is
-    validated independently. Pipeline parallelism already sizes the pool by
-    ``pp_size``.
+    The overlap scheduler needs a second non-PP slot set because it can
+    backfill seats before releasing the previous iteration's terminal slots.
+    Pipeline parallelism already sizes the pool by ``pp_size``.
     """
     if mapping.has_pp():
         num_micro_batches = mapping.pp_size
     else:
-        num_micro_batches = (2 if enable_overlap_headroom
-                             and not disable_overlap_scheduler else 1)
+        num_micro_batches = 1 if disable_overlap_scheduler else 2
     return max_batch_size * num_micro_batches
 
 
-def should_enable_dsv4_adp_dummy_fixes(model_type: Optional[str],
-                                       mapping: Mapping) -> bool:
-    """Gate DSv4 ADP dummy behavior while PP remains follow-up scope."""
-    return model_type == "deepseek_v4" and not mapping.has_pp()
-
-
-def should_enable_dsv4_overlap_headroom(
-        model_type: Optional[str], spec_config: Optional[SpeculativeConfig],
-        mapping: Mapping, disable_overlap_scheduler: bool) -> bool:
-    """Gate extra sequence slots to the validated DSv4 MTP overlap path."""
-    return (should_enable_dsv4_adp_dummy_fixes(model_type, mapping)
-            and spec_config is not None
-            and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
-            and not disable_overlap_scheduler)
+def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool:
+    """Enable transactional ADP dummy handling while PP remains follow-up."""
+    return not mapping.has_pp()
 
 
 def create_py_executor_instance(
@@ -2696,8 +2679,12 @@ def create_py_executor_instance(
             enable_prefix_aware_scheduling=enable_prefix_aware_scheduling,
         )
 
-        mb_scheduler = BindMicroBatchScheduler(max_batch_size, max_num_tokens,
-                                               ctx_chunk_config)
+        mb_scheduler = BindMicroBatchScheduler(
+            max_batch_size,
+            max_num_tokens,
+            ctx_chunk_config,
+            no_schedule_until_state=no_schedule_until_state,
+        )
 
         reorder_policy_config = llm_args.reorder_policy_config
         if reorder_policy_config is not None:

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -355,11 +355,10 @@ def __init__(
         self.mapping = mapping
         if mapping.has_pp():
             init_pp_comm(mapping)
-        # Start with the established pool size. Once the model is loaded we
-        # selectively enable headroom for the non-PP DeepSeek-V4 overlap path.
+        # The overlap scheduler can hold two iterations' requests at once.
+        # Every model-side buffer indexed by py_seq_slot must span this pool.
         from ._util import (compute_max_num_sequences,
-                            should_enable_dsv4_adp_dummy_fixes,
-                            should_enable_dsv4_overlap_headroom)
+                            should_enable_adp_dummy_fixes)
         self.max_num_seq_slots = compute_max_num_sequences(
             mapping, self.batch_size, llm_args.disable_overlap_scheduler)
         self.dist = dist
@@ -433,6 +432,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,
@@ -443,23 +443,11 @@ def __init__(
                 setattr(self, "moe_load_balancer", moe_load_balancer)
         else:
             self.model = model
-        pretrained_config = self.model.model_config.pretrained_config
-        model_type = getattr(pretrained_config, "model_type", None)
-        # Keep the scheduler/dummy fix model-scoped, while the larger slot pool
-        # is restricted to the validated MTP overlap configuration. PP remains
-        # on its established path for follow-up changes.
-        self._enable_dsv4_adp_dummy_fixes = (should_enable_dsv4_adp_dummy_fixes(
-            model_type, mapping))
-        self._enable_dsv4_overlap_headroom = (
-            should_enable_dsv4_overlap_headroom(
-                model_type, spec_config, mapping,
-                llm_args.disable_overlap_scheduler))
-        self.max_num_seq_slots = compute_max_num_sequences(
-            mapping,
-            self.batch_size,
-            llm_args.disable_overlap_scheduler,
-            enable_overlap_headroom=self._enable_dsv4_overlap_headroom,
-        )
+        self._validate_mrope_position_delta_cache_capacity()
+        # Apply transactional dummy handling to every non-PP disaggregated ADP
+        # model. Sequence-slot capacity follows the independent overlap
+        # lifecycle invariant above.
+        self._enable_adp_dummy_fixes = should_enable_adp_dummy_fixes(mapping)
         if drafting_loop_wrapper is not None:
             self.model = drafting_loop_wrapper(self.model)
             self.model_is_wrapped = True
@@ -925,6 +913,33 @@ def set_guided_decoder(self,
             return success
         return False
 
+    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 indexing past an undersized 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.")
+
     @property
     def use_mrope(self):
         use_mrope = False
@@ -2637,11 +2652,8 @@ 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
-        # Only the scoped DeepSeek-V4 overlap path opts into larger metadata
-        # buffers. Passing None preserves the established max_num_requests
-        # fallback for every other model, including MTP-Eagle with PP.
-        num_seq_slots = (self.max_num_seq_slots
-                         if self._enable_dsv4_overlap_headroom else None)
+        # Slot-indexed metadata must span the same pool as SeqSlotManager.
+        num_seq_slots = self.max_num_seq_slots
         if no_cache:
             return get_spec_metadata(
                 spec_config,
@@ -4024,7 +4036,7 @@ def _prepare_tp_inputs(
         # that carry no MRoPE metadata at all. The cache is zero-initialized and
         # the write path only ever targets real ``py_seq_slot``s, so this slot
         # permanently reads back a zero delta.
-        mrope_dummy_seq_slot = self.max_num_tokens * self.mapping.pp_size
+        mrope_dummy_seq_slot = self.max_num_seq_slots
         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

@@ -363,7 +363,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.
 
@@ -379,20 +380,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,
@@ -1402,6 +1412,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[

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -569,8 +569,8 @@ def __init__(
         self.resource_manager = resource_manager
         self.scheduler = scheduler
         self.model_engine = model_engine
-        self._enable_dsv4_adp_dummy_fixes = getattr(
-            model_engine, "_enable_dsv4_adp_dummy_fixes", False)
+        self._enable_adp_dummy_fixes = getattr(model_engine,
+                                               "_enable_adp_dummy_fixes", False)
         self.enable_attention_dp = model_engine.enable_attention_dp
         self.dist = dist
         self.sampler = sampler
@@ -3323,7 +3323,7 @@ def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None:
         must release theirs before retrying or the fixed dummy request ID leaks
         cache resources on every skipped iteration.
         """
-        if not self._enable_dsv4_adp_dummy_fixes:
+        if not self._enable_adp_dummy_fixes:
             return
 
         dummy_request = self._pending_adp_dummy_request
@@ -5682,16 +5682,15 @@ def _check_disagg_ctx_schedulable_status(self,
     def _count_schedulable_active_requests(self) -> int:
         """Count active requests that are ready for scheduling.
 
-        The non-PP DeepSeek-V4 disaggregated ADP path mirrors the decoder
-        scheduler's state window [CONTEXT_INIT, GENERATION_TO_COMPLETE). This
-        covers generation-first context requests below the lower bound and
-        terminal requests at the upper bound. Other configurations retain the
-        established ADP behavior; PP eligibility remains follow-up scope.
+        The non-PP disaggregated ADP path uses the scheduler's state-
+        eligibility contract. This keeps decoder-only and encoder-decoder
+        boundaries and special exclusions aligned without duplicating them
+        here. PP eligibility remains follow-up scope.
 
         Returns:
             The number of active requests eligible for scheduling.
         """
-        if (not self._enable_dsv4_adp_dummy_fixes
+        if (not self._enable_adp_dummy_fixes
                 or self.kv_cache_transceiver is None):
             if self.kv_cache_transceiver is None:
                 return len(self.active_requests)
@@ -5701,12 +5700,8 @@ def _count_schedulable_active_requests(self) -> int:
                 if not (req.is_disagg_generation_init_state
                         or req.is_disagg_generation_transmission_in_progress))
 
-        schedule_from_value = LlmRequestState.CONTEXT_INIT.value
-        to_complete_value = LlmRequestState.GENERATION_TO_COMPLETE.value
-
-        return sum(
-            1 for req in self.active_requests
-            if schedule_from_value <= req.state_value < to_complete_value)
+        return sum(1 for req in self.active_requests
+                   if self.scheduler.is_request_in_schedulable_state(req))
 
     def _has_adp_dummy_kv_capacity(self,
                                    token_nums: Optional[List[int]]) -> bool:
@@ -5827,7 +5822,7 @@ def _pad_attention_dp_dummy_request(self):
                 key="attention_dp_dummy_insufficient_kv_capacity")
             return
 
-        if (not self._enable_dsv4_adp_dummy_fixes
+        if (not self._enable_adp_dummy_fixes
                 or self.kv_cache_transceiver is None):
             llm_request = self.kv_cache_manager.add_dummy_requests(
                 request_ids=dummy_request_ids,
@@ -5862,9 +5857,8 @@ def _pad_attention_dp_dummy_request(self):
         except OutOfPagesError:
             dummy_requests = None
         if not dummy_requests:
-            logger.warning(
-                "Cannot allocate DeepSeek-V4 ADP pad dummy; rank schedules "
-                "an empty batch and the fleet will retry.")
+            logger.warning("Cannot allocate ADP pad dummy; rank schedules "
+                           "an empty batch and the fleet will retry.")
             return
 
         dummy_request = dummy_requests[0]

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -775,14 +775,10 @@ def drafting_loop_wrapper(model):
     if guided_decoding_config is not None:
         with allocation_scope(ExecutorMemoryType.GUIDED_DECODER):
             if mapping.is_last_pp_rank():
-                guided_decoder_slots = (max_num_seq_slots if getattr(
-                    model_engine, "_enable_dsv4_overlap_headroom", False) else
-                                        max_batch_size)
                 kwargs = {
                     "guided_decoding_config": guided_decoding_config,
-                    # The scoped DeepSeek-V4 path follows the expanded slot
-                    # pool. Other configurations retain max_batch_size.
-                    "max_num_sequences": guided_decoder_slots,
+                    # Guided-decoder state is indexed by sequence slot.
+                    "max_num_sequences": max_num_seq_slots,
                     "vocab_size_padded": model_engine.model.vocab_size_padded,
                     "rank": mapping.rank,
                 }