feat(kv-router): add worker_type to router selection log (#7258)

Signed-off-by: PeaBrane <yanrpei@gmail.com>

lib/kv-router/src/scheduling/queue.rs

@@ -272,7 +272,7 @@ mod tests {
         let (cfg_tx, cfg_rx) = watch::channel(configs);
         std::mem::forget(cfg_tx);
 
-        let selector = Box::new(DefaultWorkerSelector::default());
+        let selector = Box::new(DefaultWorkerSelector::new(None, "test"));
         let queue = Arc::new(SchedulerQueue::new(
             Arc::clone(&slots),
             cfg_rx,

lib/kv-router/src/scheduling/selector.rs

@@ -79,15 +79,17 @@ fn softmax_sample(
 }
 
 /// Default implementation matching the Python _cost_function.
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 pub struct DefaultWorkerSelector {
     pub kv_router_config: KvRouterConfig,
+    pub worker_type: &'static str,
 }
 
 impl DefaultWorkerSelector {
-    pub fn new(kv_router_config: Option<KvRouterConfig>) -> Self {
+    pub fn new(kv_router_config: Option<KvRouterConfig>, worker_type: &'static str) -> Self {
         Self {
             kv_router_config: kv_router_config.unwrap_or_default(),
+            worker_type,
         }
     }
 }
@@ -187,6 +189,21 @@ impl<C: WorkerConfigLike> WorkerSelector<C> for DefaultWorkerSelector {
 
         let best_logit = worker_logits[&best_worker];
 
+        if self.worker_type == "decode" {
+            tracing::info!(
+                "Selected worker: worker_type={}, worker_id={} dp_rank={:?}, logit: {:.3}",
+                self.worker_type,
+                best_worker.worker_id,
+                best_worker.dp_rank,
+                best_logit,
+            );
+            return Ok(WorkerSelectionResult {
+                worker: best_worker,
+                required_blocks: request_blocks as u64,
+                overlap_blocks: overlaps.get(&best_worker).copied().unwrap_or(0),
+            });
+        }
+
         let best_overlap = *overlaps.get(&best_worker).unwrap_or(&0);
 
         let total_blocks_info = workers
@@ -203,7 +220,8 @@ impl<C: WorkerConfigLike> WorkerSelector<C> for DefaultWorkerSelector {
             .unwrap_or(0);
 
         tracing::info!(
-            "Selected worker: worker_id={} dp_rank={:?}, logit: {:.3}, cached blocks: {}, tree size: {}{}",
+            "Selected worker: worker_type={}, worker_id={} dp_rank={:?}, logit: {:.3}, cached blocks: {}, tree size: {}{}",
+            self.worker_type,
             best_worker.worker_id,
             best_worker.dp_rank,
             best_logit,

lib/llm/src/discovery/model_manager.rs

@@ -588,7 +588,7 @@ impl ModelManager {
         // Get of create runtime config watcher for this endpoint
         let workers_with_configs = self.get_or_create_runtime_config_watcher(endpoint).await?;
 
-        let selector = Box::new(DefaultWorkerSelector::new(kv_router_config));
+        let selector = Box::new(DefaultWorkerSelector::new(kv_router_config, worker_type));
         let chooser = KvRouter::new(
             endpoint.clone(),
             client,

lib/llm/src/kv_router/scheduler.rs

@@ -43,7 +43,7 @@ impl KvScheduler {
         kv_router_config: &KvRouterConfig,
         worker_type: &'static str,
     ) -> Result<Self, KvSchedulerError> {
-        let selector = selector.unwrap_or(Box::new(DefaultWorkerSelector::default()));
+        let selector = selector.unwrap_or(Box::new(DefaultWorkerSelector::new(None, worker_type)));
 
         // Get initial workers from watch receiver.
         // Caller must ensure at least one worker is present (via wait_for).

lib/mocker/src/common/protocols.rs

@@ -30,7 +30,12 @@ pub type NumBlocks = usize;
 /// For Use and Promote variants, block hashes are included for KV event publishing
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub enum MoveBlock {
-    Use(Vec<UniqueBlock>, Vec<BlockHash>, Option<Vec<Vec<u32>>>),
+    Use(
+        Vec<UniqueBlock>,
+        Vec<BlockHash>,
+        Option<Vec<Vec<u32>>>,
+        Option<UniqueBlock>,
+    ),
     Destroy(Vec<UniqueBlock>),
     Deref(Vec<UniqueBlock>),
     Promote(Uuid, SequenceHash, Option<u64>, BlockHash, Option<Vec<u32>>),

lib/mocker/src/common/sequence.rs

@@ -143,7 +143,12 @@ impl ActiveSequence {
             None
         };
 
-        Some(MoveBlock::Use(blocks, hashes, token_ids))
+        let parent = if prev_blocks > 0 {
+            Some(self.unique_blocks[prev_blocks - 1].clone())
+        } else {
+            None
+        };
+        Some(MoveBlock::Use(blocks, hashes, token_ids, parent))
     }
 
     /// Commit a successful allocation by advancing `num_allocated_tokens`.
@@ -237,7 +242,7 @@ impl ActiveSequence {
 
         let new_partial_block = UniqueBlock::default();
         self.unique_blocks.push(new_partial_block.clone());
-        signals.push(MoveBlock::Use(vec![new_partial_block], vec![], None));
+        signals.push(MoveBlock::Use(vec![new_partial_block], vec![], None, None));
         Some(signals)
     }
 

lib/mocker/src/kv_manager/vllm_backend.rs

@@ -186,12 +186,12 @@ impl KvManager {
     /// For other variants, returns the total block count (they always succeed or panic).
     pub fn process(&mut self, event: &MoveBlock) -> usize {
         match event {
-            MoveBlock::Use(hashes, local_hashes, token_ids) => {
+            MoveBlock::Use(hashes, local_hashes, token_ids, parent) => {
                 let mut blocks_stored = Vec::<u64>::new();
                 let mut stored_token_ids: Option<Vec<Vec<u32>>> =
                     token_ids.as_ref().map(|_| Vec::new());
 
-                let mut parent_block: Option<&UniqueBlock> = None;
+                let mut parent_block: Option<&UniqueBlock> = parent.as_ref();
                 let mut allocated = 0;
                 for (i, hash) in hashes.iter().enumerate() {
                     // First check if it already exists in active blocks
@@ -420,7 +420,7 @@ mod tests {
         fn use_blocks(manager: &mut KvManager, ids: Vec<u64>) -> usize {
             let blocks: Vec<_> = ids.iter().map(|&id| UniqueBlock::FullBlock(id)).collect();
             let hashes: Vec<_> = ids.into_iter().collect();
-            manager.process(&MoveBlock::Use(blocks, hashes, None))
+            manager.process(&MoveBlock::Use(blocks, hashes, None, None))
         }
 
         // First use 10 blocks (0 to 9) in a batch
@@ -447,7 +447,7 @@ mod tests {
         fn use_blocks(manager: &mut KvManager, ids: Vec<u64>) {
             let blocks: Vec<_> = ids.iter().map(|&id| UniqueBlock::FullBlock(id)).collect();
             let hashes: Vec<_> = ids.into_iter().collect();
-            manager.process(&MoveBlock::Use(blocks, hashes, None));
+            manager.process(&MoveBlock::Use(blocks, hashes, None, None));
         }
 
         // Helper function to destroy multiple blocks
@@ -561,4 +561,72 @@ mod tests {
 
         use_blocks(&mut manager, vec![13]);
     }
+
+    #[test]
+    fn test_chunked_prefill_parent_hash() {
+        use std::sync::Mutex;
+
+        use crate::common::sequence::ActiveSequence;
+
+        #[derive(Default)]
+        struct CapturingSink {
+            events: Mutex<Vec<KvCacheEvent>>,
+        }
+
+        impl KvCacheEventSink for CapturingSink {
+            fn publish(
+                &self,
+                event: KvCacheEvent,
+                _block_token_ids: Option<&[Vec<u32>]>,
+            ) -> anyhow::Result<()> {
+                self.events.lock().unwrap().push(event);
+                Ok(())
+            }
+        }
+
+        let block_size = 64;
+        let tokens: Vec<u32> = (0..512).collect(); // 8 blocks
+        let mut seq = ActiveSequence::new(tokens, 100, Some(block_size), true, false);
+
+        let sink = Arc::new(CapturingSink::default());
+        let mut manager =
+            KvManager::new_with_event_sink(256, block_size, Some(sink.clone() as _), 0);
+
+        // Chunk 1: allocate blocks 0-3
+        let signal = seq.prepare_allocation(256).unwrap();
+        manager.process(&signal);
+        seq.commit_allocation(256);
+
+        // Chunk 2: allocate blocks 4-7
+        let signal = seq.prepare_allocation(512).unwrap();
+        manager.process(&signal);
+        seq.commit_allocation(512);
+
+        let events = sink.events.lock().unwrap();
+        assert_eq!(events.len(), 2, "expected two store events");
+
+        // First event: parent_hash should be None (starts from root)
+        let KvCacheEventData::Stored(ref store1) = events[0].data else {
+            panic!("expected store event");
+        };
+        assert!(
+            store1.parent_hash.is_none(),
+            "first chunk should have no parent"
+        );
+
+        // Second event: parent_hash should be the seq_hash of block 3
+        // (the last block from the first chunk)
+        let KvCacheEventData::Stored(ref store2) = events[1].data else {
+            panic!("expected store event");
+        };
+        let expected_parent = seq.unique_blocks()[3].clone();
+        let UniqueBlock::FullBlock(expected_hash) = expected_parent else {
+            panic!("expected full block");
+        };
+        assert_eq!(
+            store2.parent_hash,
+            Some(ExternalSequenceBlockHash(expected_hash)),
+            "second chunk's parent should be block 3's seq_hash"
+        );
+    }
 }

tests/serve/test_sglang.py

@@ -139,7 +139,7 @@ sglang_configs = {
                 expected_log=[
                     r"ZMQ listener .* received batch with \d+ events \(seq=\d+(?:, [^)]*)?\)",
                     r"Event processor for worker_id \d+ processing event: Stored\(",
-                    r"Selected worker: worker_id=\d+ dp_rank=.*?, logit: ",
+                    r"Selected worker: worker_type=\w+, worker_id=\d+ dp_rank=.*?, logit: ",
                 ]
             )
         ],

tests/serve/test_trtllm.py

@@ -151,7 +151,7 @@ trtllm_configs = {
             chat_payload_default(
                 expected_log=[
                     r"Event processor for worker_id \d+ processing event: Stored\(",
-                    r"Selected worker: worker_id=\d+ dp_rank=.*?, logit: ",
+                    r"Selected worker: worker_type=\w+, worker_id=\d+ dp_rank=.*?, logit: ",
                 ]
             )
         ],

tests/serve/test_vllm.py

@@ -203,7 +203,7 @@ vllm_configs = {
                 expected_log=[
                     r"ZMQ listener .* received batch with \d+ events \(seq=\d+(?:, [^)]*)?\)",
                     r"Event processor for worker_id \d+ processing event: Stored\(",
-                    r"Selected worker: worker_id=\d+ dp_rank=.*?, logit: ",
+                    r"Selected worker: worker_type=\w+, worker_id=\d+ dp_rank=.*?, logit: ",
                 ]
             )
         ],
@@ -228,7 +228,7 @@ vllm_configs = {
                 repeat_count=3,
                 expected_log=[
                     # Verify scheduler is selecting workers with cache awareness
-                    r"Selected worker: worker_id=\d+ dp_rank=.*?, logit: ",
+                    r"Selected worker: worker_type=\w+, worker_id=\d+ dp_rank=.*?, logit: ",
                     # After first request, should see cached blocks being tracked
                     r"with \d+ cached blocks",
                 ],