From 6e865bd4844d7b824faba3d5af027f5f87324275 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sun, 22 Jun 2025 19:15:47 -0700
Subject: [PATCH 01/14] Add mmhash to reuse kvcache with mm tokens

---
 .../batch_manager/kvCacheManager.h            |  18 ++-
 .../batch_manager/kvCacheManager.cpp          | 116 +++++++++++++++++-
 .../pybind/batch_manager/kvCacheManager.cpp   |   3 +-
 .../batch_manager/kvCacheManagerTest.cpp      | 114 +++++++++++++++++
 4 files changed, 244 insertions(+), 7 deletions(-)

diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
index d0daf9e43504..55a9a1e193a7 100644
--- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
+++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
@@ -39,6 +39,7 @@
 #include <set>
 #include <unordered_map>
 #include <vector>
+#include <array>
 
 namespace tensorrt_llm::batch_manager::eviction_policy
 {
@@ -68,6 +69,9 @@ using VecUniqueTokens = tensorrt_llm::runtime::VecUniqueTokens;
 using LoraTaskIdType = tensorrt_llm::runtime::LoraTaskIdType;
 using BlocksPerWindow = std::map<SizeType32, std::tuple<SizeType32, SizeType32>>;
 
+// Type alias for multimodal hash key (hash array + start offset)
+using MmKey = std::pair<std::array<uint8_t, 32>, SizeType32>;
+
 template <typename T>
 using OptionalRef = tensorrt_llm::common::OptionalRef<T>;
 
@@ -106,6 +110,10 @@ struct BlockKey
     bool usesExtraIds = false;
     std::optional<LoraTaskIdType> loraTaskId = std::nullopt;
     VecUniqueTokens uniqueTokens;
+    
+    // Extra keys for multimodal data (similar to VLLM's approach)
+    // Each extra key is a pair of (mm_hash, start_offset_in_block)
+    std::optional<std::vector<MmKey>> extraKeys = std::nullopt;
 
     BlockKey() = default;
 
@@ -119,23 +127,25 @@ struct BlockKey
         }
     }
 
-    BlockKey(bool usesExtraIds, std::optional<LoraTaskIdType> loraTaskId, VecUniqueTokens uniqueTokens)
-        : usesExtraIds(usesExtraIds)
+    explicit BlockKey(bool usesExtraIds, std::optional<LoraTaskIdType> loraTaskId, VecUniqueTokens uniqueTokens,
+        std::optional<std::vector<MmKey>> extraKeys = std::nullopt)
+        : usesExtraIds{usesExtraIds}
         , loraTaskId{loraTaskId}
         , uniqueTokens{std::move(uniqueTokens)}
+        , extraKeys{std::move(extraKeys)}
     {
     }
 
     bool operator==(BlockKey const& other) const noexcept
     {
         return (
-            usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId && uniqueTokens == other.uniqueTokens);
+            usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId && uniqueTokens == other.uniqueTokens && extraKeys == other.extraKeys);
     }
 
     int partialMatch(BlockKey const& other) const noexcept
     {
         SizeType32 numMatched{0};
-        if (loraTaskId == other.loraTaskId)
+        if (loraTaskId == other.loraTaskId && extraKeys == other.extraKeys)
         {
             auto [matchEnd, otherMatchEnd] = std::mismatch(
                 uniqueTokens.begin(), uniqueTokens.end(), other.uniqueTokens.begin(), other.uniqueTokens.end());
diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index ba3b2a94ede6..709134f42e03 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -80,10 +80,24 @@ std::vector<BlockKey> buildBlockKeys(
     std::list<VecUniqueTokens>& blockedUniqueTokens, tensorrt_llm::batch_manager::LlmRequest const& llmRequest)
 {
     std::vector<BlockKey> blockKeys;
+    
+    SizeType32 currentTokenIdx = 0;
     for (auto& uniqueTokens : blockedUniqueTokens)
     {
+        // Generate extra keys for this block
+        auto extraKeys = generateBlockHashExtraKeys(llmRequest, currentTokenIdx, currentTokenIdx + uniqueTokens.size());
+        
         blockKeys.emplace_back(
-            llmRequest.getInputTokensExtraIds().has_value(), llmRequest.getLoraTaskId(), std::move(uniqueTokens));
+            llmRequest.getInputTokensExtraIds().has_value(), 
+            llmRequest.getLoraTaskId(), 
+            std::move(uniqueTokens),
+            std::move(extraKeys));
+            
+        currentTokenIdx += uniqueTokens.size();
+    }
+    return blockKeys;
+}
+
     }
     return blockKeys;
 }
@@ -93,6 +107,74 @@ std::vector<BlockKey> buildBlockKeys(
 namespace tensorrt_llm::batch_manager::kv_cache_manager
 {
 
+std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
+    LlmRequest const& llmRequest, 
+    SizeType32 startTokenIdx, 
+    SizeType32 endTokenIdx)
+{
+    auto const multimodalHashes = llmRequest.getMultimodalHashes();
+    auto const multimodalPositions = llmRequest.getMultimodalPositions();
+    auto const multimodalLengths = llmRequest.getMultimodalLengths();
+    
+    if (!multimodalHashes || !multimodalPositions || !multimodalLengths || 
+        multimodalHashes->empty() || multimodalPositions->empty() || multimodalLengths->empty())
+    {
+        return std::nullopt;
+    }
+    
+    if (multimodalHashes->size() != multimodalPositions->size() || 
+        multimodalPositions->size() != multimodalLengths->size())
+    {
+        TLLM_LOG_WARNING("Multimodal data arrays have mismatched sizes");
+        return std::nullopt;
+    }
+    
+    std::vector<MmKey> extraKeys;  // MmKey = std::pair<std::array<uint8_t, 32>, SizeType32>
+    
+    for (size_t i = 0; i < multimodalPositions->size(); ++i)
+    {
+        auto const& startPos = (*multimodalPositions)[i];
+        auto const& length = (*multimodalLengths)[i];
+        auto const& mmHashVector = (*multimodalHashes)[i];  // This is vector<int32> - your current format
+        
+        std::array<uint8_t, 32> mmHashArray;
+        if (mmHashVector.size() == 8)  // 256-bit hash = 8 * 32-bit integers
+        {
+            // Assuming mmHashVector[j] comes from Python's int(hex_chunk, 16)
+            // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB.
+            // And assuming the overall Blake3 output wants these bytes in order: 0x00, 0x01, 0x02, 0x03...
+            for (size_t j = 0; j < 8; ++j)
+            {
+                auto const& hashPart = mmHashVector[j]; // e.g., 0x00010203
+                // Extract bytes in Big-Endian order from hashPart and place them sequentially
+                // into mmHashArray to reconstruct the original Blake3 byte sequence.
+                mmHashArray[j * 4 + 0] = static_cast<uint8_t>((hashPart >> 24) & 0xFF); // Extract 0x00 (MSB of the int)
+                mmHashArray[j * 4 + 1] = static_cast<uint8_t>((hashPart >> 16) & 0xFF); // Extract 0x01
+                mmHashArray[j * 4 + 2] = static_cast<uint8_t>((hashPart >> 8) & 0xFF);  // Extract 0x02
+                mmHashArray[j * 4 + 3] = static_cast<uint8_t>(hashPart & 0xFF);         // Extract 0x03 (LSB of the int)
+            }
+        }
+        else
+        {
+            // TODO: maybe we should raise an error here
+            TLLM_LOG_WARNING("Multimodal hash vector has unexpected size: %zu (expected 8)", mmHashVector.size());
+            continue;  // Skip this multimodal item
+        }
+        
+        // Check if this multimodal content overlaps with the current block
+        if (endTokenIdx > startPos && startTokenIdx < startPos + length)
+        {
+            // Calculate the start offset of this multimodal content within the block
+            SizeType32 mmStartInBlock = (startPos >= startTokenIdx) ? 0 : startTokenIdx - startPos;
+            
+            // Add the multimodal hash array and its start offset in the block
+            extraKeys.emplace_back(mmHashArray, mmStartInBlock);
+        }
+    }
+    
+    return extraKeys.empty() ? std::nullopt : std::make_optional(std::move(extraKeys));
+}
+
 size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) noexcept
 {
     size_t seed = blockKey.uniqueTokens.size() ^ parentHash * UINT64_C(0xbf58476d1ce4e5b9);
@@ -122,7 +204,37 @@ size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) no
         c = c ^ (c >> 31);
         seed ^= c + 0x9e3779b9 + (seed << 6) + (seed >> 2);
     }
-    // TODO: support external hashes for multimodal
+
+    // Add extra keys for multimodal data (similar to VLLM's approach)
+    if (blockKey.extraKeys)
+    {
+        for (auto const& [mmHash, startOffset] : *blockKey.extraKeys)
+        {
+            // Hash the multimodal hash array in 32-bit chunks (more efficient)
+            for (size_t i = 0; i < 32; i += 4)
+            {
+                // Combine 4 bytes into a 32-bit word (construct as little endian order)
+                uint32_t word = static_cast<uint32_t>(mmHash[i]) |
+                                (static_cast<uint32_t>(mmHash[i+1]) << 8) |
+                                (static_cast<uint32_t>(mmHash[i+2]) << 16) |
+                                (static_cast<uint32_t>(mmHash[i+3]) << 24);
+                
+                // Mix the word into the seed
+                word = ((word >> 16) ^ word) * 0x45d9f3b;
+                word = ((word >> 16) ^ word) * 0x45d9f3b;
+                word = (word >> 16) ^ word;
+                seed ^= word + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+            }
+            
+            // Hash the start offset
+            uint64_t e = static_cast<uint64_t>(startOffset);
+            e = (e ^ (e >> 30)) * UINT64_C(0xbf58476d1ce4e5b9);
+            e = (e ^ (e >> 27)) * UINT64_C(0x94d049bb133111eb);
+            e = e ^ (e >> 31);
+            seed ^= e + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+        }
+    }
+
     return seed;
 }
 
diff --git a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp
index 255b0f8efa33..3649600fb91f 100644
--- a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp
@@ -315,7 +315,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(py::module_& m)
             py::arg("lora_task_id"), py::arg("unique_tokens"))
         .def_readonly("uses_extra_ids", &tbk::BlockKey::usesExtraIds)
         .def_readonly("lora_task_id", &tbk::BlockKey::loraTaskId)
-        .def_readonly("unique_tokens", &tbk::BlockKey::uniqueTokens);
+        .def_readonly("unique_tokens", &tbk::BlockKey::uniqueTokens)
+        .def_readonly("extra_keys", &tbk::BlockKey::extraKeys);
 
     py::class_<tbk::BlockKeyHasher>(m, "BlockKeyHasher")
         .def_static("hash", &tbk::BlockKeyHasher::hash, py::arg("block_key"), py::arg("parent_hash") = 0);
diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
index 08ab45145d53..b21a7c1dd213 100644
--- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
+++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
@@ -1034,6 +1034,120 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest)
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
 }
 
+TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
+{
+    using VecTokenExtraIds = LlmRequest::VecTokenExtraIds;
+
+    auto constexpr numLayers = 12;
+    auto constexpr numKvHeads = 6;
+    auto constexpr sizePerHead = 16;
+    auto constexpr tokensPerBlock = 4;
+    auto constexpr maxBlocksPerSeq = 4;
+    auto constexpr blocksInPrimaryPool = 16;
+    auto constexpr blocksInSecondaryPool = 0;
+    auto constexpr maxNumSequences = 8;
+    auto const stream = std::make_shared<tr::CudaStream>();
+    auto constexpr onboardBlocks = true;
+    auto constexpr numReturnSequences = 1;
+    auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq;
+    auto constexpr beamWidth = 1;
+
+    auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}};
+
+    BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow,
+        maxNumSequences, stream, maxAttentionWindow, beamWidth,
+        std::vector<BlockManager::SizeType32>{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0,
+        onboardBlocks);
+    blockManager.allocatePools(false);
+
+    EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock);
+    EXPECT_EQ(blockManager.getMaxNumBlocks(), blocksInPrimaryPool);
+    EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
+
+    SizeType32 constexpr maxNewTokens{0};
+    tr::SamplingConfig const samplingConfig{beamWidth};
+    bool constexpr isStreaming{false};
+
+    // Create multimodal hash data (256-bit hash = 8 int32 values)
+    auto multimodalHashes = std::make_shared<std::vector<std::vector<SizeType32>>>(
+        std::vector<std::vector<SizeType32>>{
+            {0x12345678, 0x90ABCDEF, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666}  // Hash 1
+        });
+    auto multimodalPositions = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2});  // Start at token 2
+    auto multimodalLengths = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{4});    // Length 4 tokens
+
+    // assume prompt id starts from 100
+    auto inputTokens = std::make_shared<VecTokens>(VecTokens{100, 101, 102, 103, 104, 105, 0, 1, 2});
+    auto const inputLength = static_cast<SizeType32>(inputTokens->size());
+    LlmRequest::RequestIdType requestId{0};
+    auto llmRequest0 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
+        std::nullopt, multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
+        false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt,
+        0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION,
+        std::nullopt, numReturnSequences);
+
+    GenerationRequest seq0{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
+
+    ///////////////////////////////////////////////////////////////////////////
+    // add request and then remove it
+    auto constexpr beamIdx = 0;
+    auto promptLen0 = llmRequest0->getNumTokens(beamIdx);
+    auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock());
+    blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow);
+    EXPECT_EQ(llmRequest0->getContextCurrentPosition(), 0);
+    EXPECT_THAT(seq0.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 1, 2}));
+    llmRequest0->addNewToken(3, beamIdx);
+    llmRequest0->addNewToken(4, beamIdx);
+    auto numTokens = llmRequest0->getNumTokens(beamIdx);
+    auto numBlocks = tc::ceilDiv(numTokens, tokensPerBlock);
+    EXPECT_EQ(numBlocks, 3);
+    EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks);
+    EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks);
+
+    // Input: [100, 101, 102, 103, 104, 105, 0, 1, 2] (9 tokens)
+    // Multimodal: starts at token 2, length 4 → [102, 103, 104, 105]
+
+    // Block 0: [100, 101, 102, 103] ← Contains multimodal (102, 103)
+    // Block 1: [104, 105, 0, 1]     ← Contains multimodal (104, 105)  
+    // Block 2: [2, 3, 4]            ← No multimodal
+    blockManager.releaseBlocks(seq0, llmRequest0);
+    EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
+    EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
+
+    ///////////////////////////////////////////////////////////////////////////
+    // new request with same tokens and same multimodal hash - should reuse
+    requestId = 1;
+    auto llmRequest1 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
+        std::nullopt, multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
+        false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt,
+        0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION,
+        std::nullopt, numReturnSequences);
+    GenerationRequest seq1{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
+
+    // should reuse blocks 0, 1 and get new block 3
+    auto promptLen1 = llmRequest1->getNumTokens(beamIdx);
+    auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock());
+    blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow);
+    EXPECT_EQ(llmRequest1->getContextCurrentPosition(), 2 * tokensPerBlock);
+    EXPECT_THAT(seq1.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 1, 3}));
+    llmRequest1->addNewToken(3, beamIdx);
+    llmRequest1->addNewToken(4, beamIdx);
+    EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks);
+    EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks);
+
+    // block 3 matches block 2 and will be freed
+    blockManager.releaseBlocks(seq1, llmRequest1);
+    EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
+    EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
+
+    ///////////////////////////////////////////////////////////////////////////
+    // TODO: 
+}
+
 TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest)
 {
     // tc::Logger::getLogger()->setLevel(tc::Logger::Level::DEBUG);

From 686bc927f6cf04df349c0939fdf6a65dac697bbb Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Tue, 24 Jun 2025 06:55:12 -0700
Subject: [PATCH 02/14] Bug fixing and adding debug prints

---
 .../batch_manager/kvCacheManager.cpp          |  73 +++----
 .../batch_manager/kvCacheManagerTest.cpp      | 201 ++++++++++++++++--
 examples/llm-api/quickstart_multimodal.py     |   1 -
 .../models/modeling_multimodal_utils.py       |  12 +-
 4 files changed, 230 insertions(+), 57 deletions(-)

diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index 709134f42e03..967eaa544800 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -76,39 +76,8 @@ std::list<std::vector<T>> chopVectorIntoBlocks(
     return blockedVectors;
 }
 
-std::vector<BlockKey> buildBlockKeys(
-    std::list<VecUniqueTokens>& blockedUniqueTokens, tensorrt_llm::batch_manager::LlmRequest const& llmRequest)
-{
-    std::vector<BlockKey> blockKeys;
-    
-    SizeType32 currentTokenIdx = 0;
-    for (auto& uniqueTokens : blockedUniqueTokens)
-    {
-        // Generate extra keys for this block
-        auto extraKeys = generateBlockHashExtraKeys(llmRequest, currentTokenIdx, currentTokenIdx + uniqueTokens.size());
-        
-        blockKeys.emplace_back(
-            llmRequest.getInputTokensExtraIds().has_value(), 
-            llmRequest.getLoraTaskId(), 
-            std::move(uniqueTokens),
-            std::move(extraKeys));
-            
-        currentTokenIdx += uniqueTokens.size();
-    }
-    return blockKeys;
-}
-
-    }
-    return blockKeys;
-}
-
-} // namespace
-
-namespace tensorrt_llm::batch_manager::kv_cache_manager
-{
-
 std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
-    LlmRequest const& llmRequest, 
+    tensorrt_llm::batch_manager::LlmRequest const& llmRequest,
     SizeType32 startTokenIdx, 
     SizeType32 endTokenIdx)
 {
@@ -117,13 +86,15 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
     auto const multimodalLengths = llmRequest.getMultimodalLengths();
     
     if (!multimodalHashes || !multimodalPositions || !multimodalLengths || 
-        multimodalHashes->empty() || multimodalPositions->empty() || multimodalLengths->empty())
+        !(*multimodalHashes) || (*multimodalHashes)->empty() || 
+        !(*multimodalPositions) || (*multimodalPositions)->empty() || 
+        !(*multimodalLengths) || (*multimodalLengths)->empty())
     {
         return std::nullopt;
     }
     
-    if (multimodalHashes->size() != multimodalPositions->size() || 
-        multimodalPositions->size() != multimodalLengths->size())
+    if ((*multimodalHashes)->size() != (*multimodalPositions)->size() || 
+        (*multimodalPositions)->size() != (*multimodalLengths)->size())
     {
         TLLM_LOG_WARNING("Multimodal data arrays have mismatched sizes");
         return std::nullopt;
@@ -131,11 +102,11 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
     
     std::vector<MmKey> extraKeys;  // MmKey = std::pair<std::array<uint8_t, 32>, SizeType32>
     
-    for (size_t i = 0; i < multimodalPositions->size(); ++i)
+    for (size_t i = 0; i < (*multimodalPositions)->size(); ++i)
     {
-        auto const& startPos = (*multimodalPositions)[i];
-        auto const& length = (*multimodalLengths)[i];
-        auto const& mmHashVector = (*multimodalHashes)[i];  // This is vector<int32> - your current format
+        auto const& startPos = (*(*multimodalPositions))[i];
+        auto const& length = (*(*multimodalLengths))[i];
+        auto const& mmHashVector = (*(*multimodalHashes))[i];  // This is vector<int32> - your current format
         
         std::array<uint8_t, 32> mmHashArray;
         if (mmHashVector.size() == 8)  // 256-bit hash = 8 * 32-bit integers
@@ -175,6 +146,30 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
     return extraKeys.empty() ? std::nullopt : std::make_optional(std::move(extraKeys));
 }
 
+std::vector<BlockKey> buildBlockKeys(
+    std::list<VecUniqueTokens>& blockedUniqueTokens, tensorrt_llm::batch_manager::LlmRequest const& llmRequest)
+{
+    std::vector<BlockKey> blockKeys;
+    
+    SizeType32 currentTokenIdx = 0;
+    for (auto& uniqueTokens : blockedUniqueTokens)
+    {
+        auto extraKeys = generateBlockHashExtraKeys(llmRequest, currentTokenIdx, currentTokenIdx + uniqueTokens.size());
+        currentTokenIdx += uniqueTokens.size();
+
+        blockKeys.emplace_back(
+            llmRequest.getInputTokensExtraIds().has_value(), 
+            llmRequest.getLoraTaskId(), 
+            std::move(uniqueTokens),
+            std::move(extraKeys));
+    }
+    return blockKeys;
+}
+
+} // namespace
+
+namespace tensorrt_llm::batch_manager::kv_cache_manager
+{
 size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) noexcept
 {
     size_t seed = blockKey.uniqueTokens.size() ^ parentHash * UINT64_C(0xbf58476d1ce4e5b9);
diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
index b21a7c1dd213..de28b485d2b8 100644
--- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
+++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
@@ -1067,26 +1067,86 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     SizeType32 constexpr maxNewTokens{0};
     tr::SamplingConfig const samplingConfig{beamWidth};
     bool constexpr isStreaming{false};
+    
+    // Helper function to print block information
+    auto printBlockInfo = [&](std::string const& prefix, KVCacheBlock::IdType blockId) {
+        auto const& block = blockManager.getBlockById(blockId, maxAttentionWindow);
+        auto const& blockKey = block->getBlockKey();
+        
+        std::cout << prefix << " Block " << blockId << ":" << std::endl;
+        std::cout << "  Hash: 0x" << std::hex << block->getHash() << std::dec << std::endl;
+        std::cout << "  Tokens: [";
+        for (size_t i = 0; i < blockKey.uniqueTokens.size(); ++i) {
+            if (i > 0) std::cout << ", ";
+            std::cout << blockKey.uniqueTokens[i].tokenId;
+        }
+        std::cout << "]" << std::endl;
+        std::cout << "  UsesExtraIds: " << (blockKey.usesExtraIds ? "true" : "false") << std::endl;
+        if (blockKey.loraTaskId) {
+            std::cout << "  LoraTaskId: " << *blockKey.loraTaskId << std::endl;
+        }
+        if (blockKey.extraKeys) {
+            std::cout << "  ExtraKeys (multimodal): " << blockKey.extraKeys->size() << " items" << std::endl;
+            for (size_t i = 0; i < blockKey.extraKeys->size(); ++i) {
+                auto const& [mmHash, startOffset] = (*blockKey.extraKeys)[i];
+                std::cout << "    Item " << i << ": offset=" << startOffset << ", hash=[";
+                for (size_t j = 0; j < 8; ++j) { // Print first 8 bytes as hex
+                    if (j > 0) std::cout << " ";
+                    std::cout << std::hex << std::setw(2) << std::setfill('0') 
+                             << static_cast<int>(mmHash[j * 4]) << std::dec;
+                }
+                std::cout << "...]" << std::endl;
+            }
+        }
+        std::cout << "  IsFull: " << (block->isFull() ? "true" : "false") << std::endl;
+        std::cout << "  HasRefs: " << (block->hasRefs() ? "true" : "false") << std::endl;
+
+    };
+
+    // Helper function to print sequence block information
+    auto printSequenceBlocks = [&](std::string const& prefix, GenerationRequest const& seq) {
+        auto const& blockIds = seq.getCacheBlockIds(maxAttentionWindow).at(0);
+        std::cout << prefix << " Sequence blocks: [";
+        for (size_t i = 0; i < blockIds.size(); ++i) {
+            if (i > 0) std::cout << ", ";
+            std::cout << blockIds[i];
+        }
+        std::cout << "]" << std::endl;
+        
+        for (auto blockId : blockIds) {
+            printBlockInfo(prefix + "  ", blockId);
+        }
+    };
 
     // Create multimodal hash data (256-bit hash = 8 int32 values)
     auto multimodalHashes = std::make_shared<std::vector<std::vector<SizeType32>>>(
         std::vector<std::vector<SizeType32>>{
-            {0x12345678, 0x90ABCDEF, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666}  // Hash 1
+            {0x12345678, -0x6F543211, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666}  // Hash 1
         });
     auto multimodalPositions = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2});  // Start at token 2
     auto multimodalLengths = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{4});    // Length 4 tokens
-
+    {
+        std::cout << "\n=== Multimodal Hash Data ===" << std::endl;
+        std::cout << "Hash: [";
+        for (size_t i = 0; i < (*multimodalHashes)[0].size(); ++i) {
+            if (i > 0) std::cout << ", ";
+            std::cout << std::hex << "0x" << (*multimodalHashes)[0][i] << std::dec;
+        }
+        std::cout << "]" << std::endl;
+        std::cout << "Position: " << (*multimodalPositions)[0] << std::endl;
+        std::cout << "Length: " << (*multimodalLengths)[0] << std::endl;
+    }
+ 
     // assume prompt id starts from 100
     auto inputTokens = std::make_shared<VecTokens>(VecTokens{100, 101, 102, 103, 104, 105, 0, 1, 2});
     auto const inputLength = static_cast<SizeType32>(inputTokens->size());
     LlmRequest::RequestIdType requestId{0};
     auto llmRequest0 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
         std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
-        std::nullopt, multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
-        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
-        false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt,
-        0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION,
-        std::nullopt, numReturnSequences);
+        multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, 
+        false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
+        LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
 
     GenerationRequest seq0{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
 
@@ -1095,9 +1155,15 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     auto constexpr beamIdx = 0;
     auto promptLen0 = llmRequest0->getNumTokens(beamIdx);
     auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock());
+    {
+        std::cout << "Adding sequence with " << promptLen0 << " tokens, " << numContextBlocks0 << " context blocks" << std::endl;
+    }
     blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow);
     EXPECT_EQ(llmRequest0->getContextCurrentPosition(), 0);
     EXPECT_THAT(seq0.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 1, 2}));
+    {
+        printSequenceBlocks("After addSequence", seq0);
+    }
     llmRequest0->addNewToken(3, beamIdx);
     llmRequest0->addNewToken(4, beamIdx);
     auto numTokens = llmRequest0->getNumTokens(beamIdx);
@@ -1112,6 +1178,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     // Block 0: [100, 101, 102, 103] ← Contains multimodal (102, 103)
     // Block 1: [104, 105, 0, 1]     ← Contains multimodal (104, 105)  
     // Block 2: [2, 3, 4]            ← No multimodal
+    std::cout << "\nReleasing blocks for reuse..." << std::endl;
     blockManager.releaseBlocks(seq0, llmRequest0);
     EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
@@ -1119,33 +1186,136 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     ///////////////////////////////////////////////////////////////////////////
     // new request with same tokens and same multimodal hash - should reuse
     requestId = 1;
+    std::cout << "\n=== Second Request (seq1) - Same Context ===" << std::endl;
     auto llmRequest1 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
         std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
-        std::nullopt, multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
-        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
-        false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt,
-        0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION,
-        std::nullopt, numReturnSequences);
+        multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, 
+        false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
+        LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
     GenerationRequest seq1{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
 
     // should reuse blocks 0, 1 and get new block 3
     auto promptLen1 = llmRequest1->getNumTokens(beamIdx);
     auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock());
+    {
+        std::cout << "Adding sequence with " << promptLen1 << " tokens, " << numContextBlocks1 << " context blocks" << std::endl;
+        std::cout << "Expected: reuse blocks 0, 1 and get new block 3" << std::endl;
+    }
     blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow);
     EXPECT_EQ(llmRequest1->getContextCurrentPosition(), 2 * tokensPerBlock);
     EXPECT_THAT(seq1.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 1, 3}));
+    {
+        printSequenceBlocks("After addSequence (reuse)", seq1);
+        // Verify reuse statistics
+        std::cout << "\n=== Reuse Statistics ===" << std::endl;
+        std::cout << "Total allocated blocks: " << blockManager.getNumAllocTotalBlocks() << std::endl;
+        std::cout << "New blocks allocated: " << blockManager.getNumAllocNewBlocks() << std::endl;
+        std::cout << "Blocks reused: " << blockManager.getNumReusedBlocks() << std::endl;
+        std::cout << "Blocks missed: " << blockManager.getNumMissedBlocks() << std::endl;
+        
+        // Verify that blocks 0 and 1 were reused (same hash)
+        auto const& block0 = blockManager.getBlockById(0, maxAttentionWindow);
+        auto const& block1 = blockManager.getBlockById(1, maxAttentionWindow);
+        auto const& block3 = blockManager.getBlockById(3, maxAttentionWindow);
+        
+        std::cout << "\n=== Block Hash Verification ===" << std::endl;
+        std::cout << "Block 0 hash: 0x" << std::hex << block0->getHash() << std::dec << std::endl;
+        std::cout << "Block 1 hash: 0x" << std::hex << block1->getHash() << std::dec << std::endl;
+        std::cout << "Block 3 hash: 0x" << std::hex << block3->getHash() << std::dec << std::endl;
+        
+        // Verify that blocks 0 and 1 have the same content (should be reused)
+        auto const& blockKey0 = block0->getBlockKey();
+        auto const& blockKey1 = block1->getBlockKey();
+        auto const& blockKey3 = block3->getBlockKey();
+        
+        // Verify multimodal data is present in blocks 0 and 1
+        std::cout << "\n=== Multimodal Data Verification ===" << std::endl;
+        std::cout << "Block 0 has multimodal data: " << (blockKey0.extraKeys && !blockKey0.extraKeys->empty() ? "YES" : "NO") << std::endl;
+        std::cout << "Block 1 has multimodal data: " << (blockKey1.extraKeys && !blockKey1.extraKeys->empty() ? "YES" : "NO") << std::endl;
+        std::cout << "Block 3 has multimodal data: " << (blockKey3.extraKeys && !blockKey3.extraKeys->empty() ? "YES" : "NO") << std::endl;
+    }
     llmRequest1->addNewToken(3, beamIdx);
     llmRequest1->addNewToken(4, beamIdx);
     EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks);
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks);
-
     // block 3 matches block 2 and will be freed
+    std::cout << "\nReleasing blocks for reuse..." << std::endl;
     blockManager.releaseBlocks(seq1, llmRequest1);
     EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
 
     ///////////////////////////////////////////////////////////////////////////
-    // TODO: 
+    // Test Case 2: Different multimodal hash 
+    requestId = 2;
+    auto multimodalHashes2 = std::make_shared<std::vector<std::vector<SizeType32>>>(
+        std::vector<std::vector<SizeType32>>{
+            {0x45678123, 0x23456789, 0x34567890, 0x12121212, 0x56565656, 0x78787878, 0x54545454, 0x67676767} // Hash 2
+        });
+    auto multimodalPositions2 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2});  // Start at token 2
+    auto multimodalLengths2 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{4});    // Length 4 tokens
+    auto llmRequest2 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
+        multimodalHashes2, multimodalPositions2, multimodalLengths2, std::nullopt, std::nullopt, std::nullopt,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, 
+        false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
+        LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
+
+    GenerationRequest seq2{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
+    // no reuse, get new blocks 5, 6, 7
+    auto promptLen2 = llmRequest2->getNumTokens(beamIdx);
+    auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock());
+    blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow);
+    EXPECT_EQ(llmRequest2->getContextCurrentPosition(), 0);
+    EXPECT_THAT(seq2.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({4, 5, 6}));
+    {
+        printSequenceBlocks("Add seq2", seq2);
+    }
+    llmRequest2->addNewToken(9, beamIdx);
+    numTokens = llmRequest2->getNumTokens(beamIdx);
+    numBlocks = tc::ceilDiv(numTokens, tokensPerBlock);
+    EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks);
+    EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks);
+
+    ///////////////////////////////////////////////////////////////////////////
+    // Test Case 3: Multiple multimodal hashes and partial reuse
+    requestId = 3;
+    auto multimodalHashes3 = std::make_shared<std::vector<std::vector<SizeType32>>>(
+        std::vector<std::vector<SizeType32>>{
+            {0x12345678, -0x6F543211, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666},  // Hash 1
+            {0x45678123, 0x23456789, 0x34567890, 0x12121212, 0x56565656, 0x78787878, 0x54545454, 0x67676767} // Hash 2
+        });
+    auto multimodalPositions3 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2, 4});  // Start at token 2 and 4
+    auto multimodalLengths3 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2, 2});    // Length 2 tokens
+    
+    auto llmRequest3 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
+        multimodalHashes3, multimodalPositions3, multimodalLengths3, std::nullopt, std::nullopt, std::nullopt,
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, 
+        false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
+        LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
+    GenerationRequest seq3{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
+    // reuse block 0, get new blocks 8, 9
+    auto promptLen3 = llmRequest3->getNumTokens(beamIdx);
+    auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock());
+    blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow);
+    EXPECT_EQ(llmRequest3->getContextCurrentPosition(), tokensPerBlock); // only reuse block 0 [100, 101, 102, 103] with same hash/offset
+    EXPECT_THAT(seq3.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 7, 8}));
+    {
+        printSequenceBlocks("Add seq3", seq3);
+    }
+    llmRequest3->addNewToken(11, beamIdx);
+    numTokens = llmRequest3->getNumTokens(beamIdx);
+    numBlocks = tc::ceilDiv(numTokens, tokensPerBlock);
+    EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks * 2);
+    EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks * 2);
+
+    // clean up
+    blockManager.releaseBlocks(seq2, llmRequest2);
+    blockManager.releaseBlocks(seq3, llmRequest3);
+    EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
+    EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
+
 }
 
 TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest)
@@ -2575,7 +2745,6 @@ TEST_P(KVCacheManagerTest, KVCacheManagerRewindTokensTest)
         EXPECT_EQ(blockManager.getNumFreeBlocks(), currentNumBlocks);
     }
 }
-
 TEST_P(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowTest)
 {
     using DType = half;
@@ -4639,3 +4808,5 @@ auto const paramValues = ::testing::Values(
     });
 
 INSTANTIATE_TEST_SUITE_P(FillKvCacheAndCompleteRequestsTest, FillKvCacheAndCompleteRequestsTest, paramValues);
+
+
diff --git a/examples/llm-api/quickstart_multimodal.py b/examples/llm-api/quickstart_multimodal.py
index 967a8636e1be..31ed52b43aa9 100644
--- a/examples/llm-api/quickstart_multimodal.py
+++ b/examples/llm-api/quickstart_multimodal.py
@@ -105,7 +105,6 @@ def parse_arguments():
     parser = add_lora_args(parser)
     args = parser.parse_args()
 
-    args.disable_kv_cache_reuse = True  # kv cache reuse does not work for multimodal, force overwrite
     if args.kv_cache_fraction is None:
         args.kv_cache_fraction = 0.6  # lower the default kv cache fraction for multimodal
 
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
index 1dc86cdd1d2a..62be3c7b890a 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
@@ -78,8 +78,16 @@ def fuse_input_embeds(
 
     input_embeds[text_token_indices, :] = text_embed.to(
         dtype=input_embeds.dtype, device=input_embeds.device)
-    input_embeds[mm_token_indices, :] = mm_embed.to(dtype=input_embeds.dtype,
-                                                    device=input_embeds.device)
+
+    # When KV cache reuse enabled, mm_token_indices may only contain partial multimodal tokens or empty
+    if mm_token_indices.numel() > 0:
+        # We need to extract only the embeddings that correspond to uncached multimodal tokens
+        if mm_token_indices.numel() < mm_embed.shape[0]:
+            # Extract the last mm_token_indices.numel() embeddings from mm_embed
+            start_idx = mm_embed.shape[0] - mm_token_indices.numel()
+            input_embeds[mm_token_indices, :] = mm_embed[start_idx:, :].to(dtype=input_embeds.dtype, device=input_embeds.device)
+        else:
+            input_embeds[mm_token_indices, :] = mm_embed.to(dtype=input_embeds.dtype, device=input_embeds.device)
 
     return None, input_embeds
 

From ef6c4b8ea8da9667b281413fc807725e92333383 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Tue, 24 Jun 2025 07:01:11 -0700
Subject: [PATCH 03/14] Remove prints

---
 .../batch_manager/kvCacheManagerTest.cpp      | 115 +-----------------
 1 file changed, 2 insertions(+), 113 deletions(-)

diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
index de28b485d2b8..af3e28ff347b 100644
--- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
+++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
@@ -1067,56 +1067,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     SizeType32 constexpr maxNewTokens{0};
     tr::SamplingConfig const samplingConfig{beamWidth};
     bool constexpr isStreaming{false};
-    
-    // Helper function to print block information
-    auto printBlockInfo = [&](std::string const& prefix, KVCacheBlock::IdType blockId) {
-        auto const& block = blockManager.getBlockById(blockId, maxAttentionWindow);
-        auto const& blockKey = block->getBlockKey();
-        
-        std::cout << prefix << " Block " << blockId << ":" << std::endl;
-        std::cout << "  Hash: 0x" << std::hex << block->getHash() << std::dec << std::endl;
-        std::cout << "  Tokens: [";
-        for (size_t i = 0; i < blockKey.uniqueTokens.size(); ++i) {
-            if (i > 0) std::cout << ", ";
-            std::cout << blockKey.uniqueTokens[i].tokenId;
-        }
-        std::cout << "]" << std::endl;
-        std::cout << "  UsesExtraIds: " << (blockKey.usesExtraIds ? "true" : "false") << std::endl;
-        if (blockKey.loraTaskId) {
-            std::cout << "  LoraTaskId: " << *blockKey.loraTaskId << std::endl;
-        }
-        if (blockKey.extraKeys) {
-            std::cout << "  ExtraKeys (multimodal): " << blockKey.extraKeys->size() << " items" << std::endl;
-            for (size_t i = 0; i < blockKey.extraKeys->size(); ++i) {
-                auto const& [mmHash, startOffset] = (*blockKey.extraKeys)[i];
-                std::cout << "    Item " << i << ": offset=" << startOffset << ", hash=[";
-                for (size_t j = 0; j < 8; ++j) { // Print first 8 bytes as hex
-                    if (j > 0) std::cout << " ";
-                    std::cout << std::hex << std::setw(2) << std::setfill('0') 
-                             << static_cast<int>(mmHash[j * 4]) << std::dec;
-                }
-                std::cout << "...]" << std::endl;
-            }
-        }
-        std::cout << "  IsFull: " << (block->isFull() ? "true" : "false") << std::endl;
-        std::cout << "  HasRefs: " << (block->hasRefs() ? "true" : "false") << std::endl;
-
-    };
-
-    // Helper function to print sequence block information
-    auto printSequenceBlocks = [&](std::string const& prefix, GenerationRequest const& seq) {
-        auto const& blockIds = seq.getCacheBlockIds(maxAttentionWindow).at(0);
-        std::cout << prefix << " Sequence blocks: [";
-        for (size_t i = 0; i < blockIds.size(); ++i) {
-            if (i > 0) std::cout << ", ";
-            std::cout << blockIds[i];
-        }
-        std::cout << "]" << std::endl;
-        
-        for (auto blockId : blockIds) {
-            printBlockInfo(prefix + "  ", blockId);
-        }
-    };
 
     // Create multimodal hash data (256-bit hash = 8 int32 values)
     auto multimodalHashes = std::make_shared<std::vector<std::vector<SizeType32>>>(
@@ -1125,18 +1075,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
         });
     auto multimodalPositions = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2});  // Start at token 2
     auto multimodalLengths = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{4});    // Length 4 tokens
-    {
-        std::cout << "\n=== Multimodal Hash Data ===" << std::endl;
-        std::cout << "Hash: [";
-        for (size_t i = 0; i < (*multimodalHashes)[0].size(); ++i) {
-            if (i > 0) std::cout << ", ";
-            std::cout << std::hex << "0x" << (*multimodalHashes)[0][i] << std::dec;
-        }
-        std::cout << "]" << std::endl;
-        std::cout << "Position: " << (*multimodalPositions)[0] << std::endl;
-        std::cout << "Length: " << (*multimodalLengths)[0] << std::endl;
-    }
- 
     // assume prompt id starts from 100
     auto inputTokens = std::make_shared<VecTokens>(VecTokens{100, 101, 102, 103, 104, 105, 0, 1, 2});
     auto const inputLength = static_cast<SizeType32>(inputTokens->size());
@@ -1155,15 +1093,9 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     auto constexpr beamIdx = 0;
     auto promptLen0 = llmRequest0->getNumTokens(beamIdx);
     auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock());
-    {
-        std::cout << "Adding sequence with " << promptLen0 << " tokens, " << numContextBlocks0 << " context blocks" << std::endl;
-    }
     blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow);
     EXPECT_EQ(llmRequest0->getContextCurrentPosition(), 0);
     EXPECT_THAT(seq0.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 1, 2}));
-    {
-        printSequenceBlocks("After addSequence", seq0);
-    }
     llmRequest0->addNewToken(3, beamIdx);
     llmRequest0->addNewToken(4, beamIdx);
     auto numTokens = llmRequest0->getNumTokens(beamIdx);
@@ -1178,7 +1110,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     // Block 0: [100, 101, 102, 103] ← Contains multimodal (102, 103)
     // Block 1: [104, 105, 0, 1]     ← Contains multimodal (104, 105)  
     // Block 2: [2, 3, 4]            ← No multimodal
-    std::cout << "\nReleasing blocks for reuse..." << std::endl;
     blockManager.releaseBlocks(seq0, llmRequest0);
     EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
@@ -1186,7 +1117,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     ///////////////////////////////////////////////////////////////////////////
     // new request with same tokens and same multimodal hash - should reuse
     requestId = 1;
-    std::cout << "\n=== Second Request (seq1) - Same Context ===" << std::endl;
     auto llmRequest1 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
         std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
         multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
@@ -1198,49 +1128,14 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     // should reuse blocks 0, 1 and get new block 3
     auto promptLen1 = llmRequest1->getNumTokens(beamIdx);
     auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock());
-    {
-        std::cout << "Adding sequence with " << promptLen1 << " tokens, " << numContextBlocks1 << " context blocks" << std::endl;
-        std::cout << "Expected: reuse blocks 0, 1 and get new block 3" << std::endl;
-    }
     blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow);
     EXPECT_EQ(llmRequest1->getContextCurrentPosition(), 2 * tokensPerBlock);
     EXPECT_THAT(seq1.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 1, 3}));
-    {
-        printSequenceBlocks("After addSequence (reuse)", seq1);
-        // Verify reuse statistics
-        std::cout << "\n=== Reuse Statistics ===" << std::endl;
-        std::cout << "Total allocated blocks: " << blockManager.getNumAllocTotalBlocks() << std::endl;
-        std::cout << "New blocks allocated: " << blockManager.getNumAllocNewBlocks() << std::endl;
-        std::cout << "Blocks reused: " << blockManager.getNumReusedBlocks() << std::endl;
-        std::cout << "Blocks missed: " << blockManager.getNumMissedBlocks() << std::endl;
-        
-        // Verify that blocks 0 and 1 were reused (same hash)
-        auto const& block0 = blockManager.getBlockById(0, maxAttentionWindow);
-        auto const& block1 = blockManager.getBlockById(1, maxAttentionWindow);
-        auto const& block3 = blockManager.getBlockById(3, maxAttentionWindow);
-        
-        std::cout << "\n=== Block Hash Verification ===" << std::endl;
-        std::cout << "Block 0 hash: 0x" << std::hex << block0->getHash() << std::dec << std::endl;
-        std::cout << "Block 1 hash: 0x" << std::hex << block1->getHash() << std::dec << std::endl;
-        std::cout << "Block 3 hash: 0x" << std::hex << block3->getHash() << std::dec << std::endl;
-        
-        // Verify that blocks 0 and 1 have the same content (should be reused)
-        auto const& blockKey0 = block0->getBlockKey();
-        auto const& blockKey1 = block1->getBlockKey();
-        auto const& blockKey3 = block3->getBlockKey();
-        
-        // Verify multimodal data is present in blocks 0 and 1
-        std::cout << "\n=== Multimodal Data Verification ===" << std::endl;
-        std::cout << "Block 0 has multimodal data: " << (blockKey0.extraKeys && !blockKey0.extraKeys->empty() ? "YES" : "NO") << std::endl;
-        std::cout << "Block 1 has multimodal data: " << (blockKey1.extraKeys && !blockKey1.extraKeys->empty() ? "YES" : "NO") << std::endl;
-        std::cout << "Block 3 has multimodal data: " << (blockKey3.extraKeys && !blockKey3.extraKeys->empty() ? "YES" : "NO") << std::endl;
-    }
     llmRequest1->addNewToken(3, beamIdx);
     llmRequest1->addNewToken(4, beamIdx);
     EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks);
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks);
     // block 3 matches block 2 and will be freed
-    std::cout << "\nReleasing blocks for reuse..." << std::endl;
     blockManager.releaseBlocks(seq1, llmRequest1);
     EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
@@ -1262,15 +1157,12 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
         LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
 
     GenerationRequest seq2{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
-    // no reuse, get new blocks 5, 6, 7
+    // no reuse, get new blocks 4, 5, 6
     auto promptLen2 = llmRequest2->getNumTokens(beamIdx);
     auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock());
     blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow);
     EXPECT_EQ(llmRequest2->getContextCurrentPosition(), 0);
     EXPECT_THAT(seq2.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({4, 5, 6}));
-    {
-        printSequenceBlocks("Add seq2", seq2);
-    }
     llmRequest2->addNewToken(9, beamIdx);
     numTokens = llmRequest2->getNumTokens(beamIdx);
     numBlocks = tc::ceilDiv(numTokens, tokensPerBlock);
@@ -1295,15 +1187,12 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
         false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
         LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
     GenerationRequest seq3{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
-    // reuse block 0, get new blocks 8, 9
+    // reuse block 0, get new blocks 7, 8
     auto promptLen3 = llmRequest3->getNumTokens(beamIdx);
     auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock());
     blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow);
     EXPECT_EQ(llmRequest3->getContextCurrentPosition(), tokensPerBlock); // only reuse block 0 [100, 101, 102, 103] with same hash/offset
     EXPECT_THAT(seq3.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 7, 8}));
-    {
-        printSequenceBlocks("Add seq3", seq3);
-    }
     llmRequest3->addNewToken(11, beamIdx);
     numTokens = llmRequest3->getNumTokens(beamIdx);
     numBlocks = tc::ceilDiv(numTokens, tokensPerBlock);

From 1bd023777d0f63c96ad07ea320957fb352d7cbe8 Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Tue, 24 Jun 2025 07:27:18 -0700
Subject: [PATCH 04/14] Formatting

---
 .../batch_manager/kvCacheManager.h            |  8 +--
 .../batch_manager/kvCacheManager.cpp          | 58 +++++++--------
 .../batch_manager/kvCacheManagerTest.cpp      | 72 ++++++++++---------
 .../models/modeling_multimodal_utils.py       |  6 +-
 4 files changed, 70 insertions(+), 74 deletions(-)

diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
index 55a9a1e193a7..9d04544570ea 100644
--- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
+++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
@@ -31,6 +31,7 @@
 #include "tensorrt_llm/runtime/worldConfig.h"
 #include <NvInferRuntime.h>
 
+#include <array>
 #include <cstdint>
 #include <limits>
 #include <list>
@@ -39,7 +40,6 @@
 #include <set>
 #include <unordered_map>
 #include <vector>
-#include <array>
 
 namespace tensorrt_llm::batch_manager::eviction_policy
 {
@@ -110,7 +110,7 @@ struct BlockKey
     bool usesExtraIds = false;
     std::optional<LoraTaskIdType> loraTaskId = std::nullopt;
     VecUniqueTokens uniqueTokens;
-    
+
     // Extra keys for multimodal data (similar to VLLM's approach)
     // Each extra key is a pair of (mm_hash, start_offset_in_block)
     std::optional<std::vector<MmKey>> extraKeys = std::nullopt;
@@ -138,8 +138,8 @@ struct BlockKey
 
     bool operator==(BlockKey const& other) const noexcept
     {
-        return (
-            usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId && uniqueTokens == other.uniqueTokens && extraKeys == other.extraKeys);
+        return (usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId
+            && uniqueTokens == other.uniqueTokens && extraKeys == other.extraKeys);
     }
 
     int partialMatch(BlockKey const& other) const noexcept
diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index 967eaa544800..1e5bc3c2ff6c 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -77,39 +77,36 @@ std::list<std::vector<T>> chopVectorIntoBlocks(
 }
 
 std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
-    tensorrt_llm::batch_manager::LlmRequest const& llmRequest,
-    SizeType32 startTokenIdx, 
-    SizeType32 endTokenIdx)
+    tensorrt_llm::batch_manager::LlmRequest const& llmRequest, SizeType32 startTokenIdx, SizeType32 endTokenIdx)
 {
     auto const multimodalHashes = llmRequest.getMultimodalHashes();
     auto const multimodalPositions = llmRequest.getMultimodalPositions();
     auto const multimodalLengths = llmRequest.getMultimodalLengths();
-    
-    if (!multimodalHashes || !multimodalPositions || !multimodalLengths || 
-        !(*multimodalHashes) || (*multimodalHashes)->empty() || 
-        !(*multimodalPositions) || (*multimodalPositions)->empty() || 
-        !(*multimodalLengths) || (*multimodalLengths)->empty())
+
+    if (!multimodalHashes || !multimodalPositions || !multimodalLengths || !(*multimodalHashes)
+        || (*multimodalHashes)->empty() || !(*multimodalPositions) || (*multimodalPositions)->empty()
+        || !(*multimodalLengths) || (*multimodalLengths)->empty())
     {
         return std::nullopt;
     }
-    
-    if ((*multimodalHashes)->size() != (*multimodalPositions)->size() || 
-        (*multimodalPositions)->size() != (*multimodalLengths)->size())
+
+    if ((*multimodalHashes)->size() != (*multimodalPositions)->size()
+        || (*multimodalPositions)->size() != (*multimodalLengths)->size())
     {
         TLLM_LOG_WARNING("Multimodal data arrays have mismatched sizes");
         return std::nullopt;
     }
-    
-    std::vector<MmKey> extraKeys;  // MmKey = std::pair<std::array<uint8_t, 32>, SizeType32>
-    
+
+    std::vector<MmKey> extraKeys; // MmKey = std::pair<std::array<uint8_t, 32>, SizeType32>
+
     for (size_t i = 0; i < (*multimodalPositions)->size(); ++i)
     {
         auto const& startPos = (*(*multimodalPositions))[i];
         auto const& length = (*(*multimodalLengths))[i];
-        auto const& mmHashVector = (*(*multimodalHashes))[i];  // This is vector<int32> - your current format
-        
+        auto const& mmHashVector = (*(*multimodalHashes))[i]; // This is vector<int32> - your current format
+
         std::array<uint8_t, 32> mmHashArray;
-        if (mmHashVector.size() == 8)  // 256-bit hash = 8 * 32-bit integers
+        if (mmHashVector.size() == 8) // 256-bit hash = 8 * 32-bit integers
         {
             // Assuming mmHashVector[j] comes from Python's int(hex_chunk, 16)
             // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB.
@@ -129,20 +126,20 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
         {
             // TODO: maybe we should raise an error here
             TLLM_LOG_WARNING("Multimodal hash vector has unexpected size: %zu (expected 8)", mmHashVector.size());
-            continue;  // Skip this multimodal item
+            continue; // Skip this multimodal item
         }
-        
+
         // Check if this multimodal content overlaps with the current block
         if (endTokenIdx > startPos && startTokenIdx < startPos + length)
         {
             // Calculate the start offset of this multimodal content within the block
             SizeType32 mmStartInBlock = (startPos >= startTokenIdx) ? 0 : startTokenIdx - startPos;
-            
+
             // Add the multimodal hash array and its start offset in the block
             extraKeys.emplace_back(mmHashArray, mmStartInBlock);
         }
     }
-    
+
     return extraKeys.empty() ? std::nullopt : std::make_optional(std::move(extraKeys));
 }
 
@@ -150,18 +147,15 @@ std::vector<BlockKey> buildBlockKeys(
     std::list<VecUniqueTokens>& blockedUniqueTokens, tensorrt_llm::batch_manager::LlmRequest const& llmRequest)
 {
     std::vector<BlockKey> blockKeys;
-    
+
     SizeType32 currentTokenIdx = 0;
     for (auto& uniqueTokens : blockedUniqueTokens)
     {
         auto extraKeys = generateBlockHashExtraKeys(llmRequest, currentTokenIdx, currentTokenIdx + uniqueTokens.size());
         currentTokenIdx += uniqueTokens.size();
 
-        blockKeys.emplace_back(
-            llmRequest.getInputTokensExtraIds().has_value(), 
-            llmRequest.getLoraTaskId(), 
-            std::move(uniqueTokens),
-            std::move(extraKeys));
+        blockKeys.emplace_back(llmRequest.getInputTokensExtraIds().has_value(), llmRequest.getLoraTaskId(),
+            std::move(uniqueTokens), std::move(extraKeys));
     }
     return blockKeys;
 }
@@ -209,18 +203,16 @@ size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) no
             for (size_t i = 0; i < 32; i += 4)
             {
                 // Combine 4 bytes into a 32-bit word (construct as little endian order)
-                uint32_t word = static_cast<uint32_t>(mmHash[i]) |
-                                (static_cast<uint32_t>(mmHash[i+1]) << 8) |
-                                (static_cast<uint32_t>(mmHash[i+2]) << 16) |
-                                (static_cast<uint32_t>(mmHash[i+3]) << 24);
-                
+                uint32_t word = static_cast<uint32_t>(mmHash[i]) | (static_cast<uint32_t>(mmHash[i + 1]) << 8)
+                    | (static_cast<uint32_t>(mmHash[i + 2]) << 16) | (static_cast<uint32_t>(mmHash[i + 3]) << 24);
+
                 // Mix the word into the seed
                 word = ((word >> 16) ^ word) * 0x45d9f3b;
                 word = ((word >> 16) ^ word) * 0x45d9f3b;
                 word = (word >> 16) ^ word;
                 seed ^= word + 0x9e3779b9 + (seed << 6) + (seed >> 2);
             }
-            
+
             // Hash the start offset
             uint64_t e = static_cast<uint64_t>(startOffset);
             e = (e ^ (e >> 30)) * UINT64_C(0xbf58476d1ce4e5b9);
diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
index af3e28ff347b..ba10a17b26db 100644
--- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
+++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
@@ -1069,12 +1069,12 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     bool constexpr isStreaming{false};
 
     // Create multimodal hash data (256-bit hash = 8 int32 values)
-    auto multimodalHashes = std::make_shared<std::vector<std::vector<SizeType32>>>(
-        std::vector<std::vector<SizeType32>>{
-            {0x12345678, -0x6F543211, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666}  // Hash 1
-        });
-    auto multimodalPositions = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2});  // Start at token 2
-    auto multimodalLengths = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{4});    // Length 4 tokens
+    auto multimodalHashes = std::make_shared<std::vector<std::vector<SizeType32>>>(std::vector<std::vector<SizeType32>>{
+        {0x12345678, -0x6F543211, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666} // Hash 1
+    });
+    auto multimodalPositions
+        = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2});                    // Start at token 2
+    auto multimodalLengths = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{4}); // Length 4 tokens
     // assume prompt id starts from 100
     auto inputTokens = std::make_shared<VecTokens>(VecTokens{100, 101, 102, 103, 104, 105, 0, 1, 2});
     auto const inputLength = static_cast<SizeType32>(inputTokens->size());
@@ -1082,9 +1082,9 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     auto llmRequest0 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
         std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
         multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
-        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, 
-        false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
-        LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt,
+        std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt,
+        std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
 
     GenerationRequest seq0{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
 
@@ -1108,7 +1108,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     // Multimodal: starts at token 2, length 4 → [102, 103, 104, 105]
 
     // Block 0: [100, 101, 102, 103] ← Contains multimodal (102, 103)
-    // Block 1: [104, 105, 0, 1]     ← Contains multimodal (104, 105)  
+    // Block 1: [104, 105, 0, 1]     ← Contains multimodal (104, 105)
     // Block 2: [2, 3, 4]            ← No multimodal
     blockManager.releaseBlocks(seq0, llmRequest0);
     EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
@@ -1120,9 +1120,9 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     auto llmRequest1 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
         std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
         multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, std::nullopt,
-        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, 
-        false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
-        LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt,
+        std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt,
+        std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
     GenerationRequest seq1{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
 
     // should reuse blocks 0, 1 and get new block 3
@@ -1141,20 +1141,21 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
 
     ///////////////////////////////////////////////////////////////////////////
-    // Test Case 2: Different multimodal hash 
+    // Test Case 2: Different multimodal hash
     requestId = 2;
-    auto multimodalHashes2 = std::make_shared<std::vector<std::vector<SizeType32>>>(
-        std::vector<std::vector<SizeType32>>{
+    auto multimodalHashes2
+        = std::make_shared<std::vector<std::vector<SizeType32>>>(std::vector<std::vector<SizeType32>>{
             {0x45678123, 0x23456789, 0x34567890, 0x12121212, 0x56565656, 0x78787878, 0x54545454, 0x67676767} // Hash 2
         });
-    auto multimodalPositions2 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2});  // Start at token 2
-    auto multimodalLengths2 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{4});    // Length 4 tokens
+    auto multimodalPositions2
+        = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2});                     // Start at token 2
+    auto multimodalLengths2 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{4}); // Length 4 tokens
     auto llmRequest2 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
         std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
         multimodalHashes2, multimodalPositions2, multimodalLengths2, std::nullopt, std::nullopt, std::nullopt,
-        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, 
-        false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
-        LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt,
+        std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt,
+        std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
 
     GenerationRequest seq2{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
     // no reuse, get new blocks 4, 5, 6
@@ -1172,26 +1173,29 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     ///////////////////////////////////////////////////////////////////////////
     // Test Case 3: Multiple multimodal hashes and partial reuse
     requestId = 3;
-    auto multimodalHashes3 = std::make_shared<std::vector<std::vector<SizeType32>>>(
-        std::vector<std::vector<SizeType32>>{
-            {0x12345678, -0x6F543211, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666},  // Hash 1
-            {0x45678123, 0x23456789, 0x34567890, 0x12121212, 0x56565656, 0x78787878, 0x54545454, 0x67676767} // Hash 2
+    auto multimodalHashes3
+        = std::make_shared<std::vector<std::vector<SizeType32>>>(std::vector<std::vector<SizeType32>>{
+            {0x12345678, -0x6F543211, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666}, // Hash 1
+            {0x45678123, 0x23456789, 0x34567890, 0x12121212, 0x56565656, 0x78787878, 0x54545454, 0x67676767}   // Hash 2
         });
-    auto multimodalPositions3 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2, 4});  // Start at token 2 and 4
-    auto multimodalLengths3 = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2, 2});    // Length 2 tokens
-    
+    auto multimodalPositions3
+        = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2, 4}); // Start at token 2 and 4
+    auto multimodalLengths3
+        = std::make_shared<std::vector<SizeType32>>(std::vector<SizeType32>{2, 2}); // Length 2 tokens
+
     auto llmRequest3 = std::make_shared<LlmRequest>(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming,
         std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
         multimodalHashes3, multimodalPositions3, multimodalLengths3, std::nullopt, std::nullopt, std::nullopt,
-        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, 
-        false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt,
-        LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
+        std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt,
+        std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt,
+        std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences);
     GenerationRequest seq3{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()};
     // reuse block 0, get new blocks 7, 8
     auto promptLen3 = llmRequest3->getNumTokens(beamIdx);
     auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock());
     blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow);
-    EXPECT_EQ(llmRequest3->getContextCurrentPosition(), tokensPerBlock); // only reuse block 0 [100, 101, 102, 103] with same hash/offset
+    EXPECT_EQ(llmRequest3->getContextCurrentPosition(),
+        tokensPerBlock); // only reuse block 0 [100, 101, 102, 103] with same hash/offset
     EXPECT_THAT(seq3.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 7, 8}));
     llmRequest3->addNewToken(11, beamIdx);
     numTokens = llmRequest3->getNumTokens(beamIdx);
@@ -1204,7 +1208,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest)
     blockManager.releaseBlocks(seq3, llmRequest3);
     EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0);
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
-
 }
 
 TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest)
@@ -2634,6 +2637,7 @@ TEST_P(KVCacheManagerTest, KVCacheManagerRewindTokensTest)
         EXPECT_EQ(blockManager.getNumFreeBlocks(), currentNumBlocks);
     }
 }
+
 TEST_P(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowTest)
 {
     using DType = half;
@@ -4697,5 +4701,3 @@ auto const paramValues = ::testing::Values(
     });
 
 INSTANTIATE_TEST_SUITE_P(FillKvCacheAndCompleteRequestsTest, FillKvCacheAndCompleteRequestsTest, paramValues);
-
-
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
index 62be3c7b890a..9344b5e63253 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
@@ -85,9 +85,11 @@ def fuse_input_embeds(
         if mm_token_indices.numel() < mm_embed.shape[0]:
             # Extract the last mm_token_indices.numel() embeddings from mm_embed
             start_idx = mm_embed.shape[0] - mm_token_indices.numel()
-            input_embeds[mm_token_indices, :] = mm_embed[start_idx:, :].to(dtype=input_embeds.dtype, device=input_embeds.device)
+            input_embeds[mm_token_indices, :] = mm_embed[start_idx:, :].to(
+                dtype=input_embeds.dtype, device=input_embeds.device)
         else:
-            input_embeds[mm_token_indices, :] = mm_embed.to(dtype=input_embeds.dtype, device=input_embeds.device)
+            input_embeds[mm_token_indices, :] = mm_embed.to(
+                dtype=input_embeds.dtype, device=input_embeds.device)
 
     return None, input_embeds
 

From 140e3069a435b03012b24967c4c5560b6098d2e4 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Tue, 24 Jun 2025 07:43:28 -0700
Subject: [PATCH 05/14] Minor update on comments

---
 .../batch_manager/kvCacheManager.cpp          | 40 +++++++------------
 examples/llm-api/quickstart_multimodal.py     |  1 +
 2 files changed, 16 insertions(+), 25 deletions(-)

diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index 1e5bc3c2ff6c..b49a209c8a33 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -103,39 +103,29 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
     {
         auto const& startPos = (*(*multimodalPositions))[i];
         auto const& length = (*(*multimodalLengths))[i];
-        auto const& mmHashVector = (*(*multimodalHashes))[i]; // This is vector<int32> - your current format
+        auto const& mmHashVector = (*(*multimodalHashes))[i];
 
         std::array<uint8_t, 32> mmHashArray;
-        if (mmHashVector.size() == 8) // 256-bit hash = 8 * 32-bit integers
+        TLLM_CHECK_WITH_INFO(mmHashVector.size() == 8, "Multimodal hash vector has unexpected size: %zu (expected 8)", mmHashVector.size());
+        
+        // mmHashVector[j] comes from Python's int(hex_chunk, 16)
+        // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB (big endian)
+        // The overall Blake3 output wants these bytes in order: 0x00, 0x01, 0x02, 0x03...
+        for (size_t j = 0; j < 8; ++j)
         {
-            // Assuming mmHashVector[j] comes from Python's int(hex_chunk, 16)
-            // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB.
-            // And assuming the overall Blake3 output wants these bytes in order: 0x00, 0x01, 0x02, 0x03...
-            for (size_t j = 0; j < 8; ++j)
-            {
-                auto const& hashPart = mmHashVector[j]; // e.g., 0x00010203
-                // Extract bytes in Big-Endian order from hashPart and place them sequentially
-                // into mmHashArray to reconstruct the original Blake3 byte sequence.
-                mmHashArray[j * 4 + 0] = static_cast<uint8_t>((hashPart >> 24) & 0xFF); // Extract 0x00 (MSB of the int)
-                mmHashArray[j * 4 + 1] = static_cast<uint8_t>((hashPart >> 16) & 0xFF); // Extract 0x01
-                mmHashArray[j * 4 + 2] = static_cast<uint8_t>((hashPart >> 8) & 0xFF);  // Extract 0x02
-                mmHashArray[j * 4 + 3] = static_cast<uint8_t>(hashPart & 0xFF);         // Extract 0x03 (LSB of the int)
-            }
-        }
-        else
-        {
-            // TODO: maybe we should raise an error here
-            TLLM_LOG_WARNING("Multimodal hash vector has unexpected size: %zu (expected 8)", mmHashVector.size());
-            continue; // Skip this multimodal item
+            auto const& hashPart = mmHashVector[j]; // e.g., 0x00010203
+            // Extract bytes in Big-Endian order from hashPart and place them sequentially into mmHashArray
+            mmHashArray[j * 4 + 0] = static_cast<uint8_t>((hashPart >> 24) & 0xFF); // Extract 0x00 (MSB of the int)
+            mmHashArray[j * 4 + 1] = static_cast<uint8_t>((hashPart >> 16) & 0xFF); // Extract 0x01
+            mmHashArray[j * 4 + 2] = static_cast<uint8_t>((hashPart >> 8) & 0xFF);  // Extract 0x02
+            mmHashArray[j * 4 + 3] = static_cast<uint8_t>(hashPart & 0xFF);         // Extract 0x03 (LSB of the int)
         }
+       
 
         // Check if this multimodal content overlaps with the current block
         if (endTokenIdx > startPos && startTokenIdx < startPos + length)
         {
-            // Calculate the start offset of this multimodal content within the block
             SizeType32 mmStartInBlock = (startPos >= startTokenIdx) ? 0 : startTokenIdx - startPos;
-
-            // Add the multimodal hash array and its start offset in the block
             extraKeys.emplace_back(mmHashArray, mmStartInBlock);
         }
     }
@@ -194,7 +184,7 @@ size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) no
         seed ^= c + 0x9e3779b9 + (seed << 6) + (seed >> 2);
     }
 
-    // Add extra keys for multimodal data (similar to VLLM's approach)
+    // Add extra keys for multimodal data mixing in external mulitmodal item hash and token offset within this sequence block
     if (blockKey.extraKeys)
     {
         for (auto const& [mmHash, startOffset] : *blockKey.extraKeys)
diff --git a/examples/llm-api/quickstart_multimodal.py b/examples/llm-api/quickstart_multimodal.py
index 31ed52b43aa9..08bb6f80ba66 100644
--- a/examples/llm-api/quickstart_multimodal.py
+++ b/examples/llm-api/quickstart_multimodal.py
@@ -105,6 +105,7 @@ def parse_arguments():
     parser = add_lora_args(parser)
     args = parser.parse_args()
 
+    args.disable_kv_cache_reuse = True # kv cache reuse does not work for multimodal, force overwrite
     if args.kv_cache_fraction is None:
         args.kv_cache_fraction = 0.6  # lower the default kv cache fraction for multimodal
 

From ed9249fe493adb2c1499b5e108e4a25af54432ff Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Tue, 24 Jun 2025 08:30:05 -0700
Subject: [PATCH 06/14] Formatting

---
 cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 9 +++++----
 examples/llm-api/quickstart_multimodal.py         | 2 +-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index b49a209c8a33..8ece83afe2ef 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -106,8 +106,9 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
         auto const& mmHashVector = (*(*multimodalHashes))[i];
 
         std::array<uint8_t, 32> mmHashArray;
-        TLLM_CHECK_WITH_INFO(mmHashVector.size() == 8, "Multimodal hash vector has unexpected size: %zu (expected 8)", mmHashVector.size());
-        
+        TLLM_CHECK_WITH_INFO(mmHashVector.size() == 8, "Multimodal hash vector has unexpected size: %zu (expected 8)",
+            mmHashVector.size());
+
         // mmHashVector[j] comes from Python's int(hex_chunk, 16)
         // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB (big endian)
         // The overall Blake3 output wants these bytes in order: 0x00, 0x01, 0x02, 0x03...
@@ -120,7 +121,6 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
             mmHashArray[j * 4 + 2] = static_cast<uint8_t>((hashPart >> 8) & 0xFF);  // Extract 0x02
             mmHashArray[j * 4 + 3] = static_cast<uint8_t>(hashPart & 0xFF);         // Extract 0x03 (LSB of the int)
         }
-       
 
         // Check if this multimodal content overlaps with the current block
         if (endTokenIdx > startPos && startTokenIdx < startPos + length)
@@ -184,7 +184,8 @@ size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) no
         seed ^= c + 0x9e3779b9 + (seed << 6) + (seed >> 2);
     }
 
-    // Add extra keys for multimodal data mixing in external mulitmodal item hash and token offset within this sequence block
+    // Add extra keys for multimodal data mixing in external mulitmodal item hash and token offset within this sequence
+    // block
     if (blockKey.extraKeys)
     {
         for (auto const& [mmHash, startOffset] : *blockKey.extraKeys)
diff --git a/examples/llm-api/quickstart_multimodal.py b/examples/llm-api/quickstart_multimodal.py
index 08bb6f80ba66..967a8636e1be 100644
--- a/examples/llm-api/quickstart_multimodal.py
+++ b/examples/llm-api/quickstart_multimodal.py
@@ -105,7 +105,7 @@ def parse_arguments():
     parser = add_lora_args(parser)
     args = parser.parse_args()
 
-    args.disable_kv_cache_reuse = True # kv cache reuse does not work for multimodal, force overwrite
+    args.disable_kv_cache_reuse = True  # kv cache reuse does not work for multimodal, force overwrite
     if args.kv_cache_fraction is None:
         args.kv_cache_fraction = 0.6  # lower the default kv cache fraction for multimodal
 

From 4c153376e54085aec20014357d13ef55e8329e10 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Tue, 15 Jul 2025 17:36:25 -0700
Subject: [PATCH 07/14] Address comment

---
 .../batch_manager/kvCacheManager.cpp            | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index 8ece83afe2ef..f113984a3057 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -111,16 +111,19 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
 
         // mmHashVector[j] comes from Python's int(hex_chunk, 16)
         // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB (big endian)
-        // The overall Blake3 output wants these bytes in order: 0x00, 0x01, 0x02, 0x03...
+        // Convert 8x 32-bit integers into a 32-byte array preserving Blake3 hash byte order
+        // Example: hashPart = 0x00010203 → mmHashArray[0:3] = [0x00, 0x01, 0x02, 0x03]
+        #define GET_NTH_BYTE(hash_part, byte_idx) static_cast<uint8_t>((hash_part >> (24 - (byte_idx) * 8)) & 0xFF)
+
         for (size_t j = 0; j < 8; ++j)
         {
-            auto const& hashPart = mmHashVector[j]; // e.g., 0x00010203
-            // Extract bytes in Big-Endian order from hashPart and place them sequentially into mmHashArray
-            mmHashArray[j * 4 + 0] = static_cast<uint8_t>((hashPart >> 24) & 0xFF); // Extract 0x00 (MSB of the int)
-            mmHashArray[j * 4 + 1] = static_cast<uint8_t>((hashPart >> 16) & 0xFF); // Extract 0x01
-            mmHashArray[j * 4 + 2] = static_cast<uint8_t>((hashPart >> 8) & 0xFF);  // Extract 0x02
-            mmHashArray[j * 4 + 3] = static_cast<uint8_t>(hashPart & 0xFF);         // Extract 0x03 (LSB of the int)
+            auto const& hashPart = mmHashVector[j];
+            for (size_t byteIdx = 0; byteIdx < 4; ++byteIdx)
+            {
+                mmHashArray[j * 4 + byteIdx] = GET_NTH_BYTE(hashPart, byteIdx);
+            }
         }
+        #undef GET_NTH_BYTE
 
         // Check if this multimodal content overlaps with the current block
         if (endTokenIdx > startPos && startTokenIdx < startPos + length)

From 3f6681a6eae4032d4ebf16c3ba29815d310408cc Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Tue, 15 Jul 2025 17:44:24 -0700
Subject: [PATCH 08/14] Improve extract uncached mm embed logic

---
 .../models/modeling_multimodal_utils.py       | 101 +++++++-
 .../_torch/models/modeling_qwen2vl.py         |   3 +-
 .../_torch/pyexecutor/model_engine.py         |  11 +-
 tensorrt_llm/inputs/multimodal.py             |  55 ++++
 .../_torch/multimodal/test_kvcache_reuse.py   | 239 ++++++++++++++++++
 5 files changed, 394 insertions(+), 15 deletions(-)
 create mode 100644 tests/unittest/_torch/multimodal/test_kvcache_reuse.py

diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
index 9344b5e63253..fa0ffff37f66 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
@@ -26,7 +26,87 @@
 from torchvision.transforms import Normalize, Resize, ToTensor
 
 from tensorrt_llm._torch.modules.embedding import Embedding
+from tensorrt_llm.inputs.multimodal import MultimodalParams
+from tensorrt_llm.logger import logger
+import numpy as np
 
+def find_uncached_mm_embeds(mm_embeds: List[torch.Tensor], multimodal_params: List[MultimodalParams]) -> torch.Tensor:
+    """
+    Find the uncached multimodal mm_embeds from multimodal_params for each batch.
+    Args:
+        - mm_embeds: List[torch.Tensor]
+        - multimodal_params: List[MultimodalParams]
+    Returns:
+        - sliced_mm_embeds: List[torch.Tensor]
+          When kv_cache reuse is disabled or model not enabled/support kv_cache reuse, return the full mm_embeds.
+    Note:
+        - Current implementation assumes chunk prefill is disabled. To support chunk prefill, we might need to slightly modify the logic (see TODO below).
+    """
+    # Current support two batching modes:
+    # 1. Pre-concatenated mm_embeds for each batch, i.e., len(mm_embeds) == 1
+    # 2. Individual mm_embeds for each multimodal param, i.e., len(mm_embeds) == len(multimodal_params)
+    if len(mm_embeds) > 1 and len(mm_embeds) != len(multimodal_params):
+        raise ValueError(f"Number of mm_embeds ({len(mm_embeds)}) does not match number of multimodal params ({len(multimodal_params)}).")
+
+    if not multimodal_params or multimodal_params[0].multimodal_runtime is None:
+        # No slicing, return the full mm_embeds
+        return mm_embeds
+
+    total_cached_mm_tokens = sum([
+        param.multimodal_runtime.num_cached_mm_tokens
+        for param in multimodal_params
+    ])
+    if total_cached_mm_tokens == 0:
+        # No cached tokens, return the full mm_embeds
+        # TODO: support chunk prefill for multimodal, then we need to extract full mm_embeds for each CHUNK
+        logger.debug("No multimodal cached tokens can be reused, return the full mm_embeds")
+        return mm_embeds
+
+    if total_cached_mm_tokens == sum([
+        param.multimodal_runtime.total_mm_tokens
+        for param in multimodal_params
+    ]):
+        # All tokens are cached, return empty list
+        logger.debug("All multimodal tokens cached, skipping vision encoder forward")
+        return []
+
+    # Partial caching, return the sliced mm_embeds
+    prefix_sum = np.cumsum([param.multimodal_runtime.total_mm_tokens for param in multimodal_params])
+    current_pos = 0
+    slices = []
+    for param in multimodal_params:
+        runtime = param.multimodal_runtime
+        if runtime.num_cached_mm_tokens > 0:
+            if runtime.num_cached_mm_tokens == runtime.total_mm_tokens:
+                if len(mm_embeds) == 1: # pre-concatenated mm_embeds, need global indices
+                    current_pos += runtime.total_mm_tokens
+                slices.append((current_pos, current_pos))
+                continue
+            num_uncached_mm_tokens = runtime.total_mm_tokens - runtime.num_cached_mm_tokens
+            # TODO: support chunk prefill to extract partial mm_embeds
+            slices.append((current_pos + runtime.num_cached_mm_tokens, current_pos + runtime.total_mm_tokens))
+        else: # == 0, no cached tokens
+            # All mm_embeds in this sequence are sliced out
+            # TODO: support chunk prefill to extract partial mm_embeds
+            slices.append((current_pos, current_pos + runtime.total_mm_tokens))
+
+
+        if len(mm_embeds) == 1: # pre-concatenated mm_embeds, need global indices
+            current_pos += runtime.total_mm_tokens
+
+    sliced_mm_embeds = []
+    if len(mm_embeds) == 1:
+        for start, end in slices:
+            sliced_mm_embeds.append(mm_embeds[0][start:end])
+    else: # slice each mm_embeds individually
+        for i, (start, end) in enumerate(slices):
+            sliced_mm_embeds.append(mm_embeds[i][start:end])
+
+    if len(mm_embeds) == 1:
+        sliced_mm_embeds = [torch.cat(sliced_mm_embeds, dim=0)]
+
+    logger.debug(f"Partial caching, return sliced_mm_embeds: {sliced_mm_embeds[0].shape}")
+    return sliced_mm_embeds
 
 def fuse_input_embeds(
     embedding_layer: Embedding,
@@ -69,6 +149,13 @@ def fuse_input_embeds(
         text_token_mask = ~mm_token_mask
     text_token_indices = torch.where(text_token_mask)[0]
     mm_token_indices = torch.where(mm_token_mask)[0]
+    if len(mm_token_indices) != mm_embed.shape[0]:
+        raise ValueError(
+            f"Multimodal token count mismatch: found {len(mm_token_indices)} image tokens in input_ids "
+            f"but received {mm_embed.shape[0]} image embeddings. "
+            f"This is likely due to KV cache reuse, chunk prefill, or other optimizations that "
+            f"cause token count mismatches within the inference batch."
+        )
 
     text_embed = embedding_layer(input_ids[text_token_indices])
     input_embeds = torch.empty(input_ids.shape[0],
@@ -78,18 +165,8 @@ def fuse_input_embeds(
 
     input_embeds[text_token_indices, :] = text_embed.to(
         dtype=input_embeds.dtype, device=input_embeds.device)
-
-    # When KV cache reuse enabled, mm_token_indices may only contain partial multimodal tokens or empty
-    if mm_token_indices.numel() > 0:
-        # We need to extract only the embeddings that correspond to uncached multimodal tokens
-        if mm_token_indices.numel() < mm_embed.shape[0]:
-            # Extract the last mm_token_indices.numel() embeddings from mm_embed
-            start_idx = mm_embed.shape[0] - mm_token_indices.numel()
-            input_embeds[mm_token_indices, :] = mm_embed[start_idx:, :].to(
-                dtype=input_embeds.dtype, device=input_embeds.device)
-        else:
-            input_embeds[mm_token_indices, :] = mm_embed.to(
-                dtype=input_embeds.dtype, device=input_embeds.device)
+    input_embeds[mm_token_indices, :] = mm_embed.to(dtype=input_embeds.dtype,
+                                                    device=input_embeds.device)
 
     return None, input_embeds
 
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 2d63a4bbf92b..889a579fc709 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -18,7 +18,7 @@
 from ..attention_backend import AttentionMetadata
 from ..model_config import ModelConfig
 from .modeling_auto import AutoModelForCausalLM
-from .modeling_multimodal_utils import fuse_input_embeds
+from .modeling_multimodal_utils import fuse_input_embeds, find_uncached_mm_embeds
 from .modeling_utils import register_auto_model
 
 DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
@@ -601,6 +601,7 @@ def forward(
             mrope_config = self._parse_and_concat_mrope_config(
                 multimodal_params, num_context_requests,
                 num_generation_requests)
+            mm_embeds = find_uncached_mm_embeds(mm_embeds, multimodal_params[:num_context_requests])
 
         if 'mrope_position_deltas' in kwargs:
             mrope_config['mrope_position_deltas'] = kwargs[
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 98eb2e870d4c..29a33b6c682e 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -21,7 +21,7 @@
 from tensorrt_llm._torch.speculative.mtp import SampleStateTensorsMTP
 from tensorrt_llm._utils import (is_trace_enabled, nvtx_range, release_gc,
                                  torch_dtype_to_str, trace_func)
-from tensorrt_llm.inputs.multimodal import MultimodalParams
+from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData
 from tensorrt_llm.logger import logger
 from tensorrt_llm.lora_manager import LoraConfig, LoraModelConfig
 from tensorrt_llm.mapping import Mapping
@@ -1143,8 +1143,15 @@ def _prepare_tp_inputs(
             num_cached_tokens_per_seq.append(past_seen_token_num)
 
             # Multimodal
+            # TODO: enable chunk prefill for multimodal (maybe need to pass prompt_tokens to MultimodalRuntimeData)
+            py_multimodal_runtime = MultimodalRuntimeData(
+                mm_token_lengths=request.multimodal_lengths,
+                mm_token_positions=request.multimodal_positions,
+                num_cached_tokens=past_seen_token_num) if request.multimodal_hashes is not None else None
+
             multimodal_params = MultimodalParams(
-                multimodal_data=request.py_multimodal_data)
+                multimodal_data=request.py_multimodal_data,
+                multimodal_runtime=py_multimodal_runtime)
             multimodal_params.to_device("multimodal_data",
                                         "cuda",
                                         pin_memory=True)
diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index a6b29a9f0183..7f9b1c1c1eed 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -81,6 +81,60 @@ def to_tensor(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             torch.tensor(self.multimodal_positions, dtype=torch.int32),
             torch.tensor(self.multimodal_lengths, dtype=torch.int32))
 
+@dataclass
+class MultimodalRuntimeData:
+    """Runtime data for tracking multimodal token caching and reuse per request sequence.
+
+    This class tracks which multimodal tokens are cached vs. need to be processed
+    for each request sequence during KV cache reuse scenarios.
+
+    Attributes:
+        num_cached_tokens: Total number of cached tokens for this sequence
+        mm_token_lengths: Length of each multimodal token chunk
+        mm_token_positions: Starting positions of each multimodal token chunk
+        prompt_tokens: Current iteration of prompt tokens for this sequence (optional). Need it for chunk prefill if enabled (#TODO)
+        num_cached_mm_tokens: Number of multimodal tokens that are cached in this iteration (computed)
+        total_mm_tokens: Total number of multimodal tokens in this sequence (computed)
+    """
+    num_cached_tokens: int
+    mm_token_lengths: List[int]
+    mm_token_positions: List[int]
+
+    # TODO: support chunk prefill for multimodal
+    # When chunk prefill is enabled, we need to pass the prompt tokens for current chunk and mask to find the included mm tokens
+    prompt_tokens: Optional[List[int]] = None
+
+    num_cached_mm_tokens: Optional[int] = None
+    total_mm_tokens: Optional[int] = None
+
+    def __post_init__(self):
+        # Validate input data
+        if len(self.mm_token_positions) != len(self.mm_token_lengths):
+            raise ValueError(f"mm_token_positions ({len(self.mm_token_positions)}) and mm_token_lengths ({len(self.mm_token_lengths)}) must have the same length")
+
+        if self.num_cached_tokens < 0:
+            raise ValueError(f"num_cached_tokens must be non-negative, got {self.num_cached_tokens}")
+
+        if any(length <= 0 for length in self.mm_token_lengths):
+            raise ValueError(f"All mm_token_lengths must be positive, got {self.mm_token_lengths}")
+
+        if any(pos < 0 for pos in self.mm_token_positions):
+            raise ValueError(f"All mm_token_positions must be non-negative, got {self.mm_token_positions}")
+
+        if self.num_cached_mm_tokens is None:
+            # Compute cached multimodal tokens based on positions and cached tokens
+            self.num_cached_mm_tokens = 0
+            for pos, length in zip(self.mm_token_positions, self.mm_token_lengths):
+                if pos + length <= self.num_cached_tokens:
+                    self.num_cached_mm_tokens += length
+                elif pos < self.num_cached_tokens:
+                    # Partial overlap - only count the cached portion
+                    self.num_cached_mm_tokens += self.num_cached_tokens - pos
+
+        assert self.num_cached_mm_tokens <= self.num_cached_tokens, \
+            f"num_cached_mm_tokens ({self.num_cached_mm_tokens}) must be less than or equal to num_cached_tokens ({self.num_cached_tokens})"
+
+        self.total_mm_tokens = sum(self.mm_token_lengths)
 
 @dataclass
 class MultimodalParams:
@@ -117,6 +171,7 @@ class MultimodalParams:
 
     multimodal_input: Optional[MultimodalInput] = None
     multimodal_data: Optional[Dict[str, Any]] = field(default_factory=dict)
+    multimodal_runtime: Optional[MultimodalRuntimeData] = None
 
     def __post_init__(self):
         """Ensure default values are properly set."""
diff --git a/tests/unittest/_torch/multimodal/test_kvcache_reuse.py b/tests/unittest/_torch/multimodal/test_kvcache_reuse.py
new file mode 100644
index 000000000000..8f2876d6f281
--- /dev/null
+++ b/tests/unittest/_torch/multimodal/test_kvcache_reuse.py
@@ -0,0 +1,239 @@
+import pytest
+import torch
+import numpy as np
+from unittest.mock import Mock
+from typing import List
+
+# Import the function to test
+from tensorrt_llm._torch.models.modeling_multimodal_utils import find_uncached_mm_embeds
+from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData
+
+
+class TestMultimodalRuntimeData:
+    """Test cases for MultimodalRuntimeData computation logic, specifically num_cached_mm_tokens."""
+
+    def test_fully_cached_multimodal_tokens(self):
+        """Test when all multimodal tokens are cached."""
+        runtime = MultimodalRuntimeData(
+            num_cached_tokens=20,
+            mm_token_lengths=[5, 8, 7],  # Total: 20 tokens
+            mm_token_positions=[0, 5, 13]  # Positions: 0-5, 5-13, 13-20
+        )
+
+        # All tokens should be cached since num_cached_tokens (20) >= all positions + lengths
+        assert runtime.num_cached_mm_tokens == 20
+        assert runtime.total_mm_tokens == 20
+
+    def test_no_cached_multimodal_tokens(self):
+        """Test when no multimodal tokens are cached."""
+        runtime = MultimodalRuntimeData(
+            num_cached_tokens=10,
+            mm_token_lengths=[5, 8, 7],  # Total: 20 tokens
+            mm_token_positions=[10, 18, 30]  # All positions > num_cached_tokens
+        )
+
+        # No multimodal tokens should be cached
+        assert runtime.num_cached_mm_tokens == 0
+        assert runtime.total_mm_tokens == 20
+
+    def test_complex_scenario_with_multiple_chunks(self):
+        """Test a complex scenario with many chunks and various caching states."""
+        runtime = MultimodalRuntimeData(
+            num_cached_tokens=30,
+            mm_token_lengths=[3, 4, 5, 6, 7, 8],  # Total: 33 tokens
+            mm_token_positions=[0, 5, 10, 15, 25, 35]  # Positions: 0-3, 5-9, 10-15, 15-21, 25-32, 35-43
+        )
+
+        # Expected caching:
+        # Chunk 0: fully cached (3 tokens)
+        # Chunk 1: fully cached (4 tokens)
+        # Chunk 2: fully cached (5 tokens)
+        # Chunk 3: fully cached (6 tokens)
+        # Chunk 4: partially cached (30-25=5 out of 7 tokens)
+        # Chunk 5: not cached
+        expected_cached = 3 + 4 + 5 + 6 + 5  # 23 tokens
+        assert runtime.num_cached_mm_tokens == expected_cached
+        assert runtime.total_mm_tokens == 33
+
+
+class TestFindUncachedMmEmbed:
+    """Focused test cases for find_uncached_mm_embeds function - testing edge cases and potential bugs."""
+
+    def create_mock_runtime(self, num_cached_mm_tokens: int, total_mm_tokens: int):
+        """Helper to create a mock MultimodalRuntimeData."""
+        runtime = Mock(spec=MultimodalRuntimeData)
+        runtime.num_cached_mm_tokens = num_cached_mm_tokens
+        runtime.total_mm_tokens = total_mm_tokens
+        return runtime
+
+    def create_multimodal_params(self, num_cached_mm_tokens: int, total_mm_tokens: int):
+        """Helper to create MultimodalParams with runtime data."""
+        runtime = self.create_mock_runtime(num_cached_mm_tokens, total_mm_tokens)
+        return MultimodalParams(multimodal_runtime=runtime)
+
+    def test_mm_embed_not_batched(self):
+        """
+        Test individual batching mode where each mm_embed corresponds to one param.
+        This tests the case where len(mm_embeds) == len(multimodal_params) > 1.
+        """
+        mm_embeds = [
+            torch.randn(10, 512),  # Batch 1: 10 tokens
+            torch.randn(15, 512),  # Batch 2: 15 tokens
+            torch.randn(8, 512)    # Batch 3: 8 tokens
+        ]
+        multimodal_params = [
+            self.create_multimodal_params(3, 10),   # 3 cached, 7 uncached
+            self.create_multimodal_params(8, 15),   # 8 cached, 7 uncached
+            self.create_multimodal_params(0, 8)     # 0 cached, 8 uncached
+        ]
+
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+
+        # Should return individual slices for each batch
+        assert len(result) == 3
+        assert result[0].shape == (7, 512)   # 10 - 3 = 7
+        assert result[1].shape == (7, 512)   # 15 - 8 = 7
+        assert result[2].shape == (8, 512)   # 8 - 0 = 8
+
+        # Verify the slices are correct
+        torch.testing.assert_close(result[0], mm_embeds[0][3:10])
+        torch.testing.assert_close(result[1], mm_embeds[1][8:15])
+        torch.testing.assert_close(result[2], mm_embeds[2][0:8])
+
+    def test_mm_embed_batched(self):
+        """
+        Test batching (concatenated) mm_embeds with fused mm_embeds for each batch.
+        This tests the case where len(mm_embeds) == 1
+        """
+        mm_embeds = [torch.randn(33, 512)]  # Pre-concatenated: 10 + 13 + 10 tokens
+        multimodal_params = [
+            self.create_multimodal_params(4, 10),   # 4 cached, 6 uncached
+            self.create_multimodal_params(7, 13),   # 7 cached, 6 uncached
+            self.create_multimodal_params(3, 10)    # 3 cached, 7 uncached
+        ]
+
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+
+        # Expected slices:
+        # Batch 1: [4:10] = 6 tokens
+        # Batch 2: [10+7:10+13] = [17:23] = 6 tokens
+        # Batch 3: [23+3:23+10] = [26:33] = 7 tokens
+        # Total: 6 + 6 + 7 = 19 tokens
+        assert len(result) == 1
+        assert result[0].shape == (19, 512)
+
+        # Verify the slices are correct
+        expected = torch.cat([
+            mm_embeds[0][4:10],    # Batch 1: 6 tokens
+            mm_embeds[0][17:23],   # Batch 2: 6 tokens
+            mm_embeds[0][26:33]    # Batch 3: 7 tokens
+        ], dim=0)
+        torch.testing.assert_close(result[0], expected)
+
+    def test_mixed_caching_with_fully_cached_batches(self):
+        """
+        Test mixed scenarios where some batches are fully cached (should be skipped).
+        """
+        mm_embeds = [torch.randn(25, 512)]  # Pre-concatenated: 8 + 9 + 8 tokens
+        multimodal_params = [
+            self.create_multimodal_params(8, 8),    # All cached - should be skipped
+            self.create_multimodal_params(3, 9),    # 3 cached, 6 uncached
+            self.create_multimodal_params(8, 8)     # All cached - should be skipped
+        ]
+
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+
+        # Only batch 2 should contribute: [8+3:8+9] = [11:17] = 6 tokens
+        assert len(result) == 1
+        assert result[0].shape == (6, 512)
+
+        # Verify the slice is correct
+        torch.testing.assert_close(result[0], mm_embeds[0][11:17])
+
+    def test_all_batches_fully_cached(self):
+        """
+        Test edge case where all batches are fully cached.
+        """
+        mm_embeds = [torch.randn(30, 512)]  # Pre-concatenated: 10 + 10 + 10 tokens
+        multimodal_params = [
+            self.create_multimodal_params(10, 10),  # All cached
+            self.create_multimodal_params(10, 10),  # All cached
+            self.create_multimodal_params(10, 10)   # All cached
+        ]
+
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+
+        # Should return empty list
+        assert result == []
+
+    def test_no_batches_cached(self):
+        """
+        Test edge case where no batches have any cached tokens.
+        """
+        mm_embeds = [torch.randn(30, 512)]  # Pre-concatenated: 10 + 10 + 10 tokens
+        multimodal_params = [
+            self.create_multimodal_params(0, 10),   # No cached
+            self.create_multimodal_params(0, 10),   # No cached
+            self.create_multimodal_params(0, 10)    # No cached
+        ]
+
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+
+        # Should return the full embeddings
+        assert result == mm_embeds
+
+    def test_error_handling_mismatched_counts(self):
+        """
+        Test error handling when mm_embeds and multimodal_params counts don't match
+        in individual batching mode.
+        """
+        mm_embeds = [torch.randn(10, 512), torch.randn(15, 512)]  # 2 embeddings
+        multimodal_params = [self.create_multimodal_params(0, 10)]  # Only 1 param
+
+        with pytest.raises(ValueError, match="Number of mm_embeds \\(2\\) does not match number of multimodal params \\(1\\)"):
+            find_uncached_mm_embeds(mm_embeds, multimodal_params)
+
+    def test_single_batch_scenarios(self):
+        """
+        Test various single batch scenarios.
+        """
+        # Single batch, no caching
+        mm_embeds = [torch.randn(20, 512)]
+        multimodal_params = [self.create_multimodal_params(0, 20)]
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+        assert result == mm_embeds
+
+        # Single batch, partial caching
+        multimodal_params = [self.create_multimodal_params(5, 20)]
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+        assert len(result) == 1
+        assert result[0].shape == (15, 512)
+        torch.testing.assert_close(result[0], mm_embeds[0][5:20])
+
+        # Single batch, all cached
+        multimodal_params = [self.create_multimodal_params(20, 20)]
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+        assert result == []
+
+    def test_different_devices(self):
+        """
+        Test with tensors on different devices (if CUDA is available).
+        """
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA not available")
+
+        # Test CPU tensors
+        mm_embeds = [torch.randn(10, 512, device='cpu')]
+        multimodal_params = [self.create_multimodal_params(3, 10)]
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+        assert result[0].device == mm_embeds[0].device
+
+        # Test CUDA tensors
+        mm_embeds = [torch.randn(10, 512, device='cuda')]
+        multimodal_params = [self.create_multimodal_params(3, 10)]
+        result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
+        assert result[0].device == mm_embeds[0].device
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])
\ No newline at end of file

From abde33038fb3edd31a3d9d2f8019d05fa68d64df Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Tue, 15 Jul 2025 18:04:08 -0700
Subject: [PATCH 09/14] Formatting

---
 .../batch_manager/kvCacheManager.cpp          |  7 +-
 .../models/modeling_multimodal_utils.py       | 51 ++++++----
 .../_torch/models/modeling_qwen2vl.py         |  6 +-
 .../_torch/pyexecutor/model_engine.py         |  3 +-
 tensorrt_llm/inputs/multimodal.py             | 21 ++++-
 .../_torch/multimodal/test_kvcache_reuse.py   | 94 +++++++++++--------
 6 files changed, 114 insertions(+), 68 deletions(-)

diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index f113984a3057..fb295bb1d7e8 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -109,12 +109,11 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
         TLLM_CHECK_WITH_INFO(mmHashVector.size() == 8, "Multimodal hash vector has unexpected size: %zu (expected 8)",
             mmHashVector.size());
 
+#define GET_NTH_BYTE(hash_part, byte_idx) static_cast<uint8_t>((hash_part >> (24 - (byte_idx) *8)) & 0xFF)
         // mmHashVector[j] comes from Python's int(hex_chunk, 16)
         // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB (big endian)
         // Convert 8x 32-bit integers into a 32-byte array preserving Blake3 hash byte order
         // Example: hashPart = 0x00010203 → mmHashArray[0:3] = [0x00, 0x01, 0x02, 0x03]
-        #define GET_NTH_BYTE(hash_part, byte_idx) static_cast<uint8_t>((hash_part >> (24 - (byte_idx) * 8)) & 0xFF)
-
         for (size_t j = 0; j < 8; ++j)
         {
             auto const& hashPart = mmHashVector[j];
@@ -123,7 +122,7 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
                 mmHashArray[j * 4 + byteIdx] = GET_NTH_BYTE(hashPart, byteIdx);
             }
         }
-        #undef GET_NTH_BYTE
+#undef GET_NTH_BYTE
 
         // Check if this multimodal content overlaps with the current block
         if (endTokenIdx > startPos && startTokenIdx < startPos + length)
@@ -187,7 +186,7 @@ size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) no
         seed ^= c + 0x9e3779b9 + (seed << 6) + (seed >> 2);
     }
 
-    // Add extra keys for multimodal data mixing in external mulitmodal item hash and token offset within this sequence
+    // Add extra keys for multimodal data mixing in external multimodal item hash and token offset within this sequence
     // block
     if (blockKey.extraKeys)
     {
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
index fa0ffff37f66..f3b65d9196a6 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
@@ -19,6 +19,7 @@
 import math
 from typing import List, Optional, Tuple
 
+import numpy as np
 import torch
 import torch.nn.functional as F
 from einops import rearrange
@@ -28,9 +29,11 @@
 from tensorrt_llm._torch.modules.embedding import Embedding
 from tensorrt_llm.inputs.multimodal import MultimodalParams
 from tensorrt_llm.logger import logger
-import numpy as np
 
-def find_uncached_mm_embeds(mm_embeds: List[torch.Tensor], multimodal_params: List[MultimodalParams]) -> torch.Tensor:
+
+def find_uncached_mm_embeds(
+        mm_embeds: List[torch.Tensor],
+        multimodal_params: List[MultimodalParams]) -> torch.Tensor:
     """
     Find the uncached multimodal mm_embeds from multimodal_params for each batch.
     Args:
@@ -46,7 +49,9 @@ def find_uncached_mm_embeds(mm_embeds: List[torch.Tensor], multimodal_params: Li
     # 1. Pre-concatenated mm_embeds for each batch, i.e., len(mm_embeds) == 1
     # 2. Individual mm_embeds for each multimodal param, i.e., len(mm_embeds) == len(multimodal_params)
     if len(mm_embeds) > 1 and len(mm_embeds) != len(multimodal_params):
-        raise ValueError(f"Number of mm_embeds ({len(mm_embeds)}) does not match number of multimodal params ({len(multimodal_params)}).")
+        raise ValueError(
+            f"Number of mm_embeds ({len(mm_embeds)}) does not match number of multimodal params ({len(multimodal_params)})."
+        )
 
     if not multimodal_params or multimodal_params[0].multimodal_runtime is None:
         # No slicing, return the full mm_embeds
@@ -59,55 +64,66 @@ def find_uncached_mm_embeds(mm_embeds: List[torch.Tensor], multimodal_params: Li
     if total_cached_mm_tokens == 0:
         # No cached tokens, return the full mm_embeds
         # TODO: support chunk prefill for multimodal, then we need to extract full mm_embeds for each CHUNK
-        logger.debug("No multimodal cached tokens can be reused, return the full mm_embeds")
+        logger.debug(
+            "No multimodal cached tokens can be reused, return the full mm_embeds"
+        )
         return mm_embeds
 
     if total_cached_mm_tokens == sum([
-        param.multimodal_runtime.total_mm_tokens
-        for param in multimodal_params
+            param.multimodal_runtime.total_mm_tokens
+            for param in multimodal_params
     ]):
         # All tokens are cached, return empty list
-        logger.debug("All multimodal tokens cached, skipping vision encoder forward")
+        logger.debug(
+            "All multimodal tokens cached, skipping vision encoder forward")
         return []
 
     # Partial caching, return the sliced mm_embeds
-    prefix_sum = np.cumsum([param.multimodal_runtime.total_mm_tokens for param in multimodal_params])
+    prefix_sum = np.cumsum([
+        param.multimodal_runtime.total_mm_tokens for param in multimodal_params
+    ])
     current_pos = 0
     slices = []
     for param in multimodal_params:
         runtime = param.multimodal_runtime
         if runtime.num_cached_mm_tokens > 0:
             if runtime.num_cached_mm_tokens == runtime.total_mm_tokens:
-                if len(mm_embeds) == 1: # pre-concatenated mm_embeds, need global indices
+                if len(
+                        mm_embeds
+                ) == 1:  # pre-concatenated mm_embeds, need global indices
                     current_pos += runtime.total_mm_tokens
                 slices.append((current_pos, current_pos))
                 continue
-            num_uncached_mm_tokens = runtime.total_mm_tokens - runtime.num_cached_mm_tokens
+            runtime.total_mm_tokens - runtime.num_cached_mm_tokens
             # TODO: support chunk prefill to extract partial mm_embeds
-            slices.append((current_pos + runtime.num_cached_mm_tokens, current_pos + runtime.total_mm_tokens))
-        else: # == 0, no cached tokens
+            slices.append((current_pos + runtime.num_cached_mm_tokens,
+                           current_pos + runtime.total_mm_tokens))
+        else:  # == 0, no cached tokens
             # All mm_embeds in this sequence are sliced out
             # TODO: support chunk prefill to extract partial mm_embeds
             slices.append((current_pos, current_pos + runtime.total_mm_tokens))
 
-
-        if len(mm_embeds) == 1: # pre-concatenated mm_embeds, need global indices
+        if len(mm_embeds
+               ) == 1:  # pre-concatenated mm_embeds, need global indices
             current_pos += runtime.total_mm_tokens
 
     sliced_mm_embeds = []
     if len(mm_embeds) == 1:
         for start, end in slices:
             sliced_mm_embeds.append(mm_embeds[0][start:end])
-    else: # slice each mm_embeds individually
+    else:  # slice each mm_embeds individually
         for i, (start, end) in enumerate(slices):
             sliced_mm_embeds.append(mm_embeds[i][start:end])
 
     if len(mm_embeds) == 1:
         sliced_mm_embeds = [torch.cat(sliced_mm_embeds, dim=0)]
 
-    logger.debug(f"Partial caching, return sliced_mm_embeds: {sliced_mm_embeds[0].shape}")
+    logger.debug(
+        f"Partial caching, return sliced_mm_embeds: {sliced_mm_embeds[0].shape}"
+    )
     return sliced_mm_embeds
 
+
 def fuse_input_embeds(
     embedding_layer: Embedding,
     input_ids: torch.IntTensor,
@@ -154,8 +170,7 @@ def fuse_input_embeds(
             f"Multimodal token count mismatch: found {len(mm_token_indices)} image tokens in input_ids "
             f"but received {mm_embed.shape[0]} image embeddings. "
             f"This is likely due to KV cache reuse, chunk prefill, or other optimizations that "
-            f"cause token count mismatches within the inference batch."
-        )
+            f"cause token count mismatches within the inference batch.")
 
     text_embed = embedding_layer(input_ids[text_token_indices])
     input_embeds = torch.empty(input_ids.shape[0],
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 889a579fc709..25a2778f8b89 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -18,7 +18,8 @@
 from ..attention_backend import AttentionMetadata
 from ..model_config import ModelConfig
 from .modeling_auto import AutoModelForCausalLM
-from .modeling_multimodal_utils import fuse_input_embeds, find_uncached_mm_embeds
+from .modeling_multimodal_utils import (find_uncached_mm_embeds,
+                                        fuse_input_embeds)
 from .modeling_utils import register_auto_model
 
 DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
@@ -601,7 +602,8 @@ def forward(
             mrope_config = self._parse_and_concat_mrope_config(
                 multimodal_params, num_context_requests,
                 num_generation_requests)
-            mm_embeds = find_uncached_mm_embeds(mm_embeds, multimodal_params[:num_context_requests])
+            mm_embeds = find_uncached_mm_embeds(
+                mm_embeds, multimodal_params[:num_context_requests])
 
         if 'mrope_position_deltas' in kwargs:
             mrope_config['mrope_position_deltas'] = kwargs[
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 29a33b6c682e..bf1ad5826de7 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1147,7 +1147,8 @@ def _prepare_tp_inputs(
             py_multimodal_runtime = MultimodalRuntimeData(
                 mm_token_lengths=request.multimodal_lengths,
                 mm_token_positions=request.multimodal_positions,
-                num_cached_tokens=past_seen_token_num) if request.multimodal_hashes is not None else None
+                num_cached_tokens=past_seen_token_num
+            ) if request.multimodal_hashes is not None else None
 
             multimodal_params = MultimodalParams(
                 multimodal_data=request.py_multimodal_data,
diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index 7f9b1c1c1eed..ed15736db599 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -81,6 +81,7 @@ def to_tensor(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             torch.tensor(self.multimodal_positions, dtype=torch.int32),
             torch.tensor(self.multimodal_lengths, dtype=torch.int32))
 
+
 @dataclass
 class MultimodalRuntimeData:
     """Runtime data for tracking multimodal token caching and reuse per request sequence.
@@ -110,21 +111,30 @@ class MultimodalRuntimeData:
     def __post_init__(self):
         # Validate input data
         if len(self.mm_token_positions) != len(self.mm_token_lengths):
-            raise ValueError(f"mm_token_positions ({len(self.mm_token_positions)}) and mm_token_lengths ({len(self.mm_token_lengths)}) must have the same length")
+            raise ValueError(
+                f"mm_token_positions ({len(self.mm_token_positions)}) and mm_token_lengths ({len(self.mm_token_lengths)}) must have the same length"
+            )
 
         if self.num_cached_tokens < 0:
-            raise ValueError(f"num_cached_tokens must be non-negative, got {self.num_cached_tokens}")
+            raise ValueError(
+                f"num_cached_tokens must be non-negative, got {self.num_cached_tokens}"
+            )
 
         if any(length <= 0 for length in self.mm_token_lengths):
-            raise ValueError(f"All mm_token_lengths must be positive, got {self.mm_token_lengths}")
+            raise ValueError(
+                f"All mm_token_lengths must be positive, got {self.mm_token_lengths}"
+            )
 
         if any(pos < 0 for pos in self.mm_token_positions):
-            raise ValueError(f"All mm_token_positions must be non-negative, got {self.mm_token_positions}")
+            raise ValueError(
+                f"All mm_token_positions must be non-negative, got {self.mm_token_positions}"
+            )
 
         if self.num_cached_mm_tokens is None:
             # Compute cached multimodal tokens based on positions and cached tokens
             self.num_cached_mm_tokens = 0
-            for pos, length in zip(self.mm_token_positions, self.mm_token_lengths):
+            for pos, length in zip(self.mm_token_positions,
+                                   self.mm_token_lengths):
                 if pos + length <= self.num_cached_tokens:
                     self.num_cached_mm_tokens += length
                 elif pos < self.num_cached_tokens:
@@ -136,6 +146,7 @@ def __post_init__(self):
 
         self.total_mm_tokens = sum(self.mm_token_lengths)
 
+
 @dataclass
 class MultimodalParams:
     """Unified container for multimodal parameters.
diff --git a/tests/unittest/_torch/multimodal/test_kvcache_reuse.py b/tests/unittest/_torch/multimodal/test_kvcache_reuse.py
index 8f2876d6f281..0eb0d5f9ca40 100644
--- a/tests/unittest/_torch/multimodal/test_kvcache_reuse.py
+++ b/tests/unittest/_torch/multimodal/test_kvcache_reuse.py
@@ -1,12 +1,13 @@
+from unittest.mock import Mock
+
 import pytest
 import torch
-import numpy as np
-from unittest.mock import Mock
-from typing import List
 
 # Import the function to test
-from tensorrt_llm._torch.models.modeling_multimodal_utils import find_uncached_mm_embeds
-from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData
+from tensorrt_llm._torch.models.modeling_multimodal_utils import \
+    find_uncached_mm_embeds
+from tensorrt_llm.inputs.multimodal import (MultimodalParams,
+                                            MultimodalRuntimeData)
 
 
 class TestMultimodalRuntimeData:
@@ -41,7 +42,9 @@ def test_complex_scenario_with_multiple_chunks(self):
         runtime = MultimodalRuntimeData(
             num_cached_tokens=30,
             mm_token_lengths=[3, 4, 5, 6, 7, 8],  # Total: 33 tokens
-            mm_token_positions=[0, 5, 10, 15, 25, 35]  # Positions: 0-3, 5-9, 10-15, 15-21, 25-32, 35-43
+            mm_token_positions=[
+                0, 5, 10, 15, 25, 35
+            ]  # Positions: 0-3, 5-9, 10-15, 15-21, 25-32, 35-43
         )
 
         # Expected caching:
@@ -59,16 +62,19 @@ def test_complex_scenario_with_multiple_chunks(self):
 class TestFindUncachedMmEmbed:
     """Focused test cases for find_uncached_mm_embeds function - testing edge cases and potential bugs."""
 
-    def create_mock_runtime(self, num_cached_mm_tokens: int, total_mm_tokens: int):
+    def create_mock_runtime(self, num_cached_mm_tokens: int,
+                            total_mm_tokens: int):
         """Helper to create a mock MultimodalRuntimeData."""
         runtime = Mock(spec=MultimodalRuntimeData)
         runtime.num_cached_mm_tokens = num_cached_mm_tokens
         runtime.total_mm_tokens = total_mm_tokens
         return runtime
 
-    def create_multimodal_params(self, num_cached_mm_tokens: int, total_mm_tokens: int):
+    def create_multimodal_params(self, num_cached_mm_tokens: int,
+                                 total_mm_tokens: int):
         """Helper to create MultimodalParams with runtime data."""
-        runtime = self.create_mock_runtime(num_cached_mm_tokens, total_mm_tokens)
+        runtime = self.create_mock_runtime(num_cached_mm_tokens,
+                                           total_mm_tokens)
         return MultimodalParams(multimodal_runtime=runtime)
 
     def test_mm_embed_not_batched(self):
@@ -79,21 +85,21 @@ def test_mm_embed_not_batched(self):
         mm_embeds = [
             torch.randn(10, 512),  # Batch 1: 10 tokens
             torch.randn(15, 512),  # Batch 2: 15 tokens
-            torch.randn(8, 512)    # Batch 3: 8 tokens
+            torch.randn(8, 512)  # Batch 3: 8 tokens
         ]
         multimodal_params = [
-            self.create_multimodal_params(3, 10),   # 3 cached, 7 uncached
-            self.create_multimodal_params(8, 15),   # 8 cached, 7 uncached
-            self.create_multimodal_params(0, 8)     # 0 cached, 8 uncached
+            self.create_multimodal_params(3, 10),  # 3 cached, 7 uncached
+            self.create_multimodal_params(8, 15),  # 8 cached, 7 uncached
+            self.create_multimodal_params(0, 8)  # 0 cached, 8 uncached
         ]
 
         result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
 
         # Should return individual slices for each batch
         assert len(result) == 3
-        assert result[0].shape == (7, 512)   # 10 - 3 = 7
-        assert result[1].shape == (7, 512)   # 15 - 8 = 7
-        assert result[2].shape == (8, 512)   # 8 - 0 = 8
+        assert result[0].shape == (7, 512)  # 10 - 3 = 7
+        assert result[1].shape == (7, 512)  # 15 - 8 = 7
+        assert result[2].shape == (8, 512)  # 8 - 0 = 8
 
         # Verify the slices are correct
         torch.testing.assert_close(result[0], mm_embeds[0][3:10])
@@ -105,11 +111,12 @@ def test_mm_embed_batched(self):
         Test batching (concatenated) mm_embeds with fused mm_embeds for each batch.
         This tests the case where len(mm_embeds) == 1
         """
-        mm_embeds = [torch.randn(33, 512)]  # Pre-concatenated: 10 + 13 + 10 tokens
+        mm_embeds = [torch.randn(33,
+                                 512)]  # Pre-concatenated: 10 + 13 + 10 tokens
         multimodal_params = [
-            self.create_multimodal_params(4, 10),   # 4 cached, 6 uncached
-            self.create_multimodal_params(7, 13),   # 7 cached, 6 uncached
-            self.create_multimodal_params(3, 10)    # 3 cached, 7 uncached
+            self.create_multimodal_params(4, 10),  # 4 cached, 6 uncached
+            self.create_multimodal_params(7, 13),  # 7 cached, 6 uncached
+            self.create_multimodal_params(3, 10)  # 3 cached, 7 uncached
         ]
 
         result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
@@ -123,11 +130,13 @@ def test_mm_embed_batched(self):
         assert result[0].shape == (19, 512)
 
         # Verify the slices are correct
-        expected = torch.cat([
-            mm_embeds[0][4:10],    # Batch 1: 6 tokens
-            mm_embeds[0][17:23],   # Batch 2: 6 tokens
-            mm_embeds[0][26:33]    # Batch 3: 7 tokens
-        ], dim=0)
+        expected = torch.cat(
+            [
+                mm_embeds[0][4:10],  # Batch 1: 6 tokens
+                mm_embeds[0][17:23],  # Batch 2: 6 tokens
+                mm_embeds[0][26:33]  # Batch 3: 7 tokens
+            ],
+            dim=0)
         torch.testing.assert_close(result[0], expected)
 
     def test_mixed_caching_with_fully_cached_batches(self):
@@ -136,9 +145,11 @@ def test_mixed_caching_with_fully_cached_batches(self):
         """
         mm_embeds = [torch.randn(25, 512)]  # Pre-concatenated: 8 + 9 + 8 tokens
         multimodal_params = [
-            self.create_multimodal_params(8, 8),    # All cached - should be skipped
-            self.create_multimodal_params(3, 9),    # 3 cached, 6 uncached
-            self.create_multimodal_params(8, 8)     # All cached - should be skipped
+            self.create_multimodal_params(8,
+                                          8),  # All cached - should be skipped
+            self.create_multimodal_params(3, 9),  # 3 cached, 6 uncached
+            self.create_multimodal_params(8,
+                                          8)  # All cached - should be skipped
         ]
 
         result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
@@ -154,11 +165,12 @@ def test_all_batches_fully_cached(self):
         """
         Test edge case where all batches are fully cached.
         """
-        mm_embeds = [torch.randn(30, 512)]  # Pre-concatenated: 10 + 10 + 10 tokens
+        mm_embeds = [torch.randn(30,
+                                 512)]  # Pre-concatenated: 10 + 10 + 10 tokens
         multimodal_params = [
             self.create_multimodal_params(10, 10),  # All cached
             self.create_multimodal_params(10, 10),  # All cached
-            self.create_multimodal_params(10, 10)   # All cached
+            self.create_multimodal_params(10, 10)  # All cached
         ]
 
         result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
@@ -170,11 +182,12 @@ def test_no_batches_cached(self):
         """
         Test edge case where no batches have any cached tokens.
         """
-        mm_embeds = [torch.randn(30, 512)]  # Pre-concatenated: 10 + 10 + 10 tokens
+        mm_embeds = [torch.randn(30,
+                                 512)]  # Pre-concatenated: 10 + 10 + 10 tokens
         multimodal_params = [
-            self.create_multimodal_params(0, 10),   # No cached
-            self.create_multimodal_params(0, 10),   # No cached
-            self.create_multimodal_params(0, 10)    # No cached
+            self.create_multimodal_params(0, 10),  # No cached
+            self.create_multimodal_params(0, 10),  # No cached
+            self.create_multimodal_params(0, 10)  # No cached
         ]
 
         result = find_uncached_mm_embeds(mm_embeds, multimodal_params)
@@ -188,9 +201,14 @@ def test_error_handling_mismatched_counts(self):
         in individual batching mode.
         """
         mm_embeds = [torch.randn(10, 512), torch.randn(15, 512)]  # 2 embeddings
-        multimodal_params = [self.create_multimodal_params(0, 10)]  # Only 1 param
-
-        with pytest.raises(ValueError, match="Number of mm_embeds \\(2\\) does not match number of multimodal params \\(1\\)"):
+        multimodal_params = [self.create_multimodal_params(0,
+                                                           10)]  # Only 1 param
+
+        with pytest.raises(
+                ValueError,
+                match=
+                "Number of mm_embeds \\(2\\) does not match number of multimodal params \\(1\\)"
+        ):
             find_uncached_mm_embeds(mm_embeds, multimodal_params)
 
     def test_single_batch_scenarios(self):
@@ -236,4 +254,4 @@ def test_different_devices(self):
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
\ No newline at end of file
+    pytest.main([__file__])

From 92d4b03e8b13a4518854493c868c0a0d35a3ddb0 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Thu, 17 Jul 2025 09:40:23 -0700
Subject: [PATCH 10/14] Add comment for integer hashing

---
 cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index fb295bb1d7e8..11f7db16e661 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -158,6 +158,9 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager
 {
 size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) noexcept
 {
+    // Hashing algorithm adapted from StackOverflow:
+    // https://stackoverflow.com/questions/664014/what-integer-hash-function-are-good-that-accepts-an-integer-hash-key
+    // Constants provide very good distribution - each input bit affects each output bit with ~50% probability.
     size_t seed = blockKey.uniqueTokens.size() ^ parentHash * UINT64_C(0xbf58476d1ce4e5b9);
 
     for (auto const& uniqueToken : blockKey.uniqueTokens)

From e5eb2bd1c28eaf1c8cc520dcc4667bdf45868450 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Thu, 17 Jul 2025 10:45:12 -0700
Subject: [PATCH 11/14] Address comments

---
 .../models/modeling_multimodal_utils.py       | 25 +++----------------
 tensorrt_llm/inputs/multimodal.py             |  7 +++---
 2 files changed, 7 insertions(+), 25 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
index f3b65d9196a6..a1252fdac124 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
@@ -19,7 +19,6 @@
 import math
 from typing import List, Optional, Tuple
 
-import numpy as np
 import torch
 import torch.nn.functional as F
 from einops import rearrange
@@ -79,32 +78,14 @@ def find_uncached_mm_embeds(
         return []
 
     # Partial caching, return the sliced mm_embeds
-    prefix_sum = np.cumsum([
-        param.multimodal_runtime.total_mm_tokens for param in multimodal_params
-    ])
     current_pos = 0
     slices = []
     for param in multimodal_params:
         runtime = param.multimodal_runtime
-        if runtime.num_cached_mm_tokens > 0:
-            if runtime.num_cached_mm_tokens == runtime.total_mm_tokens:
-                if len(
-                        mm_embeds
-                ) == 1:  # pre-concatenated mm_embeds, need global indices
-                    current_pos += runtime.total_mm_tokens
-                slices.append((current_pos, current_pos))
-                continue
-            runtime.total_mm_tokens - runtime.num_cached_mm_tokens
-            # TODO: support chunk prefill to extract partial mm_embeds
-            slices.append((current_pos + runtime.num_cached_mm_tokens,
-                           current_pos + runtime.total_mm_tokens))
-        else:  # == 0, no cached tokens
-            # All mm_embeds in this sequence are sliced out
-            # TODO: support chunk prefill to extract partial mm_embeds
-            slices.append((current_pos, current_pos + runtime.total_mm_tokens))
-
+        slices.append((current_pos + runtime.num_cached_mm_tokens,
+                       current_pos + runtime.total_mm_tokens))
         if len(mm_embeds
-               ) == 1:  # pre-concatenated mm_embeds, need global indices
+               ) == 1:  # pre-concatenated mm_embeds, need global offset
             current_pos += runtime.total_mm_tokens
 
     sliced_mm_embeds = []
diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index ed15736db599..19d55ae77448 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -141,9 +141,10 @@ def __post_init__(self):
                     # Partial overlap - only count the cached portion
                     self.num_cached_mm_tokens += self.num_cached_tokens - pos
 
-        assert self.num_cached_mm_tokens <= self.num_cached_tokens, \
-            f"num_cached_mm_tokens ({self.num_cached_mm_tokens}) must be less than or equal to num_cached_tokens ({self.num_cached_tokens})"
-
+        if self.num_cached_mm_tokens > self.num_cached_tokens:
+            raise ValueError(
+                f"num_cached_mm_tokens ({self.num_cached_mm_tokens}) must be less than or equal to "
+                f"num_cached_tokens ({self.num_cached_tokens})")
         self.total_mm_tokens = sum(self.mm_token_lengths)
 
 

From 3e0ef74a21010e7d8a4191dfffa507ead62bca96 Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Sat, 19 Jul 2025 22:55:00 -0700
Subject: [PATCH 12/14] Address comments

---
 .../batch_manager/kvCacheManager.h            |  4 +--
 .../batch_manager/kvCacheManager.cpp          | 25 +++++++++++--------
 .../pybind/batch_manager/kvCacheManager.cpp   |  3 +--
 .../models/modeling_multimodal_utils.py       |  4 +--
 4 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
index 9d04544570ea..a0234cbbe49b 100644
--- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
+++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
@@ -113,7 +113,7 @@ struct BlockKey
 
     // Extra keys for multimodal data (similar to VLLM's approach)
     // Each extra key is a pair of (mm_hash, start_offset_in_block)
-    std::optional<std::vector<MmKey>> extraKeys = std::nullopt;
+    std::vector<MmKey> extraKeys;
 
     BlockKey() = default;
 
@@ -128,7 +128,7 @@ struct BlockKey
     }
 
     explicit BlockKey(bool usesExtraIds, std::optional<LoraTaskIdType> loraTaskId, VecUniqueTokens uniqueTokens,
-        std::optional<std::vector<MmKey>> extraKeys = std::nullopt)
+        std::vector<MmKey> extraKeys = {})
         : usesExtraIds{usesExtraIds}
         , loraTaskId{loraTaskId}
         , uniqueTokens{std::move(uniqueTokens)}
diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index 11f7db16e661..649c070067b2 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -76,7 +76,11 @@ std::list<std::vector<T>> chopVectorIntoBlocks(
     return blockedVectors;
 }
 
-std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
+inline uint8_t getNthByte(SizeType32 hashPart, uint8_t byteIdx) noexcept {
+    return static_cast<uint8_t>((hashPart >> (24 - byteIdx * 8)) & 0xFF);
+}
+
+std::vector<MmKey> generateBlockHashExtraKeys(
     tensorrt_llm::batch_manager::LlmRequest const& llmRequest, SizeType32 startTokenIdx, SizeType32 endTokenIdx)
 {
     auto const multimodalHashes = llmRequest.getMultimodalHashes();
@@ -87,17 +91,19 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
         || (*multimodalHashes)->empty() || !(*multimodalPositions) || (*multimodalPositions)->empty()
         || !(*multimodalLengths) || (*multimodalLengths)->empty())
     {
-        return std::nullopt;
+        return {};
     }
 
     if ((*multimodalHashes)->size() != (*multimodalPositions)->size()
         || (*multimodalPositions)->size() != (*multimodalLengths)->size())
     {
         TLLM_LOG_WARNING("Multimodal data arrays have mismatched sizes");
-        return std::nullopt;
+        return {};
     }
 
     std::vector<MmKey> extraKeys; // MmKey = std::pair<std::array<uint8_t, 32>, SizeType32>
+    extraKeys.reserve((*multimodalPositions)->size());
+    std::array<uint8_t, 32> mmHashArray;
 
     for (size_t i = 0; i < (*multimodalPositions)->size(); ++i)
     {
@@ -105,11 +111,9 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
         auto const& length = (*(*multimodalLengths))[i];
         auto const& mmHashVector = (*(*multimodalHashes))[i];
 
-        std::array<uint8_t, 32> mmHashArray;
         TLLM_CHECK_WITH_INFO(mmHashVector.size() == 8, "Multimodal hash vector has unexpected size: %zu (expected 8)",
             mmHashVector.size());
 
-#define GET_NTH_BYTE(hash_part, byte_idx) static_cast<uint8_t>((hash_part >> (24 - (byte_idx) *8)) & 0xFF)
         // mmHashVector[j] comes from Python's int(hex_chunk, 16)
         // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB (big endian)
         // Convert 8x 32-bit integers into a 32-byte array preserving Blake3 hash byte order
@@ -117,12 +121,11 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
         for (size_t j = 0; j < 8; ++j)
         {
             auto const& hashPart = mmHashVector[j];
-            for (size_t byteIdx = 0; byteIdx < 4; ++byteIdx)
+            for (uint8_t byteIdx = 0; byteIdx < 4; ++byteIdx)
             {
-                mmHashArray[j * 4 + byteIdx] = GET_NTH_BYTE(hashPart, byteIdx);
+                mmHashArray[j * 4 + byteIdx] = getNthByte(hashPart, byteIdx);
             }
         }
-#undef GET_NTH_BYTE
 
         // Check if this multimodal content overlaps with the current block
         if (endTokenIdx > startPos && startTokenIdx < startPos + length)
@@ -132,7 +135,7 @@ std::optional<std::vector<MmKey>> generateBlockHashExtraKeys(
         }
     }
 
-    return extraKeys.empty() ? std::nullopt : std::make_optional(std::move(extraKeys));
+    return extraKeys;
 }
 
 std::vector<BlockKey> buildBlockKeys(
@@ -191,9 +194,9 @@ size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) no
 
     // Add extra keys for multimodal data mixing in external multimodal item hash and token offset within this sequence
     // block
-    if (blockKey.extraKeys)
+    if (!blockKey.extraKeys.empty())
     {
-        for (auto const& [mmHash, startOffset] : *blockKey.extraKeys)
+        for (auto const& [mmHash, startOffset] : blockKey.extraKeys)
         {
             // Hash the multimodal hash array in 32-bit chunks (more efficient)
             for (size_t i = 0; i < 32; i += 4)
diff --git a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp
index 3649600fb91f..255b0f8efa33 100644
--- a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp
@@ -315,8 +315,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(py::module_& m)
             py::arg("lora_task_id"), py::arg("unique_tokens"))
         .def_readonly("uses_extra_ids", &tbk::BlockKey::usesExtraIds)
         .def_readonly("lora_task_id", &tbk::BlockKey::loraTaskId)
-        .def_readonly("unique_tokens", &tbk::BlockKey::uniqueTokens)
-        .def_readonly("extra_keys", &tbk::BlockKey::extraKeys);
+        .def_readonly("unique_tokens", &tbk::BlockKey::uniqueTokens);
 
     py::class_<tbk::BlockKeyHasher>(m, "BlockKeyHasher")
         .def_static("hash", &tbk::BlockKeyHasher::hash, py::arg("block_key"), py::arg("parent_hash") = 0);
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
index a1252fdac124..d6387f819084 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py
@@ -150,8 +150,8 @@ def fuse_input_embeds(
         raise ValueError(
             f"Multimodal token count mismatch: found {len(mm_token_indices)} image tokens in input_ids "
             f"but received {mm_embed.shape[0]} image embeddings. "
-            f"This is likely due to KV cache reuse, chunk prefill, or other optimizations that "
-            f"cause token count mismatches within the inference batch.")
+            "This is likely due to KV cache reuse, chunk prefill, or other optimizations that "
+            "cause token count mismatches within the inference batch.")
 
     text_embed = embedding_layer(input_ids[text_token_indices])
     input_embeds = torch.empty(input_ids.shape[0],

From 810f4e1e982bc0186e52462d85cabc9ab2b819ff Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Sat, 19 Jul 2025 23:16:27 -0700
Subject: [PATCH 13/14] Formatting

---
 cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index 649c070067b2..d30ba27be3ab 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -76,7 +76,8 @@ std::list<std::vector<T>> chopVectorIntoBlocks(
     return blockedVectors;
 }
 
-inline uint8_t getNthByte(SizeType32 hashPart, uint8_t byteIdx) noexcept {
+inline uint8_t getNthByte(SizeType32 hashPart, uint8_t byteIdx) noexcept
+{
     return static_cast<uint8_t>((hashPart >> (24 - byteIdx * 8)) & 0xFF);
 }
 

From 2ca1ceacbdeffe5beedac580b3e8b361068eadbf Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Sat, 19 Jul 2025 23:28:11 -0700
Subject: [PATCH 14/14] Fix lint

---
 tensorrt_llm/_torch/pyexecutor/model_engine.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index bf1ad5826de7..90529a2dd94e 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -21,7 +21,8 @@
 from tensorrt_llm._torch.speculative.mtp import SampleStateTensorsMTP
 from tensorrt_llm._utils import (is_trace_enabled, nvtx_range, release_gc,
                                  torch_dtype_to_str, trace_func)
-from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData
+from tensorrt_llm.inputs.multimodal import (MultimodalParams,
+                                            MultimodalRuntimeData)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.lora_manager import LoraConfig, LoraModelConfig
 from tensorrt_llm.mapping import Mapping