[TRTLLM-15178][fix] Pad an empty attention-DP scheduled batch so the fleet can make forward progress

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -3334,6 +3334,10 @@ def _revert_gen_alloc(self, scheduled_batch):
         """
         if self._is_kv_manager_v2:
             for req in scheduled_batch.generation_requests:
+                # The empty-batch padding dummy joins after scheduling, so its
+                # capacity was never grown; reverting would shrink it.
+                if getattr(req, "py_skip_gen_alloc_revert", False):
+                    continue
                 self.kv_cache_manager.revert_allocate_generation(req)
 
     def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None:
@@ -3343,10 +3347,10 @@ def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None:
         decision. If one rank cannot allocate its dummy, peers that succeeded
         must release theirs before retrying or the fixed dummy request ID leaks
         cache resources on every skipped iteration.
-        """
-        if not self._enable_dsv4_adp_dummy_fixes:
-            return
 
+        Only paths that register a pending dummy are affected; the others leave
+        ``_pending_adp_dummy_request`` unset and no-op here.
+        """
         dummy_request = self._pending_adp_dummy_request
         self._pending_adp_dummy_request = None
         if dummy_request is None or can_queue:
@@ -3702,6 +3706,10 @@ def _prepare_and_schedule_batch(self):
         scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(
         )
 
+        # Must run after _schedule(): the empty scheduled batch it repairs does
+        # not exist until the capacity scheduler has returned its verdict.
+        self._pad_empty_attention_dp_batch(scheduled_batch)
+
         if self.drafter is not None and not self.use_spec_decode:
             for request in scheduled_batch.all_requests():
                 request.py_disable_speculative_decoding = True
@@ -5965,6 +5973,96 @@ def _pad_attention_dp_dummy_request(self):
         self.active_requests.append(dummy_request)
         self._pending_adp_dummy_request = dummy_request
 
+    @nvtx_range("_pad_empty_attention_dp_batch")
+    def _pad_empty_attention_dp_batch(
+            self, scheduled_batch: ScheduledRequests) -> None:
+        """Pad an empty scheduled batch so attention DP can make progress.
+
+        `_can_queue` vetoes the forward pass on every rank when any rank's
+        *scheduled* batch is empty. `_pad_attention_dp_dummy_request` cannot
+        prevent that: it guarantees every rank has an *active* request, and runs
+        before `_schedule()`, so a rank whose only active request does not fit
+        the free KV cache is left unpadded and schedules an empty batch. Peers
+        then never advance their context chunks (`_update_request_states` runs
+        under `if can_queue:`), so no KV cache is freed and the rank stays
+        starved -- silently, since nothing on that path raises or times out.
+
+        Padding is rank-local because `_schedule()` allgathers inside
+        `_balance_adp_requests`. A generation dummy is used rather than the
+        context dummy the pre-schedule path would pick, since this rank is empty
+        precisely because it is short of KV cache. Not reached under pipeline
+        parallelism, which does not call `_prepare_and_schedule_batch`.
+        """
+        if not self.enable_attention_dp:
+            return
+        if scheduled_batch is None or scheduled_batch.batch_size != 0:
+            return
+        if not self.active_requests or self.expected_num_active_requests <= 0:
+            return
+        # Unlike the pre-schedule path, this one pads a rank that already holds
+        # active requests, so it must respect the cap that
+        # `_pad_attention_dp_dummy_request` asserts against. A rank at the cap
+        # is not starved anyway.
+        if len(self.active_requests) >= self.max_num_active_requests:
+            return
+        # The fill gate suppresses forwards on purpose; dummies added then are
+        # never terminated (termination follows a forward pass).
+        if self._should_skip_dummy_for_benchmark_disagg(
+                self._count_schedulable_active_requests()):
+            return
+        # ATTENTION_DP_DUMMY_REQUEST_ID is a singleton resource.
+        if any(request.py_request_id == ATTENTION_DP_DUMMY_REQUEST_ID
+               for request in self.active_requests):
+            return
+        if not self._has_adp_dummy_kv_capacity(None):
+            logger.warning_once(
+                "attention-DP rank scheduled an empty batch and cannot afford "
+                "even a padding dummy; the forward pass is vetoed fleet-wide "
+                "this iteration",
+                key="attention_dp_empty_batch_no_pad_capacity")
+            return
+
+        dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID]
+        try:
+            # Degrade to an empty batch rather than propagate: the ranks have
+            # yet to agree on can_queue, so a rank-local raise would strand the
+            # peers in the collectives that follow.
+            dummy_requests = self.kv_cache_manager.add_dummy_requests(
+                request_ids=dummy_request_ids,
+                token_nums=None,
+                is_gen=True,
+                prepare_resource=True,
+                max_num_draft_tokens=self.max_total_draft_tokens,
+            )
+        except (OutOfPagesError, NoFreeSlotsError):
+            dummy_requests = None
+        if not dummy_requests:
+            logger.warning_once(
+                "Could not allocate the attention-DP empty-batch padding "
+                "dummy; retrying next iteration",
+                key="attention_dp_empty_batch_pad_alloc_failed")
+            return
+
+        dummy_request = dummy_requests[0]
+        spec_resource_manager = self.resource_manager.get_resource_manager(
+            ResourceManagerType.SPEC_RESOURCE_MANAGER)
+        if spec_resource_manager is not None:
+            try:
+                spec_resource_manager.add_dummy_requests(dummy_request_ids)
+            except NoFreeSlotsError:
+                self.kv_cache_manager.free_resources(dummy_request)
+                return
+
+        dummy_request.is_attention_dp_dummy = True
+        # Never grown by the V2 scheduler, so `_revert_gen_alloc` must skip it.
+        dummy_request.py_skip_gen_alloc_revert = True
+        self.active_requests.append(dummy_request)
+        scheduled_batch.generation_requests.append(dummy_request)
+        # Let `_finalize_adp_dummy_allocation` roll this back if the fleet still
+        # cannot queue: no forward runs then, so the usual dummy teardown in
+        # `_handle_responses` is not reached this iteration.
+        self._pending_adp_dummy_request = dummy_request
+
     @nvtx_range("_prepare_disagg_gen_init")
     def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests):
         if fitting_disagg_gen_init_requests:

tests/unittest/_torch/executor/test_py_executor.py

@@ -26,13 +26,18 @@
     RequestQueueItem,
 )
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig
-from tensorrt_llm._torch.pyexecutor.py_executor import DisaggTransferAdmissionController, PyExecutor
+from tensorrt_llm._torch.pyexecutor.py_executor import (
+    ATTENTION_DP_DUMMY_REQUEST_ID,
+    DisaggTransferAdmissionController,
+    PyExecutor,
+)
 from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType
 from tensorrt_llm._torch.pyexecutor.scheduler import (
     FCFSWaitingQueue,
     ScheduledRequests,
     SerializableSchedulerOutput,
 )
+from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError
 
 pytestmark = pytest.mark.cpu_only
 
@@ -1290,6 +1295,7 @@ def _make_adp_request(
     )
     req.is_dummy_request = is_dummy_request
     req.is_attention_dp_dummy = False
+    req.py_skip_gen_alloc_revert = False
     req.llm_request_type = llm_request_type
     req.py_seq_slot = None
     return req
@@ -1316,6 +1322,7 @@ def __init__(
         self.num_fetch_requests = 0
         self.active_requests = []
         self.expected_num_active_requests = 1
+        self.max_num_active_requests = 8
         self.max_total_draft_tokens = 0
         self.max_num_tokens = max_num_tokens
         self._adp_dummy_is_gen = True
@@ -1724,6 +1731,179 @@ def test_pad_dummy_no_op_when_attention_dp_disabled():
     assert stub.add_dummy_calls == []
 
 
+# ---------------------------------------------------------------------------
+# Empty *scheduled* batch padding: a rank whose only active request does not fit
+# the free KV cache is skipped by _pad_attention_dp_dummy_request, yet its empty
+# scheduled batch vetoes the fleet-wide forward pass in _can_queue.
+# ---------------------------------------------------------------------------
+def _run_pad_empty(stub, scheduled_batch):
+    for helper in (
+        "_count_schedulable_active_requests",
+        "_has_adp_dummy_kv_capacity",
+        "_should_skip_dummy_for_benchmark_disagg",
+    ):
+        setattr(stub, helper, types.MethodType(getattr(PyExecutor, helper), stub))
+    PyExecutor._pad_empty_attention_dp_batch(stub, scheduled_batch)
+
+
+def _unfittable_rank(**kwargs):
+    """A rank holding one active request the capacity scheduler could not fit."""
+    stub = _StubADPExecutor(**kwargs)
+    stub.active_requests = [_make_adp_request(_STATE_GENERATION_IN_PROGRESS)]
+    stub.expected_num_active_requests = 1
+    return stub, ScheduledRequests()
+
+
+def test_pad_empty_batch_adds_generation_dummy_to_scheduled_batch():
+    stub, scheduled_batch = _unfittable_rank()
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    assert scheduled_batch.batch_size == 1
+    assert len(scheduled_batch.generation_requests) == 1
+    dummy = scheduled_batch.generation_requests[0]
+    assert dummy.is_attention_dp_dummy is True
+    assert dummy in stub.active_requests
+    # A generation dummy, not the context dummy the pre-schedule path would
+    # pick: this rank is empty precisely because it is short of KV cache.
+    assert stub.add_dummy_calls[0]["is_gen"] is True
+    assert stub.add_dummy_calls[0]["token_nums"] is None
+
+
+def test_pad_empty_batch_no_op_when_batch_is_not_empty():
+    stub, scheduled_batch = _unfittable_rank()
+    scheduled_batch.context_requests_chunking = [
+        _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=7)
+    ]
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    assert stub.add_dummy_calls == []
+    assert scheduled_batch.generation_requests == []
+
+
+def test_pad_empty_batch_no_op_when_attention_dp_disabled():
+    stub, scheduled_batch = _unfittable_rank(enable_attention_dp=False)
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    assert stub.add_dummy_calls == []
+    assert scheduled_batch.batch_size == 0
+
+
+def test_pad_empty_batch_no_op_when_rank_has_no_active_requests():
+    # The pre-schedule path already covers this case.
+    stub, scheduled_batch = _unfittable_rank()
+    stub.active_requests = []
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    assert stub.add_dummy_calls == []
+
+
+def test_pad_empty_batch_respects_max_num_active_requests():
+    # expected_num_active_requests is capped at max_num_active_requests and
+    # _pad_attention_dp_dummy_request asserts it bounds len(active_requests),
+    # so padding a rank already at the cap would trip that assert next
+    # iteration.
+    stub, scheduled_batch = _unfittable_rank()
+    stub.max_num_active_requests = 1
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    assert stub.add_dummy_calls == []
+    assert scheduled_batch.batch_size == 0
+
+
+def test_pad_empty_batch_skips_when_dummy_already_active():
+    stub, scheduled_batch = _unfittable_rank()
+    stub.active_requests.append(
+        _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=ATTENTION_DP_DUMMY_REQUEST_ID)
+    )
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    assert stub.add_dummy_calls == []
+
+
+def test_pad_empty_batch_degrades_when_kv_cache_cannot_afford_dummy():
+    stub, scheduled_batch = _unfittable_rank()
+    stub.kv_cache_manager.get_num_available_tokens.return_value = 0
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    stub.kv_cache_manager.add_dummy_requests.assert_not_called()
+    assert scheduled_batch.batch_size == 0
+
+
+@pytest.mark.parametrize("error", [OutOfPagesError("no pages"), NoFreeSlotsError("no slots")])
+def test_pad_empty_batch_degrades_on_allocation_error(error):
+    # A rank-local raise would strand the peers in the collectives that follow.
+    stub, scheduled_batch = _unfittable_rank()
+    stub.kv_cache_manager.add_dummy_requests.side_effect = error
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    assert scheduled_batch.batch_size == 0
+    assert len(stub.active_requests) == 1
+
+
+@pytest.mark.parametrize("enable_dsv4_adp_dummy_fixes", [False, True])
+def test_pad_empty_batch_dummy_rolled_back_when_fleet_still_cannot_queue(
+    enable_dsv4_adp_dummy_fixes,
+):
+    # can_queue=False skips the forward, so the usual dummy teardown in
+    # _handle_responses is not reached. Rollback must not depend on the model.
+    stub, scheduled_batch = _unfittable_rank(
+        enable_dsv4_adp_dummy_fixes=enable_dsv4_adp_dummy_fixes
+    )
+    active_request = stub.active_requests[0]
+    spec_resource_manager = Mock()
+    stub.resource_manager.get_resource_manager.return_value = spec_resource_manager
+
+    _run_pad_empty(stub, scheduled_batch)
+    dummy = stub._pending_adp_dummy_request
+    assert dummy is not None
+
+    PyExecutor._finalize_adp_dummy_allocation(stub, False)
+
+    assert stub.active_requests == [active_request]
+    assert stub._pending_adp_dummy_request is None
+    spec_resource_manager.free_resources.assert_called_once_with(dummy)
+    stub.kv_cache_manager.free_resources.assert_called_once_with(dummy)
+
+
+def test_pad_empty_batch_dummy_kept_when_fleet_can_queue():
+    # Committed dummies are terminated after the forward by _handle_responses.
+    stub, scheduled_batch = _unfittable_rank(enable_dsv4_adp_dummy_fixes=False)
+
+    _run_pad_empty(stub, scheduled_batch)
+    dummy = stub._pending_adp_dummy_request
+
+    PyExecutor._finalize_adp_dummy_allocation(stub, True)
+
+    assert dummy in stub.active_requests
+    assert stub._pending_adp_dummy_request is None
+    stub.kv_cache_manager.free_resources.assert_not_called()
+
+
+def test_pad_empty_batch_dummy_is_excluded_from_gen_alloc_revert():
+    # The dummy joins after scheduling, so its capacity was never grown.
+    stub, scheduled_batch = _unfittable_rank()
+
+    _run_pad_empty(stub, scheduled_batch)
+
+    # A request the scheduler did grow, for contrast.
+    real_gen_request = _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=9)
+    scheduled_batch.generation_requests.append(real_gen_request)
+
+    stub._is_kv_manager_v2 = True
+    PyExecutor._revert_gen_alloc(stub, scheduled_batch)
+
+    reverted = [c.args[0] for c in stub.kv_cache_manager.revert_allocate_generation.call_args_list]
+    assert reverted == [real_gen_request]
+
+
 # ---------------------------------------------------------------------------
 # ADP-safe disagg cache error handling (#13900): all TP ranks enter _handle_errors together.
 # ---------------------------------------------------------------------------
@@ -1959,6 +2139,8 @@ def _make_one_model_mtp_executor(cls, active_requests):
         model_engine.is_spec_decode is True, so _prepare_and_schedule_batch
         takes the elif draft-token normalization branch. kv_cache_transceiver
         is None to keep the test hermetic (skips the disagg blocks).
+        enable_attention_dp is False so _pad_empty_attention_dp_batch returns
+        immediately instead of padding a batch this test does not exercise.
         """
         ex = object.__new__(PyExecutor)
         ex.drafter = None
@@ -1967,6 +2149,7 @@ def _make_one_model_mtp_executor(cls, active_requests):
         ex.kv_cache_transceiver = None
         ex.is_shutdown = False
         ex.enable_iter_perf_stats = False
+        ex.enable_attention_dp = False
         ex.active_requests = active_requests
         ex.waiting_queue = []