[None][fix] Fix KV event consumption (#6346)

Signed-off-by: jthomson04 <jwillthomson19@gmail.com>

Author: jthomson04

cpp/tensorrt_llm/batch_manager/kvCacheEventManager.cpp

@@ -152,6 +152,12 @@ std::deque<tle::KVCacheEvent> KVCacheEventManager::getEvents(std::optional<std::
 void KVCacheEventManager::flush()
 {
     auto eventQueue = std::exchange(mEventQueue, {});
+
+    if (eventQueue.empty())
+    {
+        return;
+    }
+
     std::unique_lock<std::mutex> lck(mPendingEventsMutex);
     mPendingEvents.push_back(std::move(eventQueue));
     mPendingEmptyCV.notify_one();

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -17,7 +17,8 @@
 from tensorrt_llm.runtime import ModelConfig as ModelConfigPython
 from tensorrt_llm.sampling_params import SamplingParams
 
-from ..._utils import binding_to_str_dtype, get_size_in_bytes, nvtx_range
+from ..._utils import (binding_to_str_dtype, get_size_in_bytes, mpi_rank,
+                       nvtx_range)
 from ...logger import logger
 from ...mapping import CpType, Mapping
 from .kv_cache_connector import KvCacheConnectorManager
@@ -338,6 +339,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
             'copy_on_partial_reuse': kv_cache_config.copy_on_partial_reuse,
             'kv_connector_manager': self.kv_connector_manager,
         }
+
         if self.event_buffer_max_size > 0:
             if mapping.enable_attention_dp:
                 kwargs['event_manager'] = KVCacheEventManagerCpp(
@@ -347,7 +349,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
                     attention_dp_events_gather_period_ms=self.
                     attention_dp_events_gather_period_ms,
                 )
-            else:
+            elif mpi_rank() == 0:
                 kwargs['event_manager'] = KVCacheEventManagerCpp(
                     max_kv_event_entries=self.event_buffer_max_size)
 

tensorrt_llm/executor/proxy.py

@@ -205,11 +205,12 @@ def process_res(res):
 
         return True  # success
 
-    def _iteration_result_task(self, queue: Union[FusedIpcQueue,
-                                                  IntraProcessQueue],
-                               result_singleton: IterationResult) -> bool:
-        # iteration result is not urgent, so we can sleep a bit
-        time.sleep(0.2)
+    def _iteration_result_task(self,
+                               queue: Union[FusedIpcQueue, IntraProcessQueue],
+                               result_singleton: IterationResult,
+                               urgent: bool = False) -> bool:
+        if not urgent:
+            time.sleep(0.2)
 
         try:
             data = queue.get()
@@ -267,7 +268,8 @@ def dispatch_stats_task(self) -> bool:
 
     def dispatch_kv_cache_events_task(self) -> bool:
         return self._iteration_result_task(self.kv_cache_events_queue,
-                                           self._iter_kv_events_result)
+                                           self._iter_kv_events_result,
+                                           urgent=True)
 
     def _start_dispatch_threads(self):
         if self.dispatch_result_thread is None: