[Bugfix][Conductor] Clamp query match to the queried token count

docs/source/design/conductor/conductor-architecture-design.md

@@ -68,8 +68,10 @@ Within each context, Conductor stores:
   per-instance access metadata;
 - a DP-rank set used to report rank-level hit information.
 
-The current implementation computes complete-block prefix hashes from token IDs
-and ignores trailing partial blocks during `/query`.
+The vLLM strategy computes complete-block prefix hashes from token IDs and
+ignores trailing partial blocks during `/query`. The SGLang strategies also hash
+the trailing partial block, so a match that reaches it is reported as the number
+of tokens the request actually carries, never as a whole block.
 
 ## Event flow
 

mooncake-conductor/src/prefixindex/prefix_indexer.cpp

@@ -193,6 +193,21 @@ int64_t TokensForBlocks(size_t block_count, int64_t block_size) {
     return static_cast<int64_t>(block_count) * block_size;
 }
 
+// Tokens reusable from a matched run of blocks. The SGLang chains end in a
+// partial block, so a run that reaches the last block covers only the tokens
+// the query actually carries. Reporting a whole block there would hand the
+// router reusable tokens that no indexed block holds, and the engine would
+// then skip recomputing slots whose KV belongs to some other request.
+int64_t MatchedTokens(size_t block_count, int64_t block_size,
+                      size_t queried_tokens) {
+    const int64_t queried =
+        queried_tokens >
+                static_cast<size_t>(std::numeric_limits<int64_t>::max())
+            ? std::numeric_limits<int64_t>::max()
+            : static_cast<int64_t>(queried_tokens);
+    return std::min(TokensForBlocks(block_count, block_size), queried);
+}
+
 }  // namespace
 
 RegistrationResult PrefixCacheTable::ValidateRegistration(
@@ -473,6 +488,14 @@ std::map<std::string, CacheHitResult> PrefixCacheTable::Query(
         return results;
     }
     const size_t block_count = chain->BlockCount();
+    // Reusable logical KV positions carried by this query. The bigram chains
+    // hash token pairs, so n raw tokens cover n - 1 positions and the chain is
+    // built over that length. Capping at the raw count there would still let a
+    // full match report one position no indexed block holds.
+    const size_t queried_tokens =
+        (state->profile.strategy == "sglang_bigram" && !token_ids.empty())
+            ? token_ids.size() - 1
+            : token_ids.size();
 
     // Resolve the optional filter and copy rank sets before probing. The copies
     // remain valid while the probe releases and reacquires state.mutex.
@@ -559,17 +582,20 @@ std::map<std::string, CacheHitResult> PrefixCacheTable::Query(
             advance_cursor(cursor, gpu_present);
 
             RankCacheHitResult rank_match;
-            rank_match.gpu = TokensForBlocks(cursor, context.block_size);
+            rank_match.gpu =
+                MatchedTokens(cursor, context.block_size, queried_tokens);
 
             advance_cursor(cursor, [](const BlockPresence& block) {
                 return !block.cpu_owners.empty();
             });
-            rank_match.cpu = TokensForBlocks(cursor, context.block_size);
+            rank_match.cpu =
+                MatchedTokens(cursor, context.block_size, queried_tokens);
 
             advance_cursor(cursor, [](const BlockPresence& block) {
                 return !block.disk_owners.empty();
             });
-            rank_match.disk = TokensForBlocks(cursor, context.block_size);
+            rank_match.disk =
+                MatchedTokens(cursor, context.block_size, queried_tokens);
 
             result.dp.emplace(rank, rank_match.gpu);
             result.rank_matches.emplace(rank, rank_match);

mooncake-conductor/tests/prefix_indexer_test.cpp

@@ -72,6 +72,26 @@ HashProfile PickleProfile() {
             .index_projection = "low64_be"};
 }
 
+// SGLang chains carry no Python seed root, so the resolver pins an all-zero
+// sentinel digest for the registration wire contract.
+HashProfile SglangProfile() {
+    return {.strategy = "sglang",
+            .algorithm = "sha256_raw",
+            .python_hash_seed = "0",
+            .root_digest = std::string(64, '0'),
+            .index_projection = "first64_be"};
+}
+
+// The bigram chains hash token pairs, so they carry one fewer logical position
+// than the prompt has tokens.
+HashProfile SglangBigramProfile() {
+    return {.strategy = "sglang_bigram",
+            .algorithm = "sha256_raw",
+            .python_hash_seed = "0",
+            .root_digest = std::string(64, '0'),
+            .index_projection = "first64_be"};
+}
+
 EngineRegistration Registration(const std::string& instance_id = "instance-a",
                                 int64_t dp_rank = 0) {
     const ContextKey context = TestContext();
@@ -177,6 +197,49 @@ std::vector<ProjectedPrefix> Hashes(
     return prefixes;
 }
 
+std::vector<ProjectedPrefix> SglangHashes(const std::vector<int32_t>& tokens) {
+    std::string error;
+    auto strategy = mooncake::conductor::prefixindex::CreateHashStrategy(
+        SglangProfile(), &error);
+    EXPECT_TRUE(error.empty()) << error;
+    if (!strategy) {
+        return {};
+    }
+
+    std::vector<HashBlock> blocks;
+    error = strategy->Compute(TestContext(), tokens, std::nullopt, &blocks);
+    EXPECT_TRUE(error.empty()) << error;
+
+    std::vector<ProjectedPrefix> prefixes;
+    prefixes.reserve(blocks.size());
+    for (const HashBlock& block : blocks) {
+        prefixes.push_back(block.projected);
+    }
+    return prefixes;
+}
+
+std::vector<ProjectedPrefix> SglangBigramHashes(
+    const std::vector<int32_t>& tokens) {
+    std::string error;
+    auto strategy = mooncake::conductor::prefixindex::CreateHashStrategy(
+        SglangBigramProfile(), &error);
+    EXPECT_TRUE(error.empty()) << error;
+    if (!strategy) {
+        return {};
+    }
+
+    std::vector<HashBlock> blocks;
+    error = strategy->Compute(TestContext(), tokens, std::nullopt, &blocks);
+    EXPECT_TRUE(error.empty()) << error;
+
+    std::vector<ProjectedPrefix> prefixes;
+    prefixes.reserve(blocks.size());
+    for (const HashBlock& block : blocks) {
+        prefixes.push_back(block.projected);
+    }
+    return prefixes;
+}
+
 void RegisterOrFail(PrefixCacheTable& table,
                     const EngineRegistration& registration) {
     const auto result = table.Register(registration);
@@ -612,6 +675,106 @@ TEST(Query, ExactTwoInstanceSharedCacheExample) {
     EXPECT_EQ(second.disk, 48);
 }
 
+TEST(Query, TrailingPartialBlockNeverReportsMoreThanPromptTokens) {
+    PrefixCacheTable table;
+    auto registration = Registration();
+    registration.profile = SglangProfile();
+    RegisterOrFail(table, registration);
+
+    // 40 tokens over a 16-token block size: two whole blocks plus a trailing
+    // block that only covers 8 tokens.
+    const auto tokens = Tokens(40);
+    const auto hashes = SglangHashes(tokens);
+    ASSERT_EQ(hashes.size(), 3u);
+
+    ASSERT_EQ(table.StoreGpu(Gpu({hashes[0], hashes[1], hashes[2]})), "");
+    ASSERT_EQ(table.StoreShared(Shared(hashes, StorageTier::kCpu)), "");
+    ASSERT_EQ(table.StoreShared(Shared(hashes, StorageTier::kDisk)), "");
+
+    const auto result = table.Query(TestContext(), tokens).at("instance-a");
+    EXPECT_EQ(result.longest_match_tokens, 40);
+    EXPECT_EQ(result.gpu, 40);
+    EXPECT_EQ(result.dp, (std::map<int64_t, int64_t>{{0, 40}}));
+    EXPECT_EQ(
+        result.rank_matches,
+        (std::map<int64_t, RankCacheHitResult>{{0, RankMatch(40, 40, 40)}}));
+    EXPECT_EQ(result.cpu, 40);
+    EXPECT_EQ(result.disk, 40);
+}
+
+TEST(Query, WholeBlockRunIsUnaffectedByThePromptLengthClamp) {
+    PrefixCacheTable table;
+    auto registration = Registration();
+    registration.profile = SglangProfile();
+    RegisterOrFail(table, registration);
+
+    const auto tokens = Tokens(40);
+    const auto hashes = SglangHashes(tokens);
+    ASSERT_EQ(hashes.size(), 3u);
+
+    // Only the two whole blocks are indexed; the partial tail misses.
+    ASSERT_EQ(
+        table.StoreShared(Shared({hashes[0], hashes[1]}, StorageTier::kDisk)),
+        "");
+
+    const auto result = table.Query(TestContext(), tokens).at("instance-a");
+    EXPECT_EQ(result.longest_match_tokens, 32);
+    EXPECT_EQ(result.gpu, 0);
+    EXPECT_EQ(result.cpu, 0);
+    EXPECT_EQ(result.disk, 32);
+}
+
+TEST(Query, BigramFullMatchReportsOneFewerPositionThanPromptTokens) {
+    PrefixCacheTable table;
+    auto registration = Registration();
+    registration.profile = SglangBigramProfile();
+    RegisterOrFail(table, registration);
+
+    // 40 raw tokens over a 16-token block size. The bigram chain covers 39
+    // logical positions, so it is two whole blocks plus a trailing block of 7.
+    const auto tokens = Tokens(40);
+    const auto hashes = SglangBigramHashes(tokens);
+    ASSERT_EQ(hashes.size(), 3u);
+
+    ASSERT_EQ(table.StoreGpu(Gpu({hashes[0], hashes[1], hashes[2]})), "");
+    ASSERT_EQ(table.StoreShared(Shared(hashes, StorageTier::kCpu)), "");
+    ASSERT_EQ(table.StoreShared(Shared(hashes, StorageTier::kDisk)), "");
+
+    // 39, not 40 and not the 48 a whole-block count would report.
+    const auto result = table.Query(TestContext(), tokens).at("instance-a");
+    EXPECT_EQ(result.longest_match_tokens, 39);
+    EXPECT_EQ(result.gpu, 39);
+    EXPECT_EQ(result.dp, (std::map<int64_t, int64_t>{{0, 39}}));
+    EXPECT_EQ(
+        result.rank_matches,
+        (std::map<int64_t, RankCacheHitResult>{{0, RankMatch(39, 39, 39)}}));
+    EXPECT_EQ(result.cpu, 39);
+    EXPECT_EQ(result.disk, 39);
+}
+
+TEST(Query, BigramWholeBlockRunIsUnaffectedByTheLogicalLengthClamp) {
+    PrefixCacheTable table;
+    auto registration = Registration();
+    registration.profile = SglangBigramProfile();
+    RegisterOrFail(table, registration);
+
+    const auto tokens = Tokens(40);
+    const auto hashes = SglangBigramHashes(tokens);
+    ASSERT_EQ(hashes.size(), 3u);
+
+    // Only the two whole blocks are indexed; the partial tail misses. The
+    // clamp must not pull this below the 32 positions those blocks hold.
+    ASSERT_EQ(
+        table.StoreShared(Shared({hashes[0], hashes[1]}, StorageTier::kDisk)),
+        "");
+
+    const auto result = table.Query(TestContext(), tokens).at("instance-a");
+    EXPECT_EQ(result.longest_match_tokens, 32);
+    EXPECT_EQ(result.gpu, 0);
+    EXPECT_EQ(result.cpu, 0);
+    EXPECT_EQ(result.disk, 32);
+}
+
 TEST(Query, GpuCpuAndDiskExtendOneCumulativePrefix) {
     PrefixCacheTable table;
     RegisterOrFail(table, Registration());