[https://nvbugs/6919819][fix] Give non-attention-DP overlap runs the retiring-request seat headroom

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -3917,8 +3917,17 @@ def should_enable_overlap_headroom(mapping: Mapping,
     """Gate the extra micro-batch of sequence slots.
 
     True only where a retiring request and the replacement that took its place
-    can own a seat at the same time: attention DP, non-PP, overlap-on, V2 and
-    non-hybrid.
+    can own a seat at the same time: non-PP, overlap-on, V2 and non-hybrid.
+
+    Under the overlap scheduler a request that emitted its last token keeps its
+    seat until the next iteration's ``_process_previous_batch``, which runs
+    after that iteration has already admitted new work. Admission excludes such
+    retiring requests whenever this gate is on (``ADPRouter`` on the
+    attention-DP path, ``_fetch_new_requests`` otherwise), so the replacement
+    cohort needs seats of its own. Without the extra micro-batch, a workload
+    whose requests retire after one iteration -- a context server, where every
+    request finishes in its prefill step -- can only refill the seats the
+    previous iteration did not use, and runs at roughly half its batch.
 
     Widening the pool is only safe when every ``py_seq_slot``-indexed pool is
     sized from ``compute_max_num_sequences``. Two model families size one from
@@ -3935,8 +3944,7 @@ def should_enable_overlap_headroom(mapping: Mapping,
     """
     if is_hybrid or has_mrope_delta_cache or not kv_cache_manager_is_v2:
         return False
-    return (mapping.enable_attention_dp and not mapping.has_pp()
-            and not disable_overlap_scheduler)
+    return not mapping.has_pp() and not disable_overlap_scheduler
 
 
 def validate_seq_slot_pool_covers_admission(max_num_sequences: int,

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -6717,8 +6717,14 @@ def _fetch_new_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
+            total_num_live_requests = len(active_requests)
+            total_num_active_requests = total_num_live_requests
+            if self.adp_router.exclude_retiring_requests:
+                # Retiring requests hold a seat from the overlap headroom, not
+                # one that admission hands out (see
+                # should_enable_overlap_headroom). They still count as live.
+                total_num_active_requests -= count_retiring_requests(
+                    active_requests)
             all_ranks_num_active_requests = None
             all_rank_states = None
 

tests/unittest/_torch/executor/test_py_executor.py

@@ -1173,6 +1173,7 @@ def _make_executor_with_kv_cache_manager(kv_cache_manager):
         ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager
     }
     executor.enable_attention_dp = False
+    executor.adp_router = Mock(exclude_retiring_requests=False)
     executor._kv_cache_capacity = executor._read_kv_cache_capacity()
     executor._kv_load_tracker = KvLoadTracker(kv_cache_manager, rank=0, enabled=True)
     return executor
@@ -2725,6 +2726,125 @@ def test_pad_still_warns_on_a_genuine_routable_surplus():
     assert "exceeds expected_num_active_requests" in mock_logger.warning.call_args[0][0]
 
 
+def _run_non_adp_fetch(active_requests, *, exclude_retiring_requests, max_num_active_requests=4):
+    """Drive the non-attention-DP branch of ``_fetch_new_requests``.
+
+    Returns the ``total_num_live_requests`` handed to the queue wait and the
+    ``total_num_active_requests`` that admission subtracts from its capacity.
+    """
+    stub = types.SimpleNamespace(
+        enable_attention_dp=False,
+        adp_router=Mock(exclude_retiring_requests=exclude_retiring_requests),
+        max_num_active_requests=max_num_active_requests,
+        enable_iter_perf_stats=False,
+        num_fetch_requests=0,
+        dist=Mock(cp_config={}, cp_rank=0, cp_size=1, rank=0),
+        _should_exclude_last_generation_logits=lambda: False,
+        _kv_load_tracker=KvLoadTracker(None, rank=0, enabled=False),
+        _kv_cache_capacity={},
+    )
+    seen = {}
+    stub._fetch_and_enqueue_requests = lambda waiting_queue, live: seen.update(live=live)
+
+    def _pop(waiting_queue, total_num_active_requests, all_ranks_num_active_requests):
+        seen["active"] = total_num_active_requests
+        seen["all_ranks"] = all_ranks_num_active_requests
+        return []
+
+    stub._pop_from_waiting_queue = _pop
+    with patch(
+        "tensorrt_llm._torch.pyexecutor.py_executor.merge_requests", return_value=[]
+    ) as merge:
+        assert PyExecutor._fetch_new_requests(stub, Mock(), active_requests) == []
+    merge.assert_called_once()
+    return seen
+
+
+def test_non_adp_admission_excludes_retiring_requests_with_headroom():
+    """A context server under overlap: every request retires one iteration late.
+
+    With ``max_batch_size`` 4 and the previous batch's 4 requests all retiring,
+    admission must see an empty pool so it can seat a full replacement batch in
+    the headroom (nvbug 6919819); counting the retirees capped every refill at
+    the seats the previous iteration had left unused. The retirees still count
+    as live, so the idle wait does not block while they are resident.
+    """
+    retiring = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE, request_id=i) for i in range(4)]
+
+    seen = _run_non_adp_fetch(retiring, exclude_retiring_requests=True)
+
+    assert seen == {"live": 4, "active": 0, "all_ranks": None}
+
+
+def test_non_adp_admission_counts_retiring_requests_without_headroom():
+    """Without the headroom a retiring request still holds one of the B seats."""
+    active = [
+        _make_adp_request(_STATE_GENERATION_TO_COMPLETE, request_id=1),
+        _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=2),
+    ]
+
+    seen = _run_non_adp_fetch(active, exclude_retiring_requests=False)
+
+    assert seen == {"live": 2, "active": 2, "all_ranks": None}
+
+
+def test_non_adp_admission_keeps_counting_requests_that_are_not_retiring():
+    active = [
+        _make_adp_request(LlmRequestState.CONTEXT_INIT, request_id=1),
+        _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=2),
+        _make_adp_request(_STATE_GENERATION_TO_COMPLETE, request_id=3),
+    ]
+
+    seen = _run_non_adp_fetch(active, exclude_retiring_requests=True)
+
+    assert seen == {"live": 3, "active": 2, "all_ranks": None}
+
+
+@pytest.mark.parametrize(
+    ("exclude_retiring_requests", "expected_admitted_ids"),
+    [(True, [10, 11, 12, 13]), (False, [])],
+)
+def test_non_adp_fetch_pops_a_full_batch_beside_retiring_requests(
+    exclude_retiring_requests, expected_admitted_ids
+):
+    """Run the real waiting-queue pop: 4 retiring, 4 waiting, capacity 4.
+
+    With the overlap headroom the whole waiting batch is admitted beside the
+    retiring requests; without it the retiring requests fill every seat. Either
+    way the retiring requests still count as live for the idle wait.
+    """
+    retiring = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE, request_id=i) for i in range(4)]
+    waiting_queue = FCFSWaitingQueue()
+    waiting_queue.add_requests(RequestQueueItem(10 + i) for i in range(4))
+    stub = types.SimpleNamespace(
+        enable_attention_dp=False,
+        adp_router=Mock(exclude_retiring_requests=exclude_retiring_requests),
+        max_num_active_requests=4,
+        is_benchmark_disagg=False,
+        enable_iter_perf_stats=False,
+        num_fetch_requests=0,
+        dist=Mock(cp_config={}, cp_rank=0, cp_size=1, rank=0, tp_size=1),
+        _should_exclude_last_generation_logits=lambda: False,
+        _kv_load_tracker=KvLoadTracker(None, rank=0, enabled=False),
+        _kv_cache_capacity={},
+    )
+    for name in ("_get_request_admission_capacity", "_pop_from_waiting_queue"):
+        setattr(stub, name, types.MethodType(getattr(PyExecutor, name), stub))
+    live = []
+    stub._fetch_and_enqueue_requests = lambda queue, num_live: live.append(num_live)
+
+    with patch(
+        "tensorrt_llm._torch.pyexecutor.py_executor.merge_requests",
+        side_effect=lambda new_requests, **_: list(new_requests),
+    ):
+        admitted = PyExecutor._fetch_new_requests(stub, waiting_queue, retiring)
+
+    assert [item.id for item in admitted] == expected_admitted_ids
+    assert [item.id for item in waiting_queue] == [10, 11, 12, 13][len(expected_admitted_ids) :]
+    assert live == [4]
+    assert stub.num_fetch_requests == len(expected_admitted_ids)
+
+
 def test_generic_disagg_adp_mixed_rank_states_stay_queueable():
     # The generic non-PP path must give both ranks a non-empty scheduled batch:
     # one rank schedules its real request, while the terminal-only rank

tests/unittest/_torch/executor/test_seq_slot_sizing.py

@@ -13,12 +13,15 @@
 (``no_schedule_after_state=GENERATION_TO_COMPLETE``) and the ADP router has
 already excluded them from the counts admission subtracts from its capacity --
 so a replacement cohort is admitted while the retiring one is still resident.
-``should_enable_overlap_headroom`` is the single gate for that, and it requires
-attention DP: ``ADPRouter.exclude_retiring_requests`` is read only inside the
-``if self.enable_attention_dp:`` branch of ``_fetch_new_requests``, so outside
-attention DP the else-branch subtracts ``len(active_requests)`` with retirees
-included and residency is capped at ``max_batch_size * pp_size`` however many
-seats exist. It stays off for V1 (the V1 capacity schedulers hardcode
+``should_enable_overlap_headroom`` is the single gate for that. It does not
+depend on attention DP: ``ADPRouter.exclude_retiring_requests`` follows the gate,
+and both branches of ``_fetch_new_requests`` read it -- the attention-DP branch
+through the router's rank states, the other by discounting
+``count_retiring_requests`` from ``len(active_requests)``. A context server is the
+case that needs it most: every request retires the iteration after its prefill,
+so without the headroom only the seats the previous iteration left unused can be
+refilled and the batch alternates between partial fills (nvbug 6919819 measured a
+``3, 1, 3, 1`` cadence at ``max_batch_size=4``). It stays off for V1 (the V1 capacity schedulers hardcode
 ``GENERATION_COMPLETE``) and for hybrid models (SSM state is sized from
 ``max_batch_size``). "V1" there is the manager the creator *selects*, not the
 ``use_kv_cache_manager_v2`` request: the gate reads
@@ -28,12 +31,12 @@
 indexed by ``py_seq_slot``.
 
 The **index pool** (``KVCacheManagerV2.max_admissible_sequences``) is widened by
-a deliberately *broader* predicate -- ``is_disagg or (non-PP and overlap-on)``,
-with no attention-DP term -- so it equals the seat pool under attention DP and
-runs ahead of it elsewhere. Both reasons for the surplus are real: a request
-awaiting its KV transfer holds an index lease and no seat at all (because
-``SeqSlotManager.prepare_resources`` skips ``DISAGG_GENERATION_INIT``), and a
-non-ADP overlap run needs the leases without being able to occupy the seats.
+a deliberately *broader* predicate -- ``is_disagg or (non-PP and overlap-on)`` --
+so it equals the seat pool wherever the seat gate is on and runs ahead of it
+elsewhere. The surplus is real: a request awaiting its KV transfer holds an
+index lease and no seat at all (because ``SeqSlotManager.prepare_resources``
+skips ``DISAGG_GENERATION_INIT``), and a hybrid or MRoPE-delta model keeps the
+single-micro-batch seat pool while its V2 index pool still widens.
 Hence ``validate_seq_slot_pool_covers_admission`` requires the index pool to
 *cover* the seat pool rather than to equal it. Equalising them by widening the
 seat pool is the thing not to do: surplus leases cost page-table rows, surplus
@@ -92,9 +95,11 @@
 @pytest.mark.parametrize(
     "enable_attention_dp,pp_size,disable_overlap,is_v2,is_hybrid,expected",
     [
-        # The scenario this PR exists for: attention DP, no PP, overlap on, V2.
+        # Attention DP, no PP, overlap on, V2.
         (True, 1, False, True, False, True),
-        (False, 1, False, True, False, False),
+        # The same without attention DP: an aggregated or context server
+        # (nvbug 6919819) retires requests one iteration late just the same.
+        (False, 1, False, True, False, True),
         (True, 1, True, True, False, False),
         (True, 2, False, True, False, False),
         (True, 4, False, True, False, False),
@@ -131,50 +136,46 @@ def test_overlap_headroom_gate(
 def test_overlap_headroom_gate_does_not_depend_on_disaggregation():
     """Disagg is not a reason to widen the *seat* pool, only the index pool.
 
-    The gate used to take ``cache_transceiver_config`` and return True for a
-    disaggregated server regardless of attention DP. That bought seats that could
-    never be occupied: a request in KV transfer holds an index lease and no seat,
-    and outside attention DP admission subtracts ``len(active_requests)`` so
-    residency stays at ``max_batch_size * pp_size``. Asserting on the signature
-    states the intent that a value test cannot: the parameter's mere presence is
-    what invites the conflation back.
+    The gate used to take ``cache_transceiver_config`` and return True for any
+    disaggregated server. A request in KV transfer holds an index lease and no
+    seat, so disaggregation alone never needs an extra seat; only a retiring
+    request does. Asserting on the signature states the intent that a value test
+    cannot: the parameter's mere presence is what invites the conflation back.
     """
     params = inspect.signature(should_enable_overlap_headroom).parameters
     assert "cache_transceiver_config" not in params
     assert "is_disagg" not in params
 
 
-def test_overlap_headroom_gate_requires_attention_dp():
-    """Attention DP is the only place the extra seats can be occupied.
+def test_overlap_headroom_gate_does_not_require_attention_dp():
+    """Retiring requests hold a seat for one iteration with or without attention DP.
 
-    ``ADPRouter.exclude_retiring_requests`` is what lets a replacement cohort be
-    admitted while the retiring cohort still holds its seats, and that correction
-    is read only inside the ``if self.enable_attention_dp:`` branch of
-    ``_fetch_new_requests``. The else-branch subtracts ``len(active_requests)``
-    with retirees included, so residency outside attention DP is capped at
-    ``max_batch_size * pp_size`` no matter how many seats exist. Widening the
-    pool there would allocate sampler state, ``[seats, draft_len, vocab]`` draft
-    probabilities and pinned-host block-offset tables for seats that admission
-    can never hand out.
-
-    The index pool still widens on this path, so the two pools are unequal here
-    by design -- see ``test_sizing_matches_kv_manager_admission_bound``.
+    The deferral comes from the overlap executor loop, which admits new work
+    before ``_process_previous_batch`` frees the requests that finished in the
+    previous iteration. Both branches of ``_fetch_new_requests`` discount those
+    retiring requests when ``ADPRouter.exclude_retiring_requests`` is set, so the
+    seats they hold must come from the headroom in both cases too. Gating the
+    headroom on attention DP left a single-rank context server refilling only
+    the seats its previous iteration had not used (nvbug 6919819).
     """
     kwargs = dict(kv_cache_manager_is_v2=True, is_hybrid=False)
     with_adp = Mapping(world_size=1, tp_size=1, pp_size=1, enable_attention_dp=True)
     without_adp = Mapping(world_size=1, tp_size=1, pp_size=1, enable_attention_dp=False)
 
     assert should_enable_overlap_headroom(with_adp, False, **kwargs) is True
-    assert should_enable_overlap_headroom(without_adp, False, **kwargs) is False
+    assert should_enable_overlap_headroom(without_adp, False, **kwargs) is True
+    # The overlap scheduler is what defers retirement; without it there is no
+    # retiring request to seat.
+    assert should_enable_overlap_headroom(without_adp, True, **kwargs) is False
 
 
 @pytest.mark.parametrize("enable_attention_dp", [False, True])
-def test_non_adp_seat_pool_stays_at_max_batch_size(enable_attention_dp):
-    """End-to-end: the gate's ADP term reaches the seat count itself.
+def test_overlap_seat_pool_is_doubled_with_or_without_attention_dp(enable_attention_dp):
+    """End-to-end: the gate reaches the seat count itself.
 
     Asserting on ``compute_max_num_sequences`` rather than on the predicate is
-    what pins the consequence the reviewer asked for -- ``B`` seats for a non-ADP
-    overlap run, ``2B`` only under attention DP.
+    what pins the consequence: ``2B`` seats for any non-PP overlap run on V2, so
+    a retiring cohort of ``B`` and its replacement can be resident together.
     """
     max_batch_size = 8
     mapping = Mapping(
@@ -191,7 +192,7 @@ def test_non_adp_seat_pool_stays_at_max_batch_size(enable_attention_dp):
             mapping, False, kv_cache_manager_is_v2=True
         ),
     )
-    assert seats == (2 * max_batch_size if enable_attention_dp else max_batch_size)
+    assert seats == 2 * max_batch_size
 
 
 def test_overlap_headroom_gate_excludes_mrope_delta_cache_models():
@@ -522,12 +523,12 @@ def test_sizing_matches_kv_manager_admission_bound(
     ``validate_seq_slot_pool_covers_admission`` -- and the assertion is exactly
     that validator's invariant.
 
-    The seat gate is the narrower predicate: it additionally requires attention
-    DP, because that is the only branch of ``_fetch_new_requests`` that consumes
-    ``ADPRouter.exclude_retiring_requests``. So the index pool runs ahead
-    wherever the manager widens and the seat gate does not -- a non-ADP overlap
-    run, or any disaggregated run. That direction is permitted; what must never
-    happen is the reverse, which is what shipped as nvbug 6627795.
+    The seat gate is the narrower predicate: it has no disaggregation term,
+    because a request in KV transfer holds an index lease and no seat. So the
+    index pool runs ahead wherever the manager widens and the seat gate does not
+    -- a disaggregated run without the overlap term, e.g. one with the overlap
+    scheduler disabled. That direction is permitted; what must never happen is
+    the reverse, which is what shipped as nvbug 6627795.
     """
     max_batch_size = 8
     mapping = Mapping(
@@ -592,13 +593,11 @@ def test_validator_permits_an_index_pool_above_the_seat_pool():
     so the manager's ``is_disagg`` term deliberately has no seat-pool
     counterpart. A hybrid model suppresses the seat headroom -- SSM state is
     sized from ``max_batch_size`` -- while a V2 index pool behind it still
-    widens. And a non-ADP overlap run widens the index pool while the seat gate
-    stays shut, because only the attention-DP branch of ``_fetch_new_requests``
-    can admit a replacement before its predecessor releases a seat.
+    widens. And a disaggregated run with the overlap scheduler disabled widens the
+    index pool while the seat gate stays shut.
 
-    So the gap is the *ordinary* case on plain TP, not a rare exception, and that
-    is accepted rather than engineered away: admission is bounded independently
-    at ``max_batch_size * pp_size``, so a surplus lease is never extra
+    That gap is accepted rather than engineered away: admission is bounded
+    independently by the seat pool, so a surplus lease is never extra
     concurrency, and it is the cheap direction -- spare leases cost page-table
     rows, spare seats would cost sampler state, the eagerly allocated
     ``[seats, draft_len, vocab]`` draft probabilities and the pinned-host