[https://nvbugs/6627795][fix] stop charging retiring requests against ADP admission and capacity

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -642,6 +642,7 @@ def __init__(
         self._dummy_encoder_inputs: List[MultimodalParams] = []
         self._profiling_stage_data = profiling_stage_data
         self._is_disagg = is_disagg
+        self._disable_overlap_scheduler = llm_args.disable_overlap_scheduler
         self._cache_transceiver_config = llm_args.cache_transceiver_config
         self._execution_stream = execution_stream
         self._kv_cache_manager_cls = self._get_model_kv_cache_manager_cls(
@@ -1462,6 +1463,7 @@ def _create_kv_cache_manager(
             execution_stream=self._execution_stream,
             layer_mask=spec_dec_layer_mask,
             is_disagg=self._is_disagg,
+            disable_overlap_scheduler=self._disable_overlap_scheduler,
             cold_page_codec_provider=cold_page_codec_provider,
             joint_kv_cache_reuse=self._joint_kv_cache_reuse,
         )
@@ -1668,6 +1670,7 @@ def _create_one_model_draft_kv_cache_manager(
             layer_mask=spec_dec_layer_mask,
             num_layers=num_draft_layers,
             is_disagg=self._is_disagg,
+            disable_overlap_scheduler=self._disable_overlap_scheduler,
             cold_page_codec_provider=cold_page_codec_provider,
             joint_kv_cache_reuse=self._joint_kv_cache_reuse,
         )
@@ -2044,6 +2047,7 @@ def _create_cross_kv_cache_manager(
             num_layers=num_layers,
             num_kv_heads=num_kv_heads,
             head_dim=head_dim,
+            disable_overlap_scheduler=self._disable_overlap_scheduler,
             kv_cache_type=tensorrt_llm.bindings.internal.batch_manager.
             CacheType.CROSS,
         )
@@ -2363,6 +2367,7 @@ def _create_kv_cache_manager(
         head_dim: Optional[int] = None,
         kv_cache_type=None,
         is_disagg: bool = False,
+        disable_overlap_scheduler: bool = False,
         cold_page_codec_provider: Optional[object] = None,
         joint_kv_cache_reuse: bool = False) -> KVCacheManager:
     """
@@ -2504,6 +2509,8 @@ def _create_kv_cache_manager(
         manager_extra_kwargs[
             "cold_page_codec_provider"] = cold_page_codec_provider
         manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse
+        manager_extra_kwargs[
+            "disable_overlap_scheduler"] = disable_overlap_scheduler
     if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2):
         manager_extra_kwargs["is_disagg"] = is_disagg
 
@@ -3025,26 +3032,64 @@ def create_kv_cache_compression_manager(
     return None
 
 
+def is_disagg_enabled(cache_transceiver_config) -> bool:
+    """True when a cache transceiver backend is configured."""
+    return (cache_transceiver_config is not None
+            and cache_transceiver_config.backend is not None)
+
+
 def compute_max_num_sequences(mapping: Mapping,
                               max_batch_size: int,
                               disable_overlap_scheduler: bool,
                               enable_overlap_headroom: bool = False) -> int:
     """Size the sequence-slot pool (and the sampler state it indexes).
 
-    ``enable_overlap_headroom`` is intentionally opt-in. Disaggregated
-    attention-DP 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. Pipeline parallelism already sizes the pool
-    by ``pp_size``.
+    ``enable_overlap_headroom`` is intentionally opt-in; see
+    ``should_enable_disagg_adp_overlap_headroom`` for when it is set. It buys one
+    extra micro-batch worth of slots, because a finished request's teardown is
+    deferred by one iteration and its slot is still held while the replacement
+    batch is admitted (nvbug 6627795). 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 = (2 if enable_overlap_headroom else 1)
     return max_batch_size * num_micro_batches
 
 
+def resolve_max_num_sequences(model_engine,
+                              mapping: Mapping,
+                              max_batch_size: int,
+                              llm_args,
+                              max_num_sequences: Optional[int] = None) -> int:
+    """Resolve the seat-pool size for a consumer, without re-deriving it.
+
+    Order of preference, and the order matters:
+
+    1. an explicitly supplied value -- the caller already has the number the
+       engine published;
+    2. ``model_engine.max_num_seq_slots`` -- the engine's own pool, which is
+       what every seat-keyed pool was sized against;
+    3. only then a fresh ``compute_max_num_sequences``, reusing the engine's
+       headroom gate so the fallback cannot size the pool below the engine's.
+    """
+    if max_num_sequences is not None:
+        return max_num_sequences
+    engine_seats = getattr(model_engine, "max_num_seq_slots", None)
+    if engine_seats is not None:
+        return engine_seats
+    # Engines that predate the attribute (unit-test stubs, mm-encoder-only
+    # engines): recompute, but with the same gate the engine would have used.
+    return compute_max_num_sequences(mapping,
+                                     max_batch_size,
+                                     llm_args.disable_overlap_scheduler,
+                                     enable_overlap_headroom=getattr(
+                                         model_engine,
+                                         "_enable_disagg_adp_overlap_headroom",
+                                         False))
+
+
 def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool:
     """Enable transactional ADP dummy handling while PP remains follow-up."""
     return not mapping.has_pp()
@@ -3073,11 +3118,10 @@ def should_enable_disagg_adp_overlap_headroom(
         mapping: Mapping,
         cache_transceiver_config: Optional[CacheTransceiverConfig],
         disable_overlap_scheduler: bool) -> bool:
-    """Gate extra sequence slots to non-PP disaggregated attention-DP."""
-    is_disagg = (cache_transceiver_config is not None
-                 and cache_transceiver_config.backend is not None)
-    return (mapping.enable_attention_dp and is_disagg and not mapping.has_pp()
-            and not disable_overlap_scheduler)
+    """Gate extra sequence slots to non-PP attention DP."""
+    is_disagg = is_disagg_enabled(cache_transceiver_config)
+    return (mapping.enable_attention_dp and not mapping.has_pp()
+            and (is_disagg or not disable_overlap_scheduler))
 
 
 def create_py_executor_instance(
@@ -3115,15 +3159,18 @@ def create_py_executor_instance(
 
     spec_config = model_engine.spec_config
 
-    if max_num_sequences is None:
-        max_num_sequences = compute_max_num_sequences(
-            mapping, max_batch_size, llm_args.disable_overlap_scheduler)
+    is_disagg = is_disagg_enabled(cache_transceiver_config)
+
+    max_num_sequences = resolve_max_num_sequences(
+        model_engine,
+        mapping,
+        max_batch_size,
+        llm_args,
+        max_num_sequences=max_num_sequences)
 
     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}"
     )
-    is_disagg = (cache_transceiver_config is not None
-                 and cache_transceiver_config.backend is not None)
     for key, value in llm_args.extra_resource_managers.items():
         if key in resources:
             raise ValueError(
@@ -3312,8 +3359,10 @@ 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.
-    # V1 scheduler handles overlap via two_step_lookahead, so skip the
-    # slot-pool overlap factor here.
+    # V1 scheduler handles overlap via two_step_lookahead, so the capacity
+    # scheduler's budget stays at the pipeline-depth bound and deliberately does
+    # not follow the sequence-slot pool: the overlap headroom is spare seats for
+    # leases already held, not extra admission.
     scheduler_capacity = max_batch_size * mapping.pp_size
     if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager:
         scheduler_capacity += 1
@@ -3500,22 +3549,22 @@ def create_py_executor_instance(
 
 
 def create_torch_sampler_args(
-    mapping: Mapping,
     *,
     max_seq_len: int,
-    max_batch_size: int,
     speculative_config: SpeculativeConfig,
     max_beam_width: int,
     disable_overlap_scheduler: bool,
     enable_async_worker: bool,
     enable_speculative_beam_history_d2h: bool,
-    max_num_sequences: Optional[int] = None,
+    max_num_sequences: int,
 ):
-    # The sampler's per-slot state is indexed by sequence slots, so it must
-    # be sized identically to the executor's slot pool.
-    if max_num_sequences is None:
-        max_num_sequences = compute_max_num_sequences(
-            mapping, max_batch_size, disable_overlap_scheduler)
+    # 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` is
+    # required, not optional: the old default recomputed the pool from
+    # `mapping`/`max_batch_size` *without* the overlap-headroom gate, so it could
+    # only ever produce a smaller number than the slots it indexes. Those two
+    # parameters are gone with it -- keeping them would leave the raw material
+    # for the same re-derivation lying next to the resolved value.
     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
@@ -3547,10 +3596,15 @@ def instantiate_sampler(
     enable_async_worker = (confidential_compute_enabled()
                            or llm_args.sampler_force_async_worker)
 
-    sampler_args = create_torch_sampler_args(
+    max_num_sequences = resolve_max_num_sequences(
+        engine,
         mapping,
+        max_batch_size,
+        llm_args,
+        max_num_sequences=max_num_sequences)
+
+    sampler_args = create_torch_sampler_args(
         max_seq_len=engine.max_seq_len,
-        max_batch_size=max_batch_size,
         speculative_config=speculative_config,
         max_beam_width=max_beam_width,
         disable_overlap_scheduler=llm_args.disable_overlap_scheduler,

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -931,6 +931,7 @@ def __init__(
         is_disagg: bool = False,
         enable_stats: bool = False,
         num_reserved_index_slots: int = 1,
+        disable_overlap_scheduler: bool = False,
         is_estimating_kv_cache: bool = False,
         cold_page_codec_provider: Optional[object] = None,
         joint_kv_cache_reuse: bool = False,
@@ -1452,14 +1453,18 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]:
         # (TRANS_IN_PROGRESS) and continue to hold their index slots. The 2x
         # capacity lets the next batch of active requests acquire slots without
         # waiting for the previous batch's transfers to finish.
+        needs_extra_index_slots = is_disagg or (
+            mapping.enable_attention_dp and not disable_overlap_scheduler and not mapping.has_pp()
+        )
         max_num_sequences = max_batch_size * mapping.pp_size
         assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative"
         index_mapper_capacity = (
-            max_num_sequences * (2 if is_disagg else 1) + num_reserved_index_slots
+            max_num_sequences * (2 if needs_extra_index_slots else 1) + num_reserved_index_slots
         )
         logger.info(
             f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} "
             f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, "
+            f"disable_overlap_scheduler={disable_overlap_scheduler}, "
             f"num_reserved_index_slots={num_reserved_index_slots}, "
             f"max_beam_width={max_beam_width})"
         )

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -107,7 +107,7 @@
 from .scheduler import (RequestScheduler, ScheduledRequests,
                         SerializableSchedulerOutput, WaitingQueue,
                         create_waiting_queue)
-from .scheduler.adp_router import ADPRouter
+from .scheduler.adp_router import ADPRouter, count_retiring_requests
 
 if TYPE_CHECKING:
     from ray.actor import ActorHandle
@@ -692,6 +692,8 @@ def __init__(
         # can receive the transfer-manager reference at construction time.
         self.adp_router: ADPRouter = ADPRouter.create(
             dist=self.dist,
+            has_seq_slot_headroom=getattr(
+                model_engine, "_enable_disagg_adp_overlap_headroom", False),
             kv_cache_manager=self.kv_cache_manager,
             attention_dp_config=self.llm_args.attention_dp_config,
             async_transfer_manager=self.async_transfer_manager,
@@ -5794,7 +5796,7 @@ def _validate_request(self, request: LlmRequest):
         self._validate_request_budget(request)
 
     def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue,
-                                    total_num_active_requests: int) -> None:
+                                    total_num_live_requests: int) -> None:
         """Fetch requests from request_queue and enqueue to waiting_queue."""
         # Block new requests while control requests are pending
         if len(self.control_requests) != 0:
@@ -5805,7 +5807,7 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue,
         # blocking would keep the loop from reaching the
         # `should_stop_processing` check that ends it, deadlocking shutdown()
         # on `shutdown_event`.
-        idle = (total_num_active_requests == 0 and len(waiting_queue) == 0
+        idle = (total_num_live_requests == 0 and len(waiting_queue) == 0
                 and not self.is_shutdown)
         if idle:
             # In Ray path (TLLM_DISABLE_MPI=1), use a periodic heartbeat timeout so rank 0
@@ -6013,14 +6015,16 @@ def _fetch_new_requests(
                 s.num_active_requests for s in all_rank_states
             ]
             total_num_active_requests = sum(all_ranks_num_active_requests)
+            total_num_live_requests = total_num_active_requests + sum(
+                s.num_retiring_requests for s in all_rank_states)
         else:
             total_num_active_requests = len(active_requests)
+            total_num_live_requests = total_num_active_requests
             all_ranks_num_active_requests = None
             all_rank_states = None
 
         # 2. Fetch and enqueue to waiting queue
-        self._fetch_and_enqueue_requests(waiting_queue,
-                                         total_num_active_requests)
+        self._fetch_and_enqueue_requests(waiting_queue, total_num_live_requests)
 
         # 3. Pop requests from waiting queue
         new_requests = self._pop_from_waiting_queue(
@@ -7225,7 +7229,11 @@ def _pad_attention_dp_dummy_request(self):
             return
 
         expected_num_active_requests = self.expected_num_active_requests
-        if expected_num_active_requests < len(self.active_requests):
+        num_routable_active_requests = len(self.active_requests)
+        if self.adp_router.exclude_retiring_requests:
+            num_routable_active_requests -= count_retiring_requests(
+                self.active_requests)
+        if expected_num_active_requests < num_routable_active_requests:
             # Not fatal, and not a capacity violation. The router derives this
             # value as
             #   min(max(ceil(multiplier * fair_share), max(per_rank_loads)),
@@ -7246,11 +7254,12 @@ def _pad_attention_dp_dummy_request(self):
             # event loop on every affected rank at once, leaving the survivors
             # to HangDetector-abort.
             logger.warning(
-                f"active_requests ({len(self.active_requests)}) exceeds "
+                f"routable active_requests "
+                f"({num_routable_active_requests}) exceeds "
                 f"expected_num_active_requests "
                 f"({expected_num_active_requests}); tolerating (a busy rank "
                 f"needs no attention-DP dummy).")
-            expected_num_active_requests = len(self.active_requests)
+            expected_num_active_requests = num_routable_active_requests
 
         num_active_request = self._count_schedulable_active_requests()
 

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -37,8 +37,8 @@
                            get_spec_resource_manager)
 from ..virtual_memory import scope as virtual_memory_scope
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
-                    create_py_executor_instance, instantiate_sampler, is_mla,
-                    validate_feature_combination)
+                    create_py_executor_instance, instantiate_sampler,
+                    is_disagg_enabled, is_mla, validate_feature_combination)
 from .config_utils import (is_hybrid_linear, is_minimax_m3,
                            resolve_cache_transceiver_config,
                            uses_vswa_kv_cache_layout)
@@ -756,15 +756,12 @@ def allocation_scope(current_stage: ExecutorMemoryType):
     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_disagg_adp_overlap_headroom", False)
-                                        else max_batch_size)
                 kwargs = {
                     "guided_decoding_config": guided_decoding_config,
-                    # The disaggregated attention-DP overlap path follows the
-                    # expanded slot pool. Other configurations retain
-                    # max_batch_size.
-                    "max_num_sequences": guided_decoder_slots,
+                    # The guided decoder's state is indexed by py_seq_slot
+                    # (guided_decoder.py: grammar_matchers[req.seq_slot], the
+                    # bitmask rows), so it must span the whole seat pool.
+                    "max_num_sequences": max_num_seq_slots,
                     "vocab_size_padded": model_engine.model.vocab_size_padded,
                     "rank": mapping.rank,
                 }
@@ -875,8 +872,7 @@ def allocation_scope(current_stage: ExecutorMemoryType):
     if model_engine.model.model_config.is_generation:
         #NOTE: non-generation models do not have kv cache
 
-        is_disagg = (cache_transceiver_config is not None
-                     and cache_transceiver_config.backend is not None)
+        is_disagg = is_disagg_enabled(cache_transceiver_config)
         is_hybrid = is_hybrid_linear(
             model_engine.model.model_config.pretrained_config)