[https://nvbugs/6293536][fix] order KV cache transfers with overlap scheduler

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -1183,6 +1183,9 @@ class WindowBlockManager
     //! \brief Sync internal streams used by transfer manager with buffer manager stream
     void syncTransferManagerWithBufferManager();
 
+    //! \brief Whether any host-to-device or device-to-host transfer has been scheduled and not yet synchronized
+    [[nodiscard]] bool hasPendingHostTransfers() const;
+
     //! \brief Perform per-request bookkeeping
     void refreshBlocks();
 
@@ -1825,6 +1828,9 @@ class BlockManager
     //! \brief Sync internal streams used by transfer manager with buffer manager stream
     void syncTransferManagerWithBufferManager();
 
+    //! \brief Whether any host-to-device or device-to-host transfer has been scheduled and not yet synchronized
+    [[nodiscard]] bool hasPendingHostTransfers() const;
+
     //! \brief Perform per-request bookkeeping
     void refreshBlocks();
 
@@ -2091,6 +2097,7 @@ class BaseKVCacheManager
     [[nodiscard]] virtual bool isPoolLayerFirst(SizeType32 layer_idx) const = 0;
 
     virtual void syncTransferManagerWithBufferManager() = 0;
+    [[nodiscard]] virtual bool hasPendingHostTransfers() const = 0;
     virtual void refreshBlocks() = 0;
     virtual void flushIterationEvents() = 0;
     virtual void resetReuseState() = 0;
@@ -2579,6 +2586,11 @@ class KVCacheManager : public BaseKVCacheManager
         mBlockManager.syncTransferManagerWithBufferManager();
     }
 
+    [[nodiscard]] bool hasPendingHostTransfers() const override
+    {
+        return mBlockManager.hasPendingHostTransfers();
+    }
+
     //! \brief Perform per-iteration bookkeeping
     void refreshBlocks() override
     {

cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h

@@ -66,17 +66,20 @@ class KVCacheTransferManager
         std::string const& directory = "");
 
     //! \brief Synchronize internal streams with bufferManager stream.
-    //! \details The buffer manager uses the same stream as the prefill and decode kernels. This method ensures that the
-    //! internal kernels used for offloading and onboarding will wait for prefill and decode kernels before performing
-    //! any block copies. This method must be called before the first call to
-    //! KVCacheManager::addSequenceBatch in every step.
+    //! \details This method ensures that the internal kernels used for offloading and onboarding will wait for work
+    //! previously scheduled on the buffer manager stream before performing any block copies. This method must be called
+    //! before the first call to KVCacheManager::addSequenceBatch in every step when the buffer manager stream
+    //! represents the work that block transfers depend on.
     void syncWithBufferManager();
 
     //! \brief Synchronize bufferManager stream with internal streams. This method ensures that prefill and decode
     //! kernels for next step will wait for offloading and onboarding work that has already been scheduled. This method
     //! must be called after the last call to KVCacheManager::addSequenceBatch in every step.
     void syncTransfers();
 
+    //! \brief Whether any host-to-device or device-to-host transfer has been scheduled and not yet synchronized.
+    [[nodiscard]] bool hasPendingHostTransfers() const;
+
     //! \brief Get transfer stats accumulated since last call, and reset the counters.
     [[nodiscard]] KvCacheTransferStats getAndResetTransferStats();
 
@@ -121,6 +124,7 @@ class KVCacheTransferManager
     // that identifies the raw memory blocks involved in I/O, not the block Id.
     std::unordered_map<kernels::KVCacheIndex::UnderlyingType, tr::CudaEvent> mPendingReads;
     std::unordered_map<kernels::KVCacheIndex::UnderlyingType, tr::CudaEvent> mPendingWrites;
+    bool mHasPendingHostTransfers{false};
     // Reference to parent loopback agent
     std::shared_ptr<kvc::BaseLoopbackAgent> mLoopbackAgent;
     int mDeviceId;

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -2129,6 +2129,23 @@ void WindowBlockManager::syncTransferManagerWithBufferManager()
     mTransferManager->syncWithBufferManager();
 }
 
+bool BlockManager::hasPendingHostTransfers() const
+{
+    for (auto const& [_, manager] : mWindowBlockManagers)
+    {
+        if (manager.hasPendingHostTransfers())
+        {
+            return true;
+        }
+    }
+    return false;
+}
+
+bool WindowBlockManager::hasPendingHostTransfers() const
+{
+    return mTransferManager->hasPendingHostTransfers();
+}
+
 void BlockManager::refreshBlocks()
 {
     for (auto& [_, manager] : mWindowBlockManagers)

cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp

@@ -272,8 +272,8 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst,
 // for offload/onboard/partial copy is dangerous. We have an asynchronous decoder
 // that may not synchronize or synchronize at a later point in the execution stream.
 // To avoid synchronization issues caused by changes to decoder design we rely on
-// KVCacheTransferManager::syncWithBufferManager() that ensures that internal copy streams
-// will wait for prefill and decode kernels that have already been scheduled.
+// KVCacheTransferManager::syncWithBufferManager() to ensure that internal copy streams
+// wait for the work that block transfers depend on.
 //
 // Earlier versions of this code did not account for all possible cases where a new block copy
 // needed to wait for a previously scheduled copy to finish. For instance, it is possible
@@ -320,6 +320,7 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr co
     }
 
     copyBlock(offloadedBlock, block, pools, false, numTokensToCopy, mode, directory);
+    mHasPendingHostTransfers |= !offloadedBlock->isPrimary();
 
     // Update transfer statistics — distinguish host→GPU onboard from GPU→GPU intra-device copy
     {
@@ -375,6 +376,7 @@ void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offl
     }
 
     copyBlock(block, offloadBlock, pools, true, numTokensToCopy, mode, directory);
+    mHasPendingHostTransfers |= !offloadBlock->isPrimary();
 
     // Update transfer statistics
     {
@@ -404,6 +406,7 @@ void KVCacheTransferManager::syncWithBufferManager()
     // Once we synchronize, clear our list of pending transfers.
     mPendingReads.clear();
     mPendingWrites.clear();
+    mHasPendingHostTransfers = false;
 }
 
 void KVCacheTransferManager::syncTransfers()
@@ -419,6 +422,12 @@ void KVCacheTransferManager::syncTransfers()
     // Once we synchronize, clear our list of pending transfers.
     mPendingReads.clear();
     mPendingWrites.clear();
+    mHasPendingHostTransfers = false;
+}
+
+bool KVCacheTransferManager::hasPendingHostTransfers() const
+{
+    return mHasPendingHostTransfers;
 }
 
 KvCacheTransferStats KVCacheTransferManager::getAndResetTransferStats()

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp

@@ -24,6 +24,7 @@
 #include "tensorrt_llm/runtime/torchView.h"
 
 #include <ATen/ATen.h>
+#include <cstdint>
 #include <nanobind/nanobind.h>
 #include <nanobind/ndarray.h>
 #include <nanobind/operators.h>
@@ -68,7 +69,7 @@ std::optional<tensorrt_llm::runtime::ITensor::UniquePtr> from_torch(std::optiona
 class PyKvCacheManager : public tbk::BaseKVCacheManager
 {
 public:
-    NB_TRAMPOLINE(tbk::BaseKVCacheManager, 39);
+    NB_TRAMPOLINE(tbk::BaseKVCacheManager, 40);
 
     // using BaseKVCacheManager::BaseKVCacheManager; // Inherit constructors
     void allocatePools(bool useUvm = false) override
@@ -262,6 +263,11 @@ class PyKvCacheManager : public tbk::BaseKVCacheManager
         NB_OVERRIDE_PURE(syncTransferManagerWithBufferManager);
     }
 
+    bool hasPendingHostTransfers() const override
+    {
+        NB_OVERRIDE_PURE(hasPendingHostTransfers);
+    }
+
     void refreshBlocks() override
     {
         NB_OVERRIDE_PURE(refreshBlocks);
@@ -624,6 +630,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
             nb::call_guard<nb::gil_scoped_release>())
         .def("sync_transfer_manager_with_buffer_manager", &BaseKVCacheManager::syncTransferManagerWithBufferManager,
             nb::call_guard<nb::gil_scoped_release>())
+        .def("has_pending_host_transfers", &BaseKVCacheManager::hasPendingHostTransfers,
+            nb::call_guard<nb::gil_scoped_release>())
         .def("refresh_blocks", &BaseKVCacheManager::refreshBlocks, nb::call_guard<nb::gil_scoped_release>())
         .def("get_last_block_id", &BaseKVCacheManager::getLastBlockId, nb::call_guard<nb::gil_scoped_release>())
         .def("unpin_blocks_by_id", &BaseKVCacheManager::unpinBlocksById, nb::call_guard<nb::gil_scoped_release>())

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -1168,9 +1168,13 @@ def __init__(
             pool_configurations=pool_configurations,
         )
 
-    def prepare_resources(self, scheduled_batch: ScheduledRequests):
+    def prepare_resources(
+            self,
+            scheduled_batch: ScheduledRequests,
+            forward_done_event: Optional[torch.cuda.Event] = None):
         MambaCacheManager.prepare_resources(self, scheduled_batch)
-        KVCacheManager.prepare_resources(self, scheduled_batch)
+        KVCacheManager.prepare_resources(self, scheduled_batch,
+                                         forward_done_event)
 
     def free_resources(self, request: LlmRequest, pin_on_release: bool = False):
         MambaCacheManager.free_resources(self, request)
@@ -1749,8 +1753,11 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests):
                 )
                 self.mamba_ssm_rand_seed[ctx_slots] = seed_tensor
 
-    def prepare_resources(self, scheduled_batch: ScheduledRequests):
-        super().prepare_resources(scheduled_batch)
+    def prepare_resources(
+            self,
+            scheduled_batch: ScheduledRequests,
+            forward_done_event: Optional[torch.cuda.Event] = None):
+        super().prepare_resources(scheduled_batch, forward_done_event)
         if self.local_num_mamba_layers == 0:
             return
         self._prepare_resources(scheduled_batch)

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -2612,7 +2612,10 @@ def _prefetch_for_context_requests(self) -> None:
         if candidates:
             self.kv_cache_manager.prefetch_for_context_tokens(candidates)
 
-    def _prepare_and_schedule_batch(self):
+    def _prepare_and_schedule_batch(
+        self,
+        forward_done_event: Optional[torch.cuda.Event] = None
+    ) -> Tuple[Optional[ScheduledRequests], Optional[IterationStats]]:
         new_requests = self._fetch_and_activate_new_requests()
         if self.should_stop_processing:
             return None, None
@@ -2683,7 +2686,8 @@ def _prepare_and_schedule_batch(self):
 
         if self.kv_cache_transceiver:
             # For requests that are fitting disagg gen init, also prepare resources for KV cache manager
-            self._prepare_disagg_gen_init(fitting_disagg_gen_init_requests)
+            self._prepare_disagg_gen_init(fitting_disagg_gen_init_requests,
+                                          forward_done_event)
 
             all_gen_first = self.active_requests and all(
                 req.py_disaggregated_params and req.py_disaggregated_params.
@@ -3311,7 +3315,9 @@ def _executor_loop_overlap(self):
                 if self._is_kv_manager_v2 and self._can_pause_for_rebalance():
                     self._maybe_rebalance_kv_pools()
 
-                scheduled_batch, iter_stats = self._prepare_and_schedule_batch()
+                forward_done_event = self._record_execution_stream_tail_event()
+                scheduled_batch, iter_stats = self._prepare_and_schedule_batch(
+                    forward_done_event)
                 self._handle_control_request()
 
                 if scheduled_batch is None:
@@ -3353,7 +3359,8 @@ def _executor_loop_overlap(self):
 
                     self._handle_dynamic_draft_len(scheduled_batch)
 
-                    self.resource_manager.prepare_resources(scheduled_batch)
+                    self.resource_manager.prepare_resources(
+                        scheduled_batch, forward_done_event)
 
                 if self.kv_connector_manager:
                     self.kv_connector_manager.handle_metadata()
@@ -4396,7 +4403,10 @@ def _pad_attention_dp_dummy_request(self):
             self.active_requests.append(llm_request)
 
     @nvtx_range("_prepare_disagg_gen_init")
-    def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests):
+    def _prepare_disagg_gen_init(
+            self,
+            fitting_disagg_gen_init_requests: List[LlmRequest],
+            forward_done_event: Optional[torch.cuda.Event] = None) -> None:
         if fitting_disagg_gen_init_requests:
             disagg_gen_init_to_prepare = ScheduledRequests()
             disagg_gen_init_to_prepare.context_requests_last_chunk = fitting_disagg_gen_init_requests
@@ -4408,8 +4418,15 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests):
                 if (resource_mgr_type in self.resource_manager.resource_managers
                         and self.resource_manager.
                         resource_managers[resource_mgr_type] is not None):
-                    self.resource_manager.resource_managers[
-                        resource_mgr_type].prepare_resources(
+                    resource_manager = self.resource_manager.resource_managers[
+                        resource_mgr_type]
+                    if resource_mgr_type in (
+                            ResourceManagerType.KV_CACHE_MANAGER,
+                            ResourceManagerType.DRAFT_KV_CACHE_MANAGER):
+                        resource_manager.prepare_resources(
+                            disagg_gen_init_to_prepare, forward_done_event)
+                    else:
+                        resource_manager.prepare_resources(
                             disagg_gen_init_to_prepare)
 
             # Trigger KV cache exchange for new disagg_gen_init_requests
@@ -4785,6 +4802,17 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
             self._handle_errors(error_msg)
             return None
 
+    def _record_execution_stream_tail_event(self) -> Optional[torch.cuda.Event]:
+        if self.execution_stream is None:
+            return None
+
+        # In overlap mode the next prepare_resources call can start while the
+        # previous forward is still queued on execution_stream. Record at the
+        # stream tail so KVCM transfer streams wait for that in-flight work.
+        forward_done_event = torch.cuda.Event()
+        forward_done_event.record(self.execution_stream)
+        return forward_done_event
+
     def _update_generation_requests_that_will_complete_next_iteration(
             self, generation_requests: list[LlmRequest]):
         """ Update the generation requests that will complete next iteration.

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -695,7 +695,10 @@ def _context_seq_len(self, req: LlmRequest, is_cross: bool,
             return None
         return req.prompt_len
 
-    def prepare_resources(self, scheduled_batch: ScheduledRequests):
+    def prepare_resources(
+            self,
+            scheduled_batch: ScheduledRequests,
+            forward_done_event: Optional[torch.cuda.Event] = None) -> None:
         # Cross/encoder K/V is allocated once and never grows; handle it on a
         # dedicated path so the self-attention flow below stays unconditional.
         if self.kv_cache_type == CacheTypeCpp.CROSS:
@@ -746,6 +749,9 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                     self.impl.add_token(req.py_request_id)
 
             # prefill and generation kernels wait for scheduled offload/onboard/partial copy work before launching
+            if (forward_done_event is not None
+                    and self.impl.has_pending_host_transfers()):
+                forward_done_event.synchronize()
             self.impl.refresh_blocks()
 
         # A request may change from `context_requests_chunking` to
@@ -2262,10 +2268,17 @@ def get_resource_manager(
         return self.resource_managers.get(type)
 
     @nvtx_range("prepare_resources")
-    def prepare_resources(self, scheduled_batch: ScheduledRequests):
+    def prepare_resources(
+            self,
+            scheduled_batch: ScheduledRequests,
+            forward_done_event: Optional[torch.cuda.Event] = None) -> None:
         for _, resource_manager in self.resource_managers.items():
             if hasattr(resource_manager, "prepare_resources"):
-                resource_manager.prepare_resources(scheduled_batch)
+                if isinstance(resource_manager, KVCacheManager):
+                    resource_manager.prepare_resources(scheduled_batch,
+                                                       forward_done_event)
+                else:
+                    resource_manager.prepare_resources(scheduled_batch)
 
     @nvtx_range("update_resources")
     def update_resources(