Repository navigation
feat(kv-index): router prefix hash for query-path disambiguation - #565
Conversation
Remove seq_hash verification from the PositionalIndexer query path.
Backends use different hash algorithms (SHA256 for SGLang, xxHash3 for
vLLM, custom for TRT-LLM) for their position-aware SequenceHash values.
The router cannot recompute these at query time, so matching must rely
solely on position-independent ContentHash (XXH3 of token IDs).
What changed:
- kv_index/src/event_tree.rs:
- Add all_workers() and total_worker_count() to SeqEntry for
content-hash-only worker set retrieval
- Mark SeqEntry::get() as #[cfg(test)] (only used by event tests)
- Simplify get_workers_at() to return all workers at (position,
content_hash) without seq_hash filtering
- Simplify count_workers_at() and linear_scan_drain() — remove
seq_hash parameters entirely
- Simplify jump_search_matches() — remove lazy seq_hash vector and
rolling hash computation from query path
- Mark compute_next_seq_hash() and ensure_seq_hash_computed() as
#[cfg(test)] (only used to verify rolling hash correctness)
- Add compute_request_content_hashes(tokens, block_size) public
function: chunks request tokens by backend block size, computes
XXH3 ContentHash per full block, discards trailing partial chunks
- Fix deadlock in test_seq_entry_single_to_multi_upgrade: scope the
DashMap Ref to release shard lock before apply_removed
- Add 10 new tests: 7 for compute_request_content_hashes (basic,
partial trailing, fewer than block size, empty, exact multiple,
zero block size panic, block size 1) and 3 end-to-end tests
(store+query, partial overlap, different backends same content)
- kv_index/src/lib.rs:
- Export compute_request_content_hashes
Why: The original query path computed rolling XXH3 seq_hashes that
would never match the backend-specific seq_hashes stored from events.
Content-hash positional matching works because workers with different
prefix histories naturally diverge at different content positions and
get excluded from the active set during jump search.
Seq_hash remains used for event processing (store uses it for SeqEntry
keying to track per-worker block ownership, remove uses it for reverse
lookup in worker_blocks).
Refs: Phase 1 KV cache event-driven routing, Task 4
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
The previous commit removed seq_hash verification from the query path because backends use different hash algorithms. This introduced a correctness gap: count_workers_at() returned total workers across all seq_hash variants, so the jump optimization could produce false-positive skips and inflated scores. Fix: compute a router-side rolling prefix hash (XXH3) from ContentHashes during apply_stored, and store it as the SeqEntry key. The same rolling hash is recomputed at query time via ensure_seq_hash_computed, restoring Dynamo's exact disambiguation semantics. The backend's SequenceHash stays in worker_blocks only, used for apply_removed reverse lookup. What changed in kv_index/src/event_tree.rs: - LevelIndex value: (position, content_hash) → (position, content_hash, prefix_hash) to carry the router hash for removal - apply_stored: compute rolling prefix hash per block, use as SeqEntry key instead of backend seq_hash - apply_removed: extract prefix_hash from worker_blocks, use it for SeqEntry::remove instead of backend seq_hash - remove_or_clear_worker: same prefix_hash change - Un-cfg compute_next_seq_hash and ensure_seq_hash_computed — now used in both store and query paths - Restore get_workers_lazy (prefix-hash filtered), count_workers_at (prefix-hash filtered), linear_scan_drain (prefix-hash filtered) - SeqEntry::get() un-cfg'd — used in production query path again - Remove SeqEntry::all_workers() and total_worker_count() — no longer needed (were only used by the content-hash-only query path) - Update module doc to describe the dual-hash scheme - Rewrite test_seq_entry_single_to_multi_upgrade: Multi now requires different prefix histories (different content at earlier position), not just different backend seq_hashes - Remove 3 trivial tests (test_stored_block_is_copy, test_debug_format, test_compute_content_hash_determinism) that tested compiler guarantees Refs: Phase 1 KV cache event-driven routing, Task 4 Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review infoConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughIntroduces a dual-hash flow: backend provides per-block SequenceHash and token_ids; router computes ContentHash sequence and a rolling prefix_hash sequence used as SeqEntry keys. Adds Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Router
participant Backend
participant LevelIndex
participant SeqEntry
Client->>Router: send tokens + block_size
Router->>Router: compute_request_content_hashes(tokens, block_size)
Router->>Router: derive ContentHash sequence and rolling prefix_hashes
Router->>Backend: send block tokens / content_hashes
Backend->>Backend: compute backend SequenceHash (block-level)
Backend->>LevelIndex: apply_stored(position, seq_hash, content_hashes, prefix_hashes)
LevelIndex->>LevelIndex: store worker reverse lookup (pos, content_hash, prefix_hash)
LevelIndex->>SeqEntry: register entries keyed by prefix_hash
Client->>SeqEntry: query by target prefix_hash
SeqEntry-->>Client: return worker positions
Client->>Router: request removal
Router->>LevelIndex: apply_removed(prefix_hash)
LevelIndex->>SeqEntry: remove entries by prefix_hash
SeqEntry->>SeqEntry: delete index entries
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a dual-hash scheme to disambiguate sequences in the KV index, a crucial fix for the jump search optimization. By computing a router-side rolling prefix hash, it restores precise filtering and prevents false-positive jumps. However, a potential Denial of Service (DoS) vulnerability was identified where the application could panic and crash if provided with an invalid block_size of 0 in the compute_request_content_hashes function, which violates the rule against panics in code interacting with external systems. Additionally, a code review comment suggests refactoring duplicated logic in a match statement to improve clarity and maintainability, aligning with the rule to avoid duplication in such constructs.
| pub fn compute_request_content_hashes(tokens: &[u32], block_size: usize) -> Vec<ContentHash> { | ||
| assert!(block_size > 0, "block_size must be greater than 0"); | ||
| tokens | ||
| .chunks(block_size) | ||
| .filter(|chunk| chunk.len() == block_size) | ||
| .map(compute_content_hash) | ||
| .collect() | ||
| } |
There was a problem hiding this comment.
The function compute_request_content_hashes uses an assert! to verify that block_size is greater than 0. In Rust, assert! is always enabled, and if the condition fails, the entire process will panic and crash. Since this function is intended to be the entry point for the query path and block_size is likely derived from backend metadata or external configuration, an unexpected or malicious value of 0 could lead to a Denial of Service (DoS) by crashing the router. It is recommended to handle the zero case gracefully (e.g., by returning an empty vector or an error) instead of panicking.
| pub fn compute_request_content_hashes(tokens: &[u32], block_size: usize) -> Vec<ContentHash> { | |
| assert!(block_size > 0, "block_size must be greater than 0"); | |
| tokens | |
| .chunks(block_size) | |
| .filter(|chunk| chunk.len() == block_size) | |
| .map(compute_content_hash) | |
| .collect() | |
| } | |
| pub fn compute_request_content_hashes(tokens: &[u32], block_size: usize) -> Vec<ContentHash> { | |
| if block_size == 0 { | |
| return Vec::new(); | |
| } | |
| tokens | |
| .chunks(block_size) | |
| .filter(|chunk| chunk.len() == block_size) | |
| .map(compute_content_hash) | |
| .collect() | |
| } |
References
- Do not introduce panics in code that interacts with external systems if the upstream server does not handle the error. Instead, handle the error gracefully or propagate it appropriately.
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kv_index/src/event_tree.rs`:
- Around line 321-325: The current get_mut → drop(...) → remove(...) sequence is
racy and can delete an entry that a concurrent apply_stored just recreated;
replace it with the HashMap Entry API so deletion is atomic: use
self.index.entry((position, content_hash)) to match Occupied(entry) and call
entry.get_mut().remove(prefix_hash, worker), then only call entry.remove_entry()
(or entry.remove()) if the entry is now empty; apply the same change to the
other occurrence (the block around lines 369-373) to avoid lost concurrent
inserts.
- Around line 520-528: The shortcut jump that sets current_pos = next_pos when
count_workers_at(next_pos, ...) == active.len() assumes workers have
prefix-closed chains but apply_removed can create holes; either enforce
prefix-closure in apply_removed (only allow removing trailing blocks per worker
by SequenceHash so no intermediate gaps are possible) or, if keeping current
apply_removed behavior, add an intermediate-position validation before the jump:
for the loop that computes seq_hash and calls count_workers_at, verify that
every worker present at next_pos is also present for all intermediate positions
(or that the rolling seq_hash continuity matches per-worker SequenceHash) and
only set current_pos = next_pos when that validation succeeds; reference
functions/vars: apply_removed, SequenceHash, seq_hash, count_workers_at,
current_pos, next_pos, active.len().
| let count = self.count_workers_at( | ||
| next_pos, | ||
| content_hashes[next_pos], | ||
| &mut seq_hashes, | ||
| content_hashes, | ||
| ); | ||
|
|
||
| if num_workers_at_next == active.len() { | ||
| // All active workers still present at jump destination — skip ahead | ||
| if count == active.len() { | ||
| current_pos = next_pos; |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n kv_index/src/event_tree.rs | sed -n '500,540p'Repository: lightseekorg/smg
Length of output: 1545
🏁 Script executed:
rg -A 20 "fn count_workers_at" kv_index/src/event_tree.rsRepository: lightseekorg/smg
Length of output: 700
🏁 Script executed:
rg -A 30 "fn linear_scan_drain" kv_index/src/event_tree.rsRepository: lightseekorg/smg
Length of output: 1062
🏁 Script executed:
rg -A 50 "fn linear_scan_drain" kv_index/src/event_tree.rs | head -60Repository: lightseekorg/smg
Length of output: 1617
🏁 Script executed:
rg -B 10 "fn jump_search_matches" kv_index/src/event_tree.rs | head -20Repository: lightseekorg/smg
Length of output: 377
🏁 Script executed:
rg -B 5 "insert\|remove" kv_index/src/event_tree.rs | head -80Repository: lightseekorg/smg
Length of output: 42
🏁 Script executed:
cat -n kv_index/src/event_tree.rs | head -100Repository: lightseekorg/smg
Length of output: 4743
🏁 Script executed:
rg -A 30 "apply_removed\|remove_worker" kv_index/src/event_tree.rs | head -80Repository: lightseekorg/smg
Length of output: 42
🏁 Script executed:
rg "pub fn|fn " kv_index/src/event_tree.rs | head -40Repository: lightseekorg/smg
Length of output: 2048
🏁 Script executed:
rg -A 20 "pub fn apply_removed" kv_index/src/event_tree.rsRepository: lightseekorg/smg
Length of output: 825
🏁 Script executed:
rg -A 40 "pub fn apply_removed" kv_index/src/event_tree.rsRepository: lightseekorg/smg
Length of output: 1510
🏁 Script executed:
rg -A 25 "fn get_workers_lazy" kv_index/src/event_tree.rsRepository: lightseekorg/smg
Length of output: 958
🏁 Script executed:
rg -A 15 "fn compute_next_seq_hash\|fn ensure_seq_hash_computed" kv_index/src/event_tree.rsRepository: lightseekorg/smg
Length of output: 42
🏁 Script executed:
rg "compute_next_seq_hash\|ensure_seq_hash" kv_index/src/event_tree.rs -A 10Repository: lightseekorg/smg
Length of output: 42
🏁 Script executed:
rg -A 10 "ensure_seq_hash_computed" kv_index/src/Repository: lightseekorg/smg
Length of output: 3821
🏁 Script executed:
cat -n kv_index/src/event_tree.rs | sed -n '385,410p'Repository: lightseekorg/smg
Length of output: 1392
The short-circuit logic assumes workers maintain prefix-closed block chains; this invariant is not enforced.
When count == active.len(), the jump skips all intermediate positions between current_pos and next_pos. This optimization assumes any worker present at next_pos must have all blocks in the prefix chain—but apply_removed removes blocks individually by SequenceHash and does not prevent "holes" (e.g., a worker losing position 5 while retaining position 10).
Since the prefix hash (seq_hash) is computed as a rolling XXH3 across the entire sequence up to each position, a missing intermediate block corrupts the hash chain, potentially causing workers with incomplete prefixes to be treated as having continuous overlaps and inflating scores.
Enforce the invariant: either prevent apply_removed from creating gaps in a worker's block range, or add an intermediate position check in the loop to validate prefix closure before jumping.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@kv_index/src/event_tree.rs` around lines 520 - 528, The shortcut jump that
sets current_pos = next_pos when count_workers_at(next_pos, ...) == active.len()
assumes workers have prefix-closed chains but apply_removed can create holes;
either enforce prefix-closure in apply_removed (only allow removing trailing
blocks per worker by SequenceHash so no intermediate gaps are possible) or, if
keeping current apply_removed behavior, add an intermediate-position validation
before the jump: for the loop that computes seq_hash and calls count_workers_at,
verify that every worker present at next_pos is also present for all
intermediate positions (or that the rolling seq_hash continuity matches
per-worker SequenceHash) and only set current_pos = next_pos when that
validation succeeds; reference functions/vars: apply_removed, SequenceHash,
seq_hash, count_workers_at, current_pos, next_pos, active.len().
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: da40c5957e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
apply_stored now returns Result<(), ApplyError> instead of () so callers (Task 5 KvEventMonitor) can distinguish success from silently dropped events. ApplyError has WorkerNotTracked and ParentBlockNotFound variants. New test coverage (50 → 63 tests): - Jump boundary divergence at positions 31/32/33/63/64/65 on 128-block sequences with jump_size=32 - Off-by-one at jump boundaries (31, 33, 63, 65, 95, 97 block sequences) - Exact jump_size multiples (32, 64, 96) - Staggered workers across jump boundaries (depths 10/20/35/64/100) - Shared prefix divergence at jump boundary (3 workers share 40 blocks) - Very long sequence (1000 blocks: full/prefix/mid-divergence queries) - Deep continuation chain (200 blocks via 20 continuations of 10) - Multiple disjoint sequences per worker - Long sequence partial removal (remove blocks 80-99 from 100) - Remove parent does not cascade (linear scan detects gap) - Long sequence clear and rebuild with different content - Interleaved 4 workers at staggered depths (25/50/75/100) Minor: remove over-comment on From<i64>, add comments for position-0 base case and jump heuristic, use fmt import consistently. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kv_index/src/event_tree.rs`:
- Around line 1829-1854: Add a note to the apply_removed method's documentation
(in event_tree.rs) explaining that deleting an intermediate/parent block does
not cascade to child blocks and can leave orphaned/unreachable entries in the
positional indexer (they remain in storage but may be unreachable via prefix
matching, especially with jump_size=1 linear scans); update the doc comment for
PositionalIndexer::apply_removed (or the apply_removed function) to mention this
limitation and its impact on memory versus correctness so callers are aware.
- compute_request_content_hashes: return empty vec instead of panicking on block_size=0 (warn via tracing) - apply_removed / remove_or_clear_worker: use DashMap entry() API for atomic modify-and-remove, fixing the TOCTOU race between get_mut → drop → remove where a concurrent insert could be lost - linear_scan_drain: consolidate duplicated drain logic with and_then, eliminating repeated drain-break blocks - find_matches: document prefix-closed chain assumption in the jump heuristic (backends evict tail blocks via LRU) - apply_removed: document orphaned entry behavior when mid-sequence blocks are removed Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
1 participant
Summary
Add router-computed rolling prefix hash for precise jump search disambiguation, and add
compute_request_content_hashes()as the query-path entry point. This completes Task 4 (Hash Computation) of Phase 1 KV cache event-driven routing.Refs: Phase 1 KV cache event-driven routing, Task 4
What changed
kv_index/src/event_tree.rs:
Dual-hash scheme for multi-backend support:
block_hash(SequenceHash) using backend-specific algorithms (SHA256 for SGLang, xxHash3 for vLLM, custom for TRT-LLM). The router cannot recompute these at query time.apply_storedand stores it as the SeqEntry key. The same rolling hash is recomputed at query time viaensure_seq_hash_computedfor precise disambiguation — matching Dynamo's exact semantics.worker_blocksonly, used forapply_removedreverse lookup.Store path changes:
LevelIndexvalue extended:(position, content_hash)→(position, content_hash, prefix_hash)to carry the router hash for removalapply_stored: computes rolling prefix hash per block, uses it as SeqEntry key instead of backend seq_hashapply_removed: extracts prefix_hash from worker_blocks, uses it forSeqEntry::removeremove_or_clear_worker: same prefix_hash changeQuery path changes:
compute_next_seq_hashandensure_seq_hash_computedpromoted from#[cfg(test)]to production — used in both store and query pathsget_workers_lazy: filters by router prefix hash (not all_workers)count_workers_at: counts only workers matching the prefix hashlinear_scan_drain: intersects active set against prefix-hash-filtered workersSeqEntry::get()promoted from#[cfg(test)]to productionSeqEntry::all_workers()andtotal_worker_count()(no longer needed)New public function:
compute_request_content_hashes(tokens, block_size): chunks request tokens by backend block size, computes XXH3 ContentHash per full block, discards trailing partial chunks. This is the query-path entry point that feeds intofind_matches().Tests:
test_seq_entry_single_to_multi_upgradefor correct prefix-hash semantics (Multi requires different prefix histories, not just different backend seq_hashes)compute_request_content_hashes+ 3 end-to-end teststest_stored_block_is_copy,test_debug_format,test_compute_content_hash_determinism)kv_index/src/lib.rs:
compute_request_content_hashesWhy
The PositionalIndexer's jump optimization relies on precise worker counts at jump destinations. Without prefix hash filtering,
count_workers_at()returned total workers across all seq_hash variants, causing false-positive jumps when unrelated workers coincidentally produced the same count as the active set — leading to inflated scores. The router prefix hash restores Dynamo's exact disambiguation: only workers sharing the same content prefix chain are counted.Test plan
cargo test -p kv-index -- event_tree— 50 tests passcargo clippy -p kv-index --all-targets -- -D warnings— cleanmake fmt— formattedSummary by CodeRabbit
New Features
Improvements
Tests