diff --git a/py_src/vllm_router/router.py b/py_src/vllm_router/router.py index 98cd3456..89c9a6da 100644 --- a/py_src/vllm_router/router.py +++ b/py_src/vllm_router/router.py @@ -76,6 +76,8 @@ class Router: If not specified, uses the main policy. Default: None decode_policy: Specific load balancing policy for decode nodes (PD mode only). If not specified, uses the main policy. Default: None + pd_kv_cache_ttl_secs: TTL in seconds for Decode-side KV metadata cached for + bidirectional vLLM P/D transfer. Default: 0 request_id_headers: List of HTTP headers to check for request IDs. If not specified, uses common defaults: ['x-request-id', 'x-correlation-id', 'x-trace-id', 'request-id']. Example: ['x-my-request-id', 'x-custom-trace-id']. Default: None diff --git a/py_src/vllm_router/router_args.py b/py_src/vllm_router/router_args.py index 1f136635..a01008de 100644 --- a/py_src/vllm_router/router_args.py +++ b/py_src/vllm_router/router_args.py @@ -21,6 +21,7 @@ class RouterArgs: default_factory=list ) # List of (url, bootstrap_port) decode_urls: List[str] = dataclasses.field(default_factory=list) + pd_kv_cache_ttl_secs: int = 0 # Routing policy policy: str = "cache_aware" @@ -201,6 +202,12 @@ def add_cli_args( metavar=("URL",), help="Decode server URL. Can be specified multiple times.", ) + parser.add_argument( + f"--{prefix}pd-kv-cache-ttl-secs", + type=int, + default=RouterArgs.pd_kv_cache_ttl_secs, + help="TTL in seconds for Decode-side KV metadata cached for bidirectional vLLM P/D transfer.", + ) parser.add_argument( f"--{prefix}worker-startup-timeout-secs", type=int, diff --git a/src/config/types.rs b/src/config/types.rs index 0bcaf7e9..6353a876 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -84,12 +84,19 @@ pub struct RouterConfig { /// Profiling timeout in seconds (for vLLM profiling endpoints) #[serde(default = "default_profile_timeout_secs")] pub profile_timeout_secs: u64, + /// TTL for Decode-side KV metadata cached by vLLM P/D router. + #[serde(default = "default_pd_kv_cache_ttl_secs")] + pub pd_kv_cache_ttl_secs: u64, } fn default_profile_timeout_secs() -> u64 { 10 } +fn default_pd_kv_cache_ttl_secs() -> u64 { + 0 +} + fn default_history_backend() -> HistoryBackend { HistoryBackend::Memory } @@ -491,6 +498,7 @@ impl Default for RouterConfig { history_backend: default_history_backend(), enable_profiling: false, profile_timeout_secs: default_profile_timeout_secs(), + pd_kv_cache_ttl_secs: default_pd_kv_cache_ttl_secs(), } } } @@ -1063,6 +1071,7 @@ mod tests { history_backend: default_history_backend(), enable_profiling: false, profile_timeout_secs: default_profile_timeout_secs(), + pd_kv_cache_ttl_secs: default_pd_kv_cache_ttl_secs(), }; assert!(config.mode.is_pd_mode()); @@ -1131,6 +1140,7 @@ mod tests { history_backend: default_history_backend(), enable_profiling: false, profile_timeout_secs: default_profile_timeout_secs(), + pd_kv_cache_ttl_secs: default_pd_kv_cache_ttl_secs(), }; assert!(!config.mode.is_pd_mode()); @@ -1195,6 +1205,7 @@ mod tests { history_backend: default_history_backend(), enable_profiling: false, profile_timeout_secs: default_profile_timeout_secs(), + pd_kv_cache_ttl_secs: default_pd_kv_cache_ttl_secs(), }; assert!(config.has_service_discovery()); diff --git a/src/lib.rs b/src/lib.rs index 0c3d53c0..362a0180 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -104,6 +104,8 @@ struct Router { // OpenTelemetry tracing enable_trace: bool, otlp_traces_endpoint: Option, + // vLLM P/D Decode -> Prefill KV metadata cache + pd_kv_cache_ttl_secs: u64, } impl Router { @@ -253,6 +255,7 @@ impl Router { history_backend: config::HistoryBackend::Memory, enable_profiling: false, // Profiling disabled in Python binding by default profile_timeout_secs: 10, // Default profiling timeout + pd_kv_cache_ttl_secs: self.pd_kv_cache_ttl_secs, }) } } @@ -327,6 +330,8 @@ impl Router { // Tracing defaults enable_trace = false, otlp_traces_endpoint = None, + // vLLM P/D defaults + pd_kv_cache_ttl_secs = 0, ))] #[allow(clippy::too_many_arguments)] fn new( @@ -390,6 +395,7 @@ impl Router { tokenizer_path: Option, enable_trace: bool, otlp_traces_endpoint: Option, + pd_kv_cache_ttl_secs: u64, ) -> PyResult { // Determine connection mode from worker URLs let mut all_urls = worker_urls.clone(); @@ -470,6 +476,7 @@ impl Router { tokenizer_path, enable_trace, otlp_traces_endpoint, + pd_kv_cache_ttl_secs, }) } diff --git a/src/main.rs b/src/main.rs index d2ee914d..85c34792 100644 --- a/src/main.rs +++ b/src/main.rs @@ -134,6 +134,10 @@ struct CliArgs { #[arg(long, default_value_t = false)] vllm_pd_disaggregation: bool, + /// TTL in seconds for Decode-side KV metadata cached for bidirectional vLLM P/D transfer + #[arg(long, default_value_t = 0)] + pd_kv_cache_ttl_secs: u64, + /// ZMQ service discovery address for vLLM P2P NCCL coordination (e.g., "0.0.0.0:30001") /// Required for --vllm-pd-disaggregation mode. Workers register their HTTP and ZMQ addresses here. #[arg(long)] @@ -680,6 +684,7 @@ impl CliArgs { }, enable_profiling: self.profile, profile_timeout_secs: 10, // Default profiling timeout + pd_kv_cache_ttl_secs: self.pd_kv_cache_ttl_secs, }) } diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index a7bd585d..f65ed667 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -22,11 +22,17 @@ use axum::{ use serde_json::{json, Value}; use std::collections::HashMap; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use tokio::sync::Mutex; use tracing::{debug, error, info, warn}; use uuid::Uuid; +#[derive(Debug, Clone)] +struct CachedDecodeKv { + kv_transfer_params: Value, + expires_at: Instant, +} + /// vLLM PD Router that extends PDRouter with vLLM-specific request handling #[derive(Debug)] pub struct VllmPDRouter { @@ -48,6 +54,10 @@ pub struct VllmPDRouter { profiling_tasks: Arc>>, /// Intra-node data parallel size for DP-aware routing (automatically enabled when > 1) intra_node_data_parallel_size: usize, + /// Decode-side KV metadata cached for a later prefill request in the same conversation. + decode_kv_cache: Arc>>, + /// TTL for decode_kv_cache entries. Must be lower than vLLM's NIXL abort timeout. + decode_kv_cache_ttl: Duration, } impl VllmPDRouter { @@ -60,6 +70,154 @@ impl VllmPDRouter { ) } + fn conversation_key(headers: Option<&HeaderMap>) -> Option { + headers + .and_then(|headers| headers.get("x-conversation-id")) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) + } + + fn bidirectional_kv_enabled(&self) -> bool { + !self.decode_kv_cache_ttl.is_zero() + } + + fn evict_expired_decode_kv_entries(cache: &mut HashMap, now: Instant) { + cache.retain(|_, entry| entry.expires_at > now); + } + + async fn take_decode_kv_for_conversation( + &self, + conversation_key: Option<&str>, + ) -> Option { + if self.decode_kv_cache_ttl.is_zero() { + return None; + } + + let conversation_key = match conversation_key { + Some(key) => key, + None => { + debug!("PD KV cache disabled for request: no conversation key"); + return None; + } + }; + + let now = Instant::now(); + let mut cache = self.decode_kv_cache.lock().await; + Self::evict_expired_decode_kv_entries(&mut cache, now); + + match cache.remove(conversation_key) { + Some(entry) if entry.expires_at > now => { + debug!("PD KV cache hit for conversation {}", conversation_key); + Some(entry.kv_transfer_params) + } + Some(_) => { + debug!( + "PD KV cache entry expired for conversation {}", + conversation_key + ); + None + } + None => { + debug!("PD KV cache miss for conversation {}", conversation_key); + None + } + } + } + + fn default_kv_params_for_prefill() -> Value { + json!({ + "do_remote_decode": true, + "do_remote_prefill": false, + "remote_engine_id": serde_json::Value::Null, + "remote_block_ids": serde_json::Value::Null, + "remote_host": serde_json::Value::Null, + "remote_port": serde_json::Value::Null + }) + } + + fn kv_params_for_prefill(mut cached_decode_kv: Value) -> Value { + if let Some(obj) = cached_decode_kv.as_object_mut() { + obj.insert("do_remote_decode".to_string(), json!(true)); + obj.insert("do_remote_prefill".to_string(), json!(false)); + return cached_decode_kv; + } + + Self::default_kv_params_for_prefill() + } + + fn extract_decode_kv_params(decode_body: &[u8], is_streaming: bool) -> Option { + if !is_streaming { + return serde_json::from_slice::(decode_body) + .ok() + .and_then(|json| json.get("kv_transfer_params").cloned()) + .filter(Value::is_object); + } + + let body = std::str::from_utf8(decode_body).ok()?; + let mut latest = None; + for line in body.lines() { + let Some(data) = line.strip_prefix("data:") else { + continue; + }; + let data = data.trim(); + if data.is_empty() || data == "[DONE]" { + continue; + } + if let Ok(chunk) = serde_json::from_str::(data) { + if let Some(params) = chunk + .get("kv_transfer_params") + .filter(|params| params.is_object()) + { + latest = Some(params.clone()); + } + } + } + latest + } + + async fn cache_decode_kv_for_conversation( + &self, + conversation_key: Option<&str>, + decode_body: &[u8], + is_streaming: bool, + status: reqwest::StatusCode, + ) { + if self.decode_kv_cache_ttl.is_zero() || !status.is_success() { + return; + } + + let Some(conversation_key) = conversation_key else { + return; + }; + let Some(kv_transfer_params) = Self::extract_decode_kv_params(decode_body, is_streaming) + else { + debug!( + "No decode-side kv_transfer_params found for conversation {}", + conversation_key + ); + return; + }; + + let now = Instant::now(); + let expires_at = now + self.decode_kv_cache_ttl; + let mut cache = self.decode_kv_cache.lock().await; + Self::evict_expired_decode_kv_entries(&mut cache, now); + cache.insert( + conversation_key.to_string(), + CachedDecodeKv { + kv_transfer_params, + expires_at, + }, + ); + debug!( + "Cached decode-side KV metadata for conversation {} for {} seconds", + conversation_key, + self.decode_kv_cache_ttl.as_secs() + ); + } + /// Get ZMQ address for a worker URL using service discovery fn get_zmq_address(&self, http_url: &str, service_type: ServiceType) -> String { // Extract just the host:port from the URL @@ -427,6 +585,20 @@ impl VllmPDRouter { "Generated vLLM request ID for P2P coordination: {}", request_id ); + let bidirectional_kv_enabled = self.bidirectional_kv_enabled(); + let mut kv_transfer_params = Self::default_kv_params_for_prefill(); + let conversation_key = if bidirectional_kv_enabled { + let conversation_key = Self::conversation_key(headers); + if let Some(cached_decode_kv) = self + .take_decode_kv_for_conversation(conversation_key.as_deref()) + .await + { + kv_transfer_params = Self::kv_params_for_prefill(cached_decode_kv); + } + conversation_key + } else { + None + }; // DO NOT add P2P metadata to internal request_id - let vLLM generate clean internal IDs // The P2P metadata will be sent in X-Request-Id header instead @@ -434,16 +606,7 @@ impl VllmPDRouter { // Prepare prefill request (max_tokens=1 to force prefill-only mode) let mut prefill_request = Self::prepare_prefill_request(request_json.clone(), path); - // Add kv_transfer_params for NixlConnector support at top level - // This enables the prefill instance to prepare for remote decode - prefill_request["kv_transfer_params"] = json!({ - "do_remote_decode": true, - "do_remote_prefill": false, - "remote_engine_id": serde_json::Value::Null, - "remote_block_ids": serde_json::Value::Null, - "remote_host": serde_json::Value::Null, - "remote_port": serde_json::Value::Null - }); + prefill_request["kv_transfer_params"] = kv_transfer_params; debug!("Added kv_transfer_params to prefill request for NixlConnector support"); @@ -674,6 +837,16 @@ impl VllmPDRouter { .await .map_err(|e| format!("Failed to read decode response: {}", e))?; + if bidirectional_kv_enabled { + self.cache_decode_kv_for_conversation( + conversation_key.as_deref(), + &decode_body, + is_streaming, + status, + ) + .await; + } + // Per-run billing metrics on success. if let Some(rid) = run_id { if status.is_success() { @@ -727,6 +900,16 @@ impl VllmPDRouter { .await .map_err(|e| format!("Failed to read decode response: {}", e))?; + if bidirectional_kv_enabled { + self.cache_decode_kv_for_conversation( + conversation_key.as_deref(), + &decode_body, + is_streaming, + status, + ) + .await; + } + // Per-run billing metrics on success. This branch also // handles `stream=true` requests, in which case the buffered // body is SSE-framed rather than a single JSON object — the @@ -792,6 +975,20 @@ impl VllmPDRouter { self.get_zmq_address(prefill_worker.base_url(), ServiceType::Prefill); let decode_zmq_addr = self.get_zmq_address(decode_worker.base_url(), ServiceType::Decode); let request_id = Self::generate_vllm_request_id(&prefill_zmq_addr, &decode_zmq_addr); + let bidirectional_kv_enabled = self.bidirectional_kv_enabled(); + let mut kv_transfer_params = Self::default_kv_params_for_prefill(); + let conversation_key = if bidirectional_kv_enabled { + let conversation_key = Self::conversation_key(headers); + if let Some(cached_decode_kv) = self + .take_decode_kv_for_conversation(conversation_key.as_deref()) + .await + { + kv_transfer_params = Self::kv_params_for_prefill(cached_decode_kv); + } + conversation_key + } else { + None + }; debug!("Generated vLLM request ID: {}", request_id); debug!("🔍 vLLM Proxy Comparison:"); @@ -805,16 +1002,7 @@ impl VllmPDRouter { // Stage 1: Prepare prefill request with max_tokens=1 and kv_transfer_params let mut prefill_request = Self::prepare_prefill_request(original_request.clone(), path); - // Add kv_transfer_params for NixlConnector support at top level - // This enables the prefill instance to prepare for remote decode - prefill_request["kv_transfer_params"] = json!({ - "do_remote_decode": true, - "do_remote_prefill": false, - "remote_engine_id": serde_json::Value::Null, - "remote_block_ids": serde_json::Value::Null, - "remote_host": serde_json::Value::Null, - "remote_port": serde_json::Value::Null - }); + prefill_request["kv_transfer_params"] = kv_transfer_params; debug!("Added kv_transfer_params to prefill request for NixlConnector support"); @@ -1088,6 +1276,16 @@ impl VllmPDRouter { ), })?; + if bidirectional_kv_enabled { + self.cache_decode_kv_for_conversation( + conversation_key.as_deref(), + &decode_body, + is_streaming, + status, + ) + .await; + } + // Per-run billing metrics on success. if let Some(rid) = run_id { if status.is_success() { @@ -1154,6 +1352,16 @@ impl VllmPDRouter { ), })?; + if bidirectional_kv_enabled { + self.cache_decode_kv_for_conversation( + conversation_key.as_deref(), + &decode_body, + is_streaming, + status, + ) + .await; + } + // Per-run billing metrics on success. This branch also // handles `stream=true` requests, in which case the buffered // body is SSE-framed rather than a single JSON object — the @@ -1232,6 +1440,8 @@ impl VllmPDRouter { profile_timeout_secs: ctx.router_config.profile_timeout_secs, profiling_tasks: Arc::new(Mutex::new(HashMap::new())), intra_node_data_parallel_size: ctx.router_config.intra_node_data_parallel_size, + decode_kv_cache: Arc::new(Mutex::new(HashMap::new())), + decode_kv_cache_ttl: Duration::from_secs(ctx.router_config.pd_kv_cache_ttl_secs), }) } else { // Direct URL mode (same as PDRouter) @@ -1274,6 +1484,8 @@ impl VllmPDRouter { profile_timeout_secs: ctx.router_config.profile_timeout_secs, profiling_tasks: Arc::new(Mutex::new(HashMap::new())), intra_node_data_parallel_size: ctx.router_config.intra_node_data_parallel_size, + decode_kv_cache: Arc::new(Mutex::new(HashMap::new())), + decode_kv_cache_ttl: Duration::from_secs(ctx.router_config.pd_kv_cache_ttl_secs), }) } } diff --git a/tests/api_endpoints_test.rs b/tests/api_endpoints_test.rs index 21f5ab21..bd400757 100644 --- a/tests/api_endpoints_test.rs +++ b/tests/api_endpoints_test.rs @@ -66,6 +66,7 @@ impl TestContext { history_backend: vllm_router_rs::config::HistoryBackend::Memory, enable_profiling: false, profile_timeout_secs: 30, + pd_kv_cache_ttl_secs: 0, }; Self::new_with_config(config, worker_configs).await @@ -1453,6 +1454,7 @@ mod error_tests { history_backend: vllm_router_rs::config::HistoryBackend::Memory, enable_profiling: false, profile_timeout_secs: 30, + pd_kv_cache_ttl_secs: 0, }; let ctx = TestContext::new_with_config( @@ -1816,6 +1818,7 @@ mod pd_mode_tests { history_backend: vllm_router_rs::config::HistoryBackend::Memory, enable_profiling: false, profile_timeout_secs: 30, + pd_kv_cache_ttl_secs: 0, }; // Create app context @@ -1983,6 +1986,7 @@ mod request_id_tests { history_backend: vllm_router_rs::config::HistoryBackend::Memory, enable_profiling: false, profile_timeout_secs: 30, + pd_kv_cache_ttl_secs: 0, }; let ctx = TestContext::new_with_config( diff --git a/tests/test_dp_routing.rs b/tests/test_dp_routing.rs index 3e4d32fa..ed18c9ea 100644 --- a/tests/test_dp_routing.rs +++ b/tests/test_dp_routing.rs @@ -281,6 +281,7 @@ mod dp_e2e_tests { history_backend: vllm_router_rs::config::HistoryBackend::Memory, enable_profiling: false, profile_timeout_secs: 30, + pd_kv_cache_ttl_secs: 0, } } @@ -330,6 +331,7 @@ mod dp_e2e_tests { history_backend: vllm_router_rs::config::HistoryBackend::Memory, enable_profiling: false, profile_timeout_secs: 30, + pd_kv_cache_ttl_secs: 0, } } diff --git a/tests/test_pd_routing.rs b/tests/test_pd_routing.rs index 891f6f8b..bd234de2 100644 --- a/tests/test_pd_routing.rs +++ b/tests/test_pd_routing.rs @@ -196,6 +196,7 @@ mod test_pd_routing { history_backend: vllm_router_rs::config::HistoryBackend::Memory, enable_profiling: false, profile_timeout_secs: 30, + pd_kv_cache_ttl_secs: 0, }; // Router creation will fail due to health checks, but config should be valid