Skip to content

refactor(kvcache): fold prefix matching into one accumulator - #2669

Draft
nilig wants to merge 1 commit into
llm-d:mainfrom
nilig:perf/prefix-matcher-a
Draft

refactor(kvcache): fold prefix matching into one accumulator#2669
nilig wants to merge 1 commit into
llm-d:mainfrom
nilig:perf/prefix-matcher-a

Conversation

@nilig

@nilig nilig commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

What type of PR is this?
/kind cleanup

What this PR does / why we need it:
Longest-prefix matching is implemented three times: LongestPrefixScorer over a materialized Lookup for the weighted score, and the precise-prefix producer's matchedBlockCount and matchedBlockCountByTier for the block and per-tier counts it publishes, re-walking the lookup result once per endpoint. No single place defines the contract, and the re-walk grows with fleet size.

Indexer.MatchBlockKeys becomes the one entry point. A pooled accumulator implements the rules once (candidates are the pods holding the first key, a chain ends at the pod's first missing key, duplicate entries at a key take the highest weight, every tier tracks its own chain) and folds each key's entries in one call. ScoreTokens and LongestPrefixScorer.Score are projections of it; the producer only aggregates per prompt and publishes, all or nothing. KVBlockScorer and NewTracedScorer stay as deprecated compatibility surfaces; existing score_tokens span attributes keep their meaning, and the matcher's own match_block_keys span reports the longest chain.

Ownership after this change: kvblock owns storage and reads; kvcache owns the one matcher, its fallback, and match telemetry; the producer owns prompt aggregation and publication. Index reads are unchanged: this PR materializes Lookup. A follow-up adds an ordered index walk that feeds the same accumulator without materializing.

M4 Max, n=6, ScoreTokens on 1,034 blocks held by 40 endpoints with eight rank entries each (BenchmarkScoreTokensKermit): 294.5ms and 4.6 GiB per request on main, 67.7ms and 116.5 MiB here.

Which issue(s) this PR fixes:
Part of #2390

Release note:

The precise prefix-cache producer and `kvcache.Indexer.ScoreTokens` share one prefix matcher (`Indexer.MatchBlockKeys`), which emits a `match_block_keys` span. `ScoreTokens` no longer emits `compute_scores`; `kvcache.KVBlockScorer` and `kvcache.NewTracedScorer` are deprecated.

Test plan:

  • matcher parity with the replaced scorer and producer counters on 500 random fixtures (filters, gaps, duplicate pod and rank entries, unconfigured tiers, speculative entries)
  • producer aggregation across prompts, multimodal match, candidate filtering, error and cancellation publish nothing
  • score_tokens telemetry keeps its meaning (span test)
  • make presubmit

Longest-prefix matching is spread over implementations that each read
the same rules: LongestPrefixScorer over a materialized Lookup for the
weighted score, and the precise-prefix producer's matchedBlockCount and
matchedBlockCountByTier for the block and per-tier counts it publishes,
re-walking the lookup result once per endpoint. No single place defines
the contract, so nothing can test it in isolation, and the per-endpoint
re-walk grows with fleet size.

Indexer.MatchBlockKeys is the one entry point. A pooled prefixAccumulator
implements the rules once (candidates are the pods holding the first
key, a chain ends at the pod's first missing key, duplicate entries at a
key take the highest weight, every tier tracks its own chain) and folds
each key's entries in one call, skipping repeated rank entries of one
endpoint. ScoreTokens and LongestPrefixScorer.Score are projections of
it; the producer aggregates per prompt and publishes all or nothing.
KVBlockScorer and NewTracedScorer stay as deprecated compatibility
surfaces; ScoreTokens no longer scores through them. A property test
checks the accumulator against the three replaced algorithms on random
fixtures.

Index reads are unchanged: the matcher materializes Lookup. The
score_tokens span keeps blocks_found and block_hit_ratio as the count of
requested keys held by any candidate; the matcher's own span reports the
longest chain. M4 Max, n=6, ScoreTokens on the Kermit shape (1,034
blocks held by 40 endpoints with eight rank entries each): 294.5ms and
4.6GiB per request before, 67.7ms and 116.5MiB after.

Signed-off-by: nilig <nili.ifergan@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/epp area/kvcache area/scheduling kind/cleanup size/XXL Denotes a PR that changes 1000+ lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant