From 92b1b82f8afab0ac0378ba0abf64e5469016f457 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:13:49 -0700 Subject: [PATCH 1/4] [None][fix] Pad an empty attention-DP scheduled batch so the fleet can make forward progress `_can_queue` vetoes the forward pass on every attention-DP rank when any one rank's *scheduled* batch is empty: tp_batch_sizes = self.dist.tp_allgather(scheduled_batch.batch_size) can_queue = 0 not in tp_batch_sizes `_pad_attention_dp_dummy_request` is meant to make that unreachable, but it guarantees a different invariant -- every rank has at least one *active* request -- and it necessarily runs before `_schedule()`, so the capacity scheduler's verdict cannot inform it. A rank whose only active request does not fit the free KV cache therefore counts as one active request, receives no padding dummy, and schedules an empty batch. A rank-local cache shortage then stalls the whole fleet: peers never execute the context chunks they did schedule, because `_update_request_states` -- the only caller of `move_to_next_context_chunk()` -- runs under `if can_queue:`. No chunk completes, so no KV cache is released, so the starved rank stays starved. Nothing on that path raises or times out, so the stall is silent: clients see zero errors and zero progress. On a disaggregated context server, where the starved rank is waiting on its own in-flight cache transfers to release the blocks it needs, the state can persist indefinitely. Fix: after `_schedule()`, a rank left with an empty batch appends a generation dummy to the scheduled batch. Rank-local by design -- `_schedule()` performs a `tp_allgather` inside `_balance_adp_requests`, so re-scheduling on only the empty ranks would desynchronize that collective. A generation dummy is used rather than the context dummy the pre-schedule path would pick, because a rank that is empty precisely because it is short of KV cache must not be asked for `max_num_tokens` worth of it. Every allocation failure degrades to today's behaviour rather than propagating, since all ranks have yet to agree on `can_queue` and a rank-local raise would strand the peers in the collectives that follow. The pipeline-parallel loop does not call `_prepare_and_schedule_batch` and is unaffected. Verified on 4-node disaggregated attention-DP context servers on GB200 and GB300: the stall reproduced deterministically without the change, and with it the empty-batch condition still occurs and the fleet drives through every occurrence. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 126 +++++++++++++++ .../_torch/executor/test_py_executor.py | 147 +++++++++++++++++- 2 files changed, 272 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2fa5a5d0607b..a1b1d7831644 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3334,6 +3334,11 @@ 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 the batch after + # scheduling, so the V2 scheduler never grew its capacity. + # Reverting it would shrink a real allocation. + 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: @@ -3702,6 +3707,12 @@ def _prepare_and_schedule_batch(self): scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( ) + # Must run AFTER _schedule(): the condition it repairs -- an empty + # scheduled batch on one attention-DP rank, which vetoes the forward + # pass on all of them -- 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 +5976,121 @@ 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: + """Guarantee the invariant `_can_queue` actually depends on. + + `_can_queue` vetoes the forward pass on **every** attention-DP rank when + any rank's *scheduled* batch is empty:: + + tp_batch_sizes = self.dist.tp_allgather(scheduled_batch.batch_size) + can_queue = 0 not in tp_batch_sizes + + `_pad_attention_dp_dummy_request` is meant to prevent that, but it + guarantees a different invariant -- "every rank has at least one + *active* request" -- and it necessarily runs *before* `_schedule()`, so + the capacity scheduler's verdict cannot inform it. A rank whose only + active request does not fit the free KV cache therefore counts as one + active request, is given no dummy, and schedules an empty batch. + + A rank-local shortage then stalls the whole fleet: the empty rank vetoes + `can_queue` everywhere, so peers never execute the context chunks they + *did* schedule (`_update_request_states`, the only caller of + `move_to_next_context_chunk()`, runs under `if can_queue:`). No chunk + completes, so no KV cache is released, so the starved rank stays + starved. Nothing on that path raises or times out, so the stall is + silent: clients see zero errors and zero progress until enough cache is + freed elsewhere -- on a disaggregated context server, where the rank is + waiting on its own in-flight cache transfers, that may never happen. + + The padding is deliberately **rank-local**: `_schedule()` performs a + `tp_allgather` inside `_balance_adp_requests`, so re-scheduling on only + the empty ranks would desynchronize that collective. A *generation* + dummy is used rather than the context dummy the pre-schedule path would + pick, because a rank that is empty precisely because it is short of KV + cache must not be asked for `max_num_tokens` worth of it. + + Not reached by the pipeline-parallel loop, 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 + # `expected_num_active_requests` is capped at max_num_active_requests, + # and `_pad_attention_dp_dummy_request` asserts it bounds + # len(active_requests). Unlike that path -- which only pads a rank with + # no active requests at all -- this one adds to a rank that already has + # some, so it must not push the count past the cap. A rank at the cap + # has ample work to schedule and is not the starvation case repaired + # here. + 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: + # prepare_resource=True also takes a sequence slot, which can raise + # NoFreeSlotsError. Both failures must degrade to today's behaviour + # -- an empty batch this iteration -- rather than propagate: all + # ranks have yet to agree on can_queue, so a rank-local raise here + # would strand its 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 + # This dummy bypassed the capacity scheduler, so the V2 scheduler never + # grew its KV cache capacity and `_revert_gen_alloc` must leave it alone. + dummy_request.py_skip_gen_alloc_revert = True + self.active_requests.append(dummy_request) + scheduled_batch.generation_requests.append(dummy_request) + if self._enable_dsv4_adp_dummy_fixes: + # Let _finalize_adp_dummy_allocation roll the allocation back if the + # fleet still cannot queue (another rank may also be empty and have + # failed to pad). + 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: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index c9c05b39137f..5e735446edd2 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/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,144 @@ def test_pad_dummy_no_op_when_attention_dp_disabled(): assert stub.add_dummy_calls == [] +# --------------------------------------------------------------------------- +# Empty *scheduled* batch padding: _pad_attention_dp_dummy_request guarantees +# every rank has an active request, but _can_queue vetoes the fleet-wide +# forward pass on an empty *scheduled* batch. A rank whose only active request +# does not fit the free KV cache satisfies the former and violates the latter. +# --------------------------------------------------------------------------- +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): + # All ranks have yet to agree on can_queue, so a rank-local raise here + # would strand the peers in the collectives that follow. Degrade to + # today's behaviour instead. + 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 + + +def test_pad_empty_batch_dummy_is_excluded_from_gen_alloc_revert(): + # The dummy joins the batch after scheduling, so the V2 scheduler never + # grew its KV cache capacity; reverting it would shrink a real allocation. + 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. # --------------------------------------------------------------------------- From 0be7afb9c5e644641263f926f8f6de61bfe69ac7 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:35:42 -0700 Subject: [PATCH 2/4] [TRTLLM-15178][chore] Shorten the attention-DP empty-batch padding comments Address review feedback: the docstring and inline comments carried more detail than the code needs. Trim to the reasoning a reader must have -- why the padding runs after _schedule(), why it is rank-local, and why a generation dummy -- and leave the full analysis to the PR description. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 88 +++++++------------ .../_torch/executor/test_py_executor.py | 14 ++- 2 files changed, 35 insertions(+), 67 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a1b1d7831644..593dc9e7f349 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3334,9 +3334,8 @@ 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 the batch after - # scheduling, so the V2 scheduler never grew its capacity. - # Reverting it would shrink a real allocation. + # 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) @@ -3707,10 +3706,8 @@ def _prepare_and_schedule_batch(self): scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( ) - # Must run AFTER _schedule(): the condition it repairs -- an empty - # scheduled batch on one attention-DP rank, which vetoes the forward - # pass on all of them -- does not exist until the capacity scheduler has - # returned its verdict. + # 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: @@ -5979,40 +5976,22 @@ def _pad_attention_dp_dummy_request(self): @nvtx_range("_pad_empty_attention_dp_batch") def _pad_empty_attention_dp_batch( self, scheduled_batch: ScheduledRequests) -> None: - """Guarantee the invariant `_can_queue` actually depends on. - - `_can_queue` vetoes the forward pass on **every** attention-DP rank when - any rank's *scheduled* batch is empty:: - - tp_batch_sizes = self.dist.tp_allgather(scheduled_batch.batch_size) - can_queue = 0 not in tp_batch_sizes - - `_pad_attention_dp_dummy_request` is meant to prevent that, but it - guarantees a different invariant -- "every rank has at least one - *active* request" -- and it necessarily runs *before* `_schedule()`, so - the capacity scheduler's verdict cannot inform it. A rank whose only - active request does not fit the free KV cache therefore counts as one - active request, is given no dummy, and schedules an empty batch. - - A rank-local shortage then stalls the whole fleet: the empty rank vetoes - `can_queue` everywhere, so peers never execute the context chunks they - *did* schedule (`_update_request_states`, the only caller of - `move_to_next_context_chunk()`, runs under `if can_queue:`). No chunk - completes, so no KV cache is released, so the starved rank stays - starved. Nothing on that path raises or times out, so the stall is - silent: clients see zero errors and zero progress until enough cache is - freed elsewhere -- on a disaggregated context server, where the rank is - waiting on its own in-flight cache transfers, that may never happen. - - The padding is deliberately **rank-local**: `_schedule()` performs a - `tp_allgather` inside `_balance_adp_requests`, so re-scheduling on only - the empty ranks would desynchronize that collective. A *generation* - dummy is used rather than the context dummy the pre-schedule path would - pick, because a rank that is empty precisely because it is short of KV - cache must not be asked for `max_num_tokens` worth of it. - - Not reached by the pipeline-parallel loop, which does not call - `_prepare_and_schedule_batch`. + """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 @@ -6020,13 +5999,10 @@ def _pad_empty_attention_dp_batch( return if not self.active_requests or self.expected_num_active_requests <= 0: return - # `expected_num_active_requests` is capped at max_num_active_requests, - # and `_pad_attention_dp_dummy_request` asserts it bounds - # len(active_requests). Unlike that path -- which only pads a rank with - # no active requests at all -- this one adds to a rank that already has - # some, so it must not push the count past the cap. A rank at the cap - # has ample work to schedule and is not the starvation case repaired - # here. + # 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 @@ -6048,11 +6024,9 @@ def _pad_empty_attention_dp_batch( dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID] try: - # prepare_resource=True also takes a sequence slot, which can raise - # NoFreeSlotsError. Both failures must degrade to today's behaviour - # -- an empty batch this iteration -- rather than propagate: all - # ranks have yet to agree on can_queue, so a rank-local raise here - # would strand its peers in the collectives that follow. + # 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, @@ -6080,15 +6054,13 @@ def _pad_empty_attention_dp_batch( return dummy_request.is_attention_dp_dummy = True - # This dummy bypassed the capacity scheduler, so the V2 scheduler never - # grew its KV cache capacity and `_revert_gen_alloc` must leave it alone. + # 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) if self._enable_dsv4_adp_dummy_fixes: - # Let _finalize_adp_dummy_allocation roll the allocation back if the - # fleet still cannot queue (another rank may also be empty and have - # failed to pad). + # Let `_finalize_adp_dummy_allocation` roll this back if the fleet + # still cannot queue. self._pending_adp_dummy_request = dummy_request @nvtx_range("_prepare_disagg_gen_init") diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 5e735446edd2..9dd3b46ea669 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1732,10 +1732,9 @@ def test_pad_dummy_no_op_when_attention_dp_disabled(): # --------------------------------------------------------------------------- -# Empty *scheduled* batch padding: _pad_attention_dp_dummy_request guarantees -# every rank has an active request, but _can_queue vetoes the fleet-wide -# forward pass on an empty *scheduled* batch. A rank whose only active request -# does not fit the free KV cache satisfies the former and violates the latter. +# 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 ( @@ -1839,9 +1838,7 @@ def test_pad_empty_batch_degrades_when_kv_cache_cannot_afford_dummy(): @pytest.mark.parametrize("error", [OutOfPagesError("no pages"), NoFreeSlotsError("no slots")]) def test_pad_empty_batch_degrades_on_allocation_error(error): - # All ranks have yet to agree on can_queue, so a rank-local raise here - # would strand the peers in the collectives that follow. Degrade to - # today's behaviour instead. + # 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 @@ -1852,8 +1849,7 @@ def test_pad_empty_batch_degrades_on_allocation_error(error): def test_pad_empty_batch_dummy_is_excluded_from_gen_alloc_revert(): - # The dummy joins the batch after scheduling, so the V2 scheduler never - # grew its KV cache capacity; reverting it would shrink a real allocation. + # The dummy joins after scheduling, so its capacity was never grown. stub, scheduled_batch = _unfittable_rank() _run_pad_empty(stub, scheduled_batch) From c5442d95b1f04167dc17db81e009a3b5f16b95e5 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:20:34 -0700 Subject: [PATCH 3/4] [TRTLLM-15178][fix] Roll back the empty-batch padding dummy for every model The post-schedule dummy was only registered for rollback under _enable_dsv4_adp_dummy_fixes, so on other models a fleet-wide can_queue=False left it parked with its KV cache allocated until the next successful forward. _pending_adp_dummy_request is written in exactly two places -- the DeepSeek-V4 branch of _pad_attention_dp_dummy_request and the new padding path -- and every other flow leaves it None, where _finalize_adp_dummy_allocation already returns early. Dropping the model gate from that method is therefore a no-op for the existing paths and closes the gap for the new one. Rollback is needed because the usual teardown, which terminates every attention-DP dummy in _handle_responses, runs only under 'if can_queue:'. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 14 +++---- .../_torch/executor/test_py_executor.py | 39 +++++++++++++++++++ 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 593dc9e7f349..72f9d743b19a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3347,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: @@ -6058,10 +6058,10 @@ def _pad_empty_attention_dp_batch( dummy_request.py_skip_gen_alloc_revert = True self.active_requests.append(dummy_request) scheduled_batch.generation_requests.append(dummy_request) - if self._enable_dsv4_adp_dummy_fixes: - # Let `_finalize_adp_dummy_allocation` roll this back if the fleet - # still cannot queue. - self._pending_adp_dummy_request = 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): diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 9dd3b46ea669..e25cd9cb46c6 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1848,6 +1848,45 @@ def test_pad_empty_batch_degrades_on_allocation_error(error): 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() From 0723a36ca4889d557bddba8fb3aaf7fafcea6fdc Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Thu, 6 Aug 2026 21:59:39 -0700 Subject: [PATCH 4/4] [TRTLLM-15178][fix] Set enable_attention_dp on the one-model-MTP test stub `_pad_empty_attention_dp_batch` is called unconditionally from `_prepare_and_schedule_batch` and reads `self.enable_attention_dp` on its first line. `TestOneModelMTPDraftTokenScheduling` builds its executor with `object.__new__(PyExecutor)` and never assigned that attribute, so the new call raised: AttributeError: 'PyExecutor' object has no attribute 'enable_attention_dp' failing CPU-Generic-arm-1 in `unittest/_torch/executor`. Set it to False on the stub, matching `test_benchmark_disagg.py`'s executor stub, so the new padding path returns immediately for a test that does not exercise attention DP. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- tests/unittest/_torch/executor/test_py_executor.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index e25cd9cb46c6..fb09a827f424 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -2139,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 @@ -2147,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 = []