diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bc8872307d87..8f1a983db9d2 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3018,8 +3018,16 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): resource_mgr_type].prepare_resources( disagg_gen_init_to_prepare) - # Trigger KV cache exchange for new disagg_gen_init_requests - self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) + # Use the filtered batch after prepare_resources, not the original + # list. prepare_resources may skip add_sequence for some requests + # (e.g. reuse budget exhaustion) and remove them from the batch via + # reset_context_requests. Passing skipped requests to KV transfer + # causes unordered_map::at in getSequence because the sequence was + # never registered. Skipped requests remain in active_requests + # with DISAGG_GENERATION_INIT state and will be rescheduled. + accepted_requests = list( + disagg_gen_init_to_prepare.context_requests) + self._recv_disagg_gen_cache(accepted_requests) @nvtx_range("_prepare_disagg_gen_transmission_complete") def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): @@ -3642,14 +3650,13 @@ def _handle_responses(self): f"Request {request.py_request_id} has no avg_decoded_tokens_per_iter" ) - # If partial reuse is enabled, and the KV cache manager is not VSWA, and the PP size is 1, - # then we need to terminate the request. TODO: Remove this once disagg support from KVCache reuse - # path is fixed. - if self.enable_partial_reuse_for_disagg and not self.kv_cache_manager.is_vswa and self.dist.pp_size == 1: + # Never terminate a request while KV cache transfer is + # still in progress — the GEN worker may still be receiving + # data from the pinned blocks. The request will be + # terminated by _end_transfer_and_maybe_terminate once the + # transfer completes (or times out). + if not request.is_disagg_context_transmission_state: requests_to_terminate.append(request) - else: - if not request.is_disagg_context_transmission_state: - requests_to_terminate.append(request) else: new_active_requests.append(request)