[#8813][fix] Add missing event for block onboard for the kv cache transfer manager

Authored-by: @josephrocca
Signed-off-by: eopXD <[email protected]>

Author: eopXD

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -325,6 +325,21 @@ class KVCacheBlock
 
     size_t getHash() const;
 
+    void setPendingOnboardEvent(runtime::CudaEvent&& event)
+    {
+        mPendingOnboardEvent = std::move(event);
+    }
+
+    runtime::CudaEvent const* getPendingOnboardEvent() const
+    {
+        return mPendingOnboardEvent ? &mPendingOnboardEvent.value() : nullptr;
+    }
+
+    void clearPendingOnboardEvent()
+    {
+        mPendingOnboardEvent.reset();
+    }
+
 private:
     // Linear ID of block independent of pool
     IdType mBlockId;
@@ -365,6 +380,8 @@ class KVCacheBlock
     std::optional<std::chrono::steady_clock::time_point::duration> mExpirationTime;
     // Hash for the event manager
     size_t mHash;
+    // Possible pending event to onboard the block
+    std::optional<runtime::CudaEvent> mPendingOnboardEvent;
 };
 
 class GenerationRequest

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -1483,6 +1483,13 @@ void WindowBlockManager::addSequence(
 
 void WindowBlockManager::addBlockToBeam(BlockPtr& block, GenerationRequest& sequence, SizeType32 beamIdx)
 {
+    if (auto const* onboardEvent = block->getPendingOnboardEvent())
+    {
+        // Make sure block is onboarded before used
+        mBufferManager.getStream().wait(*onboardEvent);
+        block->clearPendingOnboardEvent();
+    }
+
     auto const requestId = sequence.getRequestId();
     block->incRefCount();
     if (sequence.getCacheBlockIds(mWindowSize).at(beamIdx).size() == 0)

cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp

@@ -224,6 +224,9 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadBlock, BlockPtr cons
         mOnboardManager.getStream().wait(mPendingOffloads[offloadBlock->getBlockId()]);
     }
     copyBlock(offloadBlock, block, pools, false, numTokensToCopy, mode, directory);
+    tr::CudaEvent onboardEvent;
+    mOnboardManager.getStream().record(onboardEvent);
+    block->setPendingOnboardEvent(std::move(onboardEvent));
 }
 
 void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offloadBlock,