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/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 fadc53c263dc..561d9a58f812 100644 --- a/lib/kv-router/src/standalone_indexer/mod.rs +++ b/lib/kv-router/src/standalone_indexer/mod.rs @@ -33,6 +33,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() + ); + } +} 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, })