feat: Run PyExecutor's inference flow to estimate max_num_tokens for kv_cache_manager (#3092)

Signed-off-by: Hui Gao <[email protected]>

Author: HuiGao-NV

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -118,7 +118,7 @@ class CacheTransceiver : public BaseCacheTransceiver
     std::unique_ptr<DataRequester> mDataRequester;
     std::map<LlmRequest*, std::future<void>> mResponderFutures;
     std::vector<std::pair<LlmRequest*, std::future<void>>> mRequesterFutures;
-    mpi::MpiComm const *mMpiGroupComm{}, *mMpiWorldComm{};
+    mpi::MpiComm const *mMpiGroupComm{nullptr}, *mMpiWorldComm{nullptr};
     std::shared_ptr<mpi::MpiComm> mMpiGroupTensorParaComm, mMpiGroupPipeParaComm, mMpiGroupDataComm,
         mMpiGroupTPInDPComm;
     executor::kv_cache::CommState const* mCommState;

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -64,6 +64,9 @@ def __init__(
         self._output = None
         self._graph = None
 
+    def __del__(self):
+        self._graph.reset()
+
     def capture(
         self,
         forward_fn: Callable[[Dict[str, Any]], torch.Tensor],

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -141,7 +141,8 @@ def executor_request_to_llm_request(req_id: int,
         guided_decoding_params=executor_request.guided_decoding_params,
         encoder_input_tokens=None,
         return_encoder_output=False,
-        client_id=executor_request.client_id,
+        client_id=executor_request.client_id
+        if executor_request.client_id is not None else req_id,
         priority=0.5,
         llm_request_type=llm_request_type,
         context_phase_params=executor_request.context_phase_params)

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -824,6 +824,14 @@ def _init_model_capacity(self):
         self._init_max_seq_len()
         self._init_max_num_tokens()
 
+    def _release_cuda_graphs(self):
+        for _, graph in self._cuda_graphs.items():
+            del graph
+        self._cuda_graphs.clear()
+        torch.cuda.empty_cache()
+        del self._cuda_graph_mem_pool
+        self._cuda_graph_mem_pool = None
+
     def get_max_num_sequences(self) -> int:
         """
         Return the maximum number of sequences that the model supports. PyExecutor need this to compute max_num_active_requests

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -142,7 +142,8 @@ def __init__(self,
                  max_batch_size: int = 8,
                  max_draft_tokens: int = 0,
                  kv_cache_transceiver: KvCacheTransceiver = None,
-                 draft_model_engine: Optional[ModelEngine] = None):
+                 draft_model_engine: Optional[ModelEngine] = None,
+                 start_worker: bool = True):
         super(PyExecutor, self).__init__()
         self.device_id = torch.cuda.current_device()
         self.global_rank = global_mpi_rank()
@@ -223,20 +224,34 @@ def __init__(self,
         self.stats = []
         self.start_times = {}
         self.new_active_requests_queue_latency_ms = 0
+        self.gather_all_responses = False
 
         self.kv_cache_transceiver = kv_cache_transceiver
         if self.dist.pp_size > 1:
-            event_loop = self._executor_loop_pp_overlap if enable_overlap_scheduler else self._executor_loop_pp
+            self.event_loop = self._executor_loop_pp_overlap if enable_overlap_scheduler else self._executor_loop_pp
         else:
-            event_loop = self._executor_loop_overlap if enable_overlap_scheduler else self._executor_loop
+            self.event_loop = self._executor_loop_overlap if enable_overlap_scheduler else self._executor_loop
 
-        if self.draft_model_engine is not None and event_loop.__name__ != self._executor_loop.__name__:
+        if self.draft_model_engine is not None and self.event_loop.__name__ != self._executor_loop.__name__:
             raise NotImplementedError(
                 "Drafting is not supported for selected executor loop. "
                 "Please disable disagg/pipeline parallelism/overlap scheduler.")
 
-        self.worker_thread = threading.Thread(target=event_loop, daemon=True)
-        self.worker_thread.start()
+        self.worker_started = False
+        self.worker_lock = threading.Lock()
+        if start_worker:
+            self.start_worker()
+
+    def start_worker(self):
+        self.worker_lock.acquire()
+        try:
+            if self.worker_started == False:
+                self.worker_thread = threading.Thread(target=self.event_loop,
+                                                      daemon=True)
+                self.worker_thread.start()
+                self.worker_started = True
+        finally:
+            self.worker_lock.release()
 
     def __enter__(self):
         return self
@@ -306,6 +321,7 @@ def shutdown(self):
             self.enqueue_lock.release()
         self.shutdown_event.wait()
         self.worker_thread.join()
+        self.worker_started = False
         for manager in self.resource_manager.resource_managers.values():
             if manager:
                 manager.shutdown()
@@ -372,6 +388,9 @@ def enqueue_request(self,
             self.enqueue_lock.release()
         return req_id
 
+    def set_gather_responses(self, gather_all_responses):
+        self.gather_all_responses = gather_all_responses
+
     @contextmanager
     def _profiler(self):
         it = -1
@@ -474,6 +493,9 @@ def _process_iter_stats(self, finished_requests, scheduled_batch,
                             iter_start_time, iter_stats):
         iter_end_time = time.time()
         iter_latency_ms = iter_end_time - iter_start_time
+        if iter_stats is None:
+            return
+
         self._append_iter_stats(
             self._update_iter_stats(iter_stats, iter_latency_ms,
                                     len(finished_requests), scheduled_batch))
@@ -725,7 +747,6 @@ def _executor_loop(self):
                     new_requests) or got_finish_signal
                 if got_finish_signal and len(self.active_requests) == 0:
                     break
-
                 if self.enable_iter_perf_stats:
                     iter_stats = self._get_init_iter_stats(
                         len(new_requests),
@@ -1097,6 +1118,7 @@ def _fetch_new_requests(self):
              self.dist.tp_size - 1) // self.dist.tp_size,
             max(self.all_ranks_num_active_requests),
         )
+
         self.has_context_request = False
         new_requests_cur_rank = []
         if new_requests != [] and new_requests[
@@ -1783,21 +1805,26 @@ def _handle_cancelled_requests(self):
 
     @nvtx_range("_enqueue_responses")
     def _enqueue_responses(self, responses: Dict[int, ExecutorResponse]):
-        if 0 not in self.dist.mapping.tp_group:
+        if 0 not in self.dist.mapping.tp_group and not self.gather_all_responses:
             return
 
         logger.debug(
             f'before gather, rank = {self.dist.rank}, responses = {responses}')
         if self.enable_attention_dp:
-            responses_list = self.dist.tp_gather(responses)
-            if self.dist.rank == 0:
+            if not self.gather_all_responses:
+                responses_list = self.dist.tp_gather(responses)
+            else:
+                responses_list = self.dist.allgather(responses)
+            if self.dist.rank == 0 or self.gather_all_responses:
                 gather_responses = {}
-                for resp in responses_list:
-                    gather_responses.update(resp)
-                responses = gather_responses
+                if responses_list is not None:
+                    for resp in responses_list:
+                        gather_responses.update(resp)
+                    responses = gather_responses
         logger.debug(
             f'after gather, rank = {self.dist.rank}, responses = {responses}')
-        if self.dist.rank == 0:
+
+        if self.dist.rank == 0 or self.gather_all_responses:
             with self.response_cv:
                 for req_id, resp in responses.items():
                     if req_id in self.responses.keys():
@@ -1817,6 +1844,7 @@ def _handle_responses(self):
         for request in self.active_requests:
             req_id = request.py_request_id
             #no responses for dummy request, and finish it
+
             if request.is_dummy == True:
                 requests_to_terminate.append(request)
                 continue
@@ -1825,7 +1853,6 @@ def _handle_responses(self):
             request.decoding_iter = request.py_decoding_iter
             response = request.create_response(False, self.dist.rank)
             request_done = False
-
             if response:
                 request_done = response.result.is_final
                 new_responses.update({req_id: response})
@@ -1840,7 +1867,6 @@ def _handle_responses(self):
         self._enqueue_responses(new_responses)
         for request in requests_to_terminate:
             self._terminate_request(request)
-
         return requests_to_terminate
 
     @nvtx_range("_terminate_ctx_finished_requests")