fix(kv-router): filter vLLM partial-tail KV events instead of warning on block-size mismatch

lib/kv-router/src/zmq_wire/mod.rs

@@ -11,7 +11,7 @@ use std::sync::Arc;
 use std::sync::atomic::AtomicU32;
 
 use rmp_serde as rmps;
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use crate::protocols::{DpRank, PlacementEvent, StorageTier, WorkerWithDpRank};
 
@@ -55,6 +55,11 @@ pub struct ZmqEventNormalizer {
     warning_count: Arc<AtomicU32>,
     group_metadata: FxHashMap<(DpRank, u32), KvCacheGroupMetadata>,
     cache_namespaces: FxHashMap<(WorkerWithDpRank, u64), CacheNamespaceState>,
+    /// Workers that have published at least one full-size block. Once a worker
+    /// has proven its block size matches, a smaller `BlockStored` event is a
+    /// partial tail (vLLM `--prefix-match-unit` below the block size), not a
+    /// misconfiguration, and is filtered quietly instead of warned about.
+    full_block_workers: FxHashSet<WorkerWithDpRank>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -81,6 +86,7 @@ pub enum ZmqEventFilterReason {
     UnknownKind,
     NonMainAttentionGroup,
     UnlearnedGroupIdx,
+    PartialBlock,
 }
 
 impl ZmqEventFilterReason {
@@ -96,6 +102,7 @@ impl ZmqEventFilterReason {
             Self::UnknownKind => "unknown_kind",
             Self::NonMainAttentionGroup => "non_main_attention_group",
             Self::UnlearnedGroupIdx => "unlearned_group_idx",
+            Self::PartialBlock => "partial_block",
         }
     }
 }
@@ -109,6 +116,7 @@ impl ZmqEventNormalizer {
             warning_count: Arc::new(AtomicU32::new(0)),
             group_metadata: FxHashMap::default(),
             cache_namespaces: FxHashMap::default(),
+            full_block_workers: FxHashSet::default(),
         }
     }
 
@@ -120,6 +128,7 @@ impl ZmqEventNormalizer {
             warning_count,
             group_metadata: FxHashMap::default(),
             cache_namespaces: FxHashMap::default(),
+            full_block_workers: FxHashSet::default(),
         }
     }
 
@@ -216,6 +225,22 @@ impl ZmqEventNormalizer {
             Some(reason) => return Err(reason),
             None => {}
         }
+        // Only blocks that survived group filtering count toward a worker's
+        // full-block history, so a filtered (non-main-attention) group cannot
+        // suppress a size-mismatch warning for the routed group. A clear resets
+        // the history: the engine behind this worker may come back with a
+        // different block size.
+        match &raw {
+            RawKvEvent::BlockStored { block_size, .. } => {
+                if let Some(reason) = self.filter_partial_block(*block_size, worker) {
+                    return Err(reason);
+                }
+            }
+            RawKvEvent::AllBlocksCleared { .. } => {
+                self.full_block_workers.remove(&worker);
+            }
+            _ => {}
+        }
         if !is_lower_tier {
             self.propagate_cache_namespace(&mut raw, worker)?;
         }
@@ -249,6 +274,41 @@ impl ZmqEventNormalizer {
         self.normalize_preprocessed(raw, event_id, worker)
     }
 
+    /// Filter partial-tail `BlockStored` events from engines whose hash unit is
+    /// smaller than the KV block size (vLLM `--prefix-match-unit`).
+    ///
+    /// vLLM publishes full blocks at the engine block size and the trailing
+    /// partial block of a prompt at the hash-unit size. The router can only
+    /// index full blocks, so conversion drops the partial one, but it does so
+    /// with a block-size-mismatch warning that users read as a misconfiguration.
+    /// Once a worker has published a full-size block its block size is known to
+    /// match, so a smaller event from that worker is a partial tail and is
+    /// filtered here at debug level. A mismatch on the first event from a worker
+    /// is left to conversion, which still warns, so a real block-size
+    /// misconfiguration stays visible.
+    fn filter_partial_block(
+        &mut self,
+        block_size: usize,
+        worker: WorkerWithDpRank,
+    ) -> Option<ZmqEventFilterReason> {
+        let kv_block_size = self.kv_block_size as usize;
+        if block_size == kv_block_size {
+            self.full_block_workers.insert(worker);
+            return None;
+        }
+        if block_size < kv_block_size && self.full_block_workers.contains(&worker) {
+            tracing::debug!(
+                block_size,
+                kv_block_size,
+                worker_id = worker.worker_id,
+                dp_rank = worker.dp_rank,
+                "Filtering partial-tail KV block event (hash unit smaller than block size)"
+            );
+            return Some(ZmqEventFilterReason::PartialBlock);
+        }
+        None
+    }
+
     fn learn_metadata(&mut self, metadata: KvCacheEventMetadata, dp_rank: DpRank) {
         let (Some(group_idx), Some(kind)) = (metadata.group_idx, metadata.kv_cache_spec_kind)
         else {

lib/kv-router/src/zmq_wire/tests.rs

@@ -1180,6 +1180,152 @@ fn cpu_event_with_full_payload_is_indexable() {
     assert_eq!(warning_count.load(Ordering::Relaxed), 0);
 }
 
+// ---- partial-tail block filtering (vLLM --prefix-match-unit) ----
+
+fn gpu_block_stored(block_hashes: &[u64], token_ids: &[u32], block_size: usize) -> RawKvEvent {
+    RawKvEvent::BlockStored {
+        block_hashes: block_hashes
+            .iter()
+            .copied()
+            .map(BlockHashValue::Unsigned)
+            .collect(),
+        parent_block_hash: None,
+        token_ids: token_ids.to_vec(),
+        block_size,
+        medium: None,
+        lora_name: None,
+        cache_namespace: None,
+        block_mm_infos: None,
+        is_eagle: None,
+        group_idx: None,
+        kv_cache_spec_kind: None,
+        kv_cache_spec_sliding_window: None,
+        locality: None,
+        ownership: None,
+        session_id: None,
+    }
+}
+
+/// vLLM with `--prefix-match-unit` below the block size publishes full blocks
+/// at the block size and the trailing partial block at the hash-unit size.
+/// Once a worker has published a full block, the undersized event is a partial
+/// tail and is filtered with its own reason instead of reaching conversion,
+/// where it would trigger the block-size-mismatch warning.
+#[test]
+fn partial_tail_block_is_filtered_after_a_full_block_from_the_same_worker() {
+    let mut normalizer = ZmqEventNormalizer::new(4);
+    let worker = WorkerWithDpRank::new(7, 0);
+
+    let full = gpu_block_stored(&[1], &[10, 11, 12, 13], 4);
+    assert!(normalizer.preprocess_with_reason(full, worker).is_ok());
+
+    let partial_tail = gpu_block_stored(&[2], &[14, 15], 2);
+    assert_eq!(
+        normalizer
+            .preprocess_with_reason(partial_tail, worker)
+            .err(),
+        Some(ZmqEventFilterReason::PartialBlock)
+    );
+
+    let next_full = gpu_block_stored(&[3], &[16, 17, 18, 19], 4);
+    assert!(normalizer.preprocess_with_reason(next_full, worker).is_ok());
+}
+
+/// An undersized block from a worker that has never published a full block is
+/// not a partial tail: it may be a real block-size misconfiguration, so it still
+/// reaches conversion and the existing warning fires.
+#[test]
+fn undersized_block_before_any_full_block_still_reaches_conversion_and_warns() {
+    let mut normalizer = ZmqEventNormalizer::new(4);
+    let worker = WorkerWithDpRank::new(7, 0);
+
+    let undersized = gpu_block_stored(&[2], &[14, 15], 2);
+    let raw = normalizer
+        .preprocess_with_reason(undersized, worker)
+        .expect("first undersized block must not be filtered as a partial tail");
+
+    let warning_count = Arc::new(AtomicU32::new(0));
+    let placement = convert_event(raw, 1, 4, worker, &warning_count, None, None).unwrap();
+    match placement.event.data {
+        KvCacheEventData::Stored(store_data) => assert!(store_data.blocks.is_empty()),
+        other => panic!("expected Stored event, got {other:?}"),
+    }
+    assert_eq!(warning_count.load(Ordering::Relaxed), 1);
+}
+
+/// The full-block proof is per worker: another worker's undersized event is not
+/// filtered just because a different worker has published full blocks.
+#[test]
+fn partial_tail_filter_is_tracked_per_worker() {
+    let mut normalizer = ZmqEventNormalizer::new(4);
+    let proven = WorkerWithDpRank::new(7, 0);
+    let unproven = WorkerWithDpRank::new(8, 0);
+
+    let full = gpu_block_stored(&[1], &[10, 11, 12, 13], 4);
+    assert!(normalizer.preprocess_with_reason(full, proven).is_ok());
+
+    let undersized = gpu_block_stored(&[2], &[14, 15], 2);
+    assert!(
+        normalizer
+            .preprocess_with_reason(undersized, unproven)
+            .is_ok()
+    );
+}
+
+/// A full-size block from a group the router filters out (here a Mamba state
+/// group) must not count as proof of the worker's block size: a later undersized
+/// main-attention block still has to reach conversion and warn.
+#[test]
+fn filtered_group_does_not_prove_block_size_for_partial_tail_filter() {
+    let mut normalizer = ZmqEventNormalizer::new(4);
+    let worker = WorkerWithDpRank::new(7, 0);
+
+    let mut full_non_main = gpu_block_stored(&[1], &[10, 11, 12, 13], 4);
+    if let RawKvEvent::BlockStored {
+        kv_cache_spec_kind,
+        group_idx,
+        ..
+    } = &mut full_non_main
+    {
+        *kv_cache_spec_kind = Some(KvCacheSpecKind::Mamba);
+        *group_idx = Some(1);
+    }
+    assert_eq!(
+        normalizer
+            .preprocess_with_reason(full_non_main, worker)
+            .err(),
+        Some(ZmqEventFilterReason::NonMainAttentionKind)
+    );
+
+    let undersized = gpu_block_stored(&[2], &[14, 15], 2);
+    assert!(
+        normalizer
+            .preprocess_with_reason(undersized, worker)
+            .is_ok()
+    );
+}
+
+/// `AllBlocksCleared` resets the worker's full-block history, so an engine
+/// that comes back with a different block size is not silently filtered.
+#[test]
+fn all_blocks_cleared_resets_partial_tail_filter_history() {
+    let mut normalizer = ZmqEventNormalizer::new(4);
+    let worker = WorkerWithDpRank::new(7, 0);
+
+    let full = gpu_block_stored(&[1], &[10, 11, 12, 13], 4);
+    assert!(normalizer.preprocess_with_reason(full, worker).is_ok());
+
+    let cleared = RawKvEvent::AllBlocksCleared { ownership: None };
+    assert!(normalizer.preprocess_with_reason(cleared, worker).is_ok());
+
+    let undersized = gpu_block_stored(&[2], &[14, 15], 2);
+    assert!(
+        normalizer
+            .preprocess_with_reason(undersized, worker)
+            .is_ok()
+    );
+}
+
 // ---- locality decode + storage-tier gating (this PR) ----
 
 /// LOCAL/REMOTE decode to their variants, any other string folds into