[None][fix] Size seq-slot pool to cover disagg-gen KV admission to mitigate hangs

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,
@@ -3039,21 +3046,36 @@ 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
+    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. 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``.
+    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)
-    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 should_enable_adp_dummy_fixes(mapping: Mapping) -> bool:
@@ -3085,10 +3107,28 @@ def should_enable_disagg_adp_overlap_headroom(
         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)
+    return (mapping.enable_attention_dp
+            and is_disagg_enabled(cache_transceiver_config)
+            and not mapping.has_pp() and 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.
+    """
+    admission_bound = getattr(kv_cache_manager, "max_admissible_sequences",
+                              None)
+    if admission_bound is None 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(
@@ -3126,15 +3166,18 @@ def create_py_executor_instance(
 
     spec_config = model_engine.spec_config
 
+    is_disagg = is_disagg_enabled(cache_transceiver_config)
+
     if max_num_sequences is None:
         max_num_sequences = compute_max_num_sequences(
-            mapping, max_batch_size, llm_args.disable_overlap_scheduler)
+            mapping,
+            max_batch_size,
+            llm_args.disable_overlap_scheduler,
+            is_disagg=is_disagg)
 
     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(
@@ -3292,6 +3335,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)
 

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -1454,9 +1454,11 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]:
         # waiting for the previous batch's transfers to finish.
         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 is_disagg 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}, "

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -448,13 +448,15 @@ def __init__(
         self.mapping = mapping
         if mapping.has_pp():
             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,
+        # Both reasons the pool must exceed one forward batch are known from
+        # the runtime topology, before the model loads.
+        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,9 @@ 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)
+        # Slot-indexed metadata buffers must span the whole pool, whatever
+        # sizes it.
+        num_seq_slots = self.max_num_seq_slots
         if self.spec_metadata is not None:
             return self.spec_metadata
         self.spec_metadata = get_spec_metadata(

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -37,7 +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,
+                    compute_max_num_sequences, 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,
@@ -662,8 +663,14 @@ def allocation_scope(current_stage: ExecutorMemoryType):
     resolve_cache_transceiver_config(cache_transceiver_config)
 
     config = model_engine.model.model_config.pretrained_config
-    max_num_seq_slots = getattr(model_engine, "max_num_seq_slots",
-                                max_batch_size * getattr(mapping, "pp_size", 1))
+    # Engines that are not PyTorchModelEngine do not size the pool themselves;
+    # fall back to the same sizing function so the two never disagree.
+    max_num_seq_slots = getattr(
+        model_engine, "max_num_seq_slots", None) or compute_max_num_sequences(
+            mapping,
+            max_batch_size,
+            llm_args.disable_overlap_scheduler,
+            is_disagg=is_disagg_enabled(cache_transceiver_config))
     if is_mla(config):
         if model_engine.model.model_config.enable_flash_mla:
             tokens_per_block = 64
@@ -757,15 +764,11 @@ 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,
+                    # Indexed by py_seq_slot, so it must span the whole pool
+                    # rather than one forward batch.
+                    "max_num_sequences": max_num_seq_slots,
                     "vocab_size_padded": model_engine.model.vocab_size_padded,
                     "rank": mapping.rank,
                 }
@@ -876,8 +879,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)
 

tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py

@@ -1,3 +1,5 @@
+from tensorrt_llm.logger import logger
+
 from .llm_request import LlmRequest
 from .resource_manager import BaseResourceManager, SlotManager
 from .scheduler import ScheduledRequests
@@ -18,7 +20,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests) -> None:
         for llm_req in scheduled_batch.all_requests():
             if llm_req.is_disagg_generation_init_state:
                 logger.info(
-                    f"Skip assigning sequence slot for DISAGG_GENERATION_INIT request."
+                    "Skip assigning sequence slot for DISAGG_GENERATION_INIT request."
                 )
                 continue
             if llm_req.seq_slot is None or llm_req.is_disagg_generation_transmission_complete:

tensorrt_llm/_torch/speculative/interface.py

@@ -669,8 +669,8 @@ def batch_uses_penalty(self, value: bool) -> None:
     # Vocab size used for draft_probs buffer allocation.
     vocab_size: int = 0
     # Size of the SeqSlotManager pool. py_seq_slot values range over
-    # [0, num_seq_slots); DeepSeek-V4 overlap can use 2 * max_batch_size,
-    # larger than max_num_requests (== max_batch_size).
+    # [0, num_seq_slots), which exceeds max_num_requests (== max_batch_size)
+    # under PP, DeepSeek-V4 overlap headroom, and disaggregated serving.
     # Slot-indexed buffers (draft_probs) must span this full range.
     # 0 falls back to max_num_requests.
     num_seq_slots: int = 0
@@ -845,11 +845,8 @@ def prepare_rejection_sampling_buffers(self):
         if not self.use_rejection_sampling:
             return
 
-        # Slot-indexed buffers span the full SeqSlotManager pool: py_seq_slot
-        # can range over [0, num_seq_slots), which under DeepSeek-V4 overlap
-        # exceeds max_num_requests. Fall back to max_num_requests when the pool
-        # size is unknown (0). One extra scratch row at index ``slot_capacity``
-        # absorbs CUDA-graph dummy/padding requests (``py_seq_slot is None``).
+        # One extra scratch row past the pool absorbs CUDA-graph dummy and
+        # padding requests, whose py_seq_slot is None.
         slot_capacity = self.num_seq_slots or self.max_num_requests
         num_slot_rows = slot_capacity + 1
 

tensorrt_llm/_torch/speculative/utils.py

@@ -358,8 +358,7 @@ def _build_spec_metadata(spec_config,
                          num_seq_slots=None):
     use_rejection_sampling = getattr(spec_config, "use_rejection_sampling",
                                      False)
-    # Slot-indexed buffers (draft_probs) must span the SeqSlotManager pool;
-    # DeepSeek-V4 overlap can exceed max_num_requests.
+    # See SpecMetadata.num_seq_slots for why draft_probs must span the pool.
     num_seq_slots = (num_seq_slots
                      if num_seq_slots is not None else max_num_requests)
     vocab_size = getattr(model_config, "vocab_size", 0)

tests/integration/test_lists/test-db/l0_a10.yml

@@ -37,6 +37,7 @@ l0_a10:
   - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py
   - unittest/_torch/executor/test_error_classification.py
   - unittest/_torch/executor/test_resource_manager.py
+  - unittest/_torch/executor/test_seq_slot_sizing.py
   - unittest/_torch/moe/test_communication_factory.py
   # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no
   # test list either).