fix partialMatch

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

Author: pansicheng

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -94,8 +94,8 @@ struct BlockKey
         SizeType32 numMatched{0};
         if (loraTaskId == other.loraTaskId)
         {
-            auto [matchEnd, otherMatchEnd]
-                = std::mismatch(uniqueTokens.begin(), uniqueTokens.end(), other.uniqueTokens.begin());
+            auto [matchEnd, otherMatchEnd] = std::mismatch(
+                uniqueTokens.begin(), uniqueTokens.end(), other.uniqueTokens.begin(), other.uniqueTokens.end());
             numMatched = std::distance(uniqueTokens.begin(), matchEnd);
         }
         return numMatched;

cpp/tests/batch_manager/CMakeLists.txt

@@ -28,3 +28,4 @@ target_link_libraries(trtGptModelRealDecoderTest PRIVATE modelSpecStatic
 add_gtest(peftCacheManagerTest peftCacheManagerTest.cpp)
 add_gtest(trtEncoderModelTest trtEncoderModelTest.cpp)
 add_gtest(guidedDecoderTest guidedDecoderTest.cpp)
+add_gtest(blockKeyTest blockKeyTest.cpp)

cpp/tests/batch_manager/blockKeyTest.cpp

@@ -0,0 +1,23 @@
+#include "tensorrt_llm/batch_manager/kvCacheManager.h"
+
+#include <gtest/gtest.h>
+
+using namespace tensorrt_llm::batch_manager::kv_cache_manager;
+
+class BlockKeyTest : public ::testing::Test
+{
+};
+
+TEST_F(BlockKeyTest, PartialMatch)
+{
+    VecUniqueTokens tokens0 = {{0, 0}, {0, 0}};
+    VecUniqueTokens tokens1 = {{0, 0}};
+    BlockKey bk0(false, 0, tokens0);
+    BlockKey bk1(false, 0, tokens1);
+
+    bk1.uniqueTokens.reserve(2);
+    auto ptr = reinterpret_cast<char*>(bk1.uniqueTokens.data());
+    std::fill(ptr, ptr + bk1.uniqueTokens.capacity() * sizeof(UniqueToken), 0);
+
+    EXPECT_EQ(bk0.partialMatch(bk1), 1);
+}