[None][fix] pool-qualify KV cache transfer pending keys (#15272)

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

Author: chienchunhung

cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h

@@ -24,6 +24,11 @@ namespace kvc = tensorrt_llm::executor::kv_cache;
 
 #pragma once
 
+namespace tensorrt_llm::testing
+{
+class KVCacheTransferManagerTestAccess;
+} // namespace tensorrt_llm::testing
+
 namespace tensorrt_llm::batch_manager::kv_cache_manager
 {
 
@@ -76,10 +81,15 @@ class KVCacheTransferManager
     [[nodiscard]] KvCacheTransferStats getAndResetTransferStats();
 
 private:
+    friend class ::tensorrt_llm::testing::KVCacheTransferManagerTestAccess;
+
     //! \brief Get pointer to pool specified by cache block.
     static tr::ITensor::SharedPtr computeBlockPointer(
         BlockPtr const& block, std::vector<KVCacheBlockPool> const& pools, size_t poolIdx);
 
+    //! \brief Get pool-qualified index for pending transfer tracking.
+    [[nodiscard]] static kernels::KVCacheIndex::UnderlyingType getPendingTransferIndex(BlockPtr const& block);
+
     /*!
      * \brief The key method that copies the src block to the dst block.
      *
@@ -107,8 +117,8 @@ class KVCacheTransferManager
     runtime::BufferManager mOnboardManager;
     runtime::BufferManager mOffloadManager;
 
-    // Track reads and writes for blocks. Note that it is the memory pool index that
-    // identifies the raw memory blocks involved in I/O, not the block Id.
+    // Track reads and writes for blocks. Note that it is the pool-qualified memory pool index
+    // 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;
     // Reference to parent loopback agent

cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp

@@ -97,6 +97,12 @@ tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer(
     return blockTensor;
 }
 
+tk::KVCacheIndex::UnderlyingType KVCacheTransferManager::getPendingTransferIndex(BlockPtr const& block)
+{
+    auto const blockOffset = block->getMemoryPoolBlockIndex();
+    return block->isPrimary() ? blockOffset : blockOffset | tk::KVCacheIndex::kSecondaryPoolFlag;
+}
+
 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)
@@ -252,10 +258,9 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst,
 }
 
 //
-// Note about recording events to wait for cudaMempyAsync calls between blocks:
-// The memory copy involves raw memory blocks, which are pointed to by the
-// memory pool block index. When recording events, you must use getMemoryPoolBlockIndex()
-// as the raw memory block identifier. Using getBlockId() when recording events is wrong.
+// Note about recording events to wait for cudaMemcpyAsync calls between blocks:
+// The memory copy involves raw memory blocks, which are identified by the pool-qualified
+// memory pool block index. Using getBlockId() when recording events is wrong.
 // getBlockId() returns the logical block id, which has nothing to do with the raw memory
 // block pointers involved in a cudaMemcpy.
 //
@@ -289,22 +294,25 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr co
     std::vector<KVCacheBlockPool> const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode,
     std::string const& directory)
 {
+    auto const offloadedBlockIndex = getPendingTransferIndex(offloadedBlock);
+    auto const blockIndex = getPendingTransferIndex(block);
+
     // Wait for any pending writes before reading from offloadedBlock
-    auto offloadedBlockPendingWriteItr = mPendingWrites.find(offloadedBlock->getMemoryPoolBlockIndex());
+    auto offloadedBlockPendingWriteItr = mPendingWrites.find(offloadedBlockIndex);
     if (offloadedBlockPendingWriteItr != mPendingWrites.end())
     {
         mOnboardManager.getStream().wait(offloadedBlockPendingWriteItr->second);
         // Don't erase, we are not changing state of offloadedBlock
     }
     // Wait for any pending reads before overwriting block
-    auto blockPendingReadItr = mPendingReads.find(block->getMemoryPoolBlockIndex());
+    auto blockPendingReadItr = mPendingReads.find(blockIndex);
     if (blockPendingReadItr != mPendingReads.end())
     {
         mOnboardManager.getStream().wait(blockPendingReadItr->second);
         mPendingReads.erase(blockPendingReadItr);
     }
     // Wait for any pending writes before overwriting block
-    auto blockPendingWriteItr = mPendingWrites.find(block->getMemoryPoolBlockIndex());
+    auto blockPendingWriteItr = mPendingWrites.find(blockIndex);
     if (blockPendingWriteItr != mPendingWrites.end())
     {
         mOnboardManager.getStream().wait(blockPendingWriteItr->second);
@@ -330,33 +338,36 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr co
     }
 
     // Record new pending read from offloadedBlock
-    mPendingReads[offloadedBlock->getMemoryPoolBlockIndex()] = tr::CudaEvent();
-    mOnboardManager.getStream().record(mPendingReads[offloadedBlock->getMemoryPoolBlockIndex()]);
+    mPendingReads[offloadedBlockIndex] = tr::CudaEvent();
+    mOnboardManager.getStream().record(mPendingReads[offloadedBlockIndex]);
     // Record new pending write to block
-    mPendingWrites[block->getMemoryPoolBlockIndex()] = tr::CudaEvent();
-    mOnboardManager.getStream().record(mPendingWrites[block->getMemoryPoolBlockIndex()]);
+    mPendingWrites[blockIndex] = tr::CudaEvent();
+    mOnboardManager.getStream().record(mPendingWrites[blockIndex]);
 }
 
 void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offloadBlock,
     std::vector<KVCacheBlockPool> const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode,
     std::string const& directory)
 {
+    auto const blockIndex = getPendingTransferIndex(block);
+    auto const offloadBlockIndex = getPendingTransferIndex(offloadBlock);
+
     // Wait for any pending writes before reading from block
-    auto blockPendingWriteItr = mPendingWrites.find(block->getMemoryPoolBlockIndex());
+    auto blockPendingWriteItr = mPendingWrites.find(blockIndex);
     if (blockPendingWriteItr != mPendingWrites.end())
     {
         mOffloadManager.getStream().wait(blockPendingWriteItr->second);
         // Don't erase, we are not changing state of block
     }
     // Wait for any pending reads before overwriting offloadBlock
-    auto offloadBlockPendingReadItr = mPendingReads.find(offloadBlock->getMemoryPoolBlockIndex());
+    auto offloadBlockPendingReadItr = mPendingReads.find(offloadBlockIndex);
     if (offloadBlockPendingReadItr != mPendingReads.end())
     {
         mOffloadManager.getStream().wait(offloadBlockPendingReadItr->second);
         mPendingReads.erase(offloadBlockPendingReadItr);
     }
     // Wait for any pending writes before overwriting offloadBlock
-    auto offloadBlockPendingWriteItr = mPendingWrites.find(offloadBlock->getMemoryPoolBlockIndex());
+    auto offloadBlockPendingWriteItr = mPendingWrites.find(offloadBlockIndex);
     if (offloadBlockPendingWriteItr != mPendingWrites.end())
     {
         mOffloadManager.getStream().wait(offloadBlockPendingWriteItr->second);
@@ -373,11 +384,11 @@ void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offl
     }
 
     // Record new pending read from block
-    mPendingReads[block->getMemoryPoolBlockIndex()] = tr::CudaEvent();
-    mOffloadManager.getStream().record(mPendingReads[block->getMemoryPoolBlockIndex()]);
+    mPendingReads[blockIndex] = tr::CudaEvent();
+    mOffloadManager.getStream().record(mPendingReads[blockIndex]);
     // Record new pending write to offloadBlock
-    mPendingWrites[offloadBlock->getMemoryPoolBlockIndex()] = tr::CudaEvent();
-    mOffloadManager.getStream().record(mPendingWrites[offloadBlock->getMemoryPoolBlockIndex()]);
+    mPendingWrites[offloadBlockIndex] = tr::CudaEvent();
+    mOffloadManager.getStream().record(mPendingWrites[offloadBlockIndex]);
 }
 
 void KVCacheTransferManager::syncWithBufferManager()
@@ -390,7 +401,7 @@ void KVCacheTransferManager::syncWithBufferManager()
     mBufferManager.getStream().record(readyForOnboardEvent);
     mOnboardManager.getStream().wait(readyForOnboardEvent);
 
-    // Once we synchronize, clear our list of pending thransfers.
+    // Once we synchronize, clear our list of pending transfers.
     mPendingReads.clear();
     mPendingWrites.clear();
 }
@@ -405,7 +416,7 @@ void KVCacheTransferManager::syncTransfers()
     mOnboardManager.getStream().record(onboardEvent);
     mBufferManager.getStream().wait(onboardEvent);
 
-    // Once we synchronize, clear our list of pending thransfers.
+    // Once we synchronize, clear our list of pending transfers.
     mPendingReads.clear();
     mPendingWrites.clear();
 }

cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp

@@ -62,6 +62,39 @@ namespace tle = tensorrt_llm::executor;
 namespace tr = tensorrt_llm::runtime;
 namespace fs = std::filesystem;
 
+namespace tensorrt_llm::testing
+{
+namespace tbk = batch_manager::kv_cache_manager;
+
+class KVCacheTransferManagerTestAccess
+{
+public:
+    [[nodiscard]] static std::size_t pendingReadCount(tbk::KVCacheTransferManager const& transferManager)
+    {
+        return transferManager.mPendingReads.size();
+    }
+
+    [[nodiscard]] static std::size_t pendingWriteCount(tbk::KVCacheTransferManager const& transferManager)
+    {
+        return transferManager.mPendingWrites.size();
+    }
+
+    [[nodiscard]] static bool hasPendingReadForBlock(
+        tbk::KVCacheTransferManager const& transferManager, tbk::BlockPtr const& block)
+    {
+        return transferManager.mPendingReads.find(tbk::KVCacheTransferManager::getPendingTransferIndex(block))
+            != transferManager.mPendingReads.end();
+    }
+
+    [[nodiscard]] static bool hasPendingWriteForBlock(
+        tbk::KVCacheTransferManager const& transferManager, tbk::BlockPtr const& block)
+    {
+        return transferManager.mPendingWrites.find(tbk::KVCacheTransferManager::getPendingTransferIndex(block))
+            != transferManager.mPendingWrites.end();
+    }
+};
+} // namespace tensorrt_llm::testing
+
 using BlocksPerWindow = std::map<SizeType32, std::tuple<SizeType32, SizeType32>>;
 
 using ParamType = bool;
@@ -5013,6 +5046,49 @@ TEST_F(KVCacheManagerTest, KVCacheTransferManagerConcurrencyTest)
     }
 }
 
+TEST_F(KVCacheManagerTest, KVCacheTransferManagerPendingTransfersDistinguishPrimaryAndSecondarySlots)
+{
+    using tensorrt_llm::testing::KVCacheTransferManagerTestAccess;
+
+    auto constexpr kBlockSize = 1024;
+    auto constexpr kNumSlotsPerPool = 2;
+
+    auto bufferManager = tr::BufferManager(std::make_shared<tr::CudaStream>());
+    auto transferManager = KVCacheTransferManager(bufferManager);
+
+    auto pool = KVCacheBlockPool(0, 2, 0, 0, 0);
+    pool.primaryPtr
+        = bufferManager.gpu(tr::ITensor::makeShape({kNumSlotsPerPool, kBlockSize}), nvinfer1::DataType::kFLOAT);
+    bufferManager.setZero(*pool.primaryPtr);
+
+    pool.secondaryPtr
+        = tr::BufferManager::pinned(tr::ITensor::makeShape({kNumSlotsPerPool, kBlockSize}), nvinfer1::DataType::kFLOAT);
+
+    auto primarySlot0 = std::make_shared<KVCacheBlock>(0, tk::KVCacheIndex(0, false));
+    auto primarySlot1 = std::make_shared<KVCacheBlock>(1, tk::KVCacheIndex(1, false));
+    auto secondarySlot0 = std::make_shared<KVCacheBlock>(2, tk::KVCacheIndex(0, true));
+    auto secondarySlot1 = std::make_shared<KVCacheBlock>(3, tk::KVCacheIndex(1, true));
+
+    transferManager.offload(primarySlot0, secondarySlot1, {pool});
+
+    EXPECT_EQ(KVCacheTransferManagerTestAccess::pendingReadCount(transferManager), 1);
+    EXPECT_EQ(KVCacheTransferManagerTestAccess::pendingWriteCount(transferManager), 1);
+    EXPECT_TRUE(KVCacheTransferManagerTestAccess::hasPendingReadForBlock(transferManager, primarySlot0));
+    EXPECT_TRUE(KVCacheTransferManagerTestAccess::hasPendingWriteForBlock(transferManager, secondarySlot1));
+
+    transferManager.onboard(secondarySlot0, primarySlot1, {pool});
+
+    EXPECT_EQ(KVCacheTransferManagerTestAccess::pendingReadCount(transferManager), 2);
+    EXPECT_EQ(KVCacheTransferManagerTestAccess::pendingWriteCount(transferManager), 2);
+    EXPECT_TRUE(KVCacheTransferManagerTestAccess::hasPendingReadForBlock(transferManager, primarySlot0));
+    EXPECT_TRUE(KVCacheTransferManagerTestAccess::hasPendingReadForBlock(transferManager, secondarySlot0));
+    EXPECT_TRUE(KVCacheTransferManagerTestAccess::hasPendingWriteForBlock(transferManager, primarySlot1));
+    EXPECT_TRUE(KVCacheTransferManagerTestAccess::hasPendingWriteForBlock(transferManager, secondarySlot1));
+
+    transferManager.syncTransfers();
+    bufferManager.getStream().synchronize();
+}
+
 TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerSinkTokenLengthTest)
 {
     // TODO: Support sink attention and add coverage