minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

feat(agent sessions): attribute stored KV cache blocks to sessions (#37482)

Browse Source 
Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
This commit is contained in:
ishandhanani
2026-09-14 21:34:49 -07:00
committed by GitHub
parent c9fbe5f655
commit 3f871a246c
22 changed files with 905 additions and 451 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/sgl-kv-indexer/src/bridge.rs
+158 -74
View File
@@ -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": ..., "<field>": ...}`. 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<Option<&str>, 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<i64>, 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<i64>,
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());
}
experimental/sgl-router/src/policies/kv_events/subscriber.rs
+3 -2
View File
@@ -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) => {
experimental/sgl-router/src/policies/kv_events/wire.rs
+433 -254
View File
@@ -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<D>(deserializer: D) -> Result<Self, D::Error>
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<E: de::Error>(self, v: &str) -> Result<EventField, E> {
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<A>(self, mut seq: A) -> Result<KvCacheEvent, A::Error>
fn visit_map<A>(self, mut map: A) -> Result<KvCacheEvent, A::Error>
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<String> = None;
let mut block_hashes: Option<BoundedI64Vec> = None;
let mut parent_block_hash: Option<i64> = None;
let mut token_ids: Option<BoundedU32Vec> = None;
let mut block_size: Option<u32> = None;
let mut lora_id: Option<i64> = None;
let mut medium: Option<String> = None;
while let Some(field) = map.next_key::<EventField>()? {
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::<IgnoredAny>()?;
}
}
}
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<i64> = 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<i64> = seq.next_element()?.unwrap_or(None);
// `medium` defaults to None and may be omitted.
let medium: Option<String> = seq.next_element()?.unwrap_or(None);
while seq.next_element::<IgnoredAny>()?.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<String> = seq.next_element()?.unwrap_or(None);
while seq.next_element::<IgnoredAny>()?.is_some() {}
Ok(KvCacheEvent::BlockRemoved(BlockRemoved {
block_hashes: block_hashes.0,
medium,
}))
}
"AllBlocksCleared" => {
while seq.next_element::<IgnoredAny>()?.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<u8>, 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<u8>, key: &str) {
mp::write_str(buf, key).unwrap();
}
fn write_opt_sint(buf: &mut Vec<u8>, value: Option<i64>) {
match value {
Some(v) => {
mp::write_sint(buf, v).unwrap();
}
None => mp::write_nil(buf).unwrap(),
}
}
fn write_i64_array(buf: &mut Vec<u8>, 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<u8>, 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<u8>,
block_hashes: &[i64],
parent: Option<i64>,
write_tokens: impl FnOnce(&mut Vec<u8>),
block_size: u32,
lora_id: Option<i64>,
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<i64>,
token_ids: &[u32],
block_size: u32,
lora_id: Option<i64>,
medium: Option<&str>,
extra: &[(&str, &str)],
) -> Vec<u8> {
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<i64>,
token_ids: &[u32],
block_size: u32,
lora_id: Option<i64>,
medium: Option<&str>,
) -> Vec<u8> {
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<u8> {
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<i64>,
token_ids: &[u32],
block_size: u32,
lora_id: Option<i64>,
medium: Option<&str>,
) -> Vec<u8> {
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<u8> {
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<u8> {
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, .. } => {
experimental/sgl-router/tests/component/policies/zmq_helpers.rs
+12 -5
View File
@@ -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<i64>,
@@ -37,24 +37,31 @@ pub fn encode_block_stored_event(
block_size: u32,
) -> Vec<u8> {
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
}
python/sglang/srt/disaggregation/kv_events.py
+18 -22
View File
@@ -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]]
python/sglang/srt/mem_cache/base_prefix_cache.py
+1
View File
@@ -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
python/sglang/srt/mem_cache/events.py
+16 -27
View File
@@ -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
python/sglang/srt/mem_cache/rust_tree_core/adapter.py
+4 -9
View File
@@ -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,
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+17 -4
View File
@@ -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."""
python/sglang/srt/mem_cache/unified_radix_cache.py
+2
View File
@@ -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)
rust/sglang-radix-tree/src/python_bindings.rs
+8 -1
View File
@@ -514,6 +514,7 @@ pub struct InsertParamsBinding {
pub value: Py<PyAny>,
pub extra_key: Option<String>,
pub cache_salt: Option<String>,
pub session_id: Option<String>,
pub mamba_value: Option<Py<PyAny>>,
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<PyAny>,
extra_key: Option<String>,
cache_salt: Option<String>,
session_id: Option<String>,
prev_prefix_len: usize,
swa_evicted_seqlen: usize,
swa_branching_seqlen: Option<usize>,
@@ -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<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
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<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
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<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
block_size,
medium,
cache_salt,
session_id,
} => {
let item: Py<PyAny> = (
"block_stored",
@@ -1837,6 +1843,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
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)?;
rust/sglang-radix-tree/src/tests/components/base.rs
+1
View File
@@ -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(),
rust/sglang-radix-tree/src/tests/components/full.rs
+2
View File
@@ -75,6 +75,7 @@ fn insert(tc: &mut UnifiedTreeCore<Vec<i64>>, key: &Vec<i64>, 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());
rust/sglang-radix-tree/src/tests/components/mamba.rs
+1
View File
@@ -68,6 +68,7 @@ fn insert_params_mamba<'k>(
swa_branching_seqlen: None,
chunked: false,
priority: 0,
session_id: None,
track_adopted_ranges: false,
}
}
rust/sglang-radix-tree/src/tests/components/swa.rs
+1
View File
@@ -552,6 +552,7 @@ fn insert_params_swa<'k>(
swa_branching_seqlen: None,
chunked: false,
priority: 0,
session_id: None,
track_adopted_ranges: false,
}
}
rust/sglang-radix-tree/src/tests/unified_tree_core.rs
+50 -2
View File
@@ -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<i64>) -> MatchPrefixParams<'_, Vec<i64>> {
@@ -1866,6 +1870,7 @@ fn insert_params<'k>(key: &'k Vec<i64>, value: &[i64]) -> InsertParams<'k, Vec<i
swa_branching_seqlen: None,
chunked: false,
priority: 0,
session_id: None,
track_adopted_ranges: false,
}
}
@@ -2743,6 +2748,7 @@ fn insert_coalesces_parent_linked_block_stores() {
block_size: 2,
medium: StorageMedium::Gpu,
cache_salt: None,
session_id: None,
}]
);
// Events hash lazily even though the storage tier is off.
@@ -2758,6 +2764,32 @@ fn insert_coalesces_parent_linked_block_stores() {
assert!(tc.namespaced_event_hashes.is_empty());
}
#[test]
fn insert_attributes_stored_blocks_to_session_without_changing_hashes() {
let mut tc = events_core(2);
let key = vec![1, 2, 7, 8];
let mut params = insert_params(&key, &[10, 11, 12, 13]);
params.session_id = Some("session-a");
tc.insert(&params);
let hashes = crate::node::get_hash_str::<Vec<i64>>(&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::<Vec<(i64, i64)>>(&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,
}
}
rust/sglang-radix-tree/src/unified_tree_core.rs
+38 -7
View File
@@ -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<K: ChildKeyType> {
aligned_key_len: usize,
value: Tensor,
namespace: KeyNamespace,
session_id: Option<Arc<str>>,
prev_prefix_len: usize,
swa_evicted_seqlen: usize,
swa_branching_seqlen: Option<usize>,
@@ -1422,6 +1426,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
&leaf_value,
state.priority,
state.namespace.as_ref(),
state.session_id.as_deref(),
)
} else {
state.node_id
@@ -1698,6 +1706,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
value,
priority,
KeyNamespaceRef::new(extra_key, /* cache_salt = */ None),
/* session_id = */ None,
)
}
@@ -1898,6 +1908,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
}
/// 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<K: ChildKeyType> UnifiedTreeCore<K> {
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<Arc<str>> = 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<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
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<K: ChildKeyType> UnifiedTreeCore<K> {
/* 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<K: ChildKeyType> UnifiedTreeCore<K> {
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<A> {
block_size: usize,
medium: StorageMedium,
cache_salt: Option<Arc<str>>,
session_id: Option<Arc<str>>,
},
BlockRemoved {
block_hashes: Vec<i64>,
test/registered/unit/disaggregation/test_kv_events.py
+46 -21
View File
@@ -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__":
test/registered/unit/mem_cache/test_page_interleave_shard.py
+1
View File
@@ -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)
test/registered/unit/mem_cache/test_radix_cache_unit.py
+10 -10
View File
@@ -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)
test/registered/unit/mem_cache/test_rust_tree_core_integration.py
+4 -6
View File
@@ -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",
)
]
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+79 -7
View File
@@ -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: