[None][fix] fix disagg overlap slot headroom without MTP

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -2622,12 +2622,11 @@ def compute_max_num_sequences(mapping: Mapping,
                               enable_overlap_headroom: bool = False) -> 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``.
+    ``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``.
     """
     if mapping.has_pp():
         num_micro_batches = mapping.pp_size
@@ -2650,20 +2649,14 @@ def should_enable_dsv4_adp_dummy_fixes(model_type: Optional[str],
     return model_type in _ADP_DUMMY_FIX_MODEL_TYPES 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.
-
-    Deliberately NOT routed through ``should_enable_dsv4_adp_dummy_fixes``.
-    That gate now covers more model types, while this one doubles
-    ``max_num_sequences`` (see ``compute_max_num_sequences``) and therefore
-    changes the memory envelope; it must stay pinned to the one path it was
-    measured on.
-    """
-    return (model_type == "deepseek_v4" and not mapping.has_pp()
-            and spec_config is not None
-            and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+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)
 
 

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -386,13 +386,21 @@ 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.
+        # 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,
-                            should_enable_dsv4_adp_dummy_fixes,
-                            should_enable_dsv4_overlap_headroom)
+                            should_enable_disagg_adp_overlap_headroom,
+                            should_enable_dsv4_adp_dummy_fixes)
+        self._enable_disagg_adp_overlap_headroom = (
+            should_enable_disagg_adp_overlap_headroom(
+                mapping, llm_args.cache_transceiver_config,
+                llm_args.disable_overlap_scheduler))
         self.max_num_seq_slots = compute_max_num_sequences(
-            mapping, self.batch_size, llm_args.disable_overlap_scheduler)
+            mapping,
+            self.batch_size,
+            llm_args.disable_overlap_scheduler,
+            enable_overlap_headroom=self._enable_disagg_adp_overlap_headroom,
+        )
         self.dist = dist
         if dist is not None:
             ExpertStatistic.create(self.dist.rank)
@@ -476,21 +484,8 @@ def __init__(
             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._enable_dsv4_adp_dummy_fixes = should_enable_dsv4_adp_dummy_fixes(
+            model_type, mapping)
         if drafting_loop_wrapper is not None:
             self.model = drafting_loop_wrapper(self.model)
             self.model_is_wrapped = True
@@ -2812,11 +2807,11 @@ 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
+        # The disaggregated attention-DP 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.
+        # fallback for other configurations, including PP.
         num_seq_slots = (self.max_num_seq_slots
-                         if self._enable_dsv4_overlap_headroom else None)
+                         if self._enable_disagg_adp_overlap_headroom else None)
         if no_cache:
             return get_spec_metadata(
                 spec_config,

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -785,12 +785,13 @@ def drafting_loop_wrapper(model):
         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)
+                    model_engine, "_enable_disagg_adp_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.
+                    # The disaggregated attention-DP overlap path follows the
+                    # expanded slot pool. Other configurations retain
+                    # max_batch_size.
                     "max_num_sequences": guided_decoder_slots,
                     "vocab_size_padded": model_engine.model.vocab_size_padded,
                     "rank": mapping.rank,

tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap.yaml

@@ -1,21 +1,31 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 hostname: localhost
 model: DeepSeek-V3-Lite/fp8
 backend: pytorch
 cuda_graph_config: null
 free_gpu_memory_fraction: 0.2
+max_batch_size: 1
 context_servers:
   num_instances: 1
   tensor_parallel_size: 2
   pipeline_parallel_size: 1
   enable_attention_dp: true
   disable_overlap_scheduler: true
+  kv_cache_config:
+    use_kv_cache_manager_v2: true
   cache_transceiver_config:
-    backend: DEFAULT
+    backend: NIXL
+    transceiver_runtime: PYTHON
 generation_servers:
   num_instances: 1
   tensor_parallel_size: 2
   pipeline_parallel_size: 1
   enable_attention_dp: true
   disable_overlap_scheduler: false
+  kv_cache_config:
+    use_kv_cache_manager_v2: true
   cache_transceiver_config:
-    backend: DEFAULT
+    backend: NIXL
+    transceiver_runtime: PYTHON

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -936,7 +936,8 @@ def run_disaggregated_test(example_dir,
                            disagg_schedule_style=None,
                            post_client_test=None,
                            assert_gen_log_contains=None,
-                           perf_metrics_output_dir=None):
+                           perf_metrics_output_dir=None,
+                           server_start_timeout=300):
     """Run disaggregated test using service discovery instead of MPI.
 
     If assert_gen_log_contains is set, the generation-worker logs are captured and, after the
@@ -957,7 +958,8 @@ def run_disaggregated_test(example_dir,
         setup_disagg_cluster(config_file, model_name=model_path, env=run_env, cwd=cwd,
                              schedule_style=disagg_schedule_style,
                              save_log=assert_gen_log_contains is not None,
-                             perf_metrics_output_dir=perf_metrics_output_dir)
+                             perf_metrics_output_dir=perf_metrics_output_dir,
+                             server_start_timeout=server_start_timeout)
 
     server_host = config.get("hostname", "localhost")
 
@@ -1978,7 +1980,6 @@ def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_gen_only(
                            cwd=llm_venv.get_working_directory())
 
 
-@skip_no_hopper
 @pytest.mark.skip_less_device(4)
 @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'],
                          indirect=True)
@@ -1990,9 +1991,11 @@ def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap(
 
     run_disaggregated_test(disaggregated_example_root,
                            "deepseek_v3_lite_fp_8_attention_dp_overlap",
+                           num_iters=1,
                            env=llm_venv._new_env,
                            model_path=deepseek_v3_model_root,
-                           cwd=llm_venv.get_working_directory())
+                           cwd=llm_venv.get_working_directory(),
+                           server_start_timeout=1200)
 
 
 @skip_no_hopper

tests/integration/test_lists/waives.txt

@@ -111,7 +111,6 @@ disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_co
 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_gen_only[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322)
 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_one[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322)
 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_one_mtp[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322)
-disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322)
 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap_cuda_graph[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322)
 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_ctxpp2_gentp2_one_mtp[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322)
 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_ctxtp2ep2pp2_gentp4_one_mtp_block_reuse[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322)

tests/unittest/_torch/executor/test_seq_slot_sizing.py

@@ -1,30 +1,29 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""DeepSeek-V4 seq-slot sizing includes overlap headroom.
+"""Disaggregated attention-DP seq-slot sizing includes overlap headroom.
 
 Under the overlap scheduler, requests finished in the previous iteration
 still hold their sequence slots when the next iteration's
 prepare_resources runs, while the V2 scheduler has already dropped them
 from its budget (no_schedule_after_state=GENERATION_TO_COMPLETE) and
 backfilled their seats. Transient slot demand is therefore
-2 * max_batch_size. The headroom is intentionally limited to DeepSeek-V4;
-other models preserve their established sizing pending separate validation.
+2 * max_batch_size, regardless of whether speculative decoding is enabled.
+The headroom is selected from runtime topology rather than model architecture.
 
 compute_max_num_sequences is the single sizing implementation used both
 for the executor's SeqSlotManager pool (create_py_executor_instance) and
 for the sampler state (create_torch_sampler_args).
 """
 
-from unittest.mock import Mock
-
 import pytest
 
 from tensorrt_llm._torch.pyexecutor._util import (
     compute_max_num_sequences,
     create_torch_sampler_args,
+    should_enable_disagg_adp_overlap_headroom,
     should_enable_dsv4_adp_dummy_fixes,
-    should_enable_dsv4_overlap_headroom,
 )
+from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig
 from tensorrt_llm.mapping import Mapping
 
 SIZING_CASES = [
@@ -40,31 +39,28 @@
 
 
 @pytest.mark.parametrize(
-    "model_type,has_spec,is_mtp_one_model,pp_size,disable_overlap,expected",
+    "enable_attention_dp,is_disagg,pp_size,disable_overlap,expected",
     [
-        ("deepseek_v4", True, True, 1, False, True),
-        ("deepseek_v3", True, True, 1, False, False),
-        ("deepseek_v4", False, False, 1, False, False),
-        ("deepseek_v4", True, False, 1, False, False),
-        ("deepseek_v4", True, True, 2, False, False),
-        ("deepseek_v4", True, True, 1, True, False),
-        # The widened ADP dummy gate must not leak into the headroom gate:
-        # doubling max_num_sequences changes the memory envelope and has only
-        # been validated on the DSv4 MTP overlap path.
-        ("qwen3_5_moe", True, True, 1, False, False),
+        (True, True, 1, False, True),
+        (False, True, 1, False, False),
+        (True, False, 1, False, False),
+        (True, True, 2, False, False),
+        (True, True, 1, True, False),
     ],
 )
-def test_dsv4_overlap_headroom_gate(
-    model_type, has_spec, is_mtp_one_model, pp_size, disable_overlap, expected
+def test_disagg_adp_overlap_headroom_gate(
+    enable_attention_dp, is_disagg, pp_size, disable_overlap, expected
 ):
-    spec_config = None
-    if has_spec:
-        spec_config = Mock()
-        spec_config.spec_dec_mode.is_mtp_eagle_one_model.return_value = is_mtp_one_model
-    mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size)
+    mapping = Mapping(
+        world_size=pp_size,
+        tp_size=1,
+        pp_size=pp_size,
+        enable_attention_dp=enable_attention_dp,
+    )
+    cache_config = CacheTransceiverConfig(backend="NIXL") if is_disagg else None
 
     assert (
-        should_enable_dsv4_overlap_headroom(model_type, spec_config, mapping, disable_overlap)
+        should_enable_disagg_adp_overlap_headroom(mapping, cache_config, disable_overlap)
         is expected
     )