[https://nvbugs/5625990][fix] Fix block copy from GPU to GPU for partial reuse in the KV cache manager

`KVCacheTransferManager::onboard` only covers memory movement between
CPU and GPU and not GPU to GPU.

Use the `mBufferManager` to copy block content.

This fixes the incorrect partial block copy functionality exposed
through test case `accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype_vswa_reuse_low_memory_available_partial_reuse`

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

Author: eopXD

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -995,6 +995,10 @@ class WindowBlockManager
     double mTotalInputTokens;
     // Whether blocks that are partially matched should be reused.
     bool mEnablePartialReuse;
+    // Number of partial matched blocks reused through a copy
+    SizeType32 mCopiedReusedPartialBlocks;
+    // Number of partial matched blocks reused directly without a copy
+    SizeType32 mDirectlyReusedPartialBlocks;
     // Whether partially matched blocks that are already in use should be copied and reused.
     bool mCopyOnPartialReuse;
     // The kv cache connector manager

cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h

@@ -62,6 +62,25 @@ class KVCacheTransferManager
     static tr::ITensor::SharedPtr computeBlockPointer(
         BlockPtr const& block, std::vector<KVCacheBlockPool> const& pools, size_t poolIdx);
 
+    /*!
+     * \brief Synchronize pending onboard transfers for the given blocks.
+     *
+     * \details For the src block (offloadedBlock), we wait for any pending
+     *          writes before reading from it. For the dst block (block), we
+     *          wait for any pending reads and writes before overwriting it.
+     * \param offloadedBlock    Offloaded block (to be onboarded)
+     * \param block             Block (to be copied content onto)
+     */
+    void syncPendingOnboardTransfers(BlockPtr const& offloadedBlock, BlockPtr const& block);
+
+    /*!
+     * \brief Record pending onboard transfers for the given blocks.
+     *
+     * \param offloadedBlock    Offloaded block (to be onboarded)
+     * \param block             Block (to be copied content onto)
+     */
+    void recordPendingOnboardTransfers(BlockPtr const& offloadedBlock, BlockPtr const& block);
+
     /*!
      * \brief The key method that copies the src block to the dst block.
      *
@@ -79,6 +98,15 @@ class KVCacheTransferManager
         int numTokensToCopy = 0, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM,
         std::string const& directory = "");
 
+    /*!
+     * \brief Directly copy a block from gpu to gpu.
+     *
+     * \param src             Source block
+     * \param dst             Destination block
+     * \param pools           Pools describing memory layout for KV blocks
+     */
+    void copyBlockGPUToGPU(BlockPtr const& src, BlockPtr const& dst, std::vector<KVCacheBlockPool> const& pools);
+
     runtime::BufferManager mBufferManager;
     runtime::BufferManager mOnboardManager;
     runtime::BufferManager mOffloadManager;

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -672,6 +672,8 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind
     , mReusedTokens{0.0}
     , mTotalInputTokens{0.0}
     , mEnablePartialReuse{enablePartialReuse}
+    , mCopiedReusedPartialBlocks{0}
+    , mDirectlyReusedPartialBlocks{0}
     , mCopyOnPartialReuse{copyOnPartialReuse}
     , mKvCacheConnectorManager{std::move(kvCacheConnectorManager)}
     , mEnableIndexerKCache{enableIndexerKCache}
@@ -765,6 +767,9 @@ WindowBlockManager::~WindowBlockManager()
     TLLM_LOG_DEBUG("%s - reused tokens:                       %.0f ", mLogPrefix.c_str(), mReusedTokens);
     TLLM_LOG_DEBUG("%s - reused tokens percentage (%%):        %.2f ", mLogPrefix.c_str(),
         100.0 * mReusedTokens / mTotalInputTokens);
+    TLLM_LOG_DEBUG("%s - copied reused partial blocks:          %lu  ", mLogPrefix.c_str(), mCopiedReusedPartialBlocks);
+    TLLM_LOG_DEBUG(
+        "%s - directly reused partial blocks:        %lu  ", mLogPrefix.c_str(), mDirectlyReusedPartialBlocks);
 }
 
 bool BlockManager::verifyQueueIntegrity(SizeType32 windowSize)
@@ -1246,7 +1251,7 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector<BlockKey> const&
                     // Somebody else is using block or it is not a leaf, copy reusable tokens
                     auto newBlock = getFreeBlock(
                         sequence, matchingBlock->getPriority(), matchingBlock->getDurationMs(), mode, directory);
-                    mTransferManager->onboard(matchingBlock, newBlock, mPools, numMatched, mode, directory);
+                    mTransferManager->copyBlockGPUToGPU(matchingBlock, newBlock, mPools);
                     // TODO: (optional) Send out event
                     matchingBlock = newBlock;
                     if (blockItr != blockKeys.end())
@@ -1257,6 +1262,7 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector<BlockKey> const&
                     matchingBlock->setHash();
                     TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - Copied partially filled block %d", mLogPrefix.c_str(),
                         matchingBlockId);
+                    ++mCopiedReusedPartialBlocks;
                 }
                 else
                 {
@@ -1266,6 +1272,7 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector<BlockKey> const&
                         matchingBlock, perBlockRetentions[bi].retentionPriority, perBlockRetentions[bi].durationMs);
                     TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - Reused partially filled block %d", mLogPrefix.c_str(),
                         matchingBlockId);
+                    ++mDirectlyReusedPartialBlocks;
                 }
                 searchRoot = nullptr; // no matching needed for following blocks
             }

cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp

@@ -97,6 +97,19 @@ tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer(
     return blockTensor;
 }
 
+// Directly copy a block from gpu to gpu without using the buffer manager.
+void KVCacheTransferManager::copyGPUtoGPU(
+    BlockPtr const& src, BlockPtr const& dst, std::vector<KVCacheBlockPool> const& pools)
+{
+    for (size_t poolIdx = 0; poolIdx < pools.size(); ++poolIdx)
+    {
+        auto srcPtr = computeBlockPointer(src, pools, poolIdx);
+        auto dstPtr = computeBlockPointer(dst, pools, poolIdx);
+        mBufferManager.copy(*srcPtr, *dstPtr);
+    }
+    TLLM_LOG_DEBUG("GPU-to-GPU copy for from block %d to block %d", src->getBlockId(), dst->getBlockId());
+}
+
 void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst,
     std::vector<KVCacheBlockPool> const& pools, bool isOffload, int numTokensToCopy, executor::KvCacheTransferMode mode,
     std::string const& directory)
@@ -241,9 +254,7 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst,
 // Failing to do so will lead to corrupted blocks eventually.
 //
 
-void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr const& block,
-    std::vector<KVCacheBlockPool> const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode,
-    std::string const& directory)
+void KVCacheTransferManager::syncPendingOnboardTransfers(BlockPtr const& offloadedBlock, BlockPtr const& block)
 {
     // Wait for any pending writes before reading from offloadedBlock
     auto offloadedBlockPendingWriteItr = mPendingWrites.find(offloadedBlock->getMemoryPoolBlockIndex());
@@ -266,9 +277,10 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr co
         mOnboardManager.getStream().wait(blockPendingWriteItr->second);
         mPendingWrites.erase(blockPendingWriteItr);
     }
+}
 
-    copyBlock(offloadedBlock, block, pools, false, numTokensToCopy, mode, directory);
-
+void KVCacheTransferManager::recordPendingOnboardTransfers(BlockPtr const& offloadedBlock, BlockPtr const& block)
+{
     // Record new pending read from offloadedBlock
     mPendingReads[offloadedBlock->getMemoryPoolBlockIndex()] = tr::CudaEvent();
     mOnboardManager.getStream().record(mPendingReads[offloadedBlock->getMemoryPoolBlockIndex()]);
@@ -277,6 +289,23 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr co
     mOnboardManager.getStream().record(mPendingWrites[block->getMemoryPoolBlockIndex()]);
 }
 
+void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr const& block,
+    std::vector<KVCacheBlockPool> const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode,
+    std::string const& directory)
+{
+    syncPendingOnboardTransfers(offloadedBlock, block);
+    copyBlock(offloadedBlock, block, pools, false, numTokensToCopy, mode, directory);
+    recordPendingOnboardTransfers(offloadedBlock, block);
+}
+
+void KVCacheTransferManager::copyBlockGPUToGPU(
+    BlockPtr const& offloadedBlock, BlockPtr const& block, std::vector<KVCacheBlockPool> const& pools)
+{
+    syncPendingOnboardTransfers(offloadedBlock, block);
+    copyGPUtoGPU(offloadedBlock, block, pools);
+    recordPendingOnboardTransfers(offloadedBlock, block);
+}
+
 void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offloadBlock,
     std::vector<KVCacheBlockPool> const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode,
     std::string const& directory)

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -1133,10 +1133,6 @@ def test_auto_dtype_vswa_reuse_low_memory_available_no_partial_reuse(self):
             task = MMLU(self.MODEL_NAME)
             task.evaluate(llm)
 
-    @pytest.mark.skip(
-        reason=
-        "Currently failing due to accuracy drop, https://nvbugspro.nvidia.com/bug/5625990"
-    )
     def test_auto_dtype_vswa_reuse_low_memory_available_partial_reuse(self):
         # NOTE: Test with VSWA kv cache config.
         kv_cache_config = KvCacheConfig(