diff --git a/components/src/dynamo/vllm/tests/test_vllm_kv_events_api.py b/components/src/dynamo/vllm/tests/test_vllm_kv_events_api.py index 9863f2f065fc..4e2d4a6b800c 100755 --- a/components/src/dynamo/vllm/tests/test_vllm_kv_events_api.py +++ b/components/src/dynamo/vllm/tests/test_vllm_kv_events_api.py @@ -46,6 +46,10 @@ ] +def _has_group_idx(event_cls): + return "group_idx" in event_cls.__struct_fields__ + + class TestVllmKvEventsApi: """Test vLLM KV events API compatibility.""" @@ -61,10 +65,11 @@ def test_block_stored_fields(self): 6. medium 7. lora_name (added in vLLM 0.14.0) 8. extra_keys (added in vLLM 0.17.0) + 9. group_idx (added for hybrid KV cache groups; optional for older vLLM) If vLLM adds/removes/reorders fields, this test will fail. """ - expected_fields = ( + expected_fields = [ "block_hashes", "parent_block_hash", "token_ids", @@ -73,7 +78,10 @@ def test_block_stored_fields(self): "medium", "lora_name", "extra_keys", - ) + ] + if _has_group_idx(BlockStored): + expected_fields.append("group_idx") + expected_fields = tuple(expected_fields) actual_fields = BlockStored.__struct_fields__ assert actual_fields == expected_fields, ( @@ -88,10 +96,13 @@ def test_block_stored_fields(self): def test_block_removed_fields(self): """Verify BlockRemoved has expected fields in expected order.""" - expected_fields = ( + expected_fields = [ "block_hashes", "medium", - ) + ] + if _has_group_idx(BlockRemoved): + expected_fields.append("group_idx") + expected_fields = tuple(expected_fields) actual_fields = BlockRemoved.__struct_fields__ assert actual_fields == expected_fields, ( @@ -158,16 +169,19 @@ def test_block_stored_serialization_format(self): """ import msgspec - event = BlockStored( - block_hashes=[123, 456], - parent_block_hash=789, - token_ids=[1, 2, 3, 4], - block_size=16, - lora_id=None, - medium="GPU", - lora_name=None, - extra_keys=None, - ) + event_kwargs = { + "block_hashes": [123, 456], + "parent_block_hash": 789, + "token_ids": [1, 2, 3, 4], + "block_size": 16, + "lora_id": None, + "medium": "GPU", + "lora_name": None, + "extra_keys": None, + } + if _has_group_idx(BlockStored): + event_kwargs["group_idx"] = 0 + event = BlockStored(**event_kwargs) encoded = msgspec.msgpack.encode(event) decoded = msgspec.msgpack.decode(encoded) @@ -178,9 +192,9 @@ def test_block_stored_serialization_format(self): decoded[0] == "BlockStored" ), f"Expected tag 'BlockStored', got {decoded[0]}" - # Verify field count (tag + 8 fields = 9 elements) - assert len(decoded) == 9, ( - f"Expected 9 elements (tag + 8 fields), got {len(decoded)}.\n" + expected_len = 10 if _has_group_idx(BlockStored) else 9 + assert len(decoded) == expected_len, ( + f"Expected {expected_len} elements, got {len(decoded)}.\n" f"Decoded: {decoded}\n" f"If field count changed, update Rust deserializers." ) @@ -194,22 +208,27 @@ def test_block_stored_serialization_format(self): assert decoded[6] == "GPU", f"medium at wrong position: {decoded[6]}" assert decoded[7] is None, f"lora_name at wrong position: {decoded[7]}" assert decoded[8] is None, f"extra_keys at wrong position: {decoded[8]}" + if _has_group_idx(BlockStored): + assert decoded[9] == 0, f"group_idx at wrong position: {decoded[9]}" def test_block_stored_tuple_extra_keys_serialization_format(self): """Verify multimodal tuple extra_keys keep the vLLM 0.19 wire shape.""" import msgspec mm_hash = "0123456789abcdef00112233445566778899aabbccddeefffedcba9876543210" - event = BlockStored( - block_hashes=[123], - parent_block_hash=None, - token_ids=[1, 2, 3, 4], - block_size=16, - lora_id=None, - medium="GPU", - lora_name=None, - extra_keys=[((mm_hash, 7),)], - ) + event_kwargs = { + "block_hashes": [123], + "parent_block_hash": None, + "token_ids": [1, 2, 3, 4], + "block_size": 16, + "lora_id": None, + "medium": "GPU", + "lora_name": None, + "extra_keys": [((mm_hash, 7),)], + } + if _has_group_idx(BlockStored): + event_kwargs["group_idx"] = 0 + event = BlockStored(**event_kwargs) decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event)) @@ -218,3 +237,31 @@ def test_block_stored_tuple_extra_keys_serialization_format(self): "vLLM multimodal extra_keys no longer serialize as nested tuple/list " f"payloads. Decoded: {decoded[8]!r}" ) + if _has_group_idx(BlockStored): + assert decoded[9] == 0, f"group_idx at wrong position: {decoded[9]}" + + def test_block_removed_serialization_format(self): + """Verify BlockRemoved serializes to expected msgpack array format.""" + import msgspec + + event_kwargs = { + "block_hashes": [123, 456], + "medium": "GPU", + } + if _has_group_idx(BlockRemoved): + event_kwargs["group_idx"] = 0 + event = BlockRemoved(**event_kwargs) + + decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event)) + + assert decoded[0] == "BlockRemoved" + expected_len = 4 if _has_group_idx(BlockRemoved) else 3 + assert len(decoded) == expected_len, ( + f"Expected {expected_len} elements, got {len(decoded)}.\n" + f"Decoded: {decoded}\n" + f"If field count changed, update Rust deserializers." + ) + assert decoded[1] == [123, 456], f"block_hashes at wrong position: {decoded[1]}" + assert decoded[2] == "GPU", f"medium at wrong position: {decoded[2]}" + if _has_group_idx(BlockRemoved): + assert decoded[3] == 0, f"group_idx at wrong position: {decoded[3]}" diff --git a/lib/kv-router/src/standalone_indexer/listener.rs b/lib/kv-router/src/standalone_indexer/listener.rs index e0f279c99581..6ef5d2b60c99 100644 --- a/lib/kv-router/src/standalone_indexer/listener.rs +++ b/lib/kv-router/src/standalone_indexer/listener.rs @@ -154,13 +154,15 @@ impl ListenerLoop { .data_parallel_rank .map_or(dp_rank, |rank| rank.cast_unsigned()); for raw_event in batch.events { - let placement_event = convert_event( + let Some(placement_event) = convert_event( raw_event, seq, block_size, WorkerWithDpRank::new(worker_id, effective_dp_rank), warning_count, - ); + ) else { + continue; + }; if !placement_event.placement.is_local_gpu() { continue; } @@ -223,13 +225,15 @@ impl ListenerLoop { .data_parallel_rank .map_or(self.dp_rank, |rank| rank.cast_unsigned()); for raw_event in batch.events { - let placement_event = convert_event( + let Some(placement_event) = convert_event( raw_event, seq, self.block_size, WorkerWithDpRank::new(self.worker_id, effective_dp_rank), &self.warning_count, - ); + ) else { + continue; + }; if !placement_event.placement.is_local_gpu() { continue; } diff --git a/lib/kv-router/src/zmq_wire.rs b/lib/kv-router/src/zmq_wire.rs index eab56a65d526..f0e5d363db79 100644 --- a/lib/kv-router/src/zmq_wire.rs +++ b/lib/kv-router/src/zmq_wire.rs @@ -99,6 +99,7 @@ pub enum RawKvEvent { medium: Option, }, AllBlocksCleared, + Ignored, } /// Parse MM hash from extra_keys string: @@ -126,6 +127,17 @@ pub enum ExtraKeyItem { Bool(bool), } +#[derive(Debug, Deserialize)] +#[serde(untagged)] +enum BlockStoredTrailingField { + GroupIdx(u32), + BlockMmInfos(Vec>), +} + +fn is_non_main_group(group_idx: Option) -> bool { + matches!(group_idx, Some(group_idx) if group_idx != 0) +} + /// Convert vLLM BlockStored extra_keys to block-level MM infos. /// extra_keys is a list aligned with blocks: /// - None => no MM content in that block @@ -224,6 +236,7 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { let mut lora_name: Option> = None; let mut extra_keys: Option>>>> = None; let mut block_mm_infos: Option>>> = None; + let mut group_idx: Option> = None; while let Some(key) = map.next_key::()? { match key.as_str() { @@ -254,6 +267,9 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { "block_mm_infos" => { block_mm_infos = Some(map.next_value()?); } + "group_idx" => { + group_idx = Some(map.next_value()?); + } _ => { map.next_value::()?; } @@ -278,6 +294,10 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { }; let block_size = block_size.ok_or_else(|| de::Error::missing_field("block_size"))?; + let medium = medium.unwrap_or(None); + if is_non_main_group(group_idx.unwrap_or(None)) { + return Ok(RawKvEvent::Ignored); + } let block_mm_infos = block_mm_infos .unwrap_or(None) .or_else(|| extra_keys_to_block_mm_infos(extra_keys.unwrap_or(None))); @@ -286,7 +306,7 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { parent_block_hash: parent_block_hash.unwrap_or(None), token_ids: raw_token_ids, block_size, - medium: medium.unwrap_or(None), + medium, lora_name: lora_name.unwrap_or(None), block_mm_infos, is_eagle: Some(is_eagle), @@ -295,15 +315,20 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { Some("BlockRemoved") => { let block_hashes = block_hashes.ok_or_else(|| de::Error::missing_field("block_hashes"))?; + let medium = medium.unwrap_or(None); + if is_non_main_group(group_idx.unwrap_or(None)) { + return Ok(RawKvEvent::Ignored); + } Ok(RawKvEvent::BlockRemoved { block_hashes, - medium: medium.unwrap_or(None), + medium, }) } Some("AllBlocksCleared") => Ok(RawKvEvent::AllBlocksCleared), + Some("Ignored") => Ok(RawKvEvent::Ignored), Some(other) => Err(de::Error::unknown_variant( other, - &["BlockStored", "BlockRemoved", "AllBlocksCleared"], + &["BlockStored", "BlockRemoved", "AllBlocksCleared", "Ignored"], )), None => Err(de::Error::missing_field("type")), } @@ -339,11 +364,29 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { let lora_name: Option = seq.next_element()?.unwrap_or(None); let extra_keys: Option>>> = seq.next_element()?.unwrap_or(None); - let block_mm_infos: Option>> = - seq.next_element()?.unwrap_or(None); + let mut block_mm_infos: Option>> = None; + let mut group_idx: Option = None; + + for _ in 0..2 { + let trailing: Option = + seq.next_element()?.unwrap_or(None); + match trailing { + Some(BlockStoredTrailingField::GroupIdx(idx)) => { + group_idx = Some(idx); + } + Some(BlockStoredTrailingField::BlockMmInfos(infos)) => { + block_mm_infos = Some(infos); + } + None => {} + } + } while seq.next_element::()?.is_some() {} + if is_non_main_group(group_idx) { + return Ok(RawKvEvent::Ignored); + } + let block_mm_infos = block_mm_infos.or_else(|| extra_keys_to_block_mm_infos(extra_keys)); @@ -375,9 +418,14 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { .next_element()? .ok_or_else(|| de::Error::invalid_length(1, &"missing block_hashes"))?; let medium: Option = seq.next_element()?.unwrap_or(None); + let group_idx: Option = seq.next_element()?.unwrap_or(None); while seq.next_element::()?.is_some() {} + if is_non_main_group(group_idx) { + return Ok(RawKvEvent::Ignored); + } + Ok(RawKvEvent::BlockRemoved { block_hashes, medium, @@ -387,9 +435,13 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { while seq.next_element::()?.is_some() {} Ok(RawKvEvent::AllBlocksCleared) } + "Ignored" => { + while seq.next_element::()?.is_some() {} + Ok(RawKvEvent::Ignored) + } other => Err(de::Error::unknown_variant( other, - &["BlockStored", "BlockRemoved", "AllBlocksCleared"], + &["BlockStored", "BlockRemoved", "AllBlocksCleared", "Ignored"], )), } } @@ -406,12 +458,13 @@ pub fn convert_event( kv_block_size: u32, worker: WorkerWithDpRank, warning_count: &Arc, -) -> PlacementEvent { +) -> Option { let storage_tier = match &raw { RawKvEvent::BlockStored { medium, .. } | RawKvEvent::BlockRemoved { medium, .. } => { StorageTier::from_kv_medium_or_default(medium.as_deref()) } RawKvEvent::AllBlocksCleared => StorageTier::Device, + RawKvEvent::Ignored => return None, }; let dp_rank = worker.dp_rank; let event = match raw { @@ -440,7 +493,7 @@ pub fn convert_event( // Return an empty Removed instead of Cleared to avoid nuking // the worker's entire index state. An empty Removed is a no-op // in the radix tree (zero iterations, returns Ok(())). - return PlacementEvent::new( + return Some(PlacementEvent::new( Placement::local_worker(worker.worker_id, worker.dp_rank, storage_tier), KvCacheEvent { event_id, @@ -449,7 +502,7 @@ pub fn convert_event( }), dp_rank, }, - ); + )); } } @@ -498,12 +551,13 @@ pub fn convert_event( data: KvCacheEventData::Cleared, dp_rank, }, + RawKvEvent::Ignored => unreachable!("ignored events return before conversion"), }; - PlacementEvent::new( + Some(PlacementEvent::new( Placement::local_worker(worker.worker_id, worker.dp_rank, storage_tier), event, - ) + )) } pub fn create_stored_block_from_parts( @@ -607,9 +661,16 @@ mod tests { use std::sync::atomic::AtomicU32; use rmp_serde::{from_slice, to_vec}; + use rstest::rstest; use super::*; + #[derive(Clone, Copy, Debug)] + enum TestEventKind { + BlockStored, + BlockRemoved, + } + #[test] fn test_deserialize_bigram_block_stored_sequence() { let raw_event = ( @@ -640,6 +701,97 @@ mod tests { } } + fn block_stored_sequence_with_group_idx(group_idx: Option) -> Vec { + match group_idx { + Some(group_idx) => to_vec(&( + "BlockStored", + vec![BlockHashValue::Unsigned(11)], + Option::::None, + vec![10u32, 11], + 2usize, + Option::::None, + Option::::None, + Option::::None, + Option::::None, + group_idx, + )) + .unwrap(), + None => to_vec(&( + "BlockStored", + vec![BlockHashValue::Unsigned(11)], + Option::::None, + vec![10u32, 11], + 2usize, + Option::::None, + Option::::None, + Option::::None, + )) + .unwrap(), + } + } + + fn block_removed_sequence_with_group_idx(group_idx: Option) -> Vec { + match group_idx { + Some(group_idx) => to_vec(&( + "BlockRemoved", + vec![BlockHashValue::Unsigned(11)], + Option::::None, + group_idx, + )) + .unwrap(), + None => to_vec(&( + "BlockRemoved", + vec![BlockHashValue::Unsigned(11)], + Option::::None, + )) + .unwrap(), + } + } + + fn sequence_with_group_idx(event_kind: TestEventKind, group_idx: Option) -> Vec { + match event_kind { + TestEventKind::BlockStored => block_stored_sequence_with_group_idx(group_idx), + TestEventKind::BlockRemoved => block_removed_sequence_with_group_idx(group_idx), + } + } + + fn assert_parsed_event_kind(event: RawKvEvent, expected_kind: TestEventKind) { + match (event, expected_kind) { + (RawKvEvent::BlockStored { .. }, TestEventKind::BlockStored) + | (RawKvEvent::BlockRemoved { .. }, TestEventKind::BlockRemoved) => {} + (event, expected_kind) => { + panic!("expected {expected_kind:?}, got {event:?}"); + } + } + } + + #[rstest] + #[case(TestEventKind::BlockStored)] + #[case(TestEventKind::BlockRemoved)] + fn test_deserialize_sequence_accepts_main_group_idx(#[case] event_kind: TestEventKind) { + let event: RawKvEvent = from_slice(&sequence_with_group_idx(event_kind, Some(0))).unwrap(); + + assert_parsed_event_kind(event, event_kind); + } + + #[rstest] + #[case(TestEventKind::BlockStored)] + #[case(TestEventKind::BlockRemoved)] + fn test_deserialize_sequence_ignores_non_main_group_idx(#[case] event_kind: TestEventKind) { + let event: RawKvEvent = from_slice(&sequence_with_group_idx(event_kind, Some(1))).unwrap(); + + assert!(matches!(event, RawKvEvent::Ignored)); + } + + #[rstest] + #[case(TestEventKind::BlockStored)] + #[case(TestEventKind::BlockRemoved)] + fn test_deserialize_sequence_accepts_missing_group_idx(#[case] event_kind: TestEventKind) { + let event: RawKvEvent = from_slice(&sequence_with_group_idx(event_kind, None)).unwrap(); + + assert_parsed_event_kind(event, event_kind); + } + #[test] fn test_convert_event_bigram_emits_eagle_windows() { let raw_event = RawKvEvent::BlockStored { @@ -656,7 +808,7 @@ mod tests { let placement_event = convert_event(raw_event, 7, 2, WorkerWithDpRank::new(3, 0), &warning_count); - match placement_event.event.data { + match placement_event.unwrap().event.data { KvCacheEventData::Stored(store_data) => { assert_eq!(store_data.blocks.len(), 2); assert_eq!( diff --git a/lib/llm/src/block_manager/kv_consolidator/subscriber.rs b/lib/llm/src/block_manager/kv_consolidator/subscriber.rs index b515eeee4973..2f7a06285b24 100644 --- a/lib/llm/src/block_manager/kv_consolidator/subscriber.rs +++ b/lib/llm/src/block_manager/kv_consolidator/subscriber.rs @@ -246,5 +246,7 @@ fn process_event( tracing::debug!("Processing AllBlocksCleared"); tracker.handle_clear_all(); } + + RawKvEvent::Ignored => {} } } diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 745b7b92d04d..14255c746697 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -112,7 +112,8 @@ mod test_event_processing { kv_block_size, WorkerWithDpRank::from_worker_id(1), &Arc::new(AtomicU32::new(0)), - ); + ) + .unwrap(); assert!(matches!(out.event.data, KvCacheEventData::Stored(_))); } @@ -149,14 +150,16 @@ mod test_event_processing { kv_block_size, WorkerWithDpRank::from_worker_id(1), &wc, - ); + ) + .unwrap(); let lora_out = convert_event( lora_evt, 2, kv_block_size, WorkerWithDpRank::from_worker_id(1), &wc, - ); + ) + .unwrap(); let base_hash = match &base_out.event.data { KvCacheEventData::Stored(s) => s.blocks[0].tokens_hash, @@ -205,14 +208,16 @@ mod test_event_processing { kv_block_size, WorkerWithDpRank::from_worker_id(1), &wc, - ); + ) + .unwrap(); let out2 = convert_event( evt2, 2, kv_block_size, WorkerWithDpRank::from_worker_id(1), &wc, - ); + ) + .unwrap(); let hash1 = match &out1.event.data { KvCacheEventData::Stored(s) => s.blocks[0].tokens_hash, @@ -297,7 +302,8 @@ mod test_event_processing { kv_block_size, WorkerWithDpRank::from_worker_id(1), &Arc::new(AtomicU32::new(0)), - ); + ) + .unwrap(); assert!(matches!(out.event.data, KvCacheEventData::Removed(_))); } @@ -312,7 +318,8 @@ mod test_event_processing { kv_block_size, WorkerWithDpRank::from_worker_id(1), &Arc::new(AtomicU32::new(0)), - ); + ) + .unwrap(); assert!(matches!(out.event.data, KvCacheEventData::Cleared)); } diff --git a/lib/llm/src/kv_router/publisher/zmq_listener.rs b/lib/llm/src/kv_router/publisher/zmq_listener.rs index 86620c3a0836..8769cf23cca5 100644 --- a/lib/llm/src/kv_router/publisher/zmq_listener.rs +++ b/lib/llm/src/kv_router/publisher/zmq_listener.rs @@ -103,10 +103,16 @@ pub(super) async fn start_zmq_listener( let dp_rank = batch.data_parallel_rank.unwrap_or(0).cast_unsigned(); for raw_event in batch.events { + if matches!(raw_event, RawKvEvent::Ignored) { + continue; + } let event_id = next_event_id.fetch_add(1, Ordering::SeqCst); let worker = WorkerWithDpRank::new(worker_id, dp_rank); - let event = - convert_event(raw_event, event_id, kv_block_size, worker, &warning_count); + let Some(event) = + convert_event(raw_event, event_id, kv_block_size, worker, &warning_count) + else { + continue; + }; if tx.send(event).is_err() { tracing::warn!("Failed to send message to channel - receiver dropped"); break 'main String::from("channel receiver dropped"); diff --git a/lib/llm/src/mocker.rs b/lib/llm/src/mocker.rs index 47a75dac313b..a483b0ecf777 100644 --- a/lib/llm/src/mocker.rs +++ b/lib/llm/src/mocker.rs @@ -72,9 +72,11 @@ enum ZmqRawKvEvent { parent_block_hash: Option, token_ids: Vec, block_size: u32, + group_idx: u32, }, BlockRemoved { block_hashes: Vec, + group_idx: u32, }, } @@ -277,11 +279,15 @@ fn convert_to_zmq_events( parent_block_hash, token_ids, block_size, + group_idx: 0, }] } KvCacheEventData::Removed(remove_data) => { let block_hashes: Vec = remove_data.block_hashes.iter().map(|h| h.0).collect(); - vec![ZmqRawKvEvent::BlockRemoved { block_hashes }] + vec![ZmqRawKvEvent::BlockRemoved { + block_hashes, + group_idx: 0, + }] } KvCacheEventData::Cleared => vec![], }