From cbf0f87fbf556763ada1cd745474f92ce4b3420d Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Wed, 27 May 2026 13:18:00 -0700 Subject: [PATCH 1/4] Expose Dynamo KV router APIs for external routing --- components/src/dynamo/router/__main__.py | 8 +- lib/bindings/python/rust/llm/kv.rs | 131 ++++++++++++++++++++++- lib/bindings/python/src/dynamo/_core.pyi | 72 +++++++++++++ lib/llm/src/kv_router.rs | 101 +++++++++++++++++ 4 files changed, 305 insertions(+), 7 deletions(-) diff --git a/components/src/dynamo/router/__main__.py b/components/src/dynamo/router/__main__.py index 84e2a292c23b..f15dcd83d65c 100644 --- a/components/src/dynamo/router/__main__.py +++ b/components/src/dynamo/router/__main__.py @@ -137,7 +137,9 @@ async def generate(self, request): } yield llm_engine_output - async def best_worker_id(self, token_ids, router_config_override=None): + async def best_worker_id( + self, token_ids, router_config_override=None, allowed_worker_ids=None + ): """ Get the best worker ID for a given set of tokens without actually routing. @@ -150,7 +152,9 @@ async def best_worker_id(self, token_ids, router_config_override=None): raise RuntimeError("Router not initialized") (worker_id, _dp_rank, _overlap_blocks) = await self.kv_router.best_worker( - token_ids, router_config_override + token_ids, + router_config_override, + allowed_worker_ids=allowed_worker_ids, ) yield worker_id diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 9382b85fde54..240bb7489057 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use pythonize::{depythonize, pythonize}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::ffi::OsString; use std::sync::Arc; use std::sync::atomic::AtomicU32; @@ -36,6 +36,10 @@ fn depythonize_block_mm_infos(obj: &Bound<'_, PyAny>) -> PyResult>) -> Option> { + allowed_worker_ids.map(|ids| ids.into_iter().collect()) +} + #[cfg(feature = "kv-indexer")] #[derive(Parser)] #[command( @@ -1094,7 +1098,7 @@ impl KvRouter { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (token_ids, router_config_override=None, request_id=None, update_indexer=false, block_mm_infos=None, lora_name=None, routing_constraints=None))] + #[pyo3(signature = (token_ids, router_config_override=None, request_id=None, update_indexer=false, block_mm_infos=None, lora_name=None, routing_constraints=None, allowed_worker_ids=None))] fn best_worker<'p>( &self, py: Python<'p>, @@ -1105,6 +1109,7 @@ impl KvRouter { block_mm_infos: Option, lora_name: Option, routing_constraints: Option, + allowed_worker_ids: Option>, ) -> PyResult> { let router_config_override = if let Some(obj) = router_config_override { let override_config: RouterConfigOverride = @@ -1117,11 +1122,16 @@ impl KvRouter { let block_mm_infos = block_mm_infos .map(|obj| depythonize_block_mm_infos(obj.bind(py))) .transpose()?; + let allowed_worker_ids = allowed_worker_set(allowed_worker_ids); let chooser = self.inner.chooser.clone(); let update_states = request_id.is_some(); pyo3_async_runtimes::tokio::future_into_py(py, async move { + if let Some(worker_ids) = allowed_worker_ids.as_ref() { + chooser.register_workers(worker_ids); + } + let outcome = chooser .find_best_match_details( request_id.as_deref(), @@ -1134,7 +1144,7 @@ impl KvRouter { 0.0, None, None, - None, // allowed_worker_ids: pass via RoutingHints in PreprocessedRequest path + allowed_worker_ids, routing_constraints.map(Into::into).unwrap_or_default(), ) .await @@ -1195,6 +1205,108 @@ impl KvRouter { }) } + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (token_ids, router_config_override=None, allowed_worker_ids=None, block_mm_infos=None, lora_name=None))] + fn rank_workers<'p>( + &self, + py: Python<'p>, + token_ids: Vec, + router_config_override: Option, + allowed_worker_ids: Option>, + block_mm_infos: Option, + lora_name: Option, + ) -> PyResult> { + let router_config_override = if let Some(obj) = router_config_override { + let override_config: RouterConfigOverride = + depythonize(obj.bind(py)).map_err(to_pyerr)?; + Some(override_config) + } else { + None + }; + + let block_mm_infos = block_mm_infos + .map(|obj| depythonize_block_mm_infos(obj.bind(py))) + .transpose()?; + let allowed_worker_ids = allowed_worker_set(allowed_worker_ids); + + let chooser = self.inner.chooser.clone(); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let ranked = chooser + .rank_workers( + &token_ids, + block_mm_infos.as_deref(), + router_config_override.as_ref(), + lora_name, + allowed_worker_ids, + ) + .await + .map_err(to_pyerr)?; + + Python::with_gil(|py| { + pythonize(py, &ranked) + .map(|obj| obj.unbind()) + .map_err(to_pyerr) + }) + }) + } + + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (request_id, token_ids, worker_id, dp_rank=0, overlap_blocks=0, expected_output_tokens=None, router_config_override=None, block_mm_infos=None, lora_name=None))] + fn add_request<'p>( + &self, + py: Python<'p>, + request_id: String, + token_ids: Vec, + worker_id: WorkerId, + dp_rank: DpRank, + overlap_blocks: u32, + expected_output_tokens: Option, + router_config_override: Option, + block_mm_infos: Option, + lora_name: Option, + ) -> PyResult> { + let router_config_override = if let Some(obj) = router_config_override { + let override_config: RouterConfigOverride = + depythonize(obj.bind(py)).map_err(to_pyerr)?; + Some(override_config) + } else { + None + }; + + let block_mm_infos = block_mm_infos + .map(|obj| depythonize_block_mm_infos(obj.bind(py))) + .transpose()?; + + let chooser = self.inner.chooser.clone(); + let worker = WorkerWithDpRank::new(worker_id, dp_rank); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let cached_tokens = (overlap_blocks as usize) + .saturating_mul(chooser.block_size() as usize) + .min(token_ids.len()); + chooser + .add_request( + request_id, + &token_ids, + block_mm_infos.as_deref(), + cached_tokens, + expected_output_tokens, + worker, + lora_name, + router_config_override.as_ref(), + ) + .await; + Ok::<(), pyo3::PyErr>(()) + }) + } + + #[pyo3(signature = (worker_ids))] + fn register_workers(&self, worker_ids: Vec) { + let worker_ids: HashSet = worker_ids.into_iter().collect(); + self.inner.chooser.register_workers(&worker_ids); + } + /// Mark prefill as completed for a request fn mark_prefill_complete<'p>( &self, @@ -1222,21 +1334,27 @@ impl KvRouter { }) } - #[pyo3(signature = (token_ids, block_mm_infos=None, lora_name=None))] + #[pyo3(signature = (token_ids, block_mm_infos=None, lora_name=None, allowed_worker_ids=None))] fn get_potential_loads<'p>( &self, py: Python<'p>, token_ids: Vec, block_mm_infos: Option, lora_name: Option, + allowed_worker_ids: Option>, ) -> PyResult> { let block_mm_infos = block_mm_infos .map(|obj| depythonize_block_mm_infos(obj.bind(py))) .transpose()?; + let allowed_worker_ids = allowed_worker_set(allowed_worker_ids); let chooser = self.inner.chooser.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - let loads = chooser + if let Some(worker_ids) = allowed_worker_ids.as_ref() { + chooser.register_workers(worker_ids); + } + + let mut loads = chooser .get_potential_loads( &token_ids, None, @@ -1245,6 +1363,9 @@ impl KvRouter { ) .await .map_err(to_pyerr)?; + if let Some(allowed_worker_ids) = allowed_worker_ids { + loads.retain(|load| allowed_worker_ids.contains(&load.worker_id)); + } // Return loads without aggregation - each (worker_id, dp_rank) pair is a separate entry // Use pythonize to convert Vec to Python list of dicts diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 3ca0239718e0..25ec8e1ae9ed 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -2362,6 +2362,7 @@ class KvRouter: block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, routing_constraints: Optional[RoutingConstraints] = None, + allowed_worker_ids: Optional[List[int]] = None, ) -> Tuple[int, int, int]: """ Find the best matching worker for the given tokens. @@ -2379,6 +2380,11 @@ class KvRouter: block_mm_infos: Optional block-level multimodal metadata aligned to request blocks. When provided, this is used in block hash computation to enable MM-aware worker selection. + routing_constraints: Optional topology or taint constraints to apply while + selecting a worker. + allowed_worker_ids: Optional worker IDs to consider. If provided, routing is + restricted to this set and the IDs are lazily + registered for external/Ray-owned scoring. Returns: A tuple of (worker_id, dp_rank, overlap_blocks) where: @@ -2388,11 +2394,76 @@ class KvRouter: """ ... + async def rank_workers( + self, + token_ids: List[int], + router_config_override: Optional[JsonLike] = None, + allowed_worker_ids: Optional[List[int]] = None, + block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, + lora_name: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """ + Rank candidate workers for the given tokens without booking request state. + + Args: + token_ids: List of token IDs to rank workers for. + router_config_override: Optional router configuration override. + allowed_worker_ids: Optional worker IDs to consider. If provided, + ranking is restricted to this set and the IDs are + lazily registered for external/Ray-owned scoring. + block_mm_infos: Optional block-level multimodal metadata aligned to + request blocks. + lora_name: Optional LoRA adapter name used in block hash computation. + + Returns: + A list of dictionaries sorted from best to worst. Each dictionary contains: + - worker_id: The worker ID + - dp_rank: The data parallel rank + - overlap_blocks: Number of matching cached blocks + - potential_prefill_tokens: Estimated prefill tokens if routed here + - potential_decode_blocks: Estimated active decode blocks + - score: Final router score, lower is better + """ + ... + + async def add_request( + self, + request_id: str, + token_ids: List[int], + worker_id: int, + dp_rank: int = 0, + overlap_blocks: int = 0, + expected_output_tokens: Optional[int] = None, + router_config_override: Optional[JsonLike] = None, + block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, + lora_name: Optional[str] = None, + ) -> None: + """ + Book a request in router load state after an external dispatcher accepts it. + + This is the explicit lifecycle counterpart to query-only `best_worker()` or + `rank_workers()` calls. It should be paired with `mark_prefill_complete()` + and `free()`. + """ + ... + + def register_workers(self, worker_ids: List[int]) -> None: + """ + Add externally-known workers to the router's active-load tracker. + + This is additive and does not remove workers absent from `worker_ids`. + Ray or another external router remains the source of feasible replica + membership. If Dynamo already knows a worker's DP config it is reused; + otherwise the worker is treated as a single-rank cold worker. + """ + ... + async def get_potential_loads( self, token_ids: List[int], block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, + allowed_worker_ids: Optional[List[int]] = None, ) -> List[Dict[str, int]]: """ Get potential prefill and decode loads for all workers. @@ -2402,6 +2473,7 @@ class KvRouter: block_mm_infos: Optional block-level multimodal metadata aligned to request blocks. When provided, this is used in hash computation for MM-aware potential-load estimation. + allowed_worker_ids: Optional worker IDs to include in the returned load list. Returns: A list of dictionaries, each containing: diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 9831f25de6e8..5f8c50795e00 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -33,6 +33,7 @@ use dynamo_runtime::{ traits::DistributedRuntimeProvider, }; use futures::stream; +use serde::Serialize; use tracing::Instrument; use validator::Validate; @@ -91,6 +92,55 @@ pub enum FindBestMatchOutcome { }, } +#[derive(Debug, Clone, Serialize)] +pub struct RankedWorker { + pub worker_id: WorkerId, + pub dp_rank: DpRank, + pub overlap_blocks: u32, + pub potential_prefill_tokens: usize, + pub potential_decode_blocks: usize, + pub score: f64, +} + +fn rank_potential_loads( + loads: Vec, + cache_hit_estimates: &CacheHitEstimates, + block_size: u32, + prefill_load_scale: f64, + allowed_worker_ids: Option<&HashSet>, +) -> Vec { + let mut ranked: Vec<_> = loads + .into_iter() + .filter(|load| allowed_worker_ids.is_none_or(|allowed| allowed.contains(&load.worker_id))) + .map(|load| { + let worker = WorkerWithDpRank::new(load.worker_id, load.dp_rank); + let overlap_blocks = + cache_hit_for_worker(cache_hit_estimates, worker).rounded_overlap_blocks(); + let potential_prefill_blocks = + (load.potential_prefill_tokens as f64) / (block_size as f64); + let score = + prefill_load_scale * potential_prefill_blocks + load.potential_decode_blocks as f64; + + RankedWorker { + worker_id: load.worker_id, + dp_rank: load.dp_rank, + overlap_blocks, + potential_prefill_tokens: load.potential_prefill_tokens, + potential_decode_blocks: load.potential_decode_blocks, + score, + } + }) + .collect(); + + ranked.sort_by(|a, b| { + a.score + .total_cmp(&b.score) + .then_with(|| a.worker_id.cmp(&b.worker_id)) + .then_with(|| a.dp_rank.cmp(&b.dp_rank)) + }); + ranked +} + // [gluo TODO] shouldn't need to be public // this should be discovered from the component @@ -602,6 +652,57 @@ where self.scheduler.register_workers(worker_ids); } + pub async fn rank_workers( + &self, + tokens: &[u32], + block_mm_infos: Option<&[Option]>, + router_config_override: Option<&RouterConfigOverride>, + lora_name: Option, + allowed_worker_ids: Option>, + ) -> Result> { + if let Some(worker_ids) = allowed_worker_ids.as_ref() { + self.register_workers(worker_ids); + } + + let hash_options = BlockHashOptions { + block_mm_infos, + lora_name: lora_name.as_deref(), + is_eagle: Some(self.is_eagle), + }; + let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, hash_options); + log_routing_input_hashes(None, self.block_size, tokens, &block_hashes); + + let maybe_seq_hashes = self.kv_router_config.compute_seq_hashes_for_tracking( + tokens, + self.block_size, + router_config_override, + hash_options, + Some(&block_hashes), + ); + let track_prefill_tokens = self + .kv_router_config + .track_prefill_tokens(router_config_override); + let tiered_matches = self.indexer.find_matches_by_tier(block_hashes).await?; + let cache_hit_estimates = self.cache_hit_estimates_from_tiered_matches(&tiered_matches); + let loads = self.scheduler.get_potential_loads( + maybe_seq_hashes, + tokens.len(), + cache_hit_estimates.cached_tokens.clone(), + track_prefill_tokens, + ); + let prefill_load_scale = router_config_override + .and_then(|cfg| cfg.prefill_load_scale) + .unwrap_or(self.kv_router_config.prefill_load_scale); + + Ok(rank_potential_loads( + loads, + &cache_hit_estimates, + self.block_size, + prefill_load_scale, + allowed_worker_ids.as_ref(), + )) + } + #[allow(clippy::too_many_arguments)] pub async fn add_request( &self, From 23fcb59734848c89fd1f439fb0d04a7d8af850ce Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Wed, 27 May 2026 13:18:07 -0700 Subject: [PATCH 2/4] Test Dynamo worker ranking helper --- lib/llm/src/kv_router.rs | 77 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 76 insertions(+), 1 deletion(-) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 5f8c50795e00..c4cf6d14bf49 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1153,7 +1153,7 @@ where #[cfg(test)] mod tests { use super::*; - use std::collections::HashMap; + use std::collections::{HashMap, HashSet}; use async_trait::async_trait; use dynamo_kv_router::{ @@ -1554,4 +1554,79 @@ mod tests { assert!((worker.router_credit_blocks - 1.0).abs() < f64::EPSILON); } } + + fn load( + worker_id: WorkerId, + dp_rank: DpRank, + potential_prefill_tokens: usize, + potential_decode_blocks: usize, + ) -> PotentialLoad { + PotentialLoad { + worker_id, + dp_rank, + potential_prefill_tokens, + potential_decode_blocks, + } + } + + fn cache_hit_estimates( + entries: impl IntoIterator, + ) -> CacheHitEstimates { + let mut effective_overlap_blocks = HashMap::new(); + let mut cached_tokens = HashMap::new(); + for (worker, overlap_blocks, tokens) in entries { + effective_overlap_blocks.insert(worker, overlap_blocks); + cached_tokens.insert(worker, tokens); + } + CacheHitEstimates { + effective_overlap_blocks, + cached_tokens, + } + } + + #[test] + fn rank_potential_loads_orders_by_score_and_filters_allowed_workers() { + let estimates = cache_hit_estimates([ + (WorkerWithDpRank::new(1, 0), 4.0, 16), + (WorkerWithDpRank::new(2, 0), 7.0, 28), + (WorkerWithDpRank::new(3, 0), 1.0, 4), + ]); + + let ranked = rank_potential_loads( + vec![load(1, 0, 32, 4), load(2, 0, 8, 2), load(3, 0, 4, 0)], + &estimates, + 4, + 1.0, + Some(&HashSet::from([1, 2])), + ); + + assert_eq!(ranked.len(), 2); + assert_eq!(ranked[0].worker_id, 2); + assert_eq!(ranked[0].overlap_blocks, 7); + assert_eq!(ranked[0].potential_prefill_tokens, 8); + assert_eq!(ranked[0].potential_decode_blocks, 2); + assert_eq!(ranked[0].score, 4.0); + assert_eq!(ranked[1].worker_id, 1); + assert_eq!(ranked[1].score, 12.0); + } + + #[test] + fn rank_potential_loads_is_deterministic_on_score_ties() { + let estimates = cache_hit_estimates([]); + let ranked = rank_potential_loads( + vec![load(9, 0, 4, 1), load(3, 0, 4, 1), load(3, 1, 4, 1)], + &estimates, + 4, + 1.0, + None, + ); + + assert_eq!( + ranked + .into_iter() + .map(|worker| (worker.worker_id, worker.dp_rank)) + .collect::>(), + vec![(3, 0), (3, 1), (9, 0)] + ); + } } From dfa5baae6d8017ac5ebc65bd1812b2c0a0c53404 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Wed, 27 May 2026 13:18:15 -0700 Subject: [PATCH 3/4] Honor explicit worker IDs in Dynamo KV publisher --- lib/llm/src/kv_router/publisher/mod.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/lib/llm/src/kv_router/publisher/mod.rs b/lib/llm/src/kv_router/publisher/mod.rs index 83a671705e68..0600c18f3c36 100644 --- a/lib/llm/src/kv_router/publisher/mod.rs +++ b/lib/llm/src/kv_router/publisher/mod.rs @@ -64,6 +64,13 @@ fn create_kv_stream_name(component: &Component, subject: &str) -> String { .replace("_", "-") } +fn resolve_worker_id( + default_worker_id: WorkerId, + explicit_worker_id: Option, +) -> WorkerId { + explicit_worker_id.unwrap_or(default_worker_id) +} + /// Configure the source of KV events. /// Currently, only ZMQ is supported. pub enum KvEventSourceConfig { @@ -192,7 +199,7 @@ impl KvEventPublisher { .map(|ms| ms.min(MAX_BATCHING_TIMEOUT_MS)); let (tx, rx) = mpsc::unbounded_channel::(); - let worker_id = worker_id.unwrap_or_else(|| component.drt().connection_id()); + let worker_id = resolve_worker_id(component.drt().connection_id(), worker_id); let _ = KvPublisherMetrics::from_component(&component); From 429b6dafb00170eb0a45ede771cf06cb361eaadb Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Wed, 27 May 2026 13:18:17 -0700 Subject: [PATCH 4/4] Test Dynamo KV publisher worker ID resolution --- lib/llm/src/kv_router/publisher/tests.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 863ec042ae14..8c8ac76ad321 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -14,6 +14,16 @@ use std::future::Future; use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::time::Duration; +#[test] +fn resolve_worker_id_uses_explicit_worker_when_provided() { + assert_eq!(resolve_worker_id(17, Some(42)), 42); +} + +#[test] +fn resolve_worker_id_falls_back_to_connection_id() { + assert_eq!(resolve_worker_id(17, None), 17); +} + #[cfg(test)] mod test_event_processing { use super::*;