From bfaa04b8169ba63bff5a38ce47a107a6b53537c9 Mon Sep 17 00:00:00 2001 From: Shang-Pin Date: Mon, 19 Oct 2026 20:43:21 +0000 Subject: [PATCH] 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]); +}