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

[pjdurden]

Issue #4153 — incorrect external prefix-cache reuse from partial hash hits

1. Root cause

PrefixCacheTable::Query converts a matched block count into a token count
with TokensForBlocks(cursor, context.block_size), i.e. it assumes every block
in the query's hash chain covers exactly block_size tokens.

That assumption holds for the vLLM chain, whose BlockCount() is
token_ids.size() / block_size (complete blocks only). It does not hold for
the SGLang chains: SglangHashChain deliberately mirrors the engine and ends
its chain with a partial block —

block_count_ = token_ids.size() / block_size_ + (token_ids.size() % block_size_ != 0)

(see SglangHashChain in mooncake-conductor/src/prefixindex/hash_strategy.cpp
and the existing SglangHashChain.IncludesPartialFinalBlock test).

So whenever a query's prompt length is not a multiple of block_size and
the trailing partial block is present in the index, Query reported

ceil(n / block_size) * block_size  tokens

for a prompt of only n tokens. Concretely, with block_size = 16 and a
40-token prompt the conductor answered longest_match_tokens = 48 — 8 tokens
that no indexed block ever covered, and that lie past the end of the prompt
altogether.

This is the mechanism described in the issue: the hit is real at the hash
granularity, but the reported reusable length runs past the region the hash
actually pinned down. A router or scheduler that trusts that length skips
recomputing KV slots whose contents belong to whatever request wrote them,
which surfaces as out-of-order / off-topic generation. The over-report grows
with block_size, which is why it only became visible on a deployment where
hash blocks and attention blocks were not aligned, and why forcing the two to
be equal (so no partial tail is ever hit) made it disappear.

The three tiers are all affected, because gpu, cpu and disk share one
cumulative cursor and all three go through the same conversion.

2. The fix and why

mooncake-conductor/src/prefixindex/prefix_indexer.cpp

Added a small MatchedTokens(block_count, block_size, queried_tokens) helper
next to TokensForBlocks, and used it for the gpu / cpu / disk numbers.
It is TokensForBlocks clamped to the number of tokens the query actually
carries.

Why a clamp is exactly right rather than merely safe:

• vLLM chain — block_count * block_size <= n by construction, so the
  clamp never fires. Behaviour is bit-for-bit unchanged.
• SGLang chain — blocks [0, c) cover min(c * block_size, n) tokens, so
  the clamp reproduces the true covered length, including the full-match case
  where the answer should be the whole prompt.
• SGLang bigram chain — the chain covers n - 1 bigram positions. Matching
  all of them pins every one of the n tokens, so clamping to n is correct;
  a partial run still reports c * block_size, which is one token short of
  what the bigram overlap strictly proves and therefore conservative.

Under-reporting is a missed cache hit; over-reporting is silent corruption, so
where the two chains disagree the code now rounds toward the safe side.

The alternative — dropping the partial block from the SGLang chain entirely —
was rejected: including it is a deliberate, tested property that mirrors what
the engine publishes, and removing it would throw away genuine hits. The bug is
in the accounting, not in the hashing.

The design doc claimed /query "ignores trailing partial blocks", which was
already untrue for the SGLang strategies before this change. That sentence now
states the per-strategy behaviour and the token-count rule.

3. Files changed

File	Change
mooncake-conductor/src/prefixindex/prefix_indexer.cpp	New MatchedTokens helper; Query uses it for the gpu/cpu/disk token counts.
mooncake-conductor/tests/prefix_indexer_test.cpp	SglangProfile() / SglangHashes() helpers plus two regression tests.
docs/source/design/conductor/conductor-architecture-design.md	Corrected the one-sentence description of partial-block handling in /query.

4. Risk / uncertainty

• Low blast radius. The only behaviour change is a std::min, and it can
  only lower a reported number. For the vLLM strategy it is provably a no-op
  (block_count = floor(n / block_size)), so existing vLLM deployments and all
  31 pre-existing conductor unit tests are unaffected.
• Routing shifts slightly for SGLang. Instances whose match previously
  rounded up to a block boundary now report the real prompt length. That can
  change which instance a cache-aware router picks when two candidates are
  within one block of each other. This is the intended correction.
• Bigram partial runs stay conservative (c * block_size rather than
  c * block_size + 1). Deliberate: it can only cost a single token of reuse,
  never invent one. Worth a maintainer's confirmation if bigram partial-run
  precision matters to them.
• Not addressed here (out of scope for this issue): ValidateLayout still
  requires effective_block_size == context.block_size, so a publisher whose
  attention block differs from its hash block is rejected at registration
  rather than handled. If the conductor ever needs to index a engine whose
  page size exceeds the hash unit, the match would additionally need rounding
  down to the page boundary. Nothing in the current tree exercises that path.
• The issue text also mentions eviction failures and a near-full store. Those
  are separate symptoms; this change does not touch them.

5. How it was verified

The container has no glog/gtest/OpenSSL dev packages, so the full CMake build
could not run. The affected translation units were instead compiled and
executed directly, against the real OpenSSL (libcrypto.so.3, EVP path) so
the hash chain is genuine, with a stub header only for glog's two logging
macros:

g++ -std=c++20 -DCONDUCTOR_HAS_LOWLEVEL_SHA256=0 \
    -I mooncake-conductor/include -I mooncake-common/include -I <glog-stub> \
    mooncake-conductor/src/prefixindex/{hash_strategy,prefix_indexer}.cpp

1. Reproduced the bug. A driver registering an sglang profile with
   block_size = 16, storing all 3 chain hashes of a 40-token prompt, then
   querying it. Against HEAD's prefix_indexer.cpp:

   block count: 3
   full-match  gpu=48 cpu=48 disk=48 longest=48 (prompt=40)
   

   48 reported tokens for a 40-token prompt — 8 phantom tokens.

2. Confirmed the fix. Same driver against the patched file:

   full-match  gpu=40 cpu=40 disk=40 longest=40 (prompt=40)
   whole-only  gpu=0 cpu=0 disk=32 longest=32
   

3. Ran the real unit-test file. prefix_indexer_test.cpp was linked against
   a fused gtest and run:

   • patched build: 32/32 pass, including the two new tests;
   • HEAD build with the new tests applied: 31 pass, 1 fail —
     Query.TrailingPartialBlockNeverReportsMoreThanPromptTokens.

   So the new test is a genuine regression guard, and no pre-existing conductor
   test changed behaviour.

4. Formatting. pre-commit is not installed here; clang-format (the
   mooncake-code-format hook's tool) was run over both changed C++ files and
   produced no further edits. Trailing whitespace and end-of-file newlines were
   checked by hand on all three files. make html was not run for the
   docs/ change — the Sphinx toolchain is not available offline; the edit is a
   prose-only sentence replacement inside an existing paragraph, with no links,
   directives, or toctree entries touched.

[Copilot]

🟡 Changes recommended

Add sglang_bigram regression coverage and align the remaining query documentation.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This PR fixes over-reported SGLang prefix-cache matches by clamping reusable token counts to the queried prompt length.

Changes:

• Adds clamped GPU, CPU, and disk match accounting.
• Adds SGLang partial-block regression tests.
• Updates conductor design documentation.

File summaries

File	Summary
mooncake-conductor/tests/prefix_indexer_test.cpp	Adds partial-block query coverage.
mooncake-conductor/src/prefixindex/prefix_indexer.cpp	Clamps matched token counts.
docs/source/design/conductor/conductor-architecture-design.md	Documents strategy-specific partial-block handling.

Review details

Suppressed comments (1)

mooncake-conductor/tests/prefix_indexer_test.cpp:650

• The new query coverage only registers sglang; sglang_bigram has a different block-count formula (ceil((n - 1) / block_size)) but also reaches this accounting path. Add a regression query with a non-multiple prompt to verify its conservative partial-tail result, since the existing bigram test only checks hashing and would not catch a token-accounting regression.

TEST(Query, TrailingPartialBlockNeverReportsMoreThanPromptTokens) {
    PrefixCacheTable table;
    auto registration = Registration();
    registration.profile = SglangProfile();
    RegisterOrFail(table, registration);

• Files reviewed: 3/3 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[he-yufeng]

Mechanism verified by hand against the diff, no build needed: pre-fix, a query whose chain ends in a partial block reports block_count * block_size per tier, which over-reports by up to a block; TrailingPartialBlockNeverReportsMoreThanPromptTokens pins exactly that (3 blocks x 16 tokens over a 40-token prompt must read 40, not 48), and the arithmetic fails pre-fix and passes with the clamp. WholeBlockRunIsUnaffectedByThePromptLengthClamp covers the other direction (2 whole blocks still report 32). The clamp is the right upper bound: a query can never reuse more tokens than it carries, and for vLLM-shaped chains (complete blocks only) the clamp is a no-op, so existing deployments are untouched. The doc update correctly scopes the two chain shapes. LGTM on the mechanism, with the note that I verified by reading and arithmetic, not by compiling the conductor locally.

[Misak2333]

@pjdurden
Thanks for the fix!

1. LGTM for the SGLang partial-block accounting fix. I verified it with before/after unit tests: a 40-token query with a block size of 16 previously reported 48 matched tokens; this patch correctly returns 40.
2. The sglang_bigram strategy needs a different upper bound. When counting reusable logical KV positions, n raw tokens correspond to max(n - 1, 0) bigram positions. The result should therefore be capped at that length, while still respecting the number of matched blocks.

You can simply select the bound once in PrefixCacheTable::Query():

size_t queried_tokens = token_ids.size();
if (state->profile.strategy == "sglang_bigram" &&
queried_tokens > 0) {
--queried_tokens;
}

Still, with 40 raw tokens and a block size of 16, a full bigram match now should report 39 logical positions.

[pjdurden]

@Misak2333 thanks for verifying part 1 against real before/after runs, and you are right on part 2. Pushed as 3ccff4c4.

I checked the bound against the chain construction rather than taking it on faith, and it holds. In SglangHashChain, block_count_ for bigram is built over token_ids.size() - 1, and Compute walks each block against logical_length = token_ids_.size() - 1. So the chain genuinely covers n - 1 logical positions and the raw-token clamp was still one past what any indexed block holds.

Took your placement, selecting the bound once in Query():

const size_t queried_tokens =
    (state->profile.strategy == "sglang_bigram" && !token_ids.empty())
        ? token_ids.size() - 1
        : token_ids.size();

I used !token_ids.empty() rather than queried_tokens > 0 only so the bound reads as a property of the prompt, and it gives the same max(n - 1, 0). MatchedTokens still takes the min against the matched block count, so the block side keeps bounding it from the other direction.

Working through the numbers with the two formulas as written:

n	blocks	bound	full match	unclamped
40	3	39	39	48
17	1	16	16	16
2	1	1	1	16
1	0	0	0	0
0	0	0	0	0

which matches the 39 you predicted for the 40-token case, and degrades sanely at n < 2 where the chain has no blocks at all.

Added two tests. BigramFullMatchReportsOneFewerPositionThanPromptTokens covers your case directly, 40 raw tokens at block size 16 reporting 39 across gpu/cpu/disk and the per-rank matches. BigramWholeBlockRunIsUnaffectedByTheLogicalLengthClamp indexes only the two whole blocks and asserts 32, so the new bound cannot quietly pull whole-block matches down. Both use a new SglangBigramProfile helper alongside the existing SglangProfile.

One caveat I want to be straight about: I could not build this locally. The conductor's standalone build needs glog, JsonCpp, OpenSSL and msgpack headers, and I have no way to install them on this machine, so the table above is the arithmetic from the two formulas rather than output from a run. The logic and the expected values should be right, but the numbers in those two tests are worth a look from CI or from you before this merges.

[Chase-Rong]

@Misak2333 thanks for verifying part 1 against real before/after runs, and you are right on part 2. Pushed as 3ccff4c4.

I checked the bound against the chain construction rather than taking it on faith, and it holds. In SglangHashChain, block_count_ for bigram is built over token_ids.size() - 1, and Compute walks each block against logical_length = token_ids_.size() - 1. So the chain genuinely covers n - 1 logical positions and the raw-token clamp was still one past what any indexed block holds.

Took your placement, selecting the bound once in Query():

const size_t queried_tokens =
    (state->profile.strategy == "sglang_bigram" && !token_ids.empty())
        ? token_ids.size() - 1
        : token_ids.size();

I used !token_ids.empty() rather than queried_tokens > 0 only so the bound reads as a property of the prompt, and it gives the same max(n - 1, 0). MatchedTokens still takes the min against the matched block count, so the block side keeps bounding it from the other direction.

Working through the numbers with the two formulas as written:

n blocks bound full match unclamped
40 3 39 39 48
17 1 16 16 16
2 1 1 1 16
1 0 0 0 0
0 0 0 0 0
which matches the 39 you predicted for the 40-token case, and degrades sanely at n < 2 where the chain has no blocks at all.

Added two tests. BigramFullMatchReportsOneFewerPositionThanPromptTokens covers your case directly, 40 raw tokens at block size 16 reporting 39 across gpu/cpu/disk and the per-rank matches. BigramWholeBlockRunIsUnaffectedByTheLogicalLengthClamp indexes only the two whole blocks and asserts 32, so the new bound cannot quietly pull whole-block matches down. Both use a new SglangBigramProfile helper alongside the existing SglangProfile.

One caveat I want to be straight about: I could not build this locally. The conductor's standalone build needs glog, JsonCpp, OpenSSL and msgpack headers, and I have no way to install them on this machine, so the table above is the arithmetic from the two formulas rather than output from a run. The logic and the expected values should be right, but the numbers in those two tests are worth a look from CI or from you before this merges.

LGTM, thanks for your fixs.