feat(kv-router): add standalone selector peer sync - #10745
Merged
Conversation
Signed-off-by: PeaBrane <yanrpei@gmail.com>
github-actions
Bot
added
feat
documentation
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
Contributor
Contributor
WalkthroughIntroduces a ChangesStandalone Selection Service, Replica Sync, and Overlap Refactor
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
lib/kv-router/src/services/selection/tests.rs (1)
587-588: ⚡ Quick winMake the local-only output-block assertion time-bounded, not snapshot-based.
A single 50ms sleep plus one
assert_core_loadcan miss a delayed replica event, so this test may still pass if output-block replication is accidentally queued later. Assert the remote core stays unchanged for a short bounded window.Suggested test tightening
- tokio::time::sleep(Duration::from_millis(50)).await; - assert_core_load(&core_b, 1, 1, 4); + assert_core_load_stays(&core_b, 1, 1, 4, Duration::from_millis(100)).await;Add a helper near the existing load helpers:
async fn assert_core_load_stays( core: &SelectionCore, expected_requests: usize, expected_blocks: usize, expected_tokens: usize, duration: Duration, ) { let deadline = tokio::time::Instant::now() + duration; loop { assert_core_load(core, expected_requests, expected_blocks, expected_tokens); if tokio::time::Instant::now() >= deadline { break; } tokio::time::sleep(Duration::from_millis(10)).await; } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/kv-router/src/services/selection/tests.rs` around lines 587 - 588, Replace the snapshot-based assertion with a time-bounded assertion that ensures the core state remains stable over a duration window. Add a new helper function assert_core_load_stays that takes a SelectionCore reference, expected values for requests, blocks, and tokens, and a Duration parameter. This helper should loop from the current time until the deadline (current time plus the duration), repeatedly calling the existing assert_core_load function and sleeping briefly between iterations. Then replace the current tokio::time::sleep call followed by a single assert_core_load call for core_b with a call to assert_core_load_stays passing the expected values (1, 1, 4) and an appropriate duration to catch delayed replica events.lib/kv-router/src/services/selection/scoring.rs (1)
7-10: ⚡ Quick winUse the hot-path map type for
mooncake_summaries.
mooncake_summariesis a private selection input keyed by numericWorkerId; use the project’sFxHashMapconvention here, while keeping standard maps for text-keyed wire responses.As per coding guidelines, "
lib/kv-router/**/*.rs: UseFxHashMap/FxHashSetwhen possible for internal numeric keys and hot paths in lib/kv-router."Also applies to: 14-19
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/kv-router/src/services/selection/scoring.rs` around lines 7 - 10, The mooncake_summaries map uses a standard map type for numeric WorkerId keys on a hot path, which should use the faster FxHashMap per project conventions for numeric keys in lib/kv-router. Replace the map type used for mooncake_summaries with FxHashMap while keeping standard maps for any text-keyed wire responses. This change should be applied to the mooncake_summaries declaration and any similar internal numeric-keyed maps in the same scope.Source: Coding guidelines
lib/kv-router/src/scheduling/overlap.rs (2)
4-15: ⚡ Quick winUse
FxHashMapfor new internal numeric worker-keyed hot-path maps. The shared issue is that the new overlap/selection maps are keyed by numeric worker identifiers inlib/kv-router, but use standardHashMap.
lib/kv-router/src/scheduling/overlap.rs#L4-L15: changeCacheHitEstimates.effective_overlap_blocksandCacheHitEstimates.cached_tokensto the project-standardFxHashMaptype.lib/kv-router/src/services/selection/scoring.rs#L7-L19: changeOverlapInputs.mooncake_summariesto the sameFxHashMaptype.As per coding guidelines, "
lib/kv-router/**/*.rs: UseFxHashMap/FxHashSetwhen possible for internal numeric keys and hot paths in lib/kv-router."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/kv-router/src/scheduling/overlap.rs` around lines 4 - 15, Replace standard HashMap with FxHashMap for numeric worker-keyed hot-path maps in both affected locations. In lib/kv-router/src/scheduling/overlap.rs (lines 4-15), change the type of CacheHitEstimates.effective_overlap_blocks and CacheHitEstimates.cached_tokens fields from HashMap to FxHashMap. In lib/kv-router/src/services/selection/scoring.rs (lines 7-19), change the type of OverlapInputs.mooncake_summaries field from HashMap to FxHashMap. Ensure FxHashMap is imported from the appropriate crate (typically rustc_hash or similar) in both files.Source: Coding guidelines
4-4: ⚡ Quick winUse
FxHashMapfor worker-keyed overlap estimates.These new maps are keyed by internal numeric
WorkerWithDpRankvalues on the overlap/scheduling path, so they should follow thelib/kv-routerhot-path collection guidance instead of standardHashMap.As per coding guidelines, "
lib/kv-router/**/*.rs: UseFxHashMap/FxHashSetwhen possible for internal numeric keys and hot paths in lib/kv-router."Also applies to: 13-15, 23-24
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/kv-router/src/scheduling/overlap.rs` at line 4, Replace the HashMap import with FxHashMap from the appropriate module at the import statement. Then, replace all usages of HashMap with FxHashMap throughout the overlap.rs file, particularly in the worker-keyed overlap estimate structures. Since these maps use internal numeric WorkerWithDpRank keys on the hot scheduling path, they must follow lib/kv-router guidelines by using FxHashMap instead of the standard HashMap for better performance. The comment indicates this change applies at multiple locations (lines 13-15 and 23-24 in addition to the import), so ensure all HashMap references keyed by WorkerWithDpRank or used in hot-path code are updated to FxHashMap.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/components/router/standalone-selection.md`:
- Around line 170-178: The reservation example in the `/reservations` request
body shows effective_prefill_tokens set to 4, but the `/select` response example
returns effective_prefill_tokens: 384. Since the documentation instructs to pass
the exact value from the `/select` response into the `/reservations` call,
update the effective_prefill_tokens value in this reservation example from 4 to
384 to demonstrate the correct forwarding of the value and maintain consistency
across the examples.
In `@docs/components/router/standalone-slot-tracker.md`:
- Around line 77-79: Remove the NOTE paragraph about output-block updates not
being replica-synchronized from the standalone-slot-tracker.md file (the text
spanning lines 77-79 that begins with "NOTE: Output-block updates remain
local..."). This note is redundant since the file already states that slot
tracker excludes output-block updates. If the bandwidth caveat is considered
important documentation, move it to the standalone-selection.md file instead,
but the primary action is to delete this redundant paragraph from the current
location.
---
Nitpick comments:
In `@lib/kv-router/src/scheduling/overlap.rs`:
- Around line 4-15: Replace standard HashMap with FxHashMap for numeric
worker-keyed hot-path maps in both affected locations. In
lib/kv-router/src/scheduling/overlap.rs (lines 4-15), change the type of
CacheHitEstimates.effective_overlap_blocks and CacheHitEstimates.cached_tokens
fields from HashMap to FxHashMap. In
lib/kv-router/src/services/selection/scoring.rs (lines 7-19), change the type of
OverlapInputs.mooncake_summaries field from HashMap to FxHashMap. Ensure
FxHashMap is imported from the appropriate crate (typically rustc_hash or
similar) in both files.
- Line 4: Replace the HashMap import with FxHashMap from the appropriate module
at the import statement. Then, replace all usages of HashMap with FxHashMap
throughout the overlap.rs file, particularly in the worker-keyed overlap
estimate structures. Since these maps use internal numeric WorkerWithDpRank keys
on the hot scheduling path, they must follow lib/kv-router guidelines by using
FxHashMap instead of the standard HashMap for better performance. The comment
indicates this change applies at multiple locations (lines 13-15 and 23-24 in
addition to the import), so ensure all HashMap references keyed by
WorkerWithDpRank or used in hot-path code are updated to FxHashMap.
In `@lib/kv-router/src/services/selection/scoring.rs`:
- Around line 7-10: The mooncake_summaries map uses a standard map type for
numeric WorkerId keys on a hot path, which should use the faster FxHashMap per
project conventions for numeric keys in lib/kv-router. Replace the map type used
for mooncake_summaries with FxHashMap while keeping standard maps for any
text-keyed wire responses. This change should be applied to the
mooncake_summaries declaration and any similar internal numeric-keyed maps in
the same scope.
In `@lib/kv-router/src/services/selection/tests.rs`:
- Around line 587-588: Replace the snapshot-based assertion with a time-bounded
assertion that ensures the core state remains stable over a duration window. Add
a new helper function assert_core_load_stays that takes a SelectionCore
reference, expected values for requests, blocks, and tokens, and a Duration
parameter. This helper should loop from the current time until the deadline
(current time plus the duration), repeatedly calling the existing
assert_core_load function and sleeping briefly between iterations. Then replace
the current tokio::time::sleep call followed by a single assert_core_load call
for core_b with a call to assert_core_load_stays passing the expected values (1,
1, 4) and an appropriate duration to catch delayed replica events.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3b352a0e-7224-4be7-a3e4-e30957a42591
📒 Files selected for processing (32)
docs/components/router/README.mddocs/components/router/router-guide.mddocs/components/router/standalone-selection.mddocs/components/router/standalone-slot-tracker.mddocs/index.ymllib/bindings/python/rust/llm/kv.rslib/kv-router/src/protocols.rslib/kv-router/src/scheduling/local.rslib/kv-router/src/scheduling/mod.rslib/kv-router/src/scheduling/overlap.rslib/kv-router/src/scheduling/prefill_load.rslib/kv-router/src/scheduling/queue.rslib/kv-router/src/scheduling/selector.rslib/kv-router/src/scheduling/types.rslib/kv-router/src/sequences/multi_worker.rslib/kv-router/src/services/indexer/server.rslib/kv-router/src/services/mod.rslib/kv-router/src/services/overlap.rslib/kv-router/src/services/replica_sync.rslib/kv-router/src/services/selection/README.mdlib/kv-router/src/services/selection/core/mod.rslib/kv-router/src/services/selection/input.rslib/kv-router/src/services/selection/scoring.rslib/kv-router/src/services/selection/server.rslib/kv-router/src/services/selection/tests.rslib/kv-router/src/services/selection/types.rslib/kv-router/src/services/slot_tracker/mod.rslib/kv-router/src/services/slot_tracker/registry.rslib/kv-router/src/services/slot_tracker/server.rslib/kv-router/src/services/zmq.rslib/llm/src/kv_router.rslib/llm/src/kv_router/scheduler_inputs.rs
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
ishandhanani
approved these changes
Jun 16, 2026
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
2 participants
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.
Summary
/dumpendpoint.--replica-sync-portand optional startup--replica-sync-peers; add/replica_sync/register_peer,/replica_sync/deregister_peer, and/replica_sync/peersfor dynamic in-memory membership./selectand/select_and_reserveto return Mooncake-style token overlap summaries andeffective_prefill_tokens; remove the redundant publiccached_tokensandeffective_overlap_blocksfields./reservationsto accept validatedeffective_prefill_tokensas the authoritative prefill-load hint while preserving the existing fallback when omitted.Validation
cargo test --no-default-features --features standalone-selection,standalone-slot-tracker services::cargo clippy --no-deps --all-targets --features standalone-selection,standalone-slot-tracker -- -D warningscargo test --no-default-features --features standalone-indexer services::indexercargo clippy --no-deps --all-targets --features standalone-indexer -- -D warningsselect-serviceandslot-trackercargo fmt --checkgit diff --check