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

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -121,6 +121,13 @@ def _resolve_disagg_transceiver_route(
     return backend, runtime
 
 
+def is_disagg_enabled(
+        cache_transceiver_config: Optional[CacheTransceiverConfig]) -> bool:
+    """Whether this executor participates in disaggregated serving."""
+    return (cache_transceiver_config is not None
+            and cache_transceiver_config.backend is not None)
+
+
 def get_kv_cache_manager_cls(
         model_config: ModelConfig,
         kv_cache_config: KvCacheConfig,
@@ -642,6 +649,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 +1470,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,
             kv_events_config=None
             if estimating_kv_cache or model_engine.is_draft_model else
             self._llm_args.kv_cache_config.kv_events_config,
@@ -1682,6 +1691,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,
         )
@@ -2058,6 +2068,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,
         )
@@ -2377,6 +2388,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,
         kv_events_config: Optional[KVEventsConfig] = None,
         joint_kv_cache_reuse: bool = False) -> KVCacheManager:
@@ -2520,6 +2532,8 @@ def _create_kv_cache_manager(
             "cold_page_codec_provider"] = cold_page_codec_provider
         manager_extra_kwargs["kv_events_config"] = kv_events_config
         manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse
+        manager_extra_kwargs[
+            "disable_overlap_scheduler"] = disable_overlap_scheduler
         # V2 builds the block-reuse cache key of a multimodal token run from
         # the vocabulary size. Resolve it here rather than per-branch: the
         # manager needs it whenever block reuse can meet multimodal input,
@@ -3059,21 +3073,70 @@ def create_kv_cache_compression_manager(
 def compute_max_num_sequences(mapping: Mapping,
                               max_batch_size: int,
                               disable_overlap_scheduler: bool,
-                              enable_overlap_headroom: bool = False) -> int:
+                              enable_overlap_headroom: bool = False,
+                              is_disagg: 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``.
+    The pool must seat every request the admission path can let through. On a
+    disaggregated generation server that bound is KVCacheManagerV2's
+    IndexMapper, sized at twice max_num_sequences so a batch in KV transfer
+    can overlap a batch that is generating. Seats below that bound let a
+    request be admitted that cannot be seated once its transfer lands, and
+    add_slot then raises on the executor's event-loop thread, killing the rank
+    mid-collective.
+
+    enable_overlap_headroom covers a different case: attention DP can backfill
+    seats before the overlap scheduler releases the previous iteration's
+    terminal slots (nvbug 6627795). See
+    should_enable_disagg_adp_overlap_headroom for when it is set. 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)
-    return max_batch_size * num_micro_batches
+    num_seats = max_batch_size * num_micro_batches
+    if is_disagg:
+        # max() rather than another multiplication: the disagg and
+        # overlap-headroom factors both cover one extra set of in-flight
+        # sequences, so they overlap rather than compose.
+        num_seats = max(num_seats, max_batch_size * mapping.pp_size * 2)
+    return num_seats
+
+
+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),
+        is_disagg=is_disagg_enabled(
+            getattr(llm_args, "cache_transceiver_config", None)))
 
 
 def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool:
@@ -3104,11 +3167,36 @@ 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."""
+    return (mapping.enable_attention_dp and not mapping.has_pp()
+            and (is_disagg_enabled(cache_transceiver_config)
+                 or not disable_overlap_scheduler))
+
+
+def validate_seq_slot_pool_covers_admission(max_num_sequences: int,
+                                            kv_cache_manager) -> None:
+    """Fail at startup if the seat pool is smaller than what admission allows.
+
+    Otherwise the shortfall stays invisible until a request that cannot be
+    seated arrives, and it then surfaces as a hang rather than an error.
+    No-op for managers that publish no admission bound.
+
+    ``isinstance`` rather than ``is not None``: the V1/C++ manager does not
+    publish the attribute at all, and neither do the ``Mock`` stubs that stand
+    in for a cache manager in other modules' tests -- a ``Mock`` auto-creates
+    it, so an ``is None`` opt-in would fall through to the comparison and raise
+    ``TypeError`` instead of skipping the check.
+    """
+    admission_bound = getattr(kv_cache_manager, "max_admissible_sequences",
+                              None)
+    if not isinstance(admission_bound,
+                      int) or max_num_sequences >= admission_bound:
+        return
+    raise ValueError(
+        f"Sequence-slot pool ({max_num_sequences} seats) is smaller than the "
+        f"number of sequences {type(kv_cache_manager).__name__} can admit "
+        f"({admission_bound}). Requests would be admitted that cannot be "
+        "seated; see compute_max_num_sequences.")
 
 
 def create_py_executor_instance(
@@ -3146,15 +3234,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,6 +3403,7 @@ def create_py_executor_instance(
         if isinstance(model_engine, PyTorchModelEngine):
             model_engine._init_cuda_graph_lora_manager(lora_config)
 
+    validate_seq_slot_pool_covers_admission(max_num_sequences, kv_cache_manager)
     resources[ResourceManagerType.SEQ_SLOT_MANAGER] = SeqSlotManager(
         max_num_sequences)
 
@@ -3343,8 +3435,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
@@ -3531,22 +3625,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
@@ -3578,10 +3672,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

@@ -968,6 +968,7 @@ def __init__(
         is_disagg: bool = False,
         enable_stats: bool = False,
         num_reserved_index_slots: int = 1,
+        disable_overlap_scheduler: bool = False,
         kv_events_config: Optional[KVEventsConfig] = None,
         is_estimating_kv_cache: bool = False,
         cold_page_codec_provider: Optional[object] = None,
@@ -1525,14 +1526,20 @@ 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
-        )
+        # Admission bound for real requests, excluding the reserved slots that
+        # only ever hold persistent dummies. Published so the sequence-slot pool
+        # can be checked against it at startup.
+        self.max_admissible_sequences = max_num_sequences * (2 if needs_extra_index_slots else 1)
+        index_mapper_capacity = self.max_admissible_sequences + 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/model_engine.py

@@ -450,11 +450,13 @@ def __init__(
             init_pp_comm(mapping)
         # Disaggregated attention-DP can backfill a batch before the overlap
         # scheduler releases the previous batch's terminal sequence slots.
-        from ._util import (compute_max_num_sequences,
+        from ._util import (compute_max_num_sequences, is_disagg_enabled,
                             should_enable_adp_dummy_fixes,
                             should_enable_disagg_adp_overlap_headroom,
                             should_enable_non_overlap_adp_forward_intent,
                             should_enable_scheduler_aware_adp_dummy)
+        self._is_disagg = is_disagg_enabled(
+            getattr(llm_args, "cache_transceiver_config", None))
         self._enable_disagg_adp_overlap_headroom = (
             should_enable_disagg_adp_overlap_headroom(
                 mapping, llm_args.cache_transceiver_config,
@@ -465,6 +467,7 @@ def __init__(
             self.batch_size,
             llm_args.disable_overlap_scheduler,
             enable_overlap_headroom=self._enable_disagg_adp_overlap_headroom,
+            is_disagg=self._is_disagg,
         )
         self.dist = dist
         if dist is not None:
@@ -1132,8 +1135,7 @@ def _initialize_no_kv_cache_runner(
             mm_encoder_cache_enabled=self._mm_encoder_cache_enabled,
             spec_config=self.spec_config,
             is_draft_model=self.is_draft_model,
-            num_seq_slots=(self.max_num_seq_slots if
-                           self._enable_disagg_adp_overlap_headroom else None),
+            num_seq_slots=self.max_num_seq_slots,
             original_max_draft_len=self.original_max_draft_len,
             original_max_total_draft_tokens=(
                 self.original_max_total_draft_tokens),
@@ -3709,11 +3711,6 @@ def forward_multimodal_encoder_items(
     def _set_up_spec_metadata(
             self, spec_resource_manager: Optional[BaseResourceManager]):
         spec_config = self.spec_config if self.enable_spec_decode else None
-        # The disaggregated attention-DP overlap path opts into larger metadata
-        # buffers. Passing None preserves the established max_num_requests
-        # fallback for other configurations, including PP.
-        num_seq_slots = (self.max_num_seq_slots
-                         if self._enable_disagg_adp_overlap_headroom else None)
         if self.spec_metadata is not None:
             return self.spec_metadata
         self.spec_metadata = get_spec_metadata(
@@ -3724,7 +3721,7 @@ def _set_up_spec_metadata(
             spec_resource_manager=spec_resource_manager,
             is_draft_model=self.is_draft_model,
             max_seq_len=self.max_seq_len,
-            num_seq_slots=num_seq_slots)
+            num_seq_slots=self.max_num_seq_slots)
         return self.spec_metadata
 
     def cleanup(self) -> None: