From c8328f44613441a1430071e40ff423eaee1bc9f6 Mon Sep 17 00:00:00 2001 From: Temirulan Mussayev Date: Thu, 15 Oct 2026 21:22:26 +0000 Subject: [PATCH 1/2] kv-router: opt-in refcount + SWA-aware hits for hybrid-attention models (Step-5) New --hybrid-swa-hits mode for the standalone indexer (default off; other models unchanged). For engines that interleave sliding-window groups with full attention and run MTP as EAGLE (Step-5-Preview: 4x SWA-512 + 1 full-attention group, block 64), the FA-only, set-semantics tree mispredicts the engine's prefix hit (DEEP-309: continuations 0/20 within 64 tokens, first reuse of a new scaffold junction over-predicted by 3.4K-30.7K, 9-12% under-prediction after eviction). The mode adds a per-worker oracle fed from the same ZMQ stream: - per-copy refcounts for every group; a Removed for a block with another cached copy is held back from the radix tree - sliding-window groups (filtered by the main path) are tracked by engine block hash, which vLLM shares across groups - /query reports min(tree, oracle) per worker, where the oracle predicts max{k <= min(chain, (n-1)/bs) : every SWA group holds [k-need, k)} - mtp, need = ceil((window-1)/bs) + mtp (9 for Step-5) Flags: --hybrid-swa-hits, --mtp-drop-blocks (default 1), --swa-need-blocks. Port of the reference model that matched engine cached_tokens on 56/56 parity and 60/60 eviction probes in DEEP-309. Co-Authored-By: Claude Opus 5.5 (1M context) --- lib/bindings/python/rust/llm/kv.rs | 21 + .../src/standalone_indexer/hybrid.rs | 452 ++++++++++++++++++ .../src/standalone_indexer/listener.rs | 44 +- lib/kv-router/src/standalone_indexer/mod.rs | 28 ++ .../src/standalone_indexer/server.rs | 18 +- lib/kv-router/src/zmq_wire/mod.rs | 29 ++ 6 files changed, 589 insertions(+), 3 deletions(-) create mode 100644 lib/kv-router/src/standalone_indexer/hybrid.rs diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 85a033527092..f0ca08f6fc18 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -96,6 +96,24 @@ struct KvIndexerCli { #[arg(long, default_value_t = false)] enable_logging: bool, + /// Opt-in for hybrid-attention engines (e.g. Step-5: SWA groups + full + /// attention + MTP): refcount stored block copies, track sliding-window + /// groups, and answer /query with the SWA-aware hit the engine actually + /// serves. Requires --block-size. Default off: other models are unchanged. + #[arg(long, default_value_t = false)] + hybrid_swa_hits: bool, + + /// Blocks dropped from predicted hits for EAGLE/MTP speculative decoding. + /// Only meaningful with --hybrid-swa-hits. + #[arg(long, default_value_t = 1)] + mtp_drop_blocks: u32, + + /// Contiguous sliding-window blocks a hit needs. Defaults to + /// ceil((window-1)/block_size) + mtp_drop_blocks per SWA group. + /// Only meaningful with --hybrid-swa-hits. + #[arg(long)] + swa_need_blocks: Option, + /// Kubernetes namespace to watch for engine pods. Providing this together /// with --watch-model-name (or --watch-label) enables pod auto-discovery /// (subscribe on Ready, unsubscribe on delete). @@ -193,6 +211,9 @@ where evict_retention_secs: cli.evict_retention_secs, evict_memory_threshold: cli.evict_memory_threshold, enable_logging: cli.enable_logging, + hybrid_swa_hits: cli.hybrid_swa_hits, + mtp_drop_blocks: cli.mtp_drop_blocks, + swa_need_blocks: cli.swa_need_blocks, kube_discovery, })) } diff --git a/lib/kv-router/src/standalone_indexer/hybrid.rs b/lib/kv-router/src/standalone_indexer/hybrid.rs new file mode 100644 index 000000000000..6daa196a700c --- /dev/null +++ b/lib/kv-router/src/standalone_indexer/hybrid.rs @@ -0,0 +1,452 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Opt-in prefix-hit oracle for hybrid-attention engines (`--hybrid-swa-hits`). +//! +//! The radix tree tracks only the main full-attention group with set +//! semantics. On models that interleave sliding-window (SWA) groups with full +//! attention — Step-5-Preview: four SWA-512 groups plus one full-attention +//! group, all at block size 64, with MTP run as EAGLE — that over- and +//! under-predicts what the engine actually serves: +//! +//! - SWA groups keep cached blocks only at sparse checkpoints, so the engine's +//! hit stops at the last position where every SWA group still holds the +//! blocks its window needs, often well short of the full-attention chain. +//! - When the engine's hit is shorter than the chain it recomputes and +//! re-stores the same hashes. A set-semantics index then forgets a block on +//! the first `Removed` even though another copy is still cached. +//! - EAGLE/MTP drops the last matched block, and the engine never reuses the +//! final prompt token's block (`(n - 1) / block_size` cap). +//! +//! This oracle consumes the same event stream, keeps per-copy refcounts for +//! every group, and predicts the engine's hit: +//! `k = max{k <= min(chain, (n-1)/bs) : every SWA group holds blocks [k-need, k)}`, +//! then subtracts the MTP block. Engine block hashes are shared across groups +//! (same hash for the same 64-token position), which is what makes the SWA +//! lookup by full-attention hash valid. +//! +//! Default-off: with the flag unset nothing here is constructed and the +//! listener / query paths are unchanged. + +use std::sync::OnceLock; + +use parking_lot::RwLock; +use rustc_hash::FxHashMap; + +use crate::protocols::{ + ExternalSequenceBlockHash, KvCacheEventData, LocalBlockHash, RouterEvent, WorkerWithDpRank, +}; + +/// Opt-in settings for the hybrid oracle. +#[derive(Debug, Clone, Copy)] +pub struct HybridConfig { + /// Block size of the engine's KV groups (tokens per block). + pub block_size: u32, + /// Blocks dropped from the predicted hit for EAGLE/MTP speculative + /// decoding (1 for Step-5's MTP; 0 for engines without it). + pub mtp_drop_blocks: u32, + /// Override for the contiguous SWA blocks a hit needs. `None` derives it + /// per group from the sliding window: `ceil((window - 1) / bs) + mtp_drop`. + pub swa_need_blocks: Option, +} + +#[derive(Debug, Clone, Copy)] +struct FaEntry { + parent: Option, + tokens: LocalBlockHash, + copies: u32, +} + +#[derive(Default)] +struct WorkerState { + /// Full-attention chain: (parent engine hash, local tokens hash) -> engine hash. + fa_child: FxHashMap<(Option, u64), u64>, + /// Full-attention live blocks with per-copy refcount. + fa_live: FxHashMap, + /// Per SWA group: engine hash -> copies held. + swa_live: FxHashMap>, + /// Per SWA group: sliding window in tokens (learned from events). + swa_window: FxHashMap, +} + +impl WorkerState { + fn fa_store(&mut self, parent: Option, blocks: &[(u64, LocalBlockHash)]) { + let mut parent = parent; + for &(hash, tokens) in blocks { + match self.fa_live.get_mut(&hash) { + Some(entry) => entry.copies += 1, + None => { + self.fa_live.insert( + hash, + FaEntry { + parent, + tokens, + copies: 1, + }, + ); + } + } + self.fa_child.insert((parent, tokens.0), hash); + parent = Some(hash); + } + } + + /// Returns the hashes whose LAST copy was removed (the tree may drop them). + fn fa_remove(&mut self, hashes: &[u64]) -> Vec { + let mut gone = Vec::with_capacity(hashes.len()); + for &hash in hashes { + let Some(entry) = self.fa_live.get_mut(&hash) else { + // Unknown to us: let the tree decide (it treats it as a no-op). + gone.push(hash); + continue; + }; + if entry.copies > 1 { + entry.copies -= 1; + continue; + } + let entry = self.fa_live.remove(&hash).expect("present"); + if self.fa_child.get(&(entry.parent, entry.tokens.0)) == Some(&hash) { + self.fa_child.remove(&(entry.parent, entry.tokens.0)); + } + gone.push(hash); + } + gone + } + + fn swa_store(&mut self, group: u32, window: Option, hashes: &[u64]) { + if let Some(window) = window { + self.swa_window.insert(group, window); + } + let live = self.swa_live.entry(group).or_default(); + for &hash in hashes { + *live.entry(hash).or_insert(0) += 1; + } + } + + fn swa_remove(&mut self, group: u32, hashes: &[u64]) { + let Some(live) = self.swa_live.get_mut(&group) else { + return; + }; + for &hash in hashes { + if let Some(copies) = live.get_mut(&hash) { + if *copies > 1 { + *copies -= 1; + } else { + live.remove(&hash); + } + } + } + } + + fn clear(&mut self) { + *self = WorkerState::default(); + } + + /// Predicted engine hit in blocks for a probe of `num_tokens` tokens whose + /// per-block local hashes are `local`. + fn predict(&self, local: &[LocalBlockHash], num_tokens: usize, cfg: &HybridConfig) -> u32 { + let bs = cfg.block_size as usize; + let mut chain: Vec = Vec::with_capacity(local.len()); + let mut parent = None; + for tokens in local { + let Some(&hash) = self.fa_child.get(&(parent, tokens.0)) else { + break; + }; + if !self.fa_live.contains_key(&hash) { + break; + } + chain.push(hash); + parent = Some(hash); + } + let reusable = if num_tokens == 0 { + 0 + } else { + (num_tokens - 1) / bs + }; + let cap = chain.len().min(reusable); + + let mut k = cap; + 'search: while k > 0 { + for (group, live) in &self.swa_live { + let need = cfg.swa_need_blocks.unwrap_or_else(|| { + let window = self.swa_window.get(group).copied().unwrap_or(0); + window.saturating_sub(1).div_ceil(cfg.block_size) + cfg.mtp_drop_blocks + }) as usize; + let lo = k.saturating_sub(need); + if !chain[lo..k].iter().all(|hash| live.contains_key(hash)) { + k -= 1; + continue 'search; + } + } + break; + } + (k as u32).saturating_sub(cfg.mtp_drop_blocks) + } +} + +/// Process-global oracle, keyed by worker. +pub struct HybridOracle { + cfg: HybridConfig, + workers: RwLock>, +} + +static ORACLE: OnceLock = OnceLock::new(); + +/// Install the oracle (once, at startup). No-op if already installed. +pub fn install(cfg: HybridConfig) { + let _ = ORACLE.set(HybridOracle::new(cfg)); +} + +/// The installed oracle, or `None` when `--hybrid-swa-hits` is off. +pub fn oracle() -> Option<&'static HybridOracle> { + ORACLE.get() +} + +impl HybridOracle { + pub fn new(cfg: HybridConfig) -> Self { + Self { + cfg, + workers: RwLock::new(FxHashMap::default()), + } + } + + pub fn config(&self) -> HybridConfig { + self.cfg + } + + /// Feed a main-attention event. For `Removed`, returns the event with + /// still-held copies stripped, so the radix tree only forgets a block when + /// its last copy is gone. Other events are returned unchanged. + pub fn apply_main(&self, mut event: RouterEvent) -> Option { + let worker = WorkerWithDpRank::new(event.worker_id, event.event.dp_rank); + let mut workers = self.workers.write(); + let state = workers.entry(worker).or_default(); + match &mut event.event.data { + KvCacheEventData::Stored(data) => { + let blocks: Vec<(u64, LocalBlockHash)> = data + .blocks + .iter() + .map(|b| (b.block_hash.0, b.tokens_hash)) + .collect(); + state.fa_store(data.parent_hash.map(|h| h.0), &blocks); + } + KvCacheEventData::Removed(data) => { + let hashes: Vec = data.block_hashes.iter().map(|h| h.0).collect(); + let gone = state.fa_remove(&hashes); + if gone.is_empty() { + return None; + } + data.block_hashes = gone.into_iter().map(ExternalSequenceBlockHash).collect(); + } + KvCacheEventData::Cleared => state.clear(), + } + Some(event) + } + + /// Feed a sliding-window group event that the main path filters out. + pub fn apply_swa( + &self, + worker: WorkerWithDpRank, + group: u32, + window: Option, + stored: bool, + hashes: &[u64], + ) { + let mut workers = self.workers.write(); + let state = workers.entry(worker).or_default(); + if stored { + state.swa_store(group, window, hashes); + } else { + state.swa_remove(group, hashes); + } + } + + pub fn clear(&self, worker: WorkerWithDpRank) { + if let Some(state) = self.workers.write().get_mut(&worker) { + state.clear(); + } + } + + pub fn remove_worker(&self, worker_id: u64) { + self.workers.write().retain(|w, _| w.worker_id != worker_id); + } + + /// Predicted engine hit, in blocks, for `worker`. + pub fn predict( + &self, + worker: WorkerWithDpRank, + local: &[LocalBlockHash], + num_tokens: usize, + ) -> Option { + let workers = self.workers.read(); + workers + .get(&worker) + .map(|state| state.predict(local, num_tokens, &self.cfg)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const BS: u32 = 64; + + fn cfg() -> HybridConfig { + HybridConfig { + block_size: BS, + mtp_drop_blocks: 1, + swa_need_blocks: None, + } + } + + fn lh(i: u64) -> LocalBlockHash { + LocalBlockHash(1000 + i) + } + + /// Engine hash for chain position `i` of conversation `conv`. + fn eh(conv: u64, i: u64) -> u64 { + conv * 1_000_000 + i + } + + fn store_chain(s: &mut WorkerState, conv: u64, n: u64) { + let blocks: Vec<(u64, LocalBlockHash)> = (0..n).map(|i| (eh(conv, i), lh(i))).collect(); + s.fa_store(None, &blocks); + } + + fn store_swa(s: &mut WorkerState, conv: u64, range: std::ops::Range) { + let hashes: Vec = range.map(|i| eh(conv, i)).collect(); + for g in 0..4 { + s.swa_store(g, Some(512), &hashes); + } + } + + fn probe(n_blocks: u64) -> (Vec, usize) { + ( + (0..n_blocks).map(lh).collect(), + (n_blocks as usize) * BS as usize, + ) + } + + #[test] + fn need_is_window_blocks_plus_mtp() { + // ceil(511 / 64) + 1 = 9, the reference model's SWA_NEED. + assert_eq!(511u32.div_ceil(BS) + 1, 9); + } + + #[test] + fn full_swa_coverage_predicts_chain_minus_mtp() { + let mut s = WorkerState::default(); + store_chain(&mut s, 1, 40); + store_swa(&mut s, 1, 0..40); + let (local, n) = probe(40); + // (n-1)/bs caps at 39, minus the MTP block = 38. + assert_eq!(s.predict(&local, n, &cfg()), 38); + } + + #[test] + fn swa_junction_without_checkpoint_predicts_zero() { + // DEEP-309 (b): shared scaffold, first reuse of a new junction. The + // full-attention chain covers the scaffold, but SWA groups hold only + // the original prompt's tail checkpoint, not the junction blocks. + let mut s = WorkerState::default(); + store_chain(&mut s, 1, 60); + store_swa(&mut s, 1, 51..60); // checkpoint at the original end only + let (local, n) = probe(30); // diverges at block 30 + assert_eq!(s.predict(&local, n, &cfg()), 0); + } + + #[test] + fn continuation_stops_at_previous_prompt_checkpoint() { + // DEEP-309 (a): turn N+1 extends turn N. The chain also covers turn N's + // completion, but SWA was checkpointed only at turn N's prompt end. + let mut s = WorkerState::default(); + store_chain(&mut s, 1, 50); // prompt(40) + completion(10) + store_swa(&mut s, 1, 31..40); // window ending at the prompt end + let (local, n) = probe(60); + // Largest k with SWA holding [k-9, k) is 40, minus MTP = 39. + assert_eq!(s.predict(&local, n, &cfg()), 39); + } + + #[test] + fn duplicate_copy_survives_first_removal() { + // DEEP-309 eviction: the engine re-stored blocks it recomputed; evicting + // the older copy must not forget a block the engine still holds. + let mut s = WorkerState::default(); + store_chain(&mut s, 1, 20); + store_chain(&mut s, 1, 20); // duplicate copies + store_swa(&mut s, 1, 0..20); + let gone = s.fa_remove(&[eh(1, 10)]); + assert!( + gone.is_empty(), + "first removal of a duplicated block is held back" + ); + let (local, n) = probe(20); + assert_eq!(s.predict(&local, n, &cfg()), 18); + let gone = s.fa_remove(&[eh(1, 10)]); + assert_eq!(gone, vec![eh(1, 10)]); + // Last copy gone: the chain now breaks at block 10 → k = 10, minus MTP. + assert_eq!(s.predict(&local, n, &cfg()), 9); + } + + #[test] + fn no_swa_groups_behaves_like_capped_chain() { + let mut s = WorkerState::default(); + store_chain(&mut s, 1, 10); + let (local, n) = probe(10); + assert_eq!(s.predict(&local, n, &cfg()), 8); + } + + #[test] + fn apply_main_holds_back_removal_of_duplicated_block() { + use crate::protocols::{ + KvCacheEvent, KvCacheRemoveData, KvCacheStoreData, KvCacheStoredBlockData, StorageTier, + }; + let oracle = HybridOracle::new(cfg()); + let event = |event_id, data| RouterEvent { + worker_id: 7, + storage_tier: StorageTier::Device, + event: KvCacheEvent { + event_id, + dp_rank: 0, + data, + }, + }; + let store = || { + event( + 0, + KvCacheEventData::Stored(KvCacheStoreData { + parent_hash: None, + start_position: None, + blocks: vec![KvCacheStoredBlockData { + block_hash: ExternalSequenceBlockHash(42), + tokens_hash: lh(0), + mm_extra_info: None, + }], + }), + ) + }; + let remove = || { + event( + 1, + KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: vec![ExternalSequenceBlockHash(42)], + }), + ) + }; + assert!(oracle.apply_main(store()).is_some()); + assert!(oracle.apply_main(store()).is_some()); + assert!( + oracle.apply_main(remove()).is_none(), + "one copy is still held" + ); + let passed = oracle + .apply_main(remove()) + .expect("last copy reaches the tree"); + match passed.event.data { + KvCacheEventData::Removed(d) => { + assert_eq!(d.block_hashes, vec![ExternalSequenceBlockHash(42)]) + } + _ => panic!("expected Removed"), + } + } +} diff --git a/lib/kv-router/src/standalone_indexer/listener.rs b/lib/kv-router/src/standalone_indexer/listener.rs index 3a4350a09d2a..1496053b18ba 100644 --- a/lib/kv-router/src/standalone_indexer/listener.rs +++ b/lib/kv-router/src/standalone_indexer/listener.rs @@ -11,7 +11,7 @@ use tokio_util::sync::CancellationToken; use crate::indexer::WorkerKvQueryResponse; use crate::protocols::{KvCacheEventData, RouterEvent, WorkerId, WorkerWithDpRank}; use crate::recovery::{CursorObservation, CursorState}; -use crate::zmq_wire::{ZmqEventNormalizer, decode_event_batch}; +use crate::zmq_wire::{RawKvEvent, ZmqEventNormalizer, decode_event_batch}; use super::evictions::PendingEvictions; use super::indexer::Indexer; @@ -348,6 +348,13 @@ impl ListenerLoop { // Audit-log the recovered event before the measurement filter. audit_log_event(&event, event.event.event_id, "recover"); + let event = match super::hybrid::oracle() { + Some(oracle) => match oracle.apply_main(event) { + Some(event) => event, + None => continue, + }, + None => event, + }; // Feed-layer measurement filter (same as apply_live_batch). if self.keep_evictions_intercept(&event) { @@ -405,6 +412,32 @@ impl ListenerLoop { .data_parallel_rank .map_or(self.dp_rank, |rank| rank.cast_unsigned()); for raw_event in batch.events { + // Opt-in hybrid oracle: keep the sliding-window groups the main + // path filters out (see `hybrid`). Default off. + if let Some(oracle) = super::hybrid::oracle() + && let Some((group, window)) = self + .normalizer + .sliding_window_group(&raw_event, effective_dp_rank) + { + let worker = WorkerWithDpRank::new(self.worker_id, effective_dp_rank); + match &raw_event { + RawKvEvent::BlockStored { block_hashes, .. } => { + let hashes: Vec = + block_hashes.iter().map(|h| (*h).into_u64()).collect(); + oracle.apply_swa(worker, group, window, true, &hashes); + } + RawKvEvent::BlockRemoved { block_hashes, .. } => { + let hashes: Vec = + block_hashes.iter().map(|h| (*h).into_u64()).collect(); + oracle.apply_swa(worker, group, window, false, &hashes); + } + _ => {} + } + // Still run the normalizer so it learns group metadata; + // the main path filters this event out as before. + let _ = self.normalizer.preprocess_with_reason(raw_event, worker); + continue; + } let placement_event = match self.normalizer.normalize( raw_event, seq, @@ -433,6 +466,15 @@ impl ListenerLoop { // Audit-log the event as published by the engine, before the // measurement filter below can drop it. audit_log_event(&router_event, seq, "live"); + // Opt-in hybrid oracle: refcount copies; strip removals of blocks + // that still have another cached copy before the tree sees them. + let router_event = match super::hybrid::oracle() { + Some(oracle) => match oracle.apply_main(router_event) { + Some(event) => event, + None => continue, + }, + None => router_event, + }; // Feed-layer measurement filter. if self.keep_evictions_intercept(&router_event) { continue; diff --git a/lib/kv-router/src/standalone_indexer/mod.rs b/lib/kv-router/src/standalone_indexer/mod.rs index ee71fd60c16a..b656b1fcef25 100644 --- a/lib/kv-router/src/standalone_indexer/mod.rs +++ b/lib/kv-router/src/standalone_indexer/mod.rs @@ -24,6 +24,7 @@ //! response struct in [`server`] for the exact semantics. pub mod evictions; +pub mod hybrid; pub mod indexer; pub mod listener; pub mod metrics; @@ -163,6 +164,17 @@ pub struct IndexerConfig { /// Emit verbose per-query and per-event audit logs on the `kv_audit` /// tracing target. See [`logging_enabled`]. pub enable_logging: bool, + /// Opt-in hybrid-attention hit prediction (`--hybrid-swa-hits`): refcount + /// per stored copy, track sliding-window groups, and answer `/query` with + /// the SWA-aware, MTP-adjusted hit the engine actually serves. See + /// [`hybrid`]. Requires `block_size`. + pub hybrid_swa_hits: bool, + /// Blocks dropped for EAGLE/MTP speculative decoding. Only meaningful with + /// `hybrid_swa_hits`. + pub mtp_drop_blocks: u32, + /// Override for the contiguous SWA blocks a hit needs; `None` derives it + /// from each group's sliding window. Only meaningful with `hybrid_swa_hits`. + pub swa_need_blocks: Option, } pub(super) fn validate_zmq_endpoint(endpoint: &str) -> anyhow::Result<()> { @@ -270,6 +282,22 @@ pub async fn run_server(config: IndexerConfig) -> anyhow::Result<()> { // run_server is called once, so we discard the result. let _ = KEEP_EVICTIONS.set(config.keep_evictions); let _ = ENABLE_LOGGING.set(config.enable_logging); + if config.hybrid_swa_hits { + let block_size = config + .block_size + .ok_or_else(|| anyhow::anyhow!("--hybrid-swa-hits requires --block-size"))?; + hybrid::install(hybrid::HybridConfig { + block_size, + mtp_drop_blocks: config.mtp_drop_blocks, + swa_need_blocks: config.swa_need_blocks, + }); + tracing::info!( + block_size, + mtp_drop_blocks = config.mtp_drop_blocks, + swa_need_blocks = ?config.swa_need_blocks, + "hybrid-swa-hits enabled: refcounted copies + SWA-aware /query predictions" + ); + } if config.enable_logging { tracing::info!( target: "kv_audit", diff --git a/lib/kv-router/src/standalone_indexer/server.rs b/lib/kv-router/src/standalone_indexer/server.rs index aa1e00eff6b1..28ff81c1432e 100644 --- a/lib/kv-router/src/standalone_indexer/server.rs +++ b/lib/kv-router/src/standalone_indexer/server.rs @@ -374,6 +374,7 @@ fn merge_score_responses(acc: &mut ScoreResponse, other: ScoreResponse) { async fn run_fanout_query( trees: Vec<(IndexerKey, Indexer, u32)>, mut hashes_for: impl FnMut(u32) -> Vec, + num_tokens: Option, pod_names: &HashMap, audit: QueryAudit<'_>, ) -> (StatusCode, Json) { @@ -389,8 +390,19 @@ async fn run_fanout_query( if logging && logged_hashes.is_none() { logged_hashes = Some(block_hashes.iter().map(|h| h.0).collect()); } - match indexer.find_tiered_matches(block_hashes).await { - Ok(tiered) => { + match indexer.find_tiered_matches(block_hashes.clone()).await { + Ok(mut tiered) => { + // Opt-in hybrid oracle: report the hit the engine will serve + // (SWA-aware, MTP-adjusted) instead of the raw chain length. + // Needs the probe's token count, so `/query_by_hash` is untouched. + if let (Some(oracle), Some(num_tokens)) = (super::hybrid::oracle(), num_tokens) { + for (worker, score) in tiered.device.overlap_scores.scores.iter_mut() { + if let Some(predicted) = oracle.predict(*worker, &block_hashes, num_tokens) + { + *score = predicted.min(*score); + } + } + } let response = build_score_response(&tiered, block_size, pod_names); match merged.as_mut() { Some(acc) => merge_score_responses(acc, response), @@ -473,6 +485,7 @@ async fn query( }) .clone() }, + Some(req.token_ids.len()), &state.registry.pod_names(), QueryAudit { model_name: &req.model_name, @@ -504,6 +517,7 @@ async fn query_by_hash( run_fanout_query( trees, |_| block_hashes.clone(), + None, &state.registry.pod_names(), QueryAudit { model_name: &req.model_name, diff --git a/lib/kv-router/src/zmq_wire/mod.rs b/lib/kv-router/src/zmq_wire/mod.rs index 3610db87109d..15d68523cdda 100644 --- a/lib/kv-router/src/zmq_wire/mod.rs +++ b/lib/kv-router/src/zmq_wire/mod.rs @@ -135,6 +135,35 @@ impl ZmqEventNormalizer { self.normalize_preprocessed(raw, event_id, worker) } + /// Sliding window of a learned non-main group, if the event (or the group + /// it belongs to) is a sliding-window group. Used by the opt-in hybrid + /// oracle to keep the SWA groups the main path filters out. + pub fn sliding_window_group( + &self, + raw: &RawKvEvent, + dp_rank: DpRank, + ) -> Option<(u32, Option)> { + let metadata = raw.metadata(); + let group_idx = metadata.group_idx?; + let kind = metadata.kv_cache_spec_kind.or_else(|| { + self.group_metadata + .get(&(dp_rank, group_idx)) + .map(|m| m.kind) + })?; + if !matches!( + kind, + KvCacheSpecKind::SlidingWindow | KvCacheSpecKind::SlidingWindowMla + ) { + return None; + } + let window = metadata.kv_cache_spec_sliding_window.or_else(|| { + self.group_metadata + .get(&(dp_rank, group_idx)) + .and_then(|m| m.sliding_window) + }); + Some((group_idx, window)) + } + fn learn_metadata(&mut self, metadata: KvCacheEventMetadata, dp_rank: DpRank) { let (Some(group_idx), Some(kind)) = (metadata.group_idx, metadata.kv_cache_spec_kind) else { From 9bdca6760d28d3d64b0cb24e54f77fbeb0ce6664 Mon Sep 17 00:00:00 2001 From: Temirulan Mussayev Date: Thu, 15 Oct 2026 21:55:32 +0000 Subject: [PATCH 2/2] kv-router: apply hybrid oracle to /query_by_hash (deepapi routing path) deepapi's KV routing calls /query_by_hash with xxh3 local block hashes it computes itself (seed 1337, the indexer's XXH3_SEED), not /query. With --hybrid-swa-hits the by-hash path now reports min(tree, oracle) per worker too. The probe hashes only the full blocks of a longer prompt, so the oracle treats every hashed block as reusable (n = blocks * bs + 1); with n = blocks * bs it would under-predict one block for the 63/64 of prompts that are not an exact block multiple. The override is factored into apply_hybrid() and covered by 7 by-hash tests of the DEEP-309 shapes (resend, new scaffold junction, continuation, duplicate-copy eviction, never above the tree, unknown worker). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/standalone_indexer/server.rs | 204 +++++++++++++++++- 1 file changed, 193 insertions(+), 11 deletions(-) diff --git a/lib/kv-router/src/standalone_indexer/server.rs b/lib/kv-router/src/standalone_indexer/server.rs index 28ff81c1432e..4479b7c10eee 100644 --- a/lib/kv-router/src/standalone_indexer/server.rs +++ b/lib/kv-router/src/standalone_indexer/server.rs @@ -371,10 +371,48 @@ fn merge_score_responses(acc: &mut ScoreResponse, other: ScoreResponse) { /// When `kv_audit` logging is enabled (see [`super::logging_enabled`]), emits a /// single `QUERY` line per request, before returning to the client, carrying a /// timestamp, the queried block hashes, and the full merged JSON response. +/// Probe length the opt-in hybrid oracle needs to cap a hit at what the engine +/// can reuse (`(n - 1) / block_size`). +#[derive(Clone, Copy)] +enum ProbeLen { + /// `/query`: the caller's exact token count. + Tokens(usize), + /// `/query_by_hash`: callers (deepapi's KV probe) hash only the full + /// blocks of a longer prompt, so the prompt extends past the last hashed + /// block and every hashed block is reusable: `n = blocks * bs + 1`. + FullBlockPrefix, +} + +impl ProbeLen { + fn tokens(self, num_blocks: usize, block_size: u32) -> usize { + match self { + ProbeLen::Tokens(n) => n, + ProbeLen::FullBlockPrefix => num_blocks * block_size as usize + 1, + } + } +} + +/// Replace each worker's device-tier score with the hit the engine will +/// actually serve per the hybrid oracle, never raising it above the tree's. +fn apply_hybrid( + oracle: &super::hybrid::HybridOracle, + tiered: &mut TieredMatchDetails, + block_hashes: &[LocalBlockHash], + probe_len: ProbeLen, + block_size: u32, +) { + let num_tokens = probe_len.tokens(block_hashes.len(), block_size); + for (worker, score) in tiered.device.overlap_scores.scores.iter_mut() { + if let Some(predicted) = oracle.predict(*worker, block_hashes, num_tokens) { + *score = predicted.min(*score); + } + } +} + async fn run_fanout_query( trees: Vec<(IndexerKey, Indexer, u32)>, mut hashes_for: impl FnMut(u32) -> Vec, - num_tokens: Option, + probe_len: ProbeLen, pod_names: &HashMap, audit: QueryAudit<'_>, ) -> (StatusCode, Json) { @@ -394,14 +432,9 @@ async fn run_fanout_query( Ok(mut tiered) => { // Opt-in hybrid oracle: report the hit the engine will serve // (SWA-aware, MTP-adjusted) instead of the raw chain length. - // Needs the probe's token count, so `/query_by_hash` is untouched. - if let (Some(oracle), Some(num_tokens)) = (super::hybrid::oracle(), num_tokens) { - for (worker, score) in tiered.device.overlap_scores.scores.iter_mut() { - if let Some(predicted) = oracle.predict(*worker, &block_hashes, num_tokens) - { - *score = predicted.min(*score); - } - } + // Applies to `/query` and `/query_by_hash` (deepapi routing). + if let Some(oracle) = super::hybrid::oracle() { + apply_hybrid(oracle, &mut tiered, &block_hashes, probe_len, block_size); } let response = build_score_response(&tiered, block_size, pod_names); match merged.as_mut() { @@ -485,7 +518,7 @@ async fn query( }) .clone() }, - Some(req.token_ids.len()), + ProbeLen::Tokens(req.token_ids.len()), &state.registry.pod_names(), QueryAudit { model_name: &req.model_name, @@ -517,7 +550,7 @@ async fn query_by_hash( run_fanout_query( trees, |_| block_hashes.clone(), - None, + ProbeLen::FullBlockPrefix, &state.registry.pod_names(), QueryAudit { model_name: &req.model_name, @@ -707,6 +740,155 @@ pub fn create_router(state: Arc) -> Router { #[cfg(test)] mod tests { + mod hybrid_by_hash { + use super::super::{ProbeLen, apply_hybrid}; + use crate::indexer::TieredMatchDetails; + use crate::protocols::{ + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, + KvCacheStoreData, KvCacheStoredBlockData, LocalBlockHash, RouterEvent, StorageTier, + WorkerWithDpRank, + }; + use crate::standalone_indexer::hybrid::{HybridConfig, HybridOracle}; + + const BS: u32 = 64; + const W: u64 = 7; + + fn oracle() -> HybridOracle { + HybridOracle::new(HybridConfig { + block_size: BS, + mtp_drop_blocks: 1, + swa_need_blocks: None, + }) + } + /// deepapi's per-block xxh3 local hash stand-in for position `i`. + fn lh(i: u64) -> LocalBlockHash { + LocalBlockHash(5000 + i) + } + /// Engine block hash for position `i` (shared across KV groups). + fn eh(i: u64) -> u64 { + 9_000_000 + i + } + fn store_fa(o: &HybridOracle, n: u64) { + let ev = RouterEvent { + worker_id: W, + storage_tier: StorageTier::Device, + event: KvCacheEvent { + event_id: 0, + dp_rank: 0, + data: KvCacheEventData::Stored(KvCacheStoreData { + parent_hash: None, + start_position: None, + blocks: (0..n) + .map(|i| KvCacheStoredBlockData { + block_hash: ExternalSequenceBlockHash(eh(i)), + tokens_hash: lh(i), + mm_extra_info: None, + }) + .collect(), + }), + }, + }; + assert!(o.apply_main(ev).is_some()); + } + fn remove_fa(o: &HybridOracle, i: u64) -> bool { + o.apply_main(RouterEvent { + worker_id: W, + storage_tier: StorageTier::Device, + event: KvCacheEvent { + event_id: 1, + dp_rank: 0, + data: KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: vec![ExternalSequenceBlockHash(eh(i))], + }), + }, + }) + .is_some() + } + fn store_swa(o: &HybridOracle, range: std::ops::Range) { + let hashes: Vec = range.map(eh).collect(); + for g in 0..4 { + o.apply_swa(WorkerWithDpRank::new(W, 0), g, Some(512), true, &hashes); + } + } + /// Tree score as the stock radix tree would report it (full chain). + fn by_hash(o: &HybridOracle, probe_blocks: u64, tree_blocks: u32) -> u32 { + let mut tiered = TieredMatchDetails::default(); + tiered + .device + .overlap_scores + .scores + .insert(WorkerWithDpRank::new(W, 0), tree_blocks); + let hashes: Vec = (0..probe_blocks).map(lh).collect(); + apply_hybrid(o, &mut tiered, &hashes, ProbeLen::FullBlockPrefix, BS); + tiered.device.overlap_scores.scores[&WorkerWithDpRank::new(W, 0)] + } + + #[test] + fn full_block_prefix_counts_every_hashed_block_as_reusable() { + assert_eq!(ProbeLen::FullBlockPrefix.tokens(10, BS), 641); + assert_eq!(ProbeLen::Tokens(640).tokens(10, BS), 640); + } + + #[test] + fn by_hash_resend_reports_chain_minus_mtp() { + // DEEP-309 (c): resend/extension of a long prompt: engine = chain - 1 block. + let o = oracle(); + store_fa(&o, 40); + store_swa(&o, 0..40); + assert_eq!(by_hash(&o, 40, 40), 39); + } + + #[test] + fn by_hash_new_scaffold_junction_reports_zero() { + // DEEP-309 (b): first reuse of a new junction: no SWA checkpoint there. + let o = oracle(); + store_fa(&o, 60); + store_swa(&o, 51..60); + assert_eq!(by_hash(&o, 30, 30), 0); + } + + #[test] + fn by_hash_continuation_stops_at_prompt_checkpoint() { + // DEEP-309 (a): chain covers prompt(40)+completion(10); SWA checkpoint at 40. + let o = oracle(); + store_fa(&o, 50); + store_swa(&o, 31..40); + assert_eq!(by_hash(&o, 60, 50), 39); + } + + #[test] + fn by_hash_duplicate_copy_keeps_hit_after_first_eviction() { + // DEEP-309 eviction: a re-stored duplicate survives its first Removed. + let o = oracle(); + store_fa(&o, 20); + store_fa(&o, 20); + store_swa(&o, 0..20); + assert!(!remove_fa(&o, 10), "removal held back while a copy remains"); + assert_eq!(by_hash(&o, 20, 20), 19); + } + + #[test] + fn by_hash_never_raises_tree_score() { + let o = oracle(); + store_fa(&o, 40); + store_swa(&o, 0..40); + assert_eq!(by_hash(&o, 40, 12), 12); + } + + #[test] + fn by_hash_unknown_worker_keeps_tree_score() { + let o = oracle(); + let mut tiered = TieredMatchDetails::default(); + tiered + .device + .overlap_scores + .scores + .insert(WorkerWithDpRank::new(99, 0), 8); + apply_hybrid(&o, &mut tiered, &[lh(0)], ProbeLen::FullBlockPrefix, BS); + assert_eq!(tiered.device.overlap_scores.scores[&WorkerWithDpRank::new(99, 0)], 8); + } + } + use super::*; use crate::indexer::KvIndexerInterface; use crate::standalone_indexer::indexer::create_indexer;