From 7d461346ca026c32b78239ad033016eec237877e Mon Sep 17 00:00:00 2001 From: Change72 Date: Thu, 25 Jun 2026 11:21:36 -0700 Subject: [PATCH 1/3] kv-router: route vLLM "CPU" KV events to the HostPinned tier (#10368) Cherry-pick of upstream d8a94ef28. vLLM's SimpleCPUOffloadConnector tags its events medium="CPU"; our fork only mapped "CPU_PINNED"/"CPU_TIER1", so "CPU" fell back to the Device tier. CPU-tier evictions then removed blocks the GPU still held, and CPU-only blocks looked GPU-resident. Verified on a zai-org/GLM-5.3 pod: CPU stores/removes are balanced and tagged "CPU". Conflict in zmq_wire/tests.rs resolved by keeping our tests and adding one for the medium -> tier mapping (upstream's zero-block-size CPU test dropped: our branch skips those as partial-prefix entries). (cherry picked from commit d8a94ef28e000840387658b368ba025d2ebcc15c) Claude-Session: https://claude.ai/code/session_01575gLgRp2Xk4WLR9LsiMu6 --- lib/kv-router/src/indexer/local.rs | 9 ++++-- lib/kv-router/src/indexer/lower_tier.rs | 18 +++++++++--- .../src/indexer/lower_tier_indexers.rs | 21 ++++++++++++-- lib/kv-router/src/protocols.rs | 2 +- lib/kv-router/src/zmq_wire/tests.rs | 28 ++++++++++++++++++- lib/llm/src/kv_router/indexer/mod.rs | 26 +++++++++++------ 6 files changed, 86 insertions(+), 18 deletions(-) diff --git a/lib/kv-router/src/indexer/local.rs b/lib/kv-router/src/indexer/local.rs index 43c0ae9c980f..3293ad4a86d5 100644 --- a/lib/kv-router/src/indexer/local.rs +++ b/lib/kv-router/src/indexer/local.rs @@ -213,6 +213,9 @@ pub struct LocalKvIndexer { /// This stays separate from `event_buffer` so dump wait/build state can be /// managed on the async path without holding the buffer lock across `.await`. recovery_cache: Arc, + /// Shared metrics handle, also wired into lazily created lower-tier + /// indexers so HostPinned/Disk/External traffic is counted too. + metrics: Arc, /// Maximum number of events to keep in buffer max_buffer_size: usize, // Router sets this to WORKER_KV_INDEXER_BUFFER_SIZE #[cfg(test)] @@ -230,7 +233,8 @@ impl LocalKvIndexer { max_buffer_size: usize, ) -> Self { Self { - indexer: KvIndexer::new(token, kv_block_size, metrics), + indexer: KvIndexer::new(token, kv_block_size, metrics.clone()), + metrics, lower_tier_indexers: Arc::new(Mutex::new(HashMap::new())), event_buffer: Mutex::new(VecDeque::with_capacity(max_buffer_size)), recovery_cache: Arc::new(RecoverySnapshotCache::new()), @@ -669,10 +673,11 @@ impl LocalKvIndexer { indexers .entry(storage_tier) .or_insert_with(|| { - Arc::new(ThreadPoolIndexer::new( + Arc::new(ThreadPoolIndexer::new_with_metrics( LowerTierIndexer::new(), 1, self.block_size(), + Some(self.metrics.clone()), )) }) .clone() diff --git a/lib/kv-router/src/indexer/lower_tier.rs b/lib/kv-router/src/indexer/lower_tier.rs index ef30957d685c..4f31a92b833c 100644 --- a/lib/kv-router/src/indexer/lower_tier.rs +++ b/lib/kv-router/src/indexer/lower_tier.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use dashmap::DashMap; use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; -use super::{KvIndexerMetrics, SyncIndexer, WorkerLookupStats, WorkerTask}; +use super::{EventKind, KvIndexerMetrics, SyncIndexer, WorkerLookupStats, WorkerTask}; use crate::protocols::{ ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheEventError, KvCacheStoreData, KvCacheStoredBlockData, LocalBlockHash, OverlapScores, RouterEvent, WorkerWithDpRank, @@ -516,23 +516,33 @@ impl SyncIndexer for LowerTierIndexer { fn worker( &self, event_receiver: flume::Receiver, - _metrics: Option>, + metrics: Option>, ) -> anyhow::Result<()> { let mut worker_blocks = WorkerBlockIndex::default(); + let counters = metrics.as_ref().map(|m| m.prebind()); while let Ok(task) = event_receiver.recv() { match task { WorkerTask::Event(event) => { - if let Err(error) = self.apply_event(&mut worker_blocks, event) { + let kind = EventKind::of(&event.event.data); + let result = self.apply_event(&mut worker_blocks, event); + if let Err(ref error) = result { tracing::warn!(%error, "Failed to apply lower-tier event"); } + if let Some(ref c) = counters { + c.inc(kind, result); + } } WorkerTask::EventWithAck { event, resp } => { + let kind = EventKind::of(&event.event.data); let result = self.apply_event(&mut worker_blocks, event); let applied = result.is_ok(); - if let Err(error) = result { + if let Err(ref error) = result { tracing::warn!(%error, "Failed to apply lower-tier event"); } + if let Some(ref c) = counters { + c.inc(kind, result); + } let _ = resp.send(applied); } WorkerTask::Anchor { worker, anchor } => { diff --git a/lib/kv-router/src/indexer/lower_tier_indexers.rs b/lib/kv-router/src/indexer/lower_tier_indexers.rs index 90370dd9baea..a5dfa5617ff0 100644 --- a/lib/kv-router/src/indexer/lower_tier_indexers.rs +++ b/lib/kv-router/src/indexer/lower_tier_indexers.rs @@ -18,7 +18,7 @@ use std::collections::HashMap; use std::sync::{Arc, RwLock}; use crate::indexer::{ - LowerTierContinuation, LowerTierIndexer, LowerTierMatchDetails, MatchDetails, + KvIndexerMetrics, LowerTierContinuation, LowerTierIndexer, LowerTierMatchDetails, MatchDetails, ThreadPoolIndexer, WireTieredMatchDetails, }; use crate::protocols::{LocalBlockHash, StorageTier}; @@ -27,13 +27,28 @@ use crate::protocols::{LocalBlockHash, StorageTier}; /// non-device [`StorageTier`] that has received at least one event. #[derive(Clone)] pub struct LowerTierIndexers { + metrics: Option>, num_threads: usize, block_size: u32, indexers: Arc>>>>, } impl LowerTierIndexers { + /// Metrics-less constructor for call sites without a `KvIndexerMetrics` handle. + /// Router production assembly should use [`new_with_metrics`](Self::new_with_metrics) + /// so lower-tier traffic is included in `kv_cache_events_applied`. pub fn new(num_threads: usize, block_size: u32) -> Self { + Self::new_with_metrics(num_threads, block_size, None) + } + + /// Same as [`new`](Self::new) but wires `kv_cache_events_applied` + /// counters into every lazily created per-tier indexer, matching the + /// observability of the device-tier path. + pub fn new_with_metrics( + num_threads: usize, + block_size: u32, + metrics: Option>, + ) -> Self { assert!( num_threads > 0, "lower-tier indexer threads must be non-zero" @@ -41,6 +56,7 @@ impl LowerTierIndexers { Self { num_threads, block_size, + metrics, indexers: Arc::new(RwLock::new(HashMap::new())), } } @@ -60,10 +76,11 @@ impl LowerTierIndexers { .unwrap() .entry(storage_tier) .or_insert_with(|| { - Arc::new(ThreadPoolIndexer::new( + Arc::new(ThreadPoolIndexer::new_with_metrics( LowerTierIndexer::new(), self.num_threads, self.block_size, + self.metrics.clone(), )) }) .clone() diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 9cb7761b9a82..0ca52cbab2f6 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -330,7 +330,7 @@ impl StorageTier { pub fn from_kv_medium(medium: &str) -> Option { match medium { "GPU" | "DEVICE" => Some(Self::Device), - "CPU_PINNED" | "CPU_TIER1" => Some(Self::HostPinned), + "CPU" | "CPU_PINNED" | "CPU_TIER1" => Some(Self::HostPinned), "CPU_TIER2" | "DISK" | "NVME" => Some(Self::Disk), "EXTERNAL" | "NETWORK" | "REMOTE" | "SHARED" => Some(Self::External), _ => None, diff --git a/lib/kv-router/src/zmq_wire/tests.rs b/lib/kv-router/src/zmq_wire/tests.rs index 5e3213508832..11de5f957231 100644 --- a/lib/kv-router/src/zmq_wire/tests.rs +++ b/lib/kv-router/src/zmq_wire/tests.rs @@ -8,7 +8,7 @@ use rmp_serde::{from_slice, to_vec}; use crate::protocols::{ BlockExtraInfo, BlockHashOptions, BlockMmObjectInfo, ExternalSequenceBlockHash, - KvCacheEventData, WorkerWithDpRank, compute_block_hash_for_seq, + KvCacheEventData, StorageTier, WorkerWithDpRank, compute_block_hash_for_seq, }; use super::filter::KvCacheSpecKind; @@ -641,3 +641,29 @@ fn test_convert_event_short_token_ids_keeps_parsed_blocks() { other => panic!("expected Stored event, got {other:?}"), } } + +#[test] +fn test_convert_event_cpu_medium_lands_on_host_tier() { + // vLLM's SimpleCPUOffloadConnector tags its events medium="CPU". Before + // #10368 this string was unknown and fell back to the Device tier, so a + // CPU-tier eviction removed a block the GPU still held. + for (medium, expected) in [ + (Some("CPU"), StorageTier::HostPinned), + (Some("CPU_PINNED"), StorageTier::HostPinned), + (Some("GPU"), StorageTier::Device), + (None, StorageTier::Device), + ] { + let raw_event = RawKvEvent::BlockRemoved { + block_hashes: vec![BlockHashValue::Unsigned(201)], + medium: medium.map(str::to_string), + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + }; + let warning_count = Arc::new(AtomicU32::new(0)); + let placement = convert_event(raw_event, 7, 4, WorkerWithDpRank::new(3, 0), &warning_count) + .expect("remove event converts") + .expect("remove event is kept"); + assert_eq!(placement.placement.tier, expected, "medium {medium:?}"); + } +} diff --git a/lib/llm/src/kv_router/indexer/mod.rs b/lib/llm/src/kv_router/indexer/mod.rs index 89565e3cb3d3..2142dae19fe6 100644 --- a/lib/llm/src/kv_router/indexer/mod.rs +++ b/lib/llm/src/kv_router/indexer/mod.rs @@ -116,12 +116,13 @@ impl Indexer { ConcurrentRadixTreeCompressed::new(), kv_router_config.router_event_threads as usize, block_size, - Some(kv_indexer_metrics), + Some(kv_indexer_metrics.clone()), prune_config, )), - lower_tier: LowerTierIndexers::new( + lower_tier: LowerTierIndexers::new_with_metrics( kv_router_config.router_event_threads as usize, block_size, + Some(kv_indexer_metrics), ), approx: None, primary_records_routing_decisions: true, @@ -134,10 +135,14 @@ impl Indexer { cancellation_token, None, block_size, - kv_indexer_metrics, + kv_indexer_metrics.clone(), prune_config, ), - lower_tier: LowerTierIndexers::new(1, block_size), + lower_tier: LowerTierIndexers::new_with_metrics( + 1, + block_size, + Some(kv_indexer_metrics), + ), approx: None, primary_records_routing_decisions: true, }); @@ -152,11 +157,12 @@ impl Indexer { ConcurrentRadixTreeCompressed::new(), kv_router_config.router_event_threads as usize, block_size, - Some(kv_indexer_metrics), + Some(kv_indexer_metrics.clone()), )), - lower_tier: LowerTierIndexers::new( + lower_tier: LowerTierIndexers::new_with_metrics( kv_router_config.router_event_threads as usize, block_size, + Some(kv_indexer_metrics), ), approx, primary_records_routing_decisions: false, @@ -171,10 +177,14 @@ impl Indexer { cancellation_token, None, block_size, - kv_indexer_metrics, + kv_indexer_metrics.clone(), None, ), - lower_tier: LowerTierIndexers::new(1, block_size), + lower_tier: LowerTierIndexers::new_with_metrics( + 1, + block_size, + Some(kv_indexer_metrics), + ), approx, primary_records_routing_decisions: false, }) From 11cd973221d292b776600b31c234dd456f10d4bf Mon Sep 17 00:00:00 2001 From: Thach Nguyen Date: Thu, 24 Sep 2026 00:43:07 +0000 Subject: [PATCH 2/3] standalone-indexer: promote host-only ancestors when a device store needs them With the "CPU" medium routed to the HostPinned tier (#10368), a device store whose parent the device tree only knows from the host tier is rejected with ParentBlockNotFound, and every later block of that sequence is rejected too. On zai-org/GLM-5.3 (SimpleCPUOffloadConnector) 388/388 rejected stores in a sample had a host-only parent, ~100-150/min per offload pod after the image swap. The engine can only cache a device block whose parent is resident on the device, so a device store proves its ancestor chain is there. Each listener now keeps a TierBridge: per (worker, dp_rank) sets of device and host block hashes (host with parent links), updated in event order from live and recovered events and reset on TreeDump. When a device store's parent is host-only, the bridge walks the host chain up to the nearest device block (or the root) and the listener applies that chain as a device store before the original event. A broken chain is left alone. Claude-Session: https://claude.ai/code/session_01575gLgRp2Xk4WLR9LsiMu6 --- .../src/standalone_indexer/listener.rs | 16 +- lib/kv-router/src/standalone_indexer/mod.rs | 1 + .../src/standalone_indexer/tier_bridge.rs | 293 ++++++++++++++++++ 3 files changed, 309 insertions(+), 1 deletion(-) create mode 100644 lib/kv-router/src/standalone_indexer/tier_bridge.rs diff --git a/lib/kv-router/src/standalone_indexer/listener.rs b/lib/kv-router/src/standalone_indexer/listener.rs index 135b9eb7c698..d3a52afd0bcf 100644 --- a/lib/kv-router/src/standalone_indexer/listener.rs +++ b/lib/kv-router/src/standalone_indexer/listener.rs @@ -16,6 +16,7 @@ use crate::zmq_wire::{ZmqEventNormalizer, decode_event_batch}; use super::evictions::PendingEvictions; use super::indexer::Indexer; use super::registry::ListenerRecord; +use super::tier_bridge::TierBridge; use super::zmq::{MultipartMessage, SharedSocket, connect_sub_socket, recv_multipart}; const WATERMARK_UNSET: u64 = u64::MAX; @@ -129,6 +130,7 @@ struct ListenerLoop { /// drains it through there. Untouched unless `--keep-evictions` is set. pending_evictions: Arc>, normalizer: ZmqEventNormalizer, + tier_bridge: TierBridge, messages_processed: u64, } @@ -157,6 +159,7 @@ impl ListenerLoop { watermark, pending_evictions, normalizer: ZmqEventNormalizer::new(block_size), + tier_bridge: TierBridge::new(), messages_processed: 0, } } @@ -330,6 +333,7 @@ impl ListenerLoop { // them too — replaying them against the snapshot could remove // blocks the dump says are live. self.pending_evictions.lock().clear(); + self.tier_bridge.reset(); self.indexer .remove_worker_dp_rank(self.worker_id, self.dp_rank) .await; @@ -388,7 +392,7 @@ impl ListenerLoop { /// events under the worker that was queried), so they are rewritten before /// applying. Returns the count actually applied (after the /// `keep_evictions` measurement filter). - async fn apply_recovered_events(&self, events: Vec) -> u64 { + async fn apply_recovered_events(&mut self, events: Vec) -> u64 { let mut applied = 0; for mut event in events { event.worker_id = self.worker_id; @@ -397,6 +401,11 @@ impl ListenerLoop { // Audit-log the recovered event before the measurement filter. audit_log_event(&event, event.event.event_id, "recover"); + if let Some(promoted) = self.tier_bridge.observe(&event) + && !self.keep_evictions_intercept(&promoted) + { + self.indexer.apply_event_routed(promoted).await; + } // Feed-layer measurement filter (same as apply_live_batch). if self.keep_evictions_intercept(&event) { continue; @@ -481,6 +490,11 @@ 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"); + if let Some(promoted) = self.tier_bridge.observe(&router_event) + && !self.keep_evictions_intercept(&promoted) + { + self.indexer.apply_event_routed(promoted).await; + } // 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 921f16bf78a7..03e891d845a1 100644 --- a/lib/kv-router/src/standalone_indexer/mod.rs +++ b/lib/kv-router/src/standalone_indexer/mod.rs @@ -32,6 +32,7 @@ pub mod pod_watcher; pub mod recovery; pub mod registry; pub mod server; +mod tier_bridge; mod zmq; use std::sync::{Arc, OnceLock}; diff --git a/lib/kv-router/src/standalone_indexer/tier_bridge.rs b/lib/kv-router/src/standalone_indexer/tier_bridge.rs new file mode 100644 index 000000000000..dbb2063e3ef8 --- /dev/null +++ b/lib/kv-router/src/standalone_indexer/tier_bridge.rs @@ -0,0 +1,293 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Per-listener repair for device stores whose parent is only known in the +//! host (CPU) tier. +//! +//! With CPU offload, a pod's device-tier chain can reference a parent block the +//! device index never saw stored there (seeded from a TreeDump, reloaded from +//! the CPU tier, ...). The device radix tree rejects such a store with +//! `ParentBlockNotFound` and every later block of that sequence is rejected +//! too. The engine can only cache a device block whose parent is resident on +//! the device, so a device store is proof its ancestor chain is there: this +//! bridge re-publishes that chain from the host tier as a device store first. + +use rustc_hash::{FxHashMap, FxHashSet}; + +use crate::protocols::{ + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheStoreData, + KvCacheStoredBlockData, RouterEvent, StorageTier, +}; + +struct HostBlock { + parent: Option, + block: KvCacheStoredBlockData, +} + +/// Tracks one `(worker_id, dp_rank)`'s device and host block hashes, in event +/// order, so it can synthesize the missing device ancestors of a store. +#[derive(Default)] +pub struct TierBridge { + device: FxHashSet, + host: FxHashMap, +} + +impl TierBridge { + pub fn new() -> Self { + Self::default() + } + + /// Forget everything, e.g. before a TreeDump replaces this worker's state. + pub fn reset(&mut self) { + self.device.clear(); + self.host.clear(); + } + + /// Record `event` and return a device store to apply BEFORE it, when + /// `event` is a device store whose parent is only known in the host tier. + pub fn observe(&mut self, event: &RouterEvent) -> Option { + let promotion = match (&event.event.data, event.storage_tier) { + (KvCacheEventData::Stored(store), StorageTier::Device) => store + .parent_hash + .and_then(|parent| self.promote_chain(parent)) + .map(|(anchor, blocks)| promoted_store(event, anchor, blocks)), + _ => None, + }; + if let Some(promoted) = &promotion { + self.record(promoted); + } + self.record(event); + promotion + } + + fn record(&mut self, event: &RouterEvent) { + match (&event.event.data, event.storage_tier) { + (KvCacheEventData::Stored(store), StorageTier::Device) => { + self.device + .extend(store.blocks.iter().map(|b| b.block_hash)); + } + (KvCacheEventData::Stored(store), StorageTier::HostPinned) => { + let mut parent = store.parent_hash; + for block in &store.blocks { + self.host.insert( + block.block_hash, + HostBlock { + parent, + block: block.clone(), + }, + ); + parent = Some(block.block_hash); + } + } + (KvCacheEventData::Removed(removed), StorageTier::Device) => { + for hash in &removed.block_hashes { + self.device.remove(hash); + } + } + (KvCacheEventData::Removed(removed), StorageTier::HostPinned) => { + for hash in &removed.block_hashes { + self.host.remove(hash); + } + } + (KvCacheEventData::Cleared, StorageTier::Device) => self.device.clear(), + (KvCacheEventData::Cleared, StorageTier::HostPinned) => self.host.clear(), + _ => {} + } + } + + /// Walk from `parent` up through host-only blocks until reaching a device + /// block or the root. Returns the anchor and the chain root-first, or + /// `None` when `parent` is already on the device or the chain is broken. + fn promote_chain( + &self, + parent: ExternalSequenceBlockHash, + ) -> Option<( + Option, + Vec, + )> { + if self.device.contains(&parent) { + return None; + } + let mut chain = Vec::new(); + let mut cursor = Some(parent); + while let Some(hash) = cursor { + if self.device.contains(&hash) { + break; + } + let host_block = self.host.get(&hash)?; + chain.push(host_block.block.clone()); + if chain.len() > self.host.len() { + return None; // cycle guard + } + cursor = host_block.parent; + } + chain.reverse(); + Some((cursor, chain)) + } +} + +fn promoted_store( + source: &RouterEvent, + anchor: Option, + blocks: Vec, +) -> RouterEvent { + RouterEvent::with_storage_tier( + source.worker_id, + KvCacheEvent { + event_id: source.event.event_id, + data: KvCacheEventData::Stored(KvCacheStoreData { + parent_hash: anchor, + start_position: None, + blocks, + }), + dp_rank: source.event.dp_rank, + }, + StorageTier::Device, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::{KvCacheRemoveData, LocalBlockHash}; + + fn block(hash: u64) -> KvCacheStoredBlockData { + KvCacheStoredBlockData { + block_hash: ExternalSequenceBlockHash(hash), + tokens_hash: LocalBlockHash(hash + 1000), + mm_extra_info: None, + } + } + + fn store(tier: StorageTier, parent: Option, hashes: &[u64]) -> RouterEvent { + RouterEvent::with_storage_tier( + 7, + KvCacheEvent { + event_id: 1, + data: KvCacheEventData::Stored(KvCacheStoreData { + parent_hash: parent.map(ExternalSequenceBlockHash), + start_position: None, + blocks: hashes.iter().copied().map(block).collect(), + }), + dp_rank: 0, + }, + tier, + ) + } + + fn remove(tier: StorageTier, hashes: &[u64]) -> RouterEvent { + RouterEvent::with_storage_tier( + 7, + KvCacheEvent { + event_id: 2, + data: KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: hashes + .iter() + .copied() + .map(ExternalSequenceBlockHash) + .collect(), + }), + dp_rank: 0, + }, + tier, + ) + } + + fn stored_hashes(event: &RouterEvent) -> (Option, Vec) { + match &event.event.data { + KvCacheEventData::Stored(s) => ( + s.parent_hash.map(|h| h.0), + s.blocks.iter().map(|b| b.block_hash.0).collect(), + ), + other => panic!("expected store, got {other:?}"), + } + } + + #[test] + fn device_store_with_host_only_parent_promotes_chain_to_device_anchor() { + let mut bridge = TierBridge::new(); + assert!( + bridge + .observe(&store(StorageTier::Device, None, &[1])) + .is_none() + ); + assert!( + bridge + .observe(&store(StorageTier::HostPinned, Some(1), &[2, 3])) + .is_none() + ); + + let promoted = bridge + .observe(&store(StorageTier::Device, Some(3), &[4])) + .expect("parent 3 is host-only"); + assert_eq!(promoted.storage_tier, StorageTier::Device); + assert_eq!(stored_hashes(&promoted), (Some(1), vec![2, 3])); + // Now on the device: a second child needs nothing. + assert!( + bridge + .observe(&store(StorageTier::Device, Some(4), &[5])) + .is_none() + ); + } + + #[test] + fn host_chain_reaching_the_root_is_promoted_from_the_root() { + let mut bridge = TierBridge::new(); + bridge.observe(&store(StorageTier::HostPinned, None, &[1, 2])); + let promoted = bridge + .observe(&store(StorageTier::Device, Some(2), &[3])) + .expect("chain is host-only up to the root"); + assert_eq!(stored_hashes(&promoted), (None, vec![1, 2])); + } + + #[test] + fn device_parent_or_unknown_parent_needs_no_promotion() { + let mut bridge = TierBridge::new(); + bridge.observe(&store(StorageTier::Device, None, &[1])); + bridge.observe(&store(StorageTier::HostPinned, Some(1), &[2])); + assert!( + bridge + .observe(&store(StorageTier::Device, Some(1), &[9])) + .is_none() + ); + assert!( + bridge + .observe(&store(StorageTier::Device, Some(42), &[10])) + .is_none() + ); + } + + #[test] + fn broken_host_chain_is_not_promoted() { + let mut bridge = TierBridge::new(); + bridge.observe(&store(StorageTier::HostPinned, Some(99), &[2])); + assert!( + bridge + .observe(&store(StorageTier::Device, Some(2), &[3])) + .is_none() + ); + } + + #[test] + fn removals_and_reset_are_tracked_per_tier() { + let mut bridge = TierBridge::new(); + bridge.observe(&store(StorageTier::Device, None, &[1])); + bridge.observe(&store(StorageTier::HostPinned, Some(1), &[2])); + bridge.observe(&remove(StorageTier::Device, &[1])); + // 1 left the device but is not in the host tier: chain is broken. + assert!( + bridge + .observe(&store(StorageTier::Device, Some(2), &[3])) + .is_none() + ); + + bridge.reset(); + bridge.observe(&store(StorageTier::HostPinned, None, &[5])); + bridge.observe(&remove(StorageTier::HostPinned, &[5])); + assert!( + bridge + .observe(&store(StorageTier::Device, Some(5), &[6])) + .is_none() + ); + } +} From bfaa04b8169ba63bff5a38ce47a107a6b53537c9 Mon Sep 17 00:00:00 2001 From: Shang-Pin Date: Mon, 19 Oct 2026 20:43:21 +0000 Subject: [PATCH 3/3] standalone-indexer: hash a vLLM cache_salt into its sequence's first block vLLM mixes a request's cache_salt into the hash of the sequence's first block and publishes it in that block's BlockStored extra_keys, but the normalizer kept only 64-hex MM ids. Salted and unsalted chains with the same tokens then shared tree edges while carrying different engine hashes: the tree keeps one engine hash per edge position, so the second chain's children were rejected as ParentBlockNotFound, and probes scored salted prefixes as reusable by anyone. In a BlockStored with no parent, a bare string in block 0 that is neither an MM id nor the event's LoRA name is the salt. It joins that block's MM ids as cache_salt_extra_hash(salt): a 64-hex salt parses like an MM id, any other salt is its XXH3-64. Non-root events and other blocks are unchanged, as is extra_keys_to_block_mm_infos for existing callers. deepapi's probe hashes the salt the same way under the per-model kv_probe_cache_salt flag; the tests pin vectors shared with it. Co-Authored-By: Claude Opus 5.5 (1M context) --- lib/kv-router/src/standalone_indexer/docs.md | 9 ++ lib/kv-router/src/zmq_wire/deserialize.rs | 27 +++-- lib/kv-router/src/zmq_wire/extra_keys.rs | 31 ++++- lib/kv-router/src/zmq_wire/mod.rs | 5 +- lib/kv-router/src/zmq_wire/tests.rs | 121 +++++++++++++++++++ 5 files changed, 183 insertions(+), 10 deletions(-) diff --git a/lib/kv-router/src/standalone_indexer/docs.md b/lib/kv-router/src/standalone_indexer/docs.md index c671ef235980..2c966052711d 100644 --- a/lib/kv-router/src/standalone_indexer/docs.md +++ b/lib/kv-router/src/standalone_indexer/docs.md @@ -52,6 +52,15 @@ Known gap: the engine's local-indexer TreeDumps (`/kv_recover`) carry token-only hashes, so image blocks restored at startup stay token-only until evicted. +A request's vLLM `cache_salt` is hashed the same way: vLLM appends it to the +extra keys of the sequence's first block only, so in a BlockStored with no +parent, a bare string in block 0 that is neither an MM id nor the LoRA name +joins that block's image ids as `cache_salt_extra_hash(salt)` (a 64-hex salt +parses like an image id, any other salt is its XXH3-64). Salted chains then get +their own tree path instead of sharing the unsalted one's edges. deepapi's +probe mixes the salt in under the per-model `kv_probe_cache_salt` flag. The +`/kv_recover` gap above applies to salted blocks too. + ## Data-parallel engines (`--watch-dp-size`) A vLLM engine with `--data-parallel-size N` runs N ranks per pod, each with its diff --git a/lib/kv-router/src/zmq_wire/deserialize.rs b/lib/kv-router/src/zmq_wire/deserialize.rs index d10a53a7f451..c198ce26a8bc 100644 --- a/lib/kv-router/src/zmq_wire/deserialize.rs +++ b/lib/kv-router/src/zmq_wire/deserialize.rs @@ -9,7 +9,7 @@ use serde::de::{self, IgnoredAny, MapAccess, SeqAccess, Visitor}; use crate::protocols::BlockExtraInfo; -use super::extra_keys::extra_keys_to_block_mm_infos; +use super::extra_keys::stored_extra_keys_to_block_infos; use super::filter::{BlockStoredTrailingField, KvCacheEventMetadata, KvCacheEventTrailingField}; use super::types::{BlockHashValue, ExtraKeyItem, KvTokenIds, RawKvEvent}; @@ -106,16 +106,22 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { let block_size = block_size.ok_or_else(|| de::Error::missing_field("block_size"))?; let medium = medium.unwrap_or(None); - let block_mm_infos = block_mm_infos - .unwrap_or(None) - .or_else(|| extra_keys_to_block_mm_infos(extra_keys.unwrap_or(None))); + let parent_block_hash = parent_block_hash.unwrap_or(None); + let lora_name = lora_name.unwrap_or(None); + let block_mm_infos = block_mm_infos.unwrap_or(None).or_else(|| { + stored_extra_keys_to_block_infos( + extra_keys.unwrap_or(None), + parent_block_hash.is_none(), + lora_name.as_deref(), + ) + }); Ok(RawKvEvent::BlockStored { block_hashes, - parent_block_hash: parent_block_hash.unwrap_or(None), + parent_block_hash, token_ids: raw_token_ids, block_size, medium, - lora_name: lora_name.unwrap_or(None), + lora_name, block_mm_infos, is_eagle: Some(is_eagle), group_idx: metadata.group_idx, @@ -194,8 +200,13 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { while seq.next_element::()?.is_some() {} - let block_mm_infos = - block_mm_infos.or_else(|| extra_keys_to_block_mm_infos(extra_keys)); + let block_mm_infos = block_mm_infos.or_else(|| { + stored_extra_keys_to_block_infos( + extra_keys, + parent_block_hash.is_none(), + lora_name.as_deref(), + ) + }); let (raw_token_ids, is_eagle) = normalize_token_ids(token_ids); Ok(RawKvEvent::BlockStored { diff --git a/lib/kv-router/src/zmq_wire/extra_keys.rs b/lib/kv-router/src/zmq_wire/extra_keys.rs index df33469a5660..7e1f5fe7f512 100644 --- a/lib/kv-router/src/zmq_wire/extra_keys.rs +++ b/lib/kv-router/src/zmq_wire/extra_keys.rs @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use xxhash_rust::xxh3; + use crate::protocols::{BlockExtraInfo, BlockMmObjectInfo}; use super::types::ExtraKeyItem; @@ -17,6 +19,12 @@ pub fn parse_mm_hash_from_extra_key(s: &str) -> Option { None } +/// The u64 a vLLM `cache_salt` adds to the first block of its sequence, hashed +/// like an MM id: a 64-char hex salt parses as one, any other salt is its XXH3-64. +pub fn cache_salt_extra_hash(salt: &str) -> u64 { + parse_mm_hash_from_extra_key(salt).unwrap_or_else(|| xxh3::xxh3_64(salt.as_bytes())) +} + /// Convert vLLM BlockStored extra_keys to block-level MM infos. /// extra_keys is a list aligned with blocks: /// - None => no MM content in that block @@ -25,6 +33,22 @@ pub fn parse_mm_hash_from_extra_key(s: &str) -> Option { /// start offsets (vLLM 0.19+) pub fn extra_keys_to_block_mm_infos( extra_keys: Option>>>, +) -> Option>> { + stored_extra_keys_to_block_infos(extra_keys, false, None) +} + +/// [`extra_keys_to_block_mm_infos`] for a whole BlockStored event, which also keeps +/// the request's `cache_salt`. +/// +/// vLLM appends the salt only to the extra keys of the sequence's first block +/// (`lora + mm + [cache_salt] + prompt_embeds`), so when the event starts the +/// sequence (no parent), a bare string in block 0 that is neither an MM id nor +/// the LoRA name is the salt. It joins that block's MM objects via +/// [`cache_salt_extra_hash`], so salted chains get their own tree path. +pub fn stored_extra_keys_to_block_infos( + extra_keys: Option>>>, + starts_sequence: bool, + lora_name: Option<&str>, ) -> Option>> { let extra_keys = extra_keys?; if extra_keys.is_empty() { @@ -33,11 +57,16 @@ pub fn extra_keys_to_block_mm_infos( let infos: Vec> = extra_keys .into_iter() - .map(|block_keys| { + .enumerate() + .map(|(block_idx, block_keys)| { + let salted_block = starts_sequence && block_idx == 0; let mm_objects: Vec = block_keys .unwrap_or_default() .iter() .filter_map(|key| match key { + ExtraKeyItem::Hash(s) if salted_block && Some(s.as_str()) != lora_name => { + Some(cache_salt_extra_hash(s)) + } ExtraKeyItem::Hash(hash) | ExtraKeyItem::HashWithSignedOffset((hash, _)) | ExtraKeyItem::HashWithUnsignedOffset((hash, _)) => { diff --git a/lib/kv-router/src/zmq_wire/mod.rs b/lib/kv-router/src/zmq_wire/mod.rs index ae0bb20defbf..f3d597355358 100644 --- a/lib/kv-router/src/zmq_wire/mod.rs +++ b/lib/kv-router/src/zmq_wire/mod.rs @@ -27,7 +27,10 @@ pub use convert::{ ConvertError, convert_event, create_stored_block_from_parts, create_stored_blocks, is_partial_prefix_entry, }; -pub use extra_keys::{extra_keys_to_block_mm_infos, parse_mm_hash_from_extra_key}; +pub use extra_keys::{ + cache_salt_extra_hash, extra_keys_to_block_mm_infos, parse_mm_hash_from_extra_key, + stored_extra_keys_to_block_infos, +}; pub use filter::KvCacheSpecKind; pub use types::{BlockHashValue, ExtraKeyItem, KvEventBatch, KvTokenIds, RawKvEvent}; diff --git a/lib/kv-router/src/zmq_wire/tests.rs b/lib/kv-router/src/zmq_wire/tests.rs index 11de5f957231..d5bdeb4d38bc 100644 --- a/lib/kv-router/src/zmq_wire/tests.rs +++ b/lib/kv-router/src/zmq_wire/tests.rs @@ -667,3 +667,124 @@ fn test_convert_event_cpu_medium_lands_on_host_tier() { assert_eq!(placement.placement.tier, expected, "medium {medium:?}"); } } + +#[derive(serde::Serialize)] +#[serde(untagged)] +enum WireExtraKey { + Bare(String), + WithOffset((String, i64)), +} + +fn stored_tokens_hashes( + parent: Option, + lora_name: Option<&str>, + extra_keys: Vec>>, + bare_keys: Vec>>, +) -> Vec { + let extra_keys: Vec>> = extra_keys + .into_iter() + .zip(bare_keys) + .map(|(tuples, bare)| { + let mut items: Vec = bare + .unwrap_or_default() + .into_iter() + .map(WireExtraKey::Bare) + .collect(); + items.extend(tuples.unwrap_or_default().into_iter().map(WireExtraKey::WithOffset)); + (!items.is_empty()).then_some(items) + }) + .collect(); + let encoded = to_vec(&( + "BlockStored", + vec![BlockHashValue::Unsigned(21), BlockHashValue::Unsigned(22)], + parent.map(BlockHashValue::Unsigned), + vec![10u32, 11, 12, 13], + 2usize, + Option::::None, + Option::::None, + lora_name.map(str::to_string), + Some(extra_keys), + )) + .unwrap(); + let raw_event: RawKvEvent = from_slice(&encoded).unwrap(); + let warning_count = Arc::new(AtomicU32::new(0)); + let placement = convert_event(raw_event, 7, 2, WorkerWithDpRank::new(3, 0), &warning_count) + .expect("block size matches") + .expect("store is kept"); + match placement.event.data { + KvCacheEventData::Stored(store_data) => { + store_data.blocks.iter().map(|b| b.tokens_hash.0).collect() + } + other => panic!("expected Stored event, got {other:?}"), + } +} + +// Shared with deepapi's probe tests (tests/test_kv_probe_cache_salt.py). +const PLAIN_10_11: u64 = 7369860969341005925; +const SALTED_10_11: u64 = 10774232232830598228; +const PLAIN_12_13: u64 = 2046086924859723728; +const SALTED_IMAGE_10_11: u64 = 1881935480998609139; + +#[test] +fn test_cache_salt_extra_hash() { + assert_eq!(cache_salt_extra_hash("tenant-a"), 7159698576361518977); + assert_eq!(cache_salt_extra_hash(&"ab".repeat(32)), 0xabababababababab); +} + +#[test] +fn test_cache_salt_hashes_into_first_block_only() { + let salted = stored_tokens_hashes( + None, + None, + vec![None, None], + vec![Some(vec!["tenant-a".into()]), None], + ); + assert_eq!(salted, vec![SALTED_10_11, PLAIN_12_13]); + + let plain = stored_tokens_hashes(None, None, vec![None, None], vec![None, None]); + assert_eq!(plain, vec![PLAIN_10_11, PLAIN_12_13]); +} + +#[test] +fn test_cache_salt_ignored_when_event_continues_a_sequence() { + let hashes = stored_tokens_hashes( + Some(9), + None, + vec![None, None], + vec![Some(vec!["tenant-a".into()]), None], + ); + assert_eq!(hashes, vec![PLAIN_10_11, PLAIN_12_13]); +} + +#[test] +fn test_lora_name_is_not_a_cache_salt() { + let lora_only = stored_tokens_hashes( + None, + Some("adapter"), + vec![None, None], + vec![Some(vec!["adapter".into()]), Some(vec!["adapter".into()])], + ); + let lora_seed_only = compute_block_hash_for_seq( + &[10, 11, 12, 13], + 2, + BlockHashOptions { + lora_name: Some("adapter"), + ..Default::default() + }, + ); + assert_eq!( + lora_only, + lora_seed_only.iter().map(|h| h.0).collect::>() + ); +} + +#[test] +fn test_cache_salt_combines_with_image_in_first_block() { + let hashes = stored_tokens_hashes( + None, + None, + vec![Some(vec![("ab".repeat(32), 0)]), None], + vec![Some(vec!["tenant-a".into()]), None], + ); + assert_eq!(hashes, vec![SALTED_IMAGE_10_11, PLAIN_12_13]); +}