diff --git a/lib/kv-router/src/standalone_indexer/docs.md b/lib/kv-router/src/standalone_indexer/docs.md index 22a63086f621..1819407421d6 100644 --- a/lib/kv-router/src/standalone_indexer/docs.md +++ b/lib/kv-router/src/standalone_indexer/docs.md @@ -40,6 +40,18 @@ TreeDump is tens of MB serialized inside the engine process; with the previous 10 s timeout and no gate, a fleet-wide (re)subscription lost about a third of its recoveries on 39 DP=2 pods. +## Image blocks hash from tokens only + +vLLM attaches each image's identifier to a stored block in `extra_keys`, and the +shared ZMQ normalizer would mix it into that block's tokens hash. The standalone +indexer's queriers do not: deepapi's probe hashes plain token ids and the +engine's local-indexer TreeDumps (`/kv_recover`) carry token-only hashes. With +the image hash mixed in, every query stopped matching at a conversation's first +image block. The listener therefore builds its normalizer with +`with_plain_mm_hashing()`. Trade-off: two prompts that share text but carry +different images at the same position are indexed as the same prefix, so the +indexer can over-report a hit the engine will not give. + ## 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/standalone_indexer/listener.rs b/lib/kv-router/src/standalone_indexer/listener.rs index 135b9eb7c698..0a8fd841a0b2 100644 --- a/lib/kv-router/src/standalone_indexer/listener.rs +++ b/lib/kv-router/src/standalone_indexer/listener.rs @@ -156,7 +156,7 @@ impl ListenerLoop { http_client, watermark, pending_evictions, - normalizer: ZmqEventNormalizer::new(block_size), + normalizer: ZmqEventNormalizer::new(block_size).with_plain_mm_hashing(), messages_processed: 0, } } diff --git a/lib/kv-router/src/zmq_wire/mod.rs b/lib/kv-router/src/zmq_wire/mod.rs index ae0bb20defbf..ecc276a685eb 100644 --- a/lib/kv-router/src/zmq_wire/mod.rs +++ b/lib/kv-router/src/zmq_wire/mod.rs @@ -42,6 +42,7 @@ pub struct ZmqEventNormalizer { kv_block_size: u32, warning_count: Arc, group_metadata: FxHashMap<(DpRank, u32), KvCacheGroupMetadata>, + hash_mm: bool, } #[derive(Debug, Clone, Copy)] @@ -77,6 +78,7 @@ impl ZmqEventNormalizer { kv_block_size, warning_count: Arc::new(AtomicU32::new(0)), group_metadata: FxHashMap::default(), + hash_mm: true, } } @@ -85,9 +87,19 @@ impl ZmqEventNormalizer { kv_block_size, warning_count, group_metadata: FxHashMap::default(), + hash_mm: true, } } + /// Hash image blocks from their tokens alone, ignoring the multimodal + /// identifiers vLLM attaches in `extra_keys`. Needed when the querier + /// hashes plain tokens (deepapi's probe, the engine's local-indexer + /// TreeDumps): otherwise every chain diverges at its first image block. + pub fn with_plain_mm_hashing(mut self) -> Self { + self.hash_mm = false; + self + } + pub fn preprocess(&mut self, raw: RawKvEvent, worker: WorkerWithDpRank) -> Option { self.preprocess_with_reason(raw, worker).ok() } @@ -117,6 +129,11 @@ impl ZmqEventNormalizer { event_id: u64, worker: WorkerWithDpRank, ) -> Result, ConvertError> { + let raw = if self.hash_mm { + raw + } else { + raw.without_mm_infos() + }; convert_event( raw, event_id, diff --git a/lib/kv-router/src/zmq_wire/tests.rs b/lib/kv-router/src/zmq_wire/tests.rs index 5e3213508832..b96256989a36 100644 --- a/lib/kv-router/src/zmq_wire/tests.rs +++ b/lib/kv-router/src/zmq_wire/tests.rs @@ -641,3 +641,69 @@ fn test_convert_event_short_token_ids_keeps_parsed_blocks() { other => panic!("expected Stored event, got {other:?}"), } } + +fn stored_image_block(tokens: Vec) -> RawKvEvent { + RawKvEvent::BlockStored { + block_hashes: vec![BlockHashValue::Unsigned(21)], + parent_block_hash: None, + block_size: tokens.len(), + token_ids: tokens, + medium: None, + lora_name: None, + block_mm_infos: Some(vec![Some(BlockExtraInfo { + mm_objects: vec![BlockMmObjectInfo { + mm_hash: 0x5083_86df_2042_9a4f, + offsets: vec![], + }], + })]), + is_eagle: None, + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + } +} + +fn stored_tokens_hash(event: Option) -> u64 { + match event.expect("stored event converts").event.data { + KvCacheEventData::Stored(data) => data.blocks[0].tokens_hash.0, + other => panic!("expected Stored, got {other:?}"), + } +} + +#[test] +fn test_plain_mm_hashing_matches_token_only_probe() { + let tokens: Vec = (0..4).map(|t| 129_264 + t).collect(); + let token_only = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + block_mm_infos: None, + lora_name: None, + is_eagle: None, + }, + )[0] + .0; + let worker = WorkerWithDpRank::new(3, 0); + + let mut plain = ZmqEventNormalizer::new(4).with_plain_mm_hashing(); + let plain_hash = stored_tokens_hash( + plain + .normalize(stored_image_block(tokens.clone()), 1, worker) + .unwrap(), + ); + assert_eq!( + plain_hash, token_only, + "plain mode must ignore the image hash" + ); + + let mut mm_aware = ZmqEventNormalizer::new(4); + let mm_hash = stored_tokens_hash( + mm_aware + .normalize(stored_image_block(tokens), 1, worker) + .unwrap(), + ); + assert_ne!( + mm_hash, token_only, + "default mode still mixes the image hash in" + ); +} diff --git a/lib/kv-router/src/zmq_wire/types.rs b/lib/kv-router/src/zmq_wire/types.rs index a0d7ff7a3975..f8a80b1be212 100644 --- a/lib/kv-router/src/zmq_wire/types.rs +++ b/lib/kv-router/src/zmq_wire/types.rs @@ -97,6 +97,37 @@ pub enum RawKvEvent { } impl RawKvEvent { + pub fn without_mm_infos(self) -> Self { + match self { + Self::BlockStored { + block_hashes, + parent_block_hash, + token_ids, + block_size, + medium, + lora_name, + block_mm_infos: _, + is_eagle, + group_idx, + kv_cache_spec_kind, + kv_cache_spec_sliding_window, + } => Self::BlockStored { + block_hashes, + parent_block_hash, + token_ids, + block_size, + medium, + lora_name, + block_mm_infos: None, + is_eagle, + group_idx, + kv_cache_spec_kind, + kv_cache_spec_sliding_window, + }, + other => other, + } + } + pub fn event_type_label(&self) -> &'static str { match self { Self::BlockStored { .. } => "stored",