refactor(kvcache): fold prefix matching into one accumulator - #2669
Draft
nilig wants to merge 1 commit into
Draft
Conversation
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>
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What type of PR is this?
/kind cleanup
What this PR does / why we need it:
Longest-prefix matching is implemented three times:
LongestPrefixScorerover a materializedLookupfor the weighted score, and the precise-prefix producer'smatchedBlockCountandmatchedBlockCountByTierfor 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.MatchBlockKeysbecomes 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.ScoreTokensandLongestPrefixScorer.Scoreare projections of it; the producer only aggregates per prompt and publishes, all or nothing.KVBlockScorerandNewTracedScorerstay as deprecated compatibility surfaces; existingscore_tokensspan attributes keep their meaning, and the matcher's ownmatch_block_keysspan 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,
ScoreTokenson 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:
Test plan:
score_tokenstelemetry keeps its meaning (span test)make presubmit