Initial push

Signed-off-by: Zhongxuan Wang <[email protected]>

Author: zhongxuanwang-nv

components/src/dynamo/vllm/handlers.py

@@ -211,7 +211,8 @@ def _build_completion_usage(request_output: RequestOutput) -> Dict[str, Any]:
             ),
             "prompt_tokens_details": (
                 {"cached_tokens": request_output.num_cached_tokens}
-                if request_output.num_cached_tokens
+                if request_output.num_cached_tokens is not None
+                and request_output.num_cached_tokens >= 0
                 else None
             ),
         }
@@ -241,10 +242,10 @@ async def generate_tokens(
                     out = {"token_ids": output.token_ids[num_output_tokens_so_far:]}
                     if output.finish_reason:
                         out["finish_reason"] = output.finish_reason
-                        out[
-                            "completion_usage"
-                        ] = BaseWorkerHandler._build_completion_usage(
-                            request_output=res
+                        out["completion_usage"] = (
+                            BaseWorkerHandler._build_completion_usage(
+                                request_output=res
+                            )
                         )
                     if output.stop_reason:
                         out["stop_reason"] = output.stop_reason
@@ -349,6 +350,9 @@ async def generate(self, request, context):
         request_id = context.id()
         logger.debug(f"Prefill Request ID: {request_id}")
 
+        # Extract overlap information from router (if present)
+        overlap_blocks = request.get("estimated_prefix_hit_num_blocks", 0)
+
         # Extract and decode multimodal data if present
         multi_modal_data = await self._extract_multimodal_data(request)
 
@@ -391,13 +395,18 @@ async def generate(self, request, context):
 
                     token_ids = res.outputs[0].token_ids if res.outputs else []
 
+                    # Build disaggregated_params with KV transfer params and router overlap
+                    disaggregated_params = {}
+                    if res.kv_transfer_params:
+                        disaggregated_params["kv_transfer_params"] = (
+                            res.kv_transfer_params
+                        )
+                    # Include router's overlap calculation for PrefillRouter
+                    disaggregated_params["overlap_blocks"] = overlap_blocks
+
                     output: Dict[str, Any] = {
                         "token_ids": list(token_ids),
-                        "disaggregated_params": (
-                            {"kv_transfer_params": res.kv_transfer_params}
-                            if res.kv_transfer_params
-                            else None
-                        ),
+                        "disaggregated_params": disaggregated_params,
                         "completion_usage": BaseWorkerHandler._build_completion_usage(
                             request_output=res
                         ),

lib/llm/src/entrypoint/input/common.rs

@@ -267,8 +267,9 @@ where
     };
 
     // Use the provided prefill chooser, or create a disabled one if not provided
-    let prefill_chooser =
-        prefill_chooser.unwrap_or_else(|| PrefillRouter::disabled(router_mode, enforce_disagg));
+    let block_size = card.kv_cache_block_size;
+    let prefill_chooser = prefill_chooser
+        .unwrap_or_else(|| PrefillRouter::disabled(router_mode, enforce_disagg, block_size));
     let prefill_op = prefill_chooser.into_operator();
 
     // Link with prefill chooser including backward edge for response flow

lib/llm/src/kv_router/prefill_router.rs

@@ -57,16 +57,18 @@ pub struct PrefillRouter {
     cancel_token: CancellationToken,
     router_mode: RouterMode,
     enforce_disagg: bool,
+    block_size: u32,
 }
 
 impl PrefillRouter {
     /// Create a disabled prefill router that will never activate (passthrough only)
-    pub fn disabled(router_mode: RouterMode, enforce_disagg: bool) -> Arc<Self> {
+    pub fn disabled(router_mode: RouterMode, enforce_disagg: bool, block_size: u32) -> Arc<Self> {
         Arc::new(Self {
             prefill_router: OnceLock::new(),
             cancel_token: CancellationToken::new(),
             router_mode,
             enforce_disagg,
+            block_size,
         })
     }
 
@@ -86,6 +88,7 @@ impl PrefillRouter {
             cancel_token: cancel_token.clone(),
             router_mode,
             enforce_disagg,
+            block_size: kv_cache_block_size,
         });
 
         // Spawn background task to wait for activation
@@ -180,7 +183,8 @@ impl PrefillRouter {
     async fn call_prefill(
         &self,
         request: SingleIn<PreprocessedRequest>,
-    ) -> Result<(PrefillResult, Option<u64>), PrefillError> {
+        _block_size: u32,
+    ) -> Result<(PrefillResult, Option<u64>, u32), PrefillError> {
         // Get the prefill router, error if not activated
         let Some(prefill_router) = self.prefill_router.get() else {
             return Err(PrefillError::NotActivated);
@@ -247,12 +251,22 @@ impl PrefillRouter {
                     .get("prefill_worker_id")
                     .and_then(|v| v.as_u64())
             });
+
+        // Extract overlap_blocks from the response (set by prefill worker)
+        let overlap_blocks = output
+            .disaggregated_params
+            .as_ref()
+            .and_then(|params| params.get("overlap_blocks"))
+            .and_then(|v| v.as_u64())
+            .unwrap_or(0) as u32;
+
         Ok((
             PrefillResult {
                 disaggregated_params,
                 prompt_tokens_details,
             },
             prefill_worker_id,
+            overlap_blocks,
         ))
     }
 }
@@ -297,7 +311,7 @@ impl
         let prefill_request = prefill_context;
 
         // Attempt prefill
-        let prefill_result = self.call_prefill(prefill_request).await;
+        let prefill_result = self.call_prefill(prefill_request, self.block_size).await;
 
         // Abort if cancelled during prefill
         if engine_ctx.is_stopped() || engine_ctx.is_killed() {
@@ -310,8 +324,28 @@ impl
 
         // Handle prefill result
         match prefill_result {
-            Ok((prefill_result, prefill_worker_id)) => {
-                tracing::debug!("Prefill succeeded, using disaggregated params for decode");
+            Ok((mut prefill_result, prefill_worker_id, overlap_blocks)) => {
+                // Prefer vLLM's actual cached_tokens over router's estimate
+                // vLLM queries the actual KV cache on the prefill worker (ground truth)
+                // Router's overlap is just a prediction based on its global state
+                let vllm_cached_tokens = prefill_result
+                    .prompt_tokens_details
+                    .as_ref()
+                    .and_then(|d| d.cached_tokens);
+                let final_cached_tokens = if let Some(vllm_value) = vllm_cached_tokens {
+                    vllm_value
+                } else {
+                    overlap_blocks * self.block_size
+                };
+
+                prefill_result.prompt_tokens_details =
+                    Some(dynamo_async_openai::types::PromptTokensDetails {
+                        cached_tokens: Some(final_cached_tokens),
+                        audio_tokens: prefill_result
+                            .prompt_tokens_details
+                            .as_ref()
+                            .and_then(|d| d.audio_tokens),
+                    });
 
                 let mut decode_req = req;
                 // Update request with prefill result