diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2fa5a5d0607b..72f9d743b19a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index c9c05b39137f..fb09a827f424 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,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 = []