From 52d285cbe7a32bb53a5148fc347617d5d40c3434 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Wed, 13 Aug 2025 16:18:03 -0700 Subject: [PATCH] Fix sequence slot allocation for attention DP Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py b/tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py index c43c9726412c..db4525dc79a4 100644 --- a/tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py @@ -16,7 +16,11 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int: def prepare_resources(self, scheduled_batch: ScheduledRequests) -> None: for llm_req in scheduled_batch.all_requests(): - if llm_req.seq_slot is None or llm_req.is_disagg_generation_transmission_complete: + # If the request is in the disagg generation init state, + # we don't need to add a slot for it. + if llm_req.seq_slot is None and not ( + llm_req.is_disagg_generation_init_state + or llm_req.is_disagg_generation_transmission_in_progress): llm_req.seq_slot = self.slot_manager.add_slot( llm_req.request_id) llm_req.py_seq_slot = llm_req.seq_slot