From 3f871a246c2e0c29655a625f00ea7f4d8d601dd0 Mon Sep 17 00:00:00 2001 From: ishandhanani <82981111+ishandhanani@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:34:49 -0700 Subject: [PATCH] feat(agent sessions): attribute stored KV cache blocks to sessions (#37482) Signed-off-by: Ishan Dhanani --- .../sgl-router/sgl-kv-indexer/src/bridge.rs | 232 ++++-- .../src/policies/kv_events/subscriber.rs | 5 +- .../sgl-router/src/policies/kv_events/wire.rs | 687 +++++++++++------- .../tests/component/policies/zmq_helpers.rs | 17 +- python/sglang/srt/disaggregation/kv_events.py | 40 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 1 + python/sglang/srt/mem_cache/events.py | 43 +- .../srt/mem_cache/rust_tree_core/adapter.py | 13 +- .../unified_cache/unified_tree_core.py | 21 +- .../srt/mem_cache/unified_radix_cache.py | 2 + rust/sglang-radix-tree/src/python_bindings.rs | 9 +- .../src/tests/components/base.rs | 1 + .../src/tests/components/full.rs | 2 + .../src/tests/components/mamba.rs | 1 + .../src/tests/components/swa.rs | 1 + .../src/tests/unified_tree_core.rs | 52 +- .../src/unified_tree_core.rs | 45 +- .../unit/disaggregation/test_kv_events.py | 67 +- .../mem_cache/test_page_interleave_shard.py | 1 + .../unit/mem_cache/test_radix_cache_unit.py | 20 +- .../test_rust_tree_core_integration.py | 10 +- .../test_unified_radix_cache_unittest.py | 86 ++- 22 files changed, 905 insertions(+), 451 deletions(-) diff --git a/experimental/sgl-router/sgl-kv-indexer/src/bridge.rs b/experimental/sgl-router/sgl-kv-indexer/src/bridge.rs index 644e6ce50..74e79c3b3 100644 --- a/experimental/sgl-router/sgl-kv-indexer/src/bridge.rs +++ b/experimental/sgl-router/sgl-kv-indexer/src/bridge.rs @@ -591,53 +591,76 @@ fn decode_event_batch_impl( } fn decode_event(event: &Value, actions: &mut EventActions) -> Result<(), BridgeError> { - let event = expect_array(event, "KV event")?; - let event_type = expect_str( - event - .first() - .ok_or_else(|| BridgeError::Decode("KV event is empty".to_string()))?, - "KV event tag", - )?; + // `KVCacheEvent` is a tagged map (`tag=True` without `array_like`). + match event { + Value::Map(entries) => decode_event_map(entries, actions), + _ => Err(BridgeError::Decode("KV event must be a map".to_string())), + } +} + +fn map_field<'a>(entries: &'a [(Value, Value)], name: &str) -> Option<&'a Value> { + entries + .iter() + .find(|(key, _)| key.as_str() == Some(name)) + .map(|(_, value)| value) +} + +fn required_map_field<'a>( + entries: &'a [(Value, Value)], + name: &str, +) -> Result<&'a Value, BridgeError> { + map_field(entries, name) + .ok_or_else(|| BridgeError::Decode(format!("KV event is missing `{name}`"))) +} + +/// Decodes one tagged-map event: `{"type": ..., "": ...}`. Optional +/// fields may be absent (msgspec omits `None` defaults) or present as nil. +/// Keys the indexer does not use (`token_ids`, `lora_id`, `cache_salt`, +/// `session_id`) are ignored. +fn decode_event_map( + entries: &[(Value, Value)], + actions: &mut EventActions, +) -> Result<(), BridgeError> { + let event_type = expect_str(required_map_field(entries, "type")?, "KV event type")?; + let medium = |field: &str| -> Result, BridgeError> { + match map_field(entries, "medium") { + Some(value) => expect_optional_str(value, field), + None => Ok(None), + } + }; match event_type { "BlockStored" => { - // At least 7 fields (the legacy schema); an 8th `component_types` - // slot appears with `--enable-kv-events-component-types`. Both - // shapes are accepted. - if event.len() < 7 { - return Err(BridgeError::Decode( - "BlockStored must have at least 7 array fields".to_string(), - )); - } - let tier = medium_to_tier(expect_optional_str(&event[6], "BlockStored.medium")?)?; - // `component_types` is the trailing slot: a list of component labels - // folded into a bitmask, or nil/absent for a legacy whole-block store. - let mask = match event.get(7) { + let tier = medium_to_tier(medium("BlockStored.medium")?)?; + let mask = match map_field(entries, "component_types") { Some(value) => decode_component_mask(value)?, None => None, }; - // The token count is only carried alongside component-aware stores, - // where the query path needs it to accumulate trailing windows. let block_size = match mask { - Some(_) => Some(decode_block_size(&event[4])?), + Some(_) => Some(decode_block_size(required_map_field( + entries, + "block_size", + )?)?), + None => None, + }; + let parent_block_hash = match map_field(entries, "parent_block_hash") { + Some(value) => decode_optional_hash(value, "BlockStored.parent_block_hash")?, None => None, }; actions.report( tier, - decode_optional_hash(&event[2], "BlockStored.parent_block_hash")?, - decode_hashes(&event[1])?, + parent_block_hash, + decode_hashes(required_map_field(entries, "block_hashes")?)?, mask, block_size, ); } "BlockRemoved" => { - if event.len() < 3 { - return Err(BridgeError::Decode( - "BlockRemoved must have 3 array fields".to_string(), - )); - } - let tier = medium_to_tier(expect_optional_str(&event[2], "BlockRemoved.medium")?)?; - actions.revoke(tier, decode_hashes(&event[1])?); + let tier = medium_to_tier(medium("BlockRemoved.medium")?)?; + actions.revoke( + tier, + decode_hashes(required_map_field(entries, "block_hashes")?)?, + ); } "AllBlocksCleared" => { actions.clear_all(); @@ -872,19 +895,31 @@ mod tests { } fn stored_with_parent(hashes: &[i64], parent: Option, medium: &str) -> Value { - Value::Array(vec![ - Value::String("BlockStored".into()), - ints(hashes), - parent.map_or(Value::Nil, Value::from), - ints(&[1]), // token_ids - Value::from(1_i64), // block_size - Value::Nil, // lora_id - Value::String(medium.into()), - ]) + stored_with_extra(hashes, parent, medium, vec![]) } - /// A component-aware `BlockStored` (8-element schema): trailing - /// `component_types` slot plus a concrete `block_size` token count. + /// A `BlockStored` map plus `extra` keys the bridge must ignore. + fn stored_with_extra( + hashes: &[i64], + parent: Option, + medium: &str, + extra: Vec<(&str, Value)>, + ) -> Value { + let mut entries = vec![ + ("type", Value::String("BlockStored".into())), + ("block_hashes", ints(hashes)), + ("parent_block_hash", parent.map_or(Value::Nil, Value::from)), + ("token_ids", ints(&[1])), + ("block_size", Value::from(1_i64)), + ("lora_id", Value::Nil), + ("medium", Value::String(medium.into())), + ]; + entries.extend(extra); + map_event(entries) + } + + /// A component-aware `BlockStored`: a `component_types` key plus a + /// concrete `block_size` token count. fn stored_c(hashes: &[i64], medium: &str, block_size: i64, components: Value) -> Value { stored_c_with_parent(hashes, None, medium, block_size, components) } @@ -896,15 +931,16 @@ mod tests { block_size: i64, components: Value, ) -> Value { - Value::Array(vec![ - Value::String("BlockStored".into()), - ints(hashes), - parent.map_or(Value::Nil, Value::from), - ints(&[1]), // token_ids - Value::from(block_size), - Value::Nil, // lora_id - Value::String(medium.into()), - components, // component_types (Nil or array of strings) + map_event(vec![ + ("type", Value::String("BlockStored".into())), + ("block_hashes", ints(hashes)), + ("parent_block_hash", parent.map_or(Value::Nil, Value::from)), + ("token_ids", ints(&[1])), + ("block_size", Value::from(block_size)), + ("lora_id", Value::Nil), + ("medium", Value::String(medium.into())), + // component_types (Nil or array of strings) + ("component_types", components), ]) } @@ -935,15 +971,73 @@ mod tests { } fn removed(hashes: &[i64], medium: &str) -> Value { - Value::Array(vec![ - Value::String("BlockRemoved".into()), - ints(hashes), - Value::String(medium.into()), + map_event(vec![ + ("type", Value::String("BlockRemoved".into())), + ("block_hashes", ints(hashes)), + ("medium", Value::String(medium.into())), ]) } fn cleared() -> Value { - Value::Array(vec![Value::String("AllBlocksCleared".into())]) + map_event(vec![("type", Value::String("AllBlocksCleared".into()))]) + } + + /// A tagged event map: `{"type": ..., field: value, ...}`. + fn map_event(entries: Vec<(&str, Value)>) -> Value { + Value::Map( + entries + .into_iter() + .map(|(key, value)| (Value::String(key.into()), value)) + .collect(), + ) + } + + fn tagged(event_type: &str) -> Value { + map_event(vec![("type", Value::String(event_type.into()))]) + } + + #[test] + fn attribution_keys_are_ignored() { + assert_eq!( + actions_of(vec![stored_with_extra( + &[2, 3], + Some(1), + "GPU", + vec![ + ("cache_salt", Value::String("tenant-a".into())), + ("session_id", Value::String("session-a".into())), + ], + )]), + vec![rep_with_parent(hbm(), Some(1), &["2", "3"])] + ); + } + + #[test] + fn undecodable_events_are_skipped_and_siblings_survive() { + assert_eq!( + actions_of(vec![ + // no `type` + map_event(vec![("block_hashes", ints(&[1]))]), + stored(&[1], "GPU"), + // no `medium` + map_event(vec![ + ("type", Value::String("BlockStored".into())), + ("block_hashes", ints(&[9])), + ]), + // pre-map publishers encoded events as tagged arrays + Value::Array(vec![ + Value::String("BlockStored".into()), + ints(&[7]), + Value::Nil, + ints(&[1]), + Value::from(1_i64), + Value::Nil, + Value::String("GPU".into()), + ]), + removed(&[5], "GPU"), + ]), + vec![rep(hbm(), &["1"]), rev(hbm(), &["5"])] + ); } /// Wrap events in a 3-element batch [ts, events, attn_dp_rank]. @@ -1151,7 +1245,7 @@ mod tests { #[test] fn batch_with_only_ignored_events_has_no_actions() { let config = test_config(vec![hbm()]); - let events = vec![Value::Array(vec![Value::String("BlockUpdated".into())])]; + let events = vec![tagged("BlockUpdated")]; assert!(request_of(&config, 0, events).actions.is_empty()); } @@ -1274,7 +1368,7 @@ mod tests { #[test] fn unknown_event_tag_is_ignored() { - let events = vec![Value::Array(vec![Value::String("BlockUpdated".into())])]; + let events = vec![tagged("BlockUpdated")]; assert!(actions_of(events).is_empty()); } @@ -1294,12 +1388,7 @@ mod tests { fn python_msgspec_mixed_batch_golden_decodes() { // Generated by msgspec.msgpack.Encoder from the authoritative Python // KVEventBatch schema in sglang.srt.disaggregation.kv_events. - let payload = golden_bytes(concat!( - "93cb405edd2f1a9fbe779397ab426c6f636b53746f72656492", - "cf0000011f71fb04cbd2c521974f2a940a141e280407a3475055", - "93ac426c6f636b52656d6f7665649264ccc8a44449534b", - "91b0416c6c426c6f636b73436c656172656402" - )); + let payload = golden_bytes("93cb405edd2f1a9fbe779387a474797065ab426c6f636b53746f726564ac626c6f636b5f68617368657392cf0000011f71fb04cbd2c521974fb1706172656e745f626c6f636b5f686173682aa9746f6b656e5f696473940a141e28aa626c6f636b5f73697a6504a76c6f72615f696407a66d656469756da347505583a474797065ac426c6f636b52656d6f766564ac626c6f636b5f6861736865739264ccc8a66d656469756da44449534b81a474797065b0416c6c426c6f636b73436c656172656402"); assert_eq!( decode_event_batch(&payload).unwrap().actions, vec![ @@ -1314,10 +1403,7 @@ mod tests { fn python_msgspec_bigram_tokens_golden_decodes() { // token_ids contains Python tuples as nested msgpack arrays; the // bridge ignores payload shape and indexes the published hashes. - let payload = golden_bytes(concat!( - "93cb3ff80000000000009197ab426c6f636b53746f726564916f", - "c092920a1492141e02c0a347505503" - )); + let payload = golden_bytes("93cb3ff80000000000009187a474797065ab426c6f636b53746f726564ac626c6f636b5f686173686573916fb1706172656e745f626c6f636b5f68617368c0a9746f6b656e5f69647392920a1492141eaa626c6f636b5f73697a6502a76c6f72615f6964c0a66d656469756da347505503"); assert_eq!( decode_event_batch(&payload).unwrap().actions, vec![rep(hbm(), &["111"])] @@ -1326,12 +1412,10 @@ mod tests { #[test] fn python_msgspec_nil_medium_golden_is_safely_skipped() { - // The Python schema permits medium=None; such events map to no - // Indexer tier, so they are isolated rather than given a placement. - let payload = golden_bytes(concat!( - "93cb00000000000000009297ab426c6f636b53746f7265649101", - "c092050602c0c093ac426c6f636b52656d6f7665649102c0c0" - )); + // The Python schema permits medium=None (the key is then omitted); + // such events map to no Indexer tier, so they are isolated rather + // than given a placement. + let payload = golden_bytes("93cb00000000000000009286a474797065ab426c6f636b53746f726564ac626c6f636b5f6861736865739101b1706172656e745f626c6f636b5f68617368c0a9746f6b656e5f696473920506aa626c6f636b5f73697a6502a76c6f72615f6964c082a474797065ac426c6f636b52656d6f766564ac626c6f636b5f6861736865739102c0"); assert!(decode_event_batch(&payload).unwrap().actions.is_empty()); } diff --git a/experimental/sgl-router/src/policies/kv_events/subscriber.rs b/experimental/sgl-router/src/policies/kv_events/subscriber.rs index 012ddf60c..7e18382c0 100644 --- a/experimental/sgl-router/src/policies/kv_events/subscriber.rs +++ b/experimental/sgl-router/src/policies/kv_events/subscriber.rs @@ -711,8 +711,9 @@ mod tests { mp::write_f64(&mut buf, ts).unwrap(); // events array length 1 mp::write_array_len(&mut buf, 1).unwrap(); - // event = ["AllBlocksCleared"] - mp::write_array_len(&mut buf, 1).unwrap(); + // Events use msgspec's tagged-map encoding: {"type": "AllBlocksCleared"}. + mp::write_map_len(&mut buf, 1).unwrap(); + mp::write_str(&mut buf, "type").unwrap(); mp::write_str(&mut buf, "AllBlocksCleared").unwrap(); match attn_dp_rank { Some(v) => { diff --git a/experimental/sgl-router/src/policies/kv_events/wire.rs b/experimental/sgl-router/src/policies/kv_events/wire.rs index be8bba6b6..134fd47cd 100644 --- a/experimental/sgl-router/src/policies/kv_events/wire.rs +++ b/experimental/sgl-router/src/policies/kv_events/wire.rs @@ -5,25 +5,23 @@ //! `msgspec.msgpack`. Two struct families are involved: //! //! * `EventBatch` (the outer payload) — declared with -//! `array_like=True, gc=False` (no tag). -//! * `KVCacheEvent` (each inner event variant) — additionally declared -//! with `tag=True`. +//! `array_like=True, gc=False` (no tag), so it is a msgpack **array** +//! `[ts, events, attn_dp_rank]`. +//! * `KVCacheEvent` (each inner event variant) — declared with +//! `omit_defaults=True, tag=True` and no `array_like`, so each event is a +//! msgpack **map** whose `type` key carries the class name and whose other +//! keys are field names. Optional fields left at `None` are omitted. This +//! is the same encoding vLLM uses for its KV events. //! -//! The combined effect on the wire: -//! -//! * Each struct is a msgpack **array** of its fields in declaration -//! order, not a map. -//! * `tag=True` on `KVCacheEvent` prepends a class-name string at index 0 -//! of each inner event array, so an event is -//! `[class_name_str, field1, field2, ...]`. The outer `EventBatch` -//! array does **not** carry a tag prefix. +//! Older publishers emitted each event as a tagged **array** +//! `[class_name_str, field1, ...]`; that shape is rejected. //! //! This module deserializes those bytes into Rust types and exposes a single //! [`decode_event_batch`] entry point. use std::fmt; -use serde::de::{self, Deserializer, IgnoredAny, SeqAccess, Visitor}; +use serde::de::{self, Deserializer, IgnoredAny, MapAccess, SeqAccess, Visitor}; use serde::Deserialize; /// Top-level batch payload published by SGLang. @@ -46,16 +44,16 @@ pub struct KvEventBatch { } /// A single KV cache event. The Python base class `KVCacheEvent` uses -/// `tag=True`, so each event on the wire is an array whose first element -/// is the class-name discriminator. +/// `tag=True`, so each event carries its class name under the `type` key. #[derive(Debug, Clone, PartialEq)] pub enum KvCacheEvent { - /// `["BlockStored", block_hashes, parent_block_hash, token_ids, - /// block_size, lora_id, medium?]`. + /// `{"type": "BlockStored", "block_hashes", "parent_block_hash", + /// "token_ids", "block_size", "lora_id", "medium"?, ...}`. Keys the + /// gateway does not route on (`cache_salt`, `session_id`) are ignored. BlockStored(BlockStored), - /// `["BlockRemoved", block_hashes, medium?]`. + /// `{"type": "BlockRemoved", "block_hashes", "medium"?}`. BlockRemoved(BlockRemoved), - /// `["AllBlocksCleared"]`. + /// `{"type": "AllBlocksCleared"}`. AllBlocksCleared, } @@ -329,9 +327,8 @@ impl<'de> Deserialize<'de> for BoundedU32Vec { } // --------------------------------------------------------------------------- -// Custom Deserialize impls — msgspec encodes these structs as msgpack arrays -// (not maps). The visitors also accept absent trailing optional fields for -// compatibility. +// Custom Deserialize impls — the batch is a msgpack array; each event is a +// tagged map. Optional fields may be absent or nil. // --------------------------------------------------------------------------- impl<'de> Deserialize<'de> for KvEventBatch { @@ -381,65 +378,104 @@ impl<'de> Deserialize<'de> for KvCacheEvent { where D: Deserializer<'de>, { + /// Map keys of a `KVCacheEvent` the gateway reads. Every other key + /// (`cache_salt`, `session_id`, future additions) is skipped. + enum EventField { + Type, + BlockHashes, + ParentBlockHash, + TokenIds, + BlockSize, + LoraId, + Medium, + Other, + } + + impl<'de> Deserialize<'de> for EventField { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct FieldVisitor; + impl<'de> Visitor<'de> for FieldVisitor { + type Value = EventField; + fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str("a KV event field name") + } + fn visit_str(self, v: &str) -> Result { + Ok(match v { + "type" => EventField::Type, + "block_hashes" => EventField::BlockHashes, + "parent_block_hash" => EventField::ParentBlockHash, + "token_ids" => EventField::TokenIds, + "block_size" => EventField::BlockSize, + "lora_id" => EventField::LoraId, + "medium" => EventField::Medium, + _ => EventField::Other, + }) + } + } + deserializer.deserialize_identifier(FieldVisitor) + } + } + struct EventVisitor; impl<'de> Visitor<'de> for EventVisitor { type Value = KvCacheEvent; fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { - f.write_str("a tagged msgpack array [class_name, ...fields]") + f.write_str("a tagged msgpack map {\"type\": class_name, ...fields}") } - fn visit_seq(self, mut seq: A) -> Result + fn visit_map(self, mut map: A) -> Result where - A: SeqAccess<'de>, + A: MapAccess<'de>, { - let tag: String = seq - .next_element()? - .ok_or_else(|| de::Error::missing_field("event tag"))?; - + let mut tag: Option = None; + let mut block_hashes: Option = None; + let mut parent_block_hash: Option = None; + let mut token_ids: Option = None; + let mut block_size: Option = None; + let mut lora_id: Option = None; + let mut medium: Option = None; + while let Some(field) = map.next_key::()? { + match field { + EventField::Type => tag = Some(map.next_value()?), + EventField::BlockHashes => block_hashes = Some(map.next_value()?), + // Optional fields may be present as nil or omitted entirely. + EventField::ParentBlockHash => parent_block_hash = map.next_value()?, + EventField::TokenIds => token_ids = Some(map.next_value()?), + EventField::BlockSize => block_size = Some(map.next_value()?), + EventField::LoraId => lora_id = map.next_value()?, + EventField::Medium => medium = map.next_value()?, + EventField::Other => { + map.next_value::()?; + } + } + } + let tag = tag.ok_or_else(|| de::Error::missing_field("type"))?; match tag.as_str() { - "BlockStored" => { - let block_hashes: BoundedI64Vec = seq - .next_element()? - .ok_or_else(|| de::Error::missing_field("block_hashes"))?; - let parent_block_hash: Option = seq.next_element()?.unwrap_or(None); - let token_ids: BoundedU32Vec = seq - .next_element()? - .ok_or_else(|| de::Error::missing_field("token_ids"))?; - let block_size: u32 = seq - .next_element()? - .ok_or_else(|| de::Error::missing_field("block_size"))?; - // `lora_id` is `Optional[int]` with no default — it's - // always emitted, but as nil when absent. - let lora_id: Option = seq.next_element()?.unwrap_or(None); - // `medium` defaults to None and may be omitted. - let medium: Option = seq.next_element()?.unwrap_or(None); - while seq.next_element::()?.is_some() {} - Ok(KvCacheEvent::BlockStored(BlockStored { - block_hashes: block_hashes.0, - parent_block_hash, - token_ids: token_ids.0, - block_size, - lora_id, - medium, - })) - } - "BlockRemoved" => { - let block_hashes: BoundedI64Vec = seq - .next_element()? - .ok_or_else(|| de::Error::missing_field("block_hashes"))?; - let medium: Option = seq.next_element()?.unwrap_or(None); - while seq.next_element::()?.is_some() {} - Ok(KvCacheEvent::BlockRemoved(BlockRemoved { - block_hashes: block_hashes.0, - medium, - })) - } - "AllBlocksCleared" => { - while seq.next_element::()?.is_some() {} - Ok(KvCacheEvent::AllBlocksCleared) - } + "BlockStored" => Ok(KvCacheEvent::BlockStored(BlockStored { + block_hashes: block_hashes + .ok_or_else(|| de::Error::missing_field("block_hashes"))? + .0, + parent_block_hash, + token_ids: token_ids + .ok_or_else(|| de::Error::missing_field("token_ids"))? + .0, + block_size: block_size + .ok_or_else(|| de::Error::missing_field("block_size"))?, + lora_id, + medium, + })), + "BlockRemoved" => Ok(KvCacheEvent::BlockRemoved(BlockRemoved { + block_hashes: block_hashes + .ok_or_else(|| de::Error::missing_field("block_hashes"))? + .0, + medium, + })), + "AllBlocksCleared" => Ok(KvCacheEvent::AllBlocksCleared), other => Err(de::Error::unknown_variant( other, &["BlockStored", "BlockRemoved", "AllBlocksCleared"], @@ -448,13 +484,13 @@ impl<'de> Deserialize<'de> for KvCacheEvent { } } - deserializer.deserialize_seq(EventVisitor) + deserializer.deserialize_map(EventVisitor) } } // --------------------------------------------------------------------------- // Tests — golden bytes are constructed via the `rmp` low-level encoder so -// they exercise the exact msgpack array layout SGLang emits, independent of +// they exercise the exact msgpack map layout SGLang emits, independent of // any Rust-side serializer. // --------------------------------------------------------------------------- @@ -464,11 +500,17 @@ mod tests { use rmp::encode as mp; - /// Encode a tagged event header `[tag, ...]` array of `total_len` - /// elements (tag included). - fn write_event_array(buf: &mut Vec, tag: &str, total_len: u32) { - mp::write_array_len(buf, total_len).unwrap(); - mp::write_str(buf, tag).unwrap(); + fn write_key(buf: &mut Vec, key: &str) { + mp::write_str(buf, key).unwrap(); + } + + fn write_opt_sint(buf: &mut Vec, value: Option) { + match value { + Some(v) => { + mp::write_sint(buf, v).unwrap(); + } + None => mp::write_nil(buf).unwrap(), + } } fn write_i64_array(buf: &mut Vec, values: &[i64]) { @@ -497,6 +539,101 @@ mod tests { } } + /// Start a tagged event map with `field_count` fields after `type`. + fn write_event_map(buf: &mut Vec, tag: &str, field_count: u32) { + mp::write_map_len(buf, field_count + 1).unwrap(); + write_key(buf, "type"); + mp::write_str(buf, tag).unwrap(); + } + + fn block_stored_field_count(medium: Option<&str>, extra: &[(&str, &str)]) -> u32 { + 5 + u32::from(medium.is_some()) + extra.len() as u32 + } + + /// Write the `BlockStored` fields that follow `type`. `parent_block_hash` + /// and `lora_id` have no default in the Python schema, so msgspec always + /// emits them (nil when unset); `medium` is omitted when `None`; `extra` + /// adds string-valued keys the gateway must ignore. + #[allow(clippy::too_many_arguments)] + fn write_block_stored_fields( + buf: &mut Vec, + block_hashes: &[i64], + parent: Option, + write_tokens: impl FnOnce(&mut Vec), + block_size: u32, + lora_id: Option, + medium: Option<&str>, + extra: &[(&str, &str)], + ) { + write_key(buf, "block_hashes"); + write_i64_array(buf, block_hashes); + write_key(buf, "parent_block_hash"); + write_opt_sint(buf, parent); + write_key(buf, "token_ids"); + write_tokens(buf); + write_key(buf, "block_size"); + mp::write_uint(buf, block_size as u64).unwrap(); + write_key(buf, "lora_id"); + write_opt_sint(buf, lora_id); + if let Some(m) = medium { + write_key(buf, "medium"); + mp::write_str(buf, m).unwrap(); + } + for (key, value) in extra { + write_key(buf, key); + mp::write_str(buf, value).unwrap(); + } + } + + /// Build a `BlockStored` map as msgspec emits it, plus `extra` keys. + #[allow(clippy::too_many_arguments)] + fn build_block_stored_bytes_with_extra( + block_hashes: &[i64], + parent: Option, + token_ids: &[u32], + block_size: u32, + lora_id: Option, + medium: Option<&str>, + extra: &[(&str, &str)], + ) -> Vec { + let mut buf = Vec::new(); + write_event_map( + &mut buf, + "BlockStored", + block_stored_field_count(medium, extra), + ); + write_block_stored_fields( + &mut buf, + block_hashes, + parent, + |b| write_u32_array(b, token_ids), + block_size, + lora_id, + medium, + extra, + ); + buf + } + + fn build_block_stored_bytes( + block_hashes: &[i64], + parent: Option, + token_ids: &[u32], + block_size: u32, + lora_id: Option, + medium: Option<&str>, + ) -> Vec { + build_block_stored_bytes_with_extra( + block_hashes, + parent, + token_ids, + block_size, + lora_id, + medium, + &[], + ) + } + /// Like `build_block_stored_bytes`, but `token_ids` is the bigram /// list-of-pairs shape that DeepSeek-V4-class models emit. fn build_block_stored_bigram_bytes( @@ -508,109 +645,39 @@ mod tests { medium: Option<&str>, ) -> Vec { let mut buf = Vec::new(); - write_event_array(&mut buf, "BlockStored", 7); - write_i64_array(&mut buf, block_hashes); - match parent { - Some(v) => { - mp::write_sint(&mut buf, v).unwrap(); - } - None => mp::write_nil(&mut buf).unwrap(), - } - write_bigram_token_array(&mut buf, token_pairs); - mp::write_uint(&mut buf, block_size as u64).unwrap(); - match lora_id { - Some(v) => { - mp::write_sint(&mut buf, v).unwrap(); - } - None => mp::write_nil(&mut buf).unwrap(), - } - match medium { - Some(s) => mp::write_str(&mut buf, s).unwrap(), - None => mp::write_nil(&mut buf).unwrap(), - } - buf - } - - /// Regression: bigram models (e.g. DeepSeek-V4-Flash) emit `token_ids` as - /// `[[t_i, t_{i+1}], ...]`. The decoder previously read `token_ids` as a - /// flat `u32` array and failed the entire batch with - /// "wrong msgpack marker FixArray(2)", silently disabling cache-aware - /// routing. It must instead accept the bigram shape (flattening the ints). - #[test] - fn decodes_block_stored_with_bigram_token_ids() { - let event = build_block_stored_bigram_bytes( - &[111_i64], - None, - &[(10, 20), (20, 30)], - 2, - None, - Some("GPU"), + write_event_map( + &mut buf, + "BlockStored", + block_stored_field_count(medium, &[]), + ); + write_block_stored_fields( + &mut buf, + block_hashes, + parent, + |b| write_bigram_token_array(b, token_pairs), + block_size, + lora_id, + medium, + &[], ); - let bytes = build_batch_bytes(1.5, &[event], Some(0), true); - - let batch = decode_event_batch(&bytes).expect("decode bigram token_ids"); - assert_eq!(batch.events.len(), 1); - match &batch.events[0] { - KvCacheEvent::BlockStored(b) => { - // routing-relevant fields decode unchanged - assert_eq!(b.block_hashes, vec![111]); - assert_eq!(b.parent_block_hash, None); - assert_eq!(b.block_size, 2); - // bigram pairs are flattened into the (informational) token vec - assert_eq!(b.token_ids, vec![10, 20, 20, 30]); - } - other => panic!("expected BlockStored, got {other:?}"), - } - } - - /// Build a full BlockStored event as msgspec would emit it (all 7 - /// elements: tag + 6 fields). `medium` may be Some/None. - fn build_block_stored_bytes( - block_hashes: &[i64], - parent: Option, - token_ids: &[u32], - block_size: u32, - lora_id: Option, - medium: Option<&str>, - ) -> Vec { - let mut buf = Vec::new(); - write_event_array(&mut buf, "BlockStored", 7); - write_i64_array(&mut buf, block_hashes); - match parent { - Some(v) => { - mp::write_sint(&mut buf, v).unwrap(); - } - None => mp::write_nil(&mut buf).unwrap(), - } - write_u32_array(&mut buf, token_ids); - mp::write_uint(&mut buf, block_size as u64).unwrap(); - match lora_id { - Some(v) => { - mp::write_sint(&mut buf, v).unwrap(); - } - None => mp::write_nil(&mut buf).unwrap(), - } - match medium { - Some(s) => mp::write_str(&mut buf, s).unwrap(), - None => mp::write_nil(&mut buf).unwrap(), - } buf } fn build_block_removed_bytes(block_hashes: &[i64], medium: Option<&str>) -> Vec { let mut buf = Vec::new(); - write_event_array(&mut buf, "BlockRemoved", 3); + write_event_map(&mut buf, "BlockRemoved", 1 + u32::from(medium.is_some())); + write_key(&mut buf, "block_hashes"); write_i64_array(&mut buf, block_hashes); - match medium { - Some(s) => mp::write_str(&mut buf, s).unwrap(), - None => mp::write_nil(&mut buf).unwrap(), + if let Some(m) = medium { + write_key(&mut buf, "medium"); + mp::write_str(&mut buf, m).unwrap(); } buf } fn build_all_blocks_cleared_bytes() -> Vec { let mut buf = Vec::new(); - write_event_array(&mut buf, "AllBlocksCleared", 1); + write_event_map(&mut buf, "AllBlocksCleared", 0); buf } @@ -641,6 +708,36 @@ mod tests { buf } + /// Regression: bigram models (e.g. DeepSeek-V4-Flash) emit `token_ids` as + /// `[[t_i, t_{i+1}], ...]`. The decoder previously read `token_ids` as a + /// flat `u32` array and failed the entire batch with + /// "wrong msgpack marker FixArray(2)", silently disabling cache-aware + /// routing. It must instead accept the bigram shape (flattening the ints). + #[test] + fn decodes_block_stored_with_bigram_token_ids() { + let event = build_block_stored_bigram_bytes( + &[111_i64], + None, + &[(10, 20), (20, 30)], + 2, + None, + Some("GPU"), + ); + let bytes = build_batch_bytes(1.5, &[event], Some(0), true); + + let batch = decode_event_batch(&bytes).expect("decode bigram token_ids"); + assert_eq!(batch.events.len(), 1); + match &batch.events[0] { + KvCacheEvent::BlockStored(b) => { + assert_eq!(b.block_hashes, vec![111]); + assert_eq!(b.parent_block_hash, None); + assert_eq!(b.block_size, 2); + assert_eq!(b.token_ids, vec![10, 20, 20, 30]); + } + other => panic!("expected BlockStored, got {other:?}"), + } + } + #[test] fn decodes_block_stored_with_all_fields() { let event = build_block_stored_bytes( @@ -666,12 +763,14 @@ mod tests { assert_eq!(b.lora_id, Some(7)); assert_eq!(b.medium.as_deref(), Some("GPU")); } - other => panic!("expected BlockStored, got {:?}", other), + other => panic!("expected BlockStored, got {other:?}"), } } + /// `parent_block_hash` and `lora_id` present as nil, `medium` key omitted + /// (msgspec `omit_defaults`). #[test] - fn decodes_block_stored_with_nil_optionals() { + fn decodes_block_stored_with_nil_and_omitted_optionals() { let event = build_block_stored_bytes(&[1, 2, 3], None, &[5, 6], 16, None, None); let bytes = build_batch_bytes(0.0, &[event], None, true); @@ -683,7 +782,37 @@ mod tests { assert_eq!(b.medium, None); assert_eq!(b.block_size, 16); } - other => panic!("unexpected variant: {:?}", other), + other => panic!("unexpected variant: {other:?}"), + } + } + + /// Attribution keys (`cache_salt`, `session_id`) and any future key the + /// gateway does not route on are skipped, not rejected. + #[test] + fn unknown_keys_are_ignored() { + let event = build_block_stored_bytes_with_extra( + &[10], + Some(1), + &[1, 2], + 2, + None, + Some("GPU"), + &[ + ("cache_salt", "tenant-a"), + ("session_id", "session-a"), + ("future_key", "x"), + ], + ); + let bytes = build_batch_bytes(0.0, &[event], None, true); + + let batch = decode_event_batch(&bytes).expect("decode"); + match &batch.events[0] { + KvCacheEvent::BlockStored(b) => { + assert_eq!(b.block_hashes, vec![10]); + assert_eq!(b.parent_block_hash, Some(1)); + assert_eq!(b.medium.as_deref(), Some("GPU")); + } + other => panic!("unexpected variant: {other:?}"), } } @@ -698,7 +827,22 @@ mod tests { assert_eq!(r.block_hashes, vec![100, 200]); assert_eq!(r.medium.as_deref(), Some("DISK")); } - other => panic!("unexpected variant: {:?}", other), + other => panic!("unexpected variant: {other:?}"), + } + } + + #[test] + fn medium_omitted_in_block_removed_decodes_as_none() { + let event = build_block_removed_bytes(&[42], None); + let bytes = build_batch_bytes(0.0, &[event], None, true); + + let batch = decode_event_batch(&bytes).expect("decode"); + match &batch.events[0] { + KvCacheEvent::BlockRemoved(r) => { + assert_eq!(r.block_hashes, vec![42]); + assert_eq!(r.medium, None); + } + other => panic!("unexpected variant: {other:?}"), } } @@ -738,48 +882,10 @@ mod tests { assert_eq!(batch.events.len(), 1); } - #[test] - fn medium_omitted_in_block_stored_decodes_as_none() { - // BlockStored with `medium` omitted entirely (omit_defaults can drop - // the trailing default-None field). 6 elements instead of 7. - let mut buf = Vec::new(); - write_event_array(&mut buf, "BlockStored", 6); - write_i64_array(&mut buf, &[1]); - mp::write_nil(&mut buf).unwrap(); // parent_block_hash - write_u32_array(&mut buf, &[1, 2]); - mp::write_uint(&mut buf, 2).unwrap(); // block_size - mp::write_nil(&mut buf).unwrap(); // lora_id - let bytes = build_batch_bytes(0.0, &[buf], None, true); - - let batch = decode_event_batch(&bytes).expect("decode"); - match &batch.events[0] { - KvCacheEvent::BlockStored(b) => assert_eq!(b.medium, None), - other => panic!("unexpected variant: {:?}", other), - } - } - - #[test] - fn medium_omitted_in_block_removed_decodes_as_none() { - // BlockRemoved with only [tag, block_hashes] (medium omitted). - let mut buf = Vec::new(); - write_event_array(&mut buf, "BlockRemoved", 2); - write_i64_array(&mut buf, &[42]); - let bytes = build_batch_bytes(0.0, &[buf], None, true); - - let batch = decode_event_batch(&bytes).expect("decode"); - match &batch.events[0] { - KvCacheEvent::BlockRemoved(r) => { - assert_eq!(r.block_hashes, vec![42]); - assert_eq!(r.medium, None); - } - other => panic!("unexpected variant: {:?}", other), - } - } - #[test] fn unknown_event_tag_is_rejected() { let mut buf = Vec::new(); - write_event_array(&mut buf, "MysteryEvent", 1); + write_event_map(&mut buf, "MysteryEvent", 0); let bytes = build_batch_bytes(0.0, &[buf], None, true); let err = decode_event_batch(&bytes).expect_err("should reject unknown variant"); @@ -790,12 +896,60 @@ mod tests { ); } - /// Golden bytes captured from the actual SGLang Python publisher - /// (`msgspec.msgpack.Encoder().encode(KVEventBatch(...))`). These - /// hex strings are produced by msgspec 0.21.1 against the schema in - /// `python/sglang/srt/disaggregation/kv_events.py` and lock down the - /// exact wire format the decoder is expected to consume. Regenerated - /// with `python -c '...msgspec.msgpack.Encoder().encode(...)'`. + #[test] + fn event_without_type_is_rejected() { + let mut buf = Vec::new(); + mp::write_map_len(&mut buf, 1).unwrap(); + write_key(&mut buf, "block_hashes"); + write_i64_array(&mut buf, &[1]); + let bytes = build_batch_bytes(0.0, &[buf], None, true); + + let err = decode_event_batch(&bytes).expect_err("missing type must fail"); + assert!(format!("{err}").contains("type"), "unexpected error: {err}"); + } + + #[test] + fn block_stored_without_block_hashes_is_rejected() { + let mut buf = Vec::new(); + write_event_map(&mut buf, "BlockStored", 2); + write_key(&mut buf, "token_ids"); + write_u32_array(&mut buf, &[1]); + write_key(&mut buf, "block_size"); + mp::write_uint(&mut buf, 1).unwrap(); + let bytes = build_batch_bytes(0.0, &[buf], None, true); + + let err = decode_event_batch(&bytes).expect_err("missing block_hashes must fail"); + assert!( + format!("{err}").contains("block_hashes"), + "unexpected error: {err}" + ); + } + + /// The pre-map publishers encoded events as tagged arrays; that shape is + /// no longer produced and is rejected rather than half-decoded. + #[test] + fn legacy_array_event_is_rejected() { + let mut buf = Vec::new(); + mp::write_array_len(&mut buf, 7).unwrap(); + mp::write_str(&mut buf, "BlockStored").unwrap(); + write_i64_array(&mut buf, &[1]); + mp::write_nil(&mut buf).unwrap(); + write_u32_array(&mut buf, &[1, 2]); + mp::write_uint(&mut buf, 2).unwrap(); + mp::write_nil(&mut buf).unwrap(); + mp::write_str(&mut buf, "GPU").unwrap(); + let bytes = build_batch_bytes(0.0, &[buf], None, true); + + assert!(matches!( + decode_event_batch(&bytes), + Err(DecodeError::Msgpack(_)) + )); + } + + /// Golden bytes captured from the SGLang Python publisher + /// (`msgspec.msgpack.Encoder().encode(KVEventBatch(...))`), msgspec 0.21.1, + /// against the schema in `python/sglang/srt/disaggregation/kv_events.py`. + /// They lock down the exact wire format the decoder consumes. mod msgspec_golden { use super::super::*; @@ -811,9 +965,7 @@ mod tests { // EventBatch(ts=123.456, events=[BlockStored([1234567890123, -987654321], // parent=42, tokens=[10,20,30,40], block_size=4, lora=7, medium="GPU")], // attn_dp_rank=2) - let bytes = hex_to_bytes( - "93cb405edd2f1a9fbe779197ab426c6f636b53746f72656492cf0000011f71fb04cbd2c521974f2a940a141e280407a347505502", - ); + let bytes = hex_to_bytes("93cb405edd2f1a9fbe779187a474797065ab426c6f636b53746f726564ac626c6f636b5f68617368657392cf0000011f71fb04cbd2c521974fb1706172656e745f626c6f636b5f686173682aa9746f6b656e5f696473940a141e28aa626c6f636b5f73697a6504a76c6f72615f696407a66d656469756da347505502"); let batch = decode_event_batch(&bytes).expect("decode msgspec golden"); assert_eq!(batch.ts, 123.456); assert_eq!(batch.attn_dp_rank, Some(2)); @@ -827,17 +979,15 @@ mod tests { assert_eq!(b.lora_id, Some(7)); assert_eq!(b.medium.as_deref(), Some("GPU")); } - other => panic!("expected BlockStored, got {:?}", other), + other => panic!("expected BlockStored, got {other:?}"), } } #[test] fn block_stored_with_nil_optionals() { // ts=0.0, BlockStored([1,2,3], parent=None, tokens=[5,6], block_size=16, - // lora=None, medium=None), attn_dp_rank=None - let bytes = hex_to_bytes( - "93cb00000000000000009197ab426c6f636b53746f72656493010203c092050610c0c0c0", - ); + // lora=None, medium=None -> key omitted), attn_dp_rank=None + let bytes = hex_to_bytes("93cb00000000000000009186a474797065ab426c6f636b53746f726564ac626c6f636b5f68617368657393010203b1706172656e745f626c6f636b5f68617368c0a9746f6b656e5f696473920506aa626c6f636b5f73697a6510a76c6f72615f6964c0c0"); let batch = decode_event_batch(&bytes).expect("decode msgspec golden"); assert_eq!(batch.attn_dp_rank, None); match &batch.events[0] { @@ -849,16 +999,14 @@ mod tests { assert_eq!(b.lora_id, None); assert_eq!(b.medium, None); } - other => panic!("unexpected: {:?}", other), + other => panic!("unexpected: {other:?}"), } } #[test] fn block_removed_with_medium() { // ts=1.0, [BlockRemoved([100, 200], medium="DISK")], attn_dp_rank=0 - let bytes = hex_to_bytes( - "93cb3ff00000000000009193ac426c6f636b52656d6f7665649264ccc8a44449534b00", - ); + let bytes = hex_to_bytes("93cb3ff00000000000009183a474797065ac426c6f636b52656d6f766564ac626c6f636b5f6861736865739264ccc8a66d656469756da44449534b00"); let batch = decode_event_batch(&bytes).expect("decode msgspec golden"); assert_eq!(batch.ts, 1.0); assert_eq!(batch.attn_dp_rank, Some(0)); @@ -867,15 +1015,16 @@ mod tests { assert_eq!(r.block_hashes, vec![100, 200]); assert_eq!(r.medium.as_deref(), Some("DISK")); } - other => panic!("unexpected: {:?}", other), + other => panic!("unexpected: {other:?}"), } } #[test] fn all_blocks_cleared() { // ts=2.0, [AllBlocksCleared()], attn_dp_rank=None - let bytes = - hex_to_bytes("93cb40000000000000009191b0416c6c426c6f636b73436c6561726564c0"); + let bytes = hex_to_bytes( + "93cb40000000000000009181a474797065b0416c6c426c6f636b73436c6561726564c0", + ); let batch = decode_event_batch(&bytes).expect("decode msgspec golden"); assert_eq!(batch.ts, 2.0); assert_eq!(batch.attn_dp_rank, None); @@ -886,9 +1035,7 @@ mod tests { #[test] fn mixed_batch() { // ts=99.0, [BlockStored, BlockRemoved, AllBlocksCleared], attn_dp_rank=3 - let bytes = hex_to_bytes( - "93cb4058c000000000009397ab426c6f636b53746f726564910a0192010202c0a347505593ac426c6f636b52656d6f7665649114c091b0416c6c426c6f636b73436c656172656403", - ); + let bytes = hex_to_bytes("93cb4058c000000000009387a474797065ab426c6f636b53746f726564ac626c6f636b5f686173686573910ab1706172656e745f626c6f636b5f6861736801a9746f6b656e5f696473920102aa626c6f636b5f73697a6502a76c6f72615f6964c0a66d656469756da347505582a474797065ac426c6f636b52656d6f766564ac626c6f636b5f686173686573911481a474797065b0416c6c426c6f636b73436c656172656403"); let batch = decode_event_batch(&bytes).expect("decode msgspec golden"); assert_eq!(batch.ts, 99.0); assert_eq!(batch.attn_dp_rank, Some(3)); @@ -902,17 +1049,58 @@ mod tests { assert_eq!(b.lora_id, None); assert_eq!(b.medium.as_deref(), Some("GPU")); } - other => panic!("unexpected: {:?}", other), + other => panic!("unexpected: {other:?}"), } match &batch.events[1] { KvCacheEvent::BlockRemoved(r) => { assert_eq!(r.block_hashes, vec![20]); assert_eq!(r.medium, None); } - other => panic!("unexpected: {:?}", other), + other => panic!("unexpected: {other:?}"), } assert!(matches!(batch.events[2], KvCacheEvent::AllBlocksCleared)); } + + #[test] + fn mixed_batch_with_attribution_keys() { + // ts=1.0, attn_dp_rank=0: three BlockStored (the second carries + // cache_salt and session_id, the third session_id only), one + // BlockRemoved, one AllBlocksCleared. + let bytes = hex_to_bytes("93cb3ff00000000000009587a474797065ab426c6f636b53746f726564ac626c6f636b5f686173686573920b0cb1706172656e745f626c6f636b5f68617368c0a9746f6b656e5f6964739401020304aa626c6f636b5f73697a6502a76c6f72615f6964c0a66d656469756da347505589a474797065ab426c6f636b53746f726564ac626c6f636b5f6861736865739115b1706172656e745f626c6f636b5f686173680ca9746f6b656e5f696473920506aa626c6f636b5f73697a6502a76c6f72615f6964c0a66d656469756da3475055aa63616368655f73616c74a874656e616e742d61aa73657373696f6e5f6964a6736573732d3188a474797065ab426c6f636b53746f726564ac626c6f636b5f686173686573911fb1706172656e745f626c6f636b5f68617368c0a9746f6b656e5f696473920708aa626c6f636b5f73697a6502a76c6f72615f6964c0a66d656469756da3475055aa73657373696f6e5f6964a6736573732d3283a474797065ac426c6f636b52656d6f766564ac626c6f636b5f686173686573910ba66d656469756da347505581a474797065b0416c6c426c6f636b73436c656172656400"); + let batch = decode_event_batch(&bytes).expect("decode msgspec golden"); + assert_eq!(batch.ts, 1.0); + assert_eq!(batch.attn_dp_rank, Some(0)); + assert_eq!(batch.events.len(), 5); + match &batch.events[0] { + KvCacheEvent::BlockStored(b) => { + assert_eq!(b.block_hashes, vec![11, 12]); + assert_eq!(b.parent_block_hash, None); + assert_eq!(b.token_ids, vec![1, 2, 3, 4]); + assert_eq!(b.block_size, 2); + assert_eq!(b.medium.as_deref(), Some("GPU")); + } + other => panic!("expected BlockStored, got {other:?}"), + } + match &batch.events[1] { + KvCacheEvent::BlockStored(b) => { + assert_eq!(b.block_hashes, vec![21]); + assert_eq!(b.parent_block_hash, Some(12)); + } + other => panic!("expected BlockStored, got {other:?}"), + } + match &batch.events[2] { + KvCacheEvent::BlockStored(b) => assert_eq!(b.block_hashes, vec![31]), + other => panic!("expected BlockStored, got {other:?}"), + } + match &batch.events[3] { + KvCacheEvent::BlockRemoved(r) => { + assert_eq!(r.block_hashes, vec![11]); + assert_eq!(r.medium.as_deref(), Some("GPU")); + } + other => panic!("expected BlockRemoved, got {other:?}"), + } + assert!(matches!(batch.events[4], KvCacheEvent::AllBlocksCleared)); + } } #[test] @@ -930,17 +1118,15 @@ mod tests { #[test] fn block_stored_with_too_many_hashes_rejected() { let claimed = (MAX_HASHES_PER_EVENT + 1) as u32; - let mut event = Vec::new(); - write_event_array(&mut event, "BlockStored", 7); + write_event_map(&mut event, "BlockStored", 1); + write_key(&mut event, "block_hashes"); // Oversize block_hashes prefix; only one real element. The // visitor's size_hint check fires before reading anything. mp::write_array_len(&mut event, claimed).unwrap(); mp::write_sint(&mut event, 0).unwrap(); // Trailing bytes are ignored — decoder errors out earlier. - let bytes = build_batch_bytes(0.0, &[event], None, true); - let err = decode_event_batch(&bytes).expect_err("oversize hashes should fail"); match err { DecodeError::PayloadTooLarge { field, len, cap } => { @@ -960,21 +1146,16 @@ mod tests { #[test] fn block_stored_oversize_token_ids_prefix_rejected() { let claimed = (MAX_TOKENS_PER_EVENT + 1) as u32; - let mut event = Vec::new(); - write_event_array(&mut event, "BlockStored", 7); - write_i64_array(&mut event, &[42_i64]); // block_hashes (small) - mp::write_nil(&mut event).unwrap(); // parent_block_hash - // Oversize token_ids: announce huge length but only write a - // single element. The visitor's size_hint check fires - // immediately and we never reach the truncated payload. + write_event_map(&mut event, "BlockStored", 2); + write_key(&mut event, "block_hashes"); + write_i64_array(&mut event, &[42_i64]); + write_key(&mut event, "token_ids"); + // Oversize token_ids: announce huge length but only write a single + // element. The visitor's size_hint check fires immediately. mp::write_array_len(&mut event, claimed).unwrap(); mp::write_uint(&mut event, 0).unwrap(); - // Trailing bytes after the truncated array are ignored — the - // decoder errors out on the size_hint check before reading them. - let bytes = build_batch_bytes(0.0, &[event], None, true); - let err = decode_event_batch(&bytes).expect_err("oversize token prefix should fail"); match err { DecodeError::PayloadTooLarge { field, len, cap } => { @@ -990,15 +1171,13 @@ mod tests { #[test] fn block_removed_with_too_many_hashes_rejected() { let claimed = (MAX_HASHES_PER_EVENT + 1) as u32; - let mut event = Vec::new(); - write_event_array(&mut event, "BlockRemoved", 3); + write_event_map(&mut event, "BlockRemoved", 1); + write_key(&mut event, "block_hashes"); mp::write_array_len(&mut event, claimed).unwrap(); mp::write_sint(&mut event, 0).unwrap(); // Trailing bytes ignored — decoder errors on the size hint. - let bytes = build_batch_bytes(0.0, &[event], None, true); - let err = decode_event_batch(&bytes).expect_err("oversize hashes should fail"); match err { DecodeError::PayloadTooLarge { field, cap, .. } => { diff --git a/experimental/sgl-router/tests/component/policies/zmq_helpers.rs b/experimental/sgl-router/tests/component/policies/zmq_helpers.rs index ec36eb2f3..9f0e81624 100644 --- a/experimental/sgl-router/tests/component/policies/zmq_helpers.rs +++ b/experimental/sgl-router/tests/component/policies/zmq_helpers.rs @@ -27,9 +27,9 @@ pub async fn make_pub_bound() -> (PubSocket, u16) { (sock, port) } -/// Encode a single `BlockStored` event in the wire format msgspec -/// emits. Layout: `["BlockStored", block_hashes, parent, token_ids, -/// block_size, lora_id, medium]`. +/// Encode a single `BlockStored` event in the wire format msgspec emits: a +/// tagged map `{"type": "BlockStored", "block_hashes", "parent_block_hash", +/// "token_ids", "block_size", "lora_id", "medium"}`. pub fn encode_block_stored_event( block_hashes: &[i64], parent: Option, @@ -37,24 +37,31 @@ pub fn encode_block_stored_event( block_size: u32, ) -> Vec { let mut buf = Vec::new(); - mp::write_array_len(&mut buf, 7).unwrap(); + mp::write_map_len(&mut buf, 7).unwrap(); + mp::write_str(&mut buf, "type").unwrap(); mp::write_str(&mut buf, "BlockStored").unwrap(); + mp::write_str(&mut buf, "block_hashes").unwrap(); mp::write_array_len(&mut buf, block_hashes.len() as u32).unwrap(); for v in block_hashes { mp::write_sint(&mut buf, *v).unwrap(); } + mp::write_str(&mut buf, "parent_block_hash").unwrap(); match parent { Some(v) => { mp::write_sint(&mut buf, v).unwrap(); } None => mp::write_nil(&mut buf).unwrap(), } + mp::write_str(&mut buf, "token_ids").unwrap(); mp::write_array_len(&mut buf, token_ids.len() as u32).unwrap(); for v in token_ids { mp::write_uint(&mut buf, *v as u64).unwrap(); } + mp::write_str(&mut buf, "block_size").unwrap(); mp::write_uint(&mut buf, block_size as u64).unwrap(); - mp::write_nil(&mut buf).unwrap(); // lora_id + mp::write_str(&mut buf, "lora_id").unwrap(); + mp::write_nil(&mut buf).unwrap(); + mp::write_str(&mut buf, "medium").unwrap(); mp::write_str(&mut buf, "GPU").unwrap(); buf } diff --git a/python/sglang/srt/disaggregation/kv_events.py b/python/sglang/srt/disaggregation/kv_events.py index 1672cbf8a..b16cff019 100644 --- a/python/sglang/srt/disaggregation/kv_events.py +++ b/python/sglang/srt/disaggregation/kv_events.py @@ -241,11 +241,20 @@ class EventBatch( class KVCacheEvent( msgspec.Struct, - array_like=True, # type: ignore[call-arg] + omit_defaults=True, # type: ignore[call-arg] gc=False, # type: ignore[call-arg] tag=True, ): - """Base class for all KV cache-related events""" + """Base class for all KV cache-related events. + + Events are tagged msgpack maps: ``type`` carries the class name and every + other key is a field name. Optional fields left at ``None`` are omitted, so + adding an optional field never changes the shape an older consumer sees. + This is the same encoding vLLM uses for its ``KVCacheEvent``, so a consumer + such as Dynamo decodes both engines with one code path. + + ``EventBatch`` stays a positional array ``[ts, events, attn_dp_rank]``. + """ class StorageMedium(str, enum.Enum): @@ -257,12 +266,6 @@ class StorageMedium(str, enum.Enum): EXTERNAL = "EXTERNAL" # L4: shared / remote pool (e.g. Mooncake) -class BlockStoredMetadata(msgspec.Struct, omit_defaults=True, gc=False): - """Typed request metadata attached to a stored KV block.""" - - cache_salt: str - - class OffloadedState(msgspec.Struct): """Decode-side offload progress for one request, keyed by Req in the manager.""" @@ -279,16 +282,13 @@ class BlockStored(KVCacheEvent): block_size: int lora_id: Optional[int] medium: Optional[str] = None - - -class BlockStoredWithMetadata(BlockStored, tag="BlockStored", kw_only=True): - """BlockStored wire extension used only when typed metadata is present. - - A separate struct keeps unsalted events at their legacy array length; an - optional field on BlockStored would still serialize a trailing null. - """ - - metadata: BlockStoredMetadata + # Salt of the request that stored these blocks. Block hashes are already + # namespaced by it; consumers index the emitted hashes rather than + # recompute them. + cache_salt: Optional[str] = None + # Session that triggered this store. Attribution only: the blocks may be + # shared with other sessions, and the hash does not depend on it. + session_id: Optional[str] = None class BlockRemoved(KVCacheEvent): @@ -301,10 +301,6 @@ class AllBlocksCleared(KVCacheEvent): class KVEventBatch(EventBatch): - # BlockStoredWithMetadata deliberately stays out of this tagged union. - # Existing typed consumers decode its shared "BlockStored" tag as the base - # type and ignore the trailing metadata; adding both types would give - # msgspec duplicate tags and make the union invalid. events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]] diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 77e152903..cb1ed7032 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -90,6 +90,7 @@ class InsertParams: # General chunked: bool = False priority: int = 0 + session_id: Optional[str] = None track_adopted_ranges: bool = False # Logical-page KV sharding: rotation base of the chain the inserted diff --git a/python/sglang/srt/mem_cache/events.py b/python/sglang/srt/mem_cache/events.py index a3a857866..c5f7c341c 100644 --- a/python/sglang/srt/mem_cache/events.py +++ b/python/sglang/srt/mem_cache/events.py @@ -24,8 +24,6 @@ from sglang.srt.disaggregation.kv_events import ( AllBlocksCleared, BlockRemoved, BlockStored, - BlockStoredMetadata, - BlockStoredWithMetadata, StorageMedium, ) from sglang.srt.mem_cache.utils import ( @@ -63,19 +61,12 @@ class KVCacheEventRecorder: return elif isinstance(tail, BlockStored) and isinstance(event, BlockStored): - tail_metadata = ( - tail.metadata if isinstance(tail, BlockStoredWithMetadata) else None - ) - event_metadata = ( - event.metadata - if isinstance(event, BlockStoredWithMetadata) - else None - ) if ( tail.medium == event.medium and tail.lora_id == event.lora_id and tail.block_size == event.block_size - and tail_metadata == event_metadata + and tail.cache_salt == event.cache_salt + and tail.session_id == event.session_id and tail.block_hashes and event.parent_block_hash == tail.block_hashes[-1] ): @@ -112,7 +103,9 @@ class KVCacheEventRecorder: return None return hash_str_to_int64(parent_hash_values[-1]) - def record_store(self, node: Any, medium=None) -> None: + def record_store( + self, node: Any, medium=None, *, session_id: Optional[str] = None + ) -> None: # One BlockStored per ``page_size`` chunk. # ``medium`` defaults to StorageMedium.GPU but callers may override # for lower-tier insertions (e.g. StorageMedium.CPU for host/L2 cache). @@ -140,22 +133,18 @@ class KVCacheEventRecorder: block_hash = hash_str_to_int64(event_hash_values[page_index]) - event_args = { - "block_hashes": [block_hash], - "parent_block_hash": parent_block_hash, - "token_ids": page_tokens, - "block_size": len(page_tokens), - "lora_id": None, - "medium": medium, - } - if node.key.cache_salt is None: - event = BlockStored(**event_args) - else: - event = BlockStoredWithMetadata( - **event_args, - metadata=BlockStoredMetadata(cache_salt=node.key.cache_salt), + self.enqueue( + BlockStored( + block_hashes=[block_hash], + parent_block_hash=parent_block_hash, + token_ids=page_tokens, + block_size=len(page_tokens), + lora_id=None, + medium=medium, + cache_salt=node.key.cache_salt, + session_id=session_id, ) - self.enqueue(event) + ) parent_block_hash = block_hash page_index += 1 diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py index ff11a09a8..9383e4916 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py +++ b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py @@ -12,8 +12,6 @@ from sglang.srt.disaggregation.kv_events import ( AllBlocksCleared, BlockRemoved, BlockStored, - BlockStoredMetadata, - BlockStoredWithMetadata, StorageMedium, ) from sglang.srt.mem_cache.base_prefix_cache import ( @@ -84,19 +82,15 @@ def _kv_event_from_tagged(event: tuple): """Build the Python KV cache event for one of the binding's tagged tuples.""" tag = event[0] if tag == "block_stored": - event_args = dict( + return BlockStored( block_hashes=event[1], parent_block_hash=event[2], token_ids=event[3], block_size=event[4], lora_id=None, medium=StorageMedium(event[5]), - ) - if event[6] is None: - return BlockStored(**event_args) - return BlockStoredWithMetadata( - **event_args, - metadata=BlockStoredMetadata(cache_salt=event[6]), + cache_salt=event[6], + session_id=event[7], ) if tag == "block_removed": return BlockRemoved(block_hashes=event[1], medium=StorageMedium(event[2])) @@ -588,6 +582,7 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface): value=value, extra_key=key.extra_key, cache_salt=key.cache_salt, + session_id=params.session_id, mamba_value=params.mamba_value, prev_prefix_len=params.prev_prefix_len, swa_evicted_seqlen=params.swa_evicted_seqlen, diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 796a39912..e452e3513 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -1161,7 +1161,11 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): node.priority = max(node.priority, state.priority) if node.evicted: - self._unevict_node_on_insert(node, state.value[:prefix_len]) + self._unevict_node_on_insert( + node, + state.value[:prefix_len], + session_id=state.params.session_id, + ) state.result.record_adopted_range( BASE_COMPONENT_TYPE, state.total_prefix_length, @@ -1235,6 +1239,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): state.key, state.value, priority=state.priority, + session_id=state.params.session_id, rotation_base=state.params.rotation_base, ) state.is_new_leaf = True @@ -1360,6 +1365,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): key: RadixKey, value: torch.Tensor, priority: int = 0, + session_id: Optional[str] = None, rotation_base: Optional[int] = None, ) -> UnifiedTreeNode: new_node = self._new_node(priority=priority) @@ -1378,11 +1384,14 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): self._update_evictable_leaf_sets(new_node) self._update_evictable_leaf_sets(parent) - self.kv_events.record_store(new_node) + self.kv_events.record_store(new_node, session_id=session_id) return new_node def _unevict_node_on_insert( - self, node: UnifiedTreeNode, fresh_value: torch.Tensor + self, + node: UnifiedTreeNode, + fresh_value: torch.Tensor, + session_id: Optional[str] = None, ) -> None: """Restore an evicted node's Full device value from fresh KV indices during insert.""" @@ -1400,7 +1409,11 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): self._update_duplicate_tracking(node) if node.parent is not None: self._update_evictable_leaf_sets(node.parent) - self.kv_events.record_store(node, medium=StorageMedium.GPU) + self.kv_events.record_store( + node, + medium=StorageMedium.GPU, + session_id=session_id, + ) def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None: """Update both device and host leaf sets for a node.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 420fa8cc1..f19d7b6c1 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -977,6 +977,7 @@ class UnifiedRadixCache(BasePrefixCache): insert_params = InsertParams( prev_prefix_len=req.kv.cache_protected_len, priority=getattr(req, "priority", 0) or 0, + session_id=req.session_id, rotation_base=req.kv_rotation_base, ) @@ -1113,6 +1114,7 @@ class UnifiedRadixCache(BasePrefixCache): prev_prefix_len=req.kv.cache_protected_len, chunked=chunked, priority=getattr(req, "priority", 0) or 0, + session_id=req.session_id, rotation_base=req.kv_rotation_base, ) effective_cache_len = len(token_ids) diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index d34a2862b..c19793ec6 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -514,6 +514,7 @@ pub struct InsertParamsBinding { pub value: Py, pub extra_key: Option, pub cache_salt: Option, + pub session_id: Option, pub mamba_value: Option>, pub prev_prefix_len: usize, pub swa_evicted_seqlen: usize, @@ -526,13 +527,14 @@ pub struct InsertParamsBinding { #[pymethods] impl InsertParamsBinding { #[new] - #[pyo3(signature = (key, value, extra_key = None, cache_salt = None, prev_prefix_len = 0, swa_evicted_seqlen = 0, swa_branching_seqlen = None, chunked = false, priority = 0, mamba_value = None, track_adopted_ranges = false))] + #[pyo3(signature = (key, value, extra_key = None, cache_salt = None, session_id = None, prev_prefix_len = 0, swa_evicted_seqlen = 0, swa_branching_seqlen = None, chunked = false, priority = 0, mamba_value = None, track_adopted_ranges = false))] fn new( py: Python<'_>, key: &Bound<'_, PyAny>, value: Py, extra_key: Option, cache_salt: Option, + session_id: Option, prev_prefix_len: usize, swa_evicted_seqlen: usize, swa_branching_seqlen: Option, @@ -546,6 +548,7 @@ impl InsertParamsBinding { value, extra_key, cache_salt, + session_id, mamba_value, prev_prefix_len, swa_evicted_seqlen, @@ -1023,6 +1026,7 @@ impl TreeCoreBinding { params.extra_key.as_deref(), params.cache_salt.as_deref(), ), + session_id: params.session_id.as_deref(), value: value.0, mamba_value, prev_prefix_len: params.prev_prefix_len, @@ -1059,6 +1063,7 @@ impl TreeCoreBinding { params.extra_key.as_deref(), params.cache_salt.as_deref(), ), + session_id: params.session_id.as_deref(), value: value.0, mamba_value, prev_prefix_len: params.prev_prefix_len, @@ -1828,6 +1833,7 @@ impl TreeCoreBinding { block_size, medium, cache_salt, + session_id, } => { let item: Py = ( "block_stored", @@ -1837,6 +1843,7 @@ impl TreeCoreBinding { block_size, medium.as_str(), cache_salt.map(|salt| salt.to_string()), + session_id.map(|session_id| session_id.to_string()), ) .into_py(py); list.append(item)?; diff --git a/rust/sglang-radix-tree/src/tests/components/base.rs b/rust/sglang-radix-tree/src/tests/components/base.rs index 270fa5d71..c835597c2 100644 --- a/rust/sglang-radix-tree/src/tests/components/base.rs +++ b/rust/sglang-radix-tree/src/tests/components/base.rs @@ -98,6 +98,7 @@ fn insert_overlap_default_consumes_nothing() { swa_branching_seqlen: None, chunked: false, priority: 0, + session_id: None, track_adopted_ranges: false, }, &mut InsertResult::default(), diff --git a/rust/sglang-radix-tree/src/tests/components/full.rs b/rust/sglang-radix-tree/src/tests/components/full.rs index 18ad7dae3..c466fd198 100644 --- a/rust/sglang-radix-tree/src/tests/components/full.rs +++ b/rust/sglang-radix-tree/src/tests/components/full.rs @@ -75,6 +75,7 @@ fn insert(tc: &mut UnifiedTreeCore>, key: &Vec, value: &[i64]) { swa_branching_seqlen: None, chunked: false, priority: 0, + session_id: None, track_adopted_ranges: false, }); } @@ -477,6 +478,7 @@ fn host_drive_is_a_noop_without_host_leaves() { swa_branching_seqlen: None, chunked: false, priority: 0, + session_id: None, track_adopted_ranges: false, }); let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); diff --git a/rust/sglang-radix-tree/src/tests/components/mamba.rs b/rust/sglang-radix-tree/src/tests/components/mamba.rs index 672f70bac..1a1821e6b 100644 --- a/rust/sglang-radix-tree/src/tests/components/mamba.rs +++ b/rust/sglang-radix-tree/src/tests/components/mamba.rs @@ -68,6 +68,7 @@ fn insert_params_mamba<'k>( swa_branching_seqlen: None, chunked: false, priority: 0, + session_id: None, track_adopted_ranges: false, } } diff --git a/rust/sglang-radix-tree/src/tests/components/swa.rs b/rust/sglang-radix-tree/src/tests/components/swa.rs index 8f53a7236..eab37178f 100644 --- a/rust/sglang-radix-tree/src/tests/components/swa.rs +++ b/rust/sglang-radix-tree/src/tests/components/swa.rs @@ -552,6 +552,7 @@ fn insert_params_swa<'k>( swa_branching_seqlen: None, chunked: false, priority: 0, + session_id: None, track_adopted_ranges: false, } } diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index 01a8fb208..bd0765ed6 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -1158,7 +1158,7 @@ fn unevict_restores_the_value_and_the_leaf_sets() { .set_device_value(p, FULL, Tensor::from_slice(&[0i64])); tc.evictable_device_leaves.add(p); let mut fresh = Tensor::from_slice(&[20i64]); - tc.unevict_node_on_insert_(c, &fresh); + tc.unevict_node_on_insert_(c, &fresh, /* session_id = */ None); assert_eq!(tc.evictable_size_(FULL), 1); assert!(tc.evictable_device_leaves.contains(c)); assert!(!tc.evictable_device_leaves.contains(p)); @@ -1187,7 +1187,11 @@ fn unevict_panics_on_a_node_that_still_has_its_value() { .unwrap(); tc.arena .set_device_value(a, FULL, Tensor::from_slice(&[0i64])); - tc.unevict_node_on_insert_(a, &Tensor::from_slice(&[1i64])); + tc.unevict_node_on_insert_( + a, + &Tensor::from_slice(&[1i64]), + /* session_id = */ None, + ); } fn match_params(key: &Vec) -> MatchPrefixParams<'_, Vec> { @@ -1866,6 +1870,7 @@ fn insert_params<'k>(key: &'k Vec, value: &[i64]) -> InsertParams<'k, Vec>(&key, None, 2); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: hashes + .iter() + .map(|hash| crate::node::hash_str_to_int64(hash)) + .collect(), + parent_block_hash: None, + token_ids: key, + block_size: 2, + medium: StorageMedium::Gpu, + cache_salt: None, + session_id: Some(Arc::from("session-a")), + }] + ); +} + #[test] fn namespaced_event_hashes_are_sparse_and_removed_with_the_node() { let mut tc = events_core(2); @@ -2828,6 +2860,7 @@ fn extra_key_nodes_publish_token_only_event_hashes() { block_size: 2, medium: StorageMedium::Gpu, cache_salt: None, + session_id: None, }] ); @@ -2920,6 +2953,7 @@ fn event_coalescing_respects_store_remove_and_clear_boundaries() { block_size: 2, medium: StorageMedium::Gpu, cache_salt: None, + session_id: None, }); assert_eq!(tc.kv_event_queue.len(), 1); // A different block size must not join the parent-linked store tail. @@ -2930,6 +2964,7 @@ fn event_coalescing_respects_store_remove_and_clear_boundaries() { block_size: 1, medium: StorageMedium::Gpu, cache_salt: None, + session_id: None, }); assert_eq!(tc.kv_event_queue.len(), 2); // Matching size and parent are still separated across media. @@ -2940,6 +2975,7 @@ fn event_coalescing_respects_store_remove_and_clear_boundaries() { block_size: 1, medium: StorageMedium::Cpu, cache_salt: None, + session_id: None, }); assert_eq!(tc.kv_event_queue.len(), 3); // Matching size and medium are still separated without the parent link. @@ -2950,6 +2986,7 @@ fn event_coalescing_respects_store_remove_and_clear_boundaries() { block_size: 1, medium: StorageMedium::Cpu, cache_salt: None, + session_id: None, }); assert_eq!(tc.kv_event_queue.len(), 4); tc.enqueue_kv_event_(KvCacheEvent::BlockRemoved { @@ -2992,6 +3029,7 @@ fn event_coalescing_respects_store_remove_and_clear_boundaries() { block_size: 2, medium: StorageMedium::Gpu, cache_salt: Some(Arc::from("tenant-a")), + session_id: None, }); tc.enqueue_kv_event_(KvCacheEvent::BlockStored { block_hashes: vec![2], @@ -3000,6 +3038,7 @@ fn event_coalescing_respects_store_remove_and_clear_boundaries() { block_size: 2, medium: StorageMedium::Gpu, cache_salt: Some(Arc::from("tenant-b")), + session_id: None, }); assert_eq!(tc.kv_event_queue.len(), 2); } @@ -3055,6 +3094,7 @@ fn bigram_insert_events_carry_pair_token_payloads() { swa_branching_seqlen: None, chunked: false, priority: 0, + session_id: None, track_adopted_ranges: false, }); let hashes = crate::node::get_hash_str::>(&key, None, 1); @@ -3070,6 +3110,7 @@ fn bigram_insert_events_carry_pair_token_payloads() { block_size: 1, medium: StorageMedium::Gpu, cache_salt: None, + session_id: None, }] ); } @@ -3093,6 +3134,7 @@ fn finish_write_through_emits_cpu_stored_events() { block_size: 1, medium: StorageMedium::Cpu, cache_salt: None, + session_id: None, }] ); } @@ -3148,6 +3190,7 @@ fn load_back_commit_emits_gpu_stored_events() { block_size: 1, medium: StorageMedium::Gpu, cache_salt: None, + session_id: None, }] ); } @@ -3167,6 +3210,7 @@ fn unevict_on_insert_emits_a_gpu_stored_event() { block_size: 1, medium: StorageMedium::Gpu, cache_salt: None, + session_id: None, }] ); } @@ -3255,6 +3299,7 @@ fn split_insert_stores_only_the_new_block_chained_to_the_split_parent() { block_size: 2, medium: StorageMedium::Gpu, cache_salt: None, + session_id: None, }] ); // The split divided the page hashes between the two fragments. @@ -3336,6 +3381,7 @@ fn finish_write_through_after_a_split_publishes_both_fragments() { block_size: 2, medium: StorageMedium::Cpu, cache_salt: None, + session_id: None, }] ); // The matching ack cleared the pending mark on both fragments. @@ -3664,6 +3710,7 @@ fn insert_host_publishes_a_host_store_event() { block_size: 2, medium: StorageMedium::Cpu, cache_salt: None, + session_id: None, }] ); } @@ -8437,6 +8484,7 @@ fn sequence_insert_params<'k>( swa_branching_seqlen: None, chunked: false, priority: 0, + session_id: None, track_adopted_ranges: false, } } diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 44ab0acff..9600aab95 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -119,6 +119,9 @@ pub struct InsertParams<'k, K: ChildKeyType> { pub key: &'k K, /// Namespace of the insert; picks the matching subtree root. pub namespace: KeyNamespaceRef<'k>, + /// Request session attributed to newly stored blocks. This is event metadata only; + /// it does not participate in tree matching or block hashing. + pub session_id: Option<&'k str>, /// Device KV indices covering the key, one row per atom. pub value: Tensor, /// Tokens of this request already cached before the insert (the duplicate @@ -211,6 +214,7 @@ pub struct InsertWalkState { aligned_key_len: usize, value: Tensor, namespace: KeyNamespace, + session_id: Option>, prev_prefix_len: usize, swa_evicted_seqlen: usize, swa_branching_seqlen: Option, @@ -1422,6 +1426,7 @@ impl UnifiedTreeCore { aligned_key_len, value: params.value.narrow(0, 0, aligned_key_len as i64), namespace: params.namespace.to_owned(), + session_id: params.session_id.map(Arc::from), prev_prefix_len: params.prev_prefix_len, swa_evicted_seqlen: params.swa_evicted_seqlen, swa_branching_seqlen: params.swa_branching_seqlen, @@ -1547,6 +1552,7 @@ impl UnifiedTreeCore { let params = InsertParams { key: &state.key, namespace: state.namespace.as_ref(), + session_id: state.session_id.as_deref(), value: state.value.shallow_clone(), prev_prefix_len: state.prev_prefix_len, swa_evicted_seqlen: state.swa_evicted_seqlen, @@ -1560,6 +1566,7 @@ impl UnifiedTreeCore { self.unevict_node_on_insert_( node_id, &state.value.narrow(0, cursor as i64, prefix_len as i64), + state.session_id.as_deref(), ); state .result @@ -1679,6 +1686,7 @@ impl UnifiedTreeCore { &leaf_value, state.priority, state.namespace.as_ref(), + state.session_id.as_deref(), ) } else { state.node_id @@ -1698,6 +1706,7 @@ impl UnifiedTreeCore { let params = InsertParams { key: &state.key, namespace: state.namespace.as_ref(), + session_id: state.session_id.as_deref(), value: state.value.shallow_clone(), prev_prefix_len: state.prev_prefix_len, swa_evicted_seqlen: state.swa_evicted_seqlen, @@ -1888,6 +1897,7 @@ impl UnifiedTreeCore { value, priority, KeyNamespaceRef::new(extra_key, /* cache_salt = */ None), + /* session_id = */ None, ) } @@ -1898,6 +1908,7 @@ impl UnifiedTreeCore { value: &Tensor, priority: i64, namespace: KeyNamespaceRef<'_>, + session_id: Option<&str>, ) -> NodeIdx_ { let page_size = self.page_size; let child_map_key = key.child_key(page_size); @@ -1921,13 +1932,18 @@ impl UnifiedTreeCore { self.update_evictable_leaf_sets_(new_node_id); self.update_evictable_leaf_sets_(parent_id); - self.record_store_event_(new_node_id, StorageMedium::Gpu); + self.record_store_event_(new_node_id, StorageMedium::Gpu, session_id); new_node_id } /// Restore an evicted node's Full device value from fresh KV indices /// during insert. - pub fn unevict_node_on_insert_(&mut self, node_id: NodeIdx_, fresh_value: &Tensor) { + pub fn unevict_node_on_insert_( + &mut self, + node_id: NodeIdx_, + fresh_value: &Tensor, + session_id: Option<&str>, + ) { self.arena .set_device_value(node_id, FULL, fresh_value.copy()); let tokens = fresh_value.size()[0] as usize; @@ -1943,7 +1959,7 @@ impl UnifiedTreeCore { if let Some(parent_id) = self.arena.node(node_id).try_parent() { self.update_evictable_leaf_sets_(parent_id); } - self.record_store_event_(node_id, StorageMedium::Gpu); + self.record_store_event_(node_id, StorageMedium::Gpu, session_id); } /// Update both device and host leaf sets for a node. @@ -2768,6 +2784,7 @@ impl UnifiedTreeCore { block_size: tail_block_size, medium: tail_medium, cache_salt: tail_cache_salt, + session_id: tail_session_id, .. }), KvCacheEvent::BlockStored { @@ -2777,10 +2794,12 @@ impl UnifiedTreeCore { block_size, medium, cache_salt, + session_id, }, ) if *tail_medium == medium && *tail_block_size == block_size && *tail_cache_salt == cache_salt + && *tail_session_id == session_id && !tail_hashes.is_empty() && parent_block_hash == tail_hashes.last().copied() => { @@ -2847,7 +2866,12 @@ impl UnifiedTreeCore { } /// Build one BlockStored per page and coalesce compatible queue neighbors. - fn record_store_event_(&mut self, node_id: NodeIdx_, medium: StorageMedium) { + fn record_store_event_( + &mut self, + node_id: NodeIdx_, + medium: StorageMedium, + session_id: Option<&str>, + ) { if !self.enable_kv_cache_events { return; } @@ -2856,6 +2880,7 @@ impl UnifiedTreeCore { self.arena.node_mut(node_id).hash_value = Some(hash_values); } let cache_salt = self.arena.node(node_id).namespace.cache_salt_arc(); + let session_id: Option> = session_id.map(Arc::from); let namespaced = self.arena.node(node_id).namespace != KeyNamespace::default(); if namespaced { self.ensure_namespaced_event_hashes_(node_id); @@ -2885,6 +2910,7 @@ impl UnifiedTreeCore { block_size: page.len(), medium, cache_salt: cache_salt.clone(), + session_id: session_id.clone(), }); parent_block_hash = Some(block_hash); }; @@ -3107,7 +3133,11 @@ impl UnifiedTreeCore { self.update_evictable_leaf_sets_(new_node_id); self.update_evictable_leaf_sets_(node_id); result.inserted_host_node = Some(self.arena.node(new_node_id).id); - self.record_store_event_(new_node_id, StorageMedium::Cpu); + self.record_store_event_( + new_node_id, + StorageMedium::Cpu, + /* session_id = */ None, + ); Ok(result) } @@ -3653,7 +3683,7 @@ impl UnifiedTreeCore { /* pool_storage_result = */ None, ); for loaded_idx in loaded_node_indices { - self.record_store_event_(loaded_idx, StorageMedium::Gpu); + self.record_store_event_(loaded_idx, StorageMedium::Gpu, /* session_id = */ None); } for (component_type, transfers) in comp_xfers { self.component_by_type_(component_type) @@ -3853,7 +3883,7 @@ impl UnifiedTreeCore { node.write_through_pending_id = None; self.update_full_coexisting_host_tracking_(node_idx); } - self.record_store_event_(node_idx, StorageMedium::Cpu); + self.record_store_event_(node_idx, StorageMedium::Cpu, /* session_id = */ None); } Ok(()) } @@ -5045,6 +5075,7 @@ pub enum KvCacheEvent { block_size: usize, medium: StorageMedium, cache_salt: Option>, + session_id: Option>, }, BlockRemoved { block_hashes: Vec, diff --git a/test/registered/unit/disaggregation/test_kv_events.py b/test/registered/unit/disaggregation/test_kv_events.py index e92a0c42b..077cda160 100644 --- a/test/registered/unit/disaggregation/test_kv_events.py +++ b/test/registered/unit/disaggregation/test_kv_events.py @@ -12,9 +12,9 @@ import unittest import msgspec from sglang.srt.disaggregation.kv_events import ( + AllBlocksCleared, + BlockRemoved, BlockStored, - BlockStoredMetadata, - BlockStoredWithMetadata, KVEventBatch, StorageMedium, ZmqEventPublisher, @@ -186,42 +186,67 @@ class TestSelectKvPublisherDpRank(CustomTestCase): class TestBlockStoredWireFormat(CustomTestCase): - def _event(self, metadata=None): - event_type = BlockStored if metadata is None else BlockStoredWithMetadata - kwargs = dict( + def _event(self, **extra): + return BlockStored( block_hashes=[123], parent_block_hash=None, token_ids=[1, 2], block_size=2, lora_id=None, medium=StorageMedium.GPU, + **extra, ) - if metadata is not None: - kwargs["metadata"] = metadata - return event_type(**kwargs) - def test_unsalted_event_keeps_legacy_array_shape(self): + def test_event_is_a_tagged_map(self): decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(self._event())) - self.assertEqual(len(decoded), 7) + self.assertIsInstance(decoded, dict) + self.assertEqual(decoded["type"], "BlockStored") + self.assertEqual( + set(decoded), + { + "type", + "block_hashes", + "parent_block_hash", + "token_ids", + "block_size", + "lora_id", + "medium", + }, + ) - def test_salted_event_appends_typed_metadata(self): - event = self._event(BlockStoredMetadata(cache_salt="tenant-a")) - encoded = msgspec.msgpack.encode(event) - decoded = msgspec.msgpack.decode(encoded) - round_tripped = msgspec.msgpack.decode(encoded, type=BlockStoredWithMetadata) - self.assertEqual(len(decoded), 8) - self.assertEqual(decoded[7], {"cache_salt": "tenant-a"}) - self.assertEqual(round_tripped.metadata.cache_salt, "tenant-a") + def test_salt_and_session_are_named_fields(self): + event = self._event(cache_salt="tenant-a", session_id="session-a") + decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event)) + self.assertEqual(decoded["cache_salt"], "tenant-a") + self.assertEqual(decoded["session_id"], "session-a") - def test_salted_event_remains_compatible_with_typed_batch_consumers(self): + def test_one_decoder_reads_a_mixed_batch(self): batch = KVEventBatch( ts=1.0, - events=[self._event(BlockStoredMetadata(cache_salt="tenant-a"))], + events=[ + self._event(), + self._event(cache_salt="tenant-a"), + self._event(session_id="session-a"), + BlockRemoved(block_hashes=[123], medium=StorageMedium.GPU), + AllBlocksCleared(), + ], ) round_tripped = msgspec.msgpack.decode( msgspec.msgpack.encode(batch), type=KVEventBatch ) - self.assertEqual(round_tripped.events[0].block_hashes, [123]) + stored = round_tripped.events[:3] + self.assertEqual([e.cache_salt for e in stored], [None, "tenant-a", None]) + self.assertEqual([e.session_id for e in stored], [None, None, "session-a"]) + self.assertIsInstance(round_tripped.events[3], BlockRemoved) + self.assertIsInstance(round_tripped.events[4], AllBlocksCleared) + + def test_batch_stays_a_positional_array_of_maps(self): + batch = KVEventBatch(ts=1.0, events=[self._event()], attn_dp_rank=0) + decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(batch)) + self.assertEqual(decoded[0], 1.0) + self.assertEqual(decoded[2], 0) + self.assertIsInstance(decoded[1][0], dict) + self.assertEqual(len(decoded), 3) if __name__ == "__main__": diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index 09e467046..f9b909db9 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -612,6 +612,7 @@ class _GraftReq: self.swa_prefix_lock_released = False self.finished_reason = None self.session = None + self.session_id = None def get_fill_ids(self): return array("q", self.fill_ids) diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index e52324432..29e94d89d 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -34,8 +34,6 @@ from sglang.srt.disaggregation.kv_events import ( AllBlocksCleared, BlockRemoved, BlockStored, - BlockStoredMetadata, - BlockStoredWithMetadata, StorageMedium, ) from sglang.srt.managers.schedule_batch import ReqKvInfo @@ -66,20 +64,17 @@ class TestKVCacheEventQueue(unittest.TestCase): medium: StorageMedium = StorageMedium.GPU, lora_id: int | None = None, cache_salt: str | None = None, + session_id: str | None = None, ) -> BlockStored: - event_args = dict( + return BlockStored( block_hashes=[block_hash], parent_block_hash=parent_block_hash, token_ids=[block_hash, block_hash + 1][:block_size], block_size=block_size, lora_id=lora_id, medium=medium, - ) - if cache_salt is None: - return BlockStored(**event_args) - return BlockStoredWithMetadata( - **event_args, - metadata=BlockStoredMetadata(cache_salt=cache_salt), + cache_salt=cache_salt, + session_id=session_id, ) def test_enqueue_coalesces_compatible_stores(self): @@ -133,6 +128,11 @@ class TestKVCacheEventQueue(unittest.TestCase): queue.enqueue(self._store(2, 1, cache_salt="tenant-b")) self.assertEqual(len(queue.take()), 2) + queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE) + queue.enqueue(self._store(1, None, session_id="session-a")) + queue.enqueue(self._store(2, 1, session_id="session-b")) + self.assertEqual(len(queue.take()), 2) + class TestRadixKey(unittest.TestCase): """Test cases for RadixKey class.""" @@ -781,7 +781,7 @@ class TestRadixCache(CustomTestCase): removed = [event for event in events if isinstance(event, BlockRemoved)] self.assertEqual(len(stored), 1) - self.assertEqual(stored[0].metadata.cache_salt, "tenant-a") + self.assertEqual(stored[0].cache_salt, "tenant-a") self.assertEqual(stored[0].parent_block_hash, None) self.assertEqual(len(stored[0].block_hashes), 2) self.assertEqual(removed[0].block_hashes, stored[0].block_hashes) diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index 688e8bbff..3a1d5f0ed 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -20,8 +20,6 @@ from sglang.srt.disaggregation.kv_events import ( AllBlocksCleared, BlockRemoved, BlockStored, - BlockStoredMetadata, - BlockStoredWithMetadata, StorageMedium, ) from sglang.srt.environ import envs @@ -971,14 +969,14 @@ def test_salted_events_match_python_hash_and_metadata_contract(): for value in mem_cache.get_hash_str(array("q", [1, 2, 7, 8]), seed, 2) ] assert core.take_events() == [ - BlockStoredWithMetadata( + BlockStored( block_hashes=hashes, parent_block_hash=None, token_ids=[1, 2, 7, 8], block_size=2, lora_id=None, medium=StorageMedium.GPU, - metadata=BlockStoredMetadata(cache_salt="tenant-a"), + cache_salt="tenant-a", ) ] @@ -1012,14 +1010,14 @@ def test_salted_eagle_events_match_the_bigram_hash_contract(): for value in mem_cache.get_hash_str(raw_tokens, seed, 2, is_bigram=True) ] assert core.take_events() == [ - BlockStoredWithMetadata( + BlockStored( block_hashes=hashes, parent_block_hash=None, token_ids=[(1, 2), (2, 3), (3, 4), (4, 5)], block_size=2, lora_id=None, medium=StorageMedium.GPU, - metadata=BlockStoredMetadata(cache_salt="tenant-a"), + cache_salt="tenant-a", ) ] diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 719324844..e6cc0fbfb 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -22,7 +22,6 @@ from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape from sglang.srt.disaggregation.kv_events import ( BlockRemoved, BlockStored, - BlockStoredWithMetadata, StorageMedium, ) from sglang.srt.environ import envs @@ -1012,11 +1011,18 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): *, extra_key=None, cache_salt=None, + session_id=None, ): key = RadixKey(array("q", tokens), extra_key=extra_key, cache_salt=cache_salt) value = allocator.alloc(len(tokens)) self.assertIsNotNone(value) - return cache.insert(InsertParams(key=key, value=value[: len(key)])) + return cache.insert( + InsertParams( + key=key, + value=value[: len(key)], + session_id=session_id, + ) + ) def _stored_events(self, cache, medium=None): events = [e for e in cache.take_events() if isinstance(e, BlockStored)] @@ -1096,8 +1102,7 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): self._insert(cache, allocator, seq, cache_salt="tenant-a") stored = self._stored_events(cache, StorageMedium.GPU) self.assertEqual(len(stored), 1) - self.assertIsInstance(stored[0], BlockStoredWithMetadata) - self.assertEqual(stored[0].metadata.cache_salt, "tenant-a") + self.assertEqual(stored[0].cache_salt, "tenant-a") salted_hashes = self._event_hashes(stored) cache.evict(EvictParams(num_tokens=len(seq))) @@ -1115,6 +1120,73 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): salted_hashes, ) + def test_session_id_is_attributed_without_changing_block_hash(self): + cache_a, allocator_a, _ = build_fixture(self.cfg, enable_kv_cache_events=True) + cache_a.take_events() + self._insert(cache_a, allocator_a, [1, 2, 3, 4], session_id="session-a") + stored_a = self._stored_events(cache_a, StorageMedium.GPU) + self.assertEqual(len(stored_a), 1) + self.assertEqual(stored_a[0].session_id, "session-a") + + cache_b, allocator_b, _ = build_fixture(self.cfg, enable_kv_cache_events=True) + cache_b.take_events() + self._insert(cache_b, allocator_b, [1, 2, 3, 4], session_id="session-b") + stored_b = self._stored_events(cache_b, StorageMedium.GPU) + self.assertEqual(stored_b[0].session_id, "session-b") + self.assertEqual(self._event_hashes(stored_a), self._event_hashes(stored_b)) + + def test_shared_prefix_hit_is_quiet_and_divergent_tails_are_attributed(self): + cache, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True) + cache.take_events() + + shared_prefix = [1, 2, 3, 4] + self._insert(cache, allocator, shared_prefix, session_id="session-a") + initial = self._stored_events(cache, StorageMedium.GPU) + self.assertEqual(len(initial), 1) + shared_parent = initial[0].block_hashes[-1] + + self._insert(cache, allocator, shared_prefix, session_id="session-b") + self.assertEqual(self._stored_events(cache, StorageMedium.GPU), []) + + self._insert( + cache, + allocator, + shared_prefix + [5, 6], + session_id="session-a", + ) + session_a_tail = self._stored_events(cache, StorageMedium.GPU) + self.assertEqual(len(session_a_tail), 1) + self.assertEqual(session_a_tail[0].parent_block_hash, shared_parent) + self.assertEqual(list(session_a_tail[0].token_ids), [5, 6]) + self.assertEqual(session_a_tail[0].session_id, "session-a") + + self._insert( + cache, + allocator, + shared_prefix + [7, 8], + session_id="session-b", + ) + session_b_tail = self._stored_events(cache, StorageMedium.GPU) + self.assertEqual(len(session_b_tail), 1) + self.assertEqual(session_b_tail[0].parent_block_hash, shared_parent) + self.assertEqual(list(session_b_tail[0].token_ids), [7, 8]) + self.assertEqual(session_b_tail[0].session_id, "session-b") + + def test_session_id_and_cache_salt_are_both_attributed(self): + cache, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True) + cache.take_events() + self._insert( + cache, + allocator, + [1, 2, 3, 4], + cache_salt="tenant-a", + session_id="session-a", + ) + stored = self._stored_events(cache, StorageMedium.GPU) + self.assertEqual(len(stored), 1) + self.assertEqual(stored[0].cache_salt, "tenant-a") + self.assertEqual(stored[0].session_id, "session-a") + def test_cache_salt_event_parentage_survives_node_split(self): cache, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True) cache.take_events() @@ -1127,8 +1199,7 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): self._insert(cache, allocator, [1, 2, 5, 6], cache_salt="tenant-a") branch = self._stored_events(cache, StorageMedium.GPU) self.assertEqual(len(branch), 1) - self.assertIsInstance(branch[0], BlockStoredWithMetadata) - self.assertEqual(branch[0].metadata.cache_salt, "tenant-a") + self.assertEqual(branch[0].cache_salt, "tenant-a") self.assertEqual(branch[0].parent_block_hash, original[0].block_hashes[0]) self.assertEqual(list(branch[0].token_ids), [5, 6]) @@ -1296,10 +1367,11 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): self.assertTrue(cache.tree_core.is_full_device_evicted(node)) self.assertTrue(cache.tree_core.is_backuped(node)) - self._insert(cache, allocator, seq) + self._insert(cache, allocator, seq, session_id="session-a") restored_gpu = self._stored_events(cache, StorageMedium.GPU) self.assertFalse(cache.tree_core.is_full_device_evicted(node)) self.assertCountEqual(self._event_hashes(restored_gpu), stored_hashes) + self.assertEqual(restored_gpu[0].session_id, "session-a") class UnifiedRadixCacheSuite: