From 7eac3b7eaf986b88bb70e5c1ef49a60c9fd51e12 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 23 Apr 2026 14:39:02 +0200 Subject: [PATCH 01/29] Initial commit Signed-off-by: simondanielsson From a6374a4268a99cb5b5856899c39c6e2185369ce4 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 23 Apr 2026 14:39:51 +0200 Subject: [PATCH 02/29] feat: first impl of async (WRITE) mode and raise if mori used without service discovery Signed-off-by: simondanielsson --- src/config/validation.rs | 62 +- src/routers/http/vllm_pd_router.rs | 687 ++++++++++++--------- src/routers/http/vllm_service_discovery.rs | 171 ++++- 3 files changed, 626 insertions(+), 294 deletions(-) diff --git a/src/config/validation.rs b/src/config/validation.rs index f51d8092..8a255f15 100644 --- a/src/config/validation.rs +++ b/src/config/validation.rs @@ -462,8 +462,18 @@ impl ConfigValidator { // All policies are now supported for both router types thanks to the unified trait design // No mode/policy restrictions needed anymore - // Check if service discovery is enabled for worker count validation - let has_service_discovery = config.discovery.as_ref().is_some_and(|d| d.enabled); + // Check if service discovery is enabled for worker count validation. + // This covers both K8s service discovery (config.discovery) and vLLM ZMQ + // service discovery (VllmPrefillDecode { discovery_address: Some(_) }). + let has_vllm_discovery = matches!( + &config.mode, + RoutingMode::VllmPrefillDecode { + discovery_address: Some(_), + .. + } + ); + let has_service_discovery = + config.discovery.as_ref().is_some_and(|d| d.enabled) || has_vllm_discovery; // Only validate worker counts if service discovery is disabled if !has_service_discovery { @@ -510,6 +520,17 @@ impl ConfigValidator { // DP-aware routing is now automatically enabled when data_parallel_size > 1 // and is compatible with service discovery + // MoRI-IO requires service discovery: ZMQ addresses are obtained via instance + // registration and are not available in direct URL mode. + if config.kv_connector == KvConnector::MoriIO && !has_service_discovery { + return Err(ConfigError::IncompatibleConfig { + reason: "MoRI-IO KV connector requires service discovery to be enabled \ + (ZMQ addresses are obtained via instance registration). Please \ + run with `--vllm-discovery-address ${address}`" + .to_string(), + }); + } + Ok(()) } @@ -848,4 +869,41 @@ mod tests { let result = ConfigValidator::validate(&config); assert!(result.is_ok()); } + + #[test] + fn test_moriio_requires_service_discovery() { + let mut config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec!["http://worker:8000".to_string()], + }, + PolicyConfig::Random, + ); + config.kv_connector = KvConnector::MoriIO; + config.discovery = None; + + let result = ConfigValidator::validate(&config); + assert!(result.is_err()); + if let Err(e) = result { + assert!(e.to_string().contains("MoRI-IO KV connector requires service discovery")); + } + } + + #[test] + fn test_moriio_with_service_discovery_is_valid() { + let mut config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec!["http://worker:8000".to_string()], + }, + PolicyConfig::Random, + ); + config.kv_connector = KvConnector::MoriIO; + config.discovery = Some(DiscoveryConfig { + enabled: true, + selector: [("app".to_string(), "vllm".to_string())].into(), + ..Default::default() + }); + + let result = ConfigValidator::validate(&config); + assert!(result.is_ok()); + } } diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 8ede26ab..22098252 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -4,7 +4,7 @@ use super::dp_utils; use super::logprobs_merge; use super::pd_router::PDRouter; use super::pd_types::{error_chain, PDRouterError}; -use super::vllm_service_discovery::{ServiceRegistry, ServiceType}; +use super::vllm_service_discovery::{MoriIOTransferMode, ServiceRegistry, ServiceType}; use crate::config::KvConnector; use crate::core::{BasicWorker, Worker, WorkerType}; use crate::metrics::RouterMetrics; @@ -145,6 +145,11 @@ impl VllmPDRouter { } /// Generate a connector-specific transfer ID for correlating prefill and decode legs. + /// Returns the MoRI-IO transfer mode if it has been set by a registration, or `None`. + fn moriio_transfer_mode(&self) -> Option { + self.service_registry.moriio_transfer_mode.get().copied() + } + /// Returns `None` for connectors that do not use a transfer_id (e.g. NIXL). fn generate_transfer_id(&self) -> Option { match self.kv_connector { @@ -171,14 +176,24 @@ impl VllmPDRouter { }) } KvConnector::MoriIO => { - json!({ - "do_remote_decode": true, - "do_remote_prefill": false, - "remote_engine_id": serde_json::Value::Null, - "remote_block_ids": serde_json::Value::Null, - "transfer_id": transfer_id.unwrap_or(""), - "remote_dp_size": self.intra_node_data_parallel_size, - }) + if matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)) { + // WRITE mode: prefill pushes blocks to decode; no do_remote_* flags needed. + json!({ + "remote_dp_size": self.intra_node_data_parallel_size, + "remote_tp_size": 1, + "transfer_id": transfer_id.unwrap_or(""), + }) + } else { + // READ mode (default): prefill waits for decode to pull blocks. + json!({ + "do_remote_decode": true, + "do_remote_prefill": false, + "remote_engine_id": serde_json::Value::Null, + "remote_block_ids": serde_json::Value::Null, + "transfer_id": transfer_id.unwrap_or(""), + "remote_dp_size": self.intra_node_data_parallel_size, + }) + } } KvConnector::Nixl => { json!({ @@ -193,6 +208,63 @@ impl VllmPDRouter { } } + /// Build decode kv_transfer_params for all connectors. Returns `None` if params cannot be + /// determined (e.g. Mooncake bootstrap info missing, or no kv_transfer_params in prefill response). + async fn build_decode_kv_params( + &self, + prefill_url: &str, + prefill_response_json: Option<&Value>, + transfer_id: Option<&str>, + prefill_dp_rank: Option, + ) -> Option { + match self.kv_connector { + KvConnector::Mooncake => { + let Some((bootstrap_addr, engine_id)) = + self.get_mooncake_info(prefill_url, prefill_dp_rank.map(|r| r as usize)).await + else { + warn!( + "No Mooncake bootstrap info for prefill {}, decode will proceed without kv_transfer_params", + prefill_url + ); + return None; + }; + Some(self.build_mooncake_decode_kv_transfer_params( + transfer_id.unwrap_or(""), + &bootstrap_addr, + &engine_id, + )) + } + KvConnector::MoriIO => { + if matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)) { + // WRITE mode: build decode params directly; decode does not need the prefill response. + let mut params = json!({ + "do_remote_decode": false, + "do_remote_prefill": true, + "remote_engine_id": serde_json::Value::Null, + "remote_block_ids": serde_json::Value::Null, + "transfer_id": transfer_id.unwrap_or(""), + "remote_dp_size": self.intra_node_data_parallel_size, + "remote_tp_size": 1, + }); + if self.intra_node_data_parallel_size > 1 { + if let Some(rank) = prefill_dp_rank { + params["remote_dp_rank"] = json!(rank); + } + } + Some(params) + } else { + // READ mode: extract params from prefill response and inject remote_dp_size. + let mut params = prefill_response_json?.get("kv_transfer_params")?.clone(); + params["remote_dp_size"] = json!(self.intra_node_data_parallel_size); + Some(params) + } + } + KvConnector::Nixl => { + Some(prefill_response_json?.get("kv_transfer_params")?.clone()) + } + } + } + /// Build kv_transfer_params for the decode request (Mooncake only). /// For NIXL, decode params come from the prefill response instead. fn build_mooncake_decode_kv_transfer_params( @@ -561,15 +633,6 @@ impl VllmPDRouter { self.kv_connector ); - let prefill_request_str = serde_json::to_string(&prefill_request) - .map_err(|e| format!("Failed to serialize prefill request: {}", e))?; - - // Stage 1: Send to prefill server with max_tokens=1 and P2P coordination header - debug!( - "Stage 1: Sending prefill-only request (max_tokens=1) to prefill server at http://{}", - prefill_http - ); - // Extract dp_rank from prefill_http if intra_node_data_parallel_size > 1 let (prefill_base_http, prefill_dp_rank) = if self.intra_node_data_parallel_size > 1 { let prefill_url = format!("http://{}", prefill_http); @@ -580,52 +643,80 @@ impl VllmPDRouter { (prefill_http.to_string(), None) }; - // Start profiling on prefill server - self.start_profiling(&format!("http://{}", prefill_base_http)) - .await; + // Extract dp_rank from decode_http if intra_node_data_parallel_size > 1 + let (decode_base_http, decode_dp_rank) = if self.intra_node_data_parallel_size > 1 { + let decode_url = format!("http://{}", decode_http); + let (base, rank) = dp_utils::parse_worker_url(&decode_url); + let base_http = base.replace("http://", "").replace("https://", ""); + (base_http, rank) + } else { + (decode_http.to_string(), None) + }; - let mut prefill_request_builder = self - .http_client - .post(format!("http://{}{}", prefill_base_http, path)) - .header("Content-Type", "application/json") - .header("X-Request-Id", &request_id); // P2P coordination metadata in header + let is_moriio_write = matches!(self.kv_connector, KvConnector::MoriIO) + && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)); - // Add X-data-parallel-rank header using shared utilities - prefill_request_builder = - dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); - if let Some(rank) = prefill_dp_rank { + let prefill_request_str = serde_json::to_string(&prefill_request) + .map_err(|e| format!("Failed to serialize prefill request: {}", e))?; + + // Stage 1: dispatch prefill. + // WRITE mode: fire-and-forget (decode does not need the prefill response). + // READ mode: await and parse — decode params come from the prefill response body. + let prefill_response_json: Option = if is_moriio_write { + let http_client = self.http_client.clone(); + let prefill_url = format!("http://{}{}", prefill_base_http, path); + let prefill_request_id = request_id.clone(); + let prefill_dp_rank_copy = prefill_dp_rank; + tokio::spawn(async move { + let mut builder = http_client + .post(&prefill_url) + .header("Content-Type", "application/json") + .header("X-Request-Id", &prefill_request_id); + builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); + match builder.body(prefill_request_str).send().await { + Ok(resp) => debug!("MoRI-IO WRITE prefill completed with status {}", resp.status()), + Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), + } + }); + None + } else { debug!( - "Added X-data-parallel-rank={} header to prefill request", - rank + "Stage 1: Sending prefill-only request (max_tokens=1) to prefill server at http://{}", + prefill_http ); - } - - let prefill_request_url = format!("http://{}{}", prefill_base_http, path); - let prefill_response = match otel_http::send_client_request( - prefill_request_builder.body(prefill_request_str), - headers, - ClientRequestOptions { - method: "POST", - url: &prefill_request_url, - route: Some(path), - request_phase: Some("prefill"), - }, - ) - .await - { - Ok(resp) => resp, - Err(e) => { - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(prefill_http); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - return Err(format!( - "Prefill request failed to {}: {}", - prefill_http, full_error - )); - } - }; + self.start_profiling(&format!("http://{}", prefill_base_http)).await; + + let mut prefill_request_builder = self + .http_client + .post(format!("http://{}{}", prefill_base_http, path)) + .header("Content-Type", "application/json") + .header("X-Request-Id", &request_id); + prefill_request_builder = + dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); + + let prefill_request_url = format!("http://{}{}", prefill_base_http, path); + let prefill_response = match otel_http::send_client_request( + prefill_request_builder.body(prefill_request_str), + headers, + ClientRequestOptions { + method: "POST", + url: &prefill_request_url, + route: Some(path), + request_phase: Some("prefill"), + }, + ) + .await + { + Ok(resp) => resp, + Err(e) => { + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(prefill_http); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + return Err(format!("Prefill request failed to {}: {}", prefill_http, full_error)); + } + }; let prefill_status = prefill_response.status(); debug!("Prefill server responded with status: {}", prefill_status); @@ -653,59 +744,29 @@ impl VllmPDRouter { debug!("Prefill response body: {}", prefill_response_text); - let prefill_response_json: Value = serde_json::from_str(&prefill_response_text) + let prefill_json: Value = serde_json::from_str(&prefill_response_text) .map_err(|e| format!("Failed to parse prefill response as JSON: {}", e))?; - // Extract kv_transfer_params from prefill response if present - let kv_transfer_params = prefill_response_json.get("kv_transfer_params").cloned(); + Some(prefill_json) + }; // end prefill_response_json - if let Some(ref params) = kv_transfer_params { + // Prepare decode request + let mut decode_request = request_json.clone(); + let prefill_url_key = format!("http://{}", prefill_base_http); + if let Some(params) = self + .build_decode_kv_params( + &prefill_url_key, + prefill_response_json.as_ref(), + transfer_id.as_deref(), + prefill_dp_rank.map(|r| r as u32), + ) + .await + { + decode_request["kv_transfer_params"] = params; debug!( - "Extracted kv_transfer_params from prefill response: {}", - serde_json::to_string_pretty(params).unwrap_or_default() + "Added kv_transfer_params to decode request for {:?} connector", + self.kv_connector ); - } else { - debug!("No kv_transfer_params found in prefill response, will proceed without them"); - } - - // Prepare decode request with kv_transfer_params from prefill response at top level - let mut decode_request = request_json.clone(); - if matches!(self.kv_connector, KvConnector::Mooncake) { - // Mooncake: set decode params proactively from bootstrap info - let prefill_url_key = format!("http://{}", prefill_base_http); - if let Some((bootstrap_addr, engine_id)) = self - .get_mooncake_info(&prefill_url_key, prefill_dp_rank) - .await - { - decode_request["kv_transfer_params"] = self - .build_mooncake_decode_kv_transfer_params( - transfer_id.as_deref().unwrap_or(""), - &bootstrap_addr, - &engine_id, - ); - debug!( - "Set Mooncake decode kv_transfer_params with bootstrap_addr={}, engine_id={}", - bootstrap_addr, engine_id - ); - } else { - warn!( - "No Mooncake bootstrap info for prefill {}, decode will proceed without kv_transfer_params", - prefill_url_key - ); - } - } else { - // NIXL and MoRI-IO: extract kv_transfer_params from prefill response - if let Some(mut params) = kv_transfer_params { - if matches!(self.kv_connector, KvConnector::MoriIO) { - // MoRI-IO decode connector needs to know how many prefill DP ranks to handshake with. - params["remote_dp_size"] = json!(self.intra_node_data_parallel_size); - } - decode_request["kv_transfer_params"] = params; - debug!( - "Added kv_transfer_params to decode request for {:?} connector", - self.kv_connector - ); - } } let decode_request_str = serde_json::to_string(&decode_request) @@ -721,16 +782,6 @@ impl VllmPDRouter { decode_http ); - // Extract dp_rank from decode_http if intra_node_data_parallel_size > 1 - let (decode_base_http, decode_dp_rank) = if self.intra_node_data_parallel_size > 1 { - let decode_url = format!("http://{}", decode_http); - let (base, rank) = dp_utils::parse_worker_url(&decode_url); - let base_http = base.replace("http://", "").replace("https://", ""); - (base_http, rank) - } else { - (decode_http.to_string(), None) - }; - // Start profiling on decode server self.start_profiling(&format!("http://{}", decode_base_http)) .await; @@ -826,8 +877,10 @@ impl VllmPDRouter { .map_err(|e| format!("Failed to parse decode response as JSON: {}", e))?; // Merge logprobs from prefill into decode response + let empty_json = Value::Null; + let prefill_json_ref = prefill_response_json.as_ref().unwrap_or(&empty_json); let merged = - logprobs_merge::merge_logprobs_in_json(&prefill_response_json, &mut decode_json); + logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); if merged { debug!("Successfully merged logprobs from prefill and decode responses"); } else { @@ -935,187 +988,189 @@ impl VllmPDRouter { let prefill_dp_rank = prefill_worker.dp_rank(); let prefill_url = prefill_worker.endpoint_url(path); - debug!( - "🚀 vLLM Stage 1 - Prefill: {} with request_id: {}", - prefill_url, request_id - ); - if let Some(rank) = prefill_dp_rank { - debug!("📤 Prefill request headers: Authorization=Bearer [REDACTED], X-Request-Id={}, X-data-parallel-rank={}", request_id, rank); + let decode_base_url = decode_worker.base_url().to_string(); + let decode_dp_rank = decode_worker.dp_rank(); + let decode_url = decode_worker.endpoint_url(path); + + let is_moriio_write = matches!(self.kv_connector, KvConnector::MoriIO) + && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)); + + // Stage 1: dispatch prefill. + // WRITE mode: fire-and-forget (decode does not need the prefill response). + // READ mode: await and parse — decode params come from the prefill response body. + let prefill_response_json: Option = if is_moriio_write { + debug!("MoRI-IO WRITE mode: dispatching prefill fire-and-forget"); + let http_client = self.pd_router.client.clone(); + let prefill_request_clone = prefill_request.clone(); + let prefill_request_id = request_id.clone(); + let prefill_url_clone = prefill_url.clone(); + let prefill_dp_rank_copy = prefill_dp_rank; + tokio::spawn(async move { + let mut builder = http_client + .post(&prefill_url_clone) + .header("Content-Type", "application/json") + .header( + "Authorization", + format!( + "Bearer {}", + std::env::var("OPENAI_API_KEY").unwrap_or_default() + ), + ) + .header("X-Request-Id", &prefill_request_id); + builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); + match builder.json(&prefill_request_clone).send().await { + Ok(resp) => debug!( + "MoRI-IO WRITE prefill completed with status {}", + resp.status() + ), + Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), + } + }); + prefill_worker.decrement_load(); + decode_worker.increment_load(); + None } else { debug!( - "📤 Prefill request headers: Authorization=Bearer [REDACTED], X-Request-Id={}", - request_id + "🚀 vLLM Stage 1 - Prefill: {} with request_id: {}", + prefill_url, request_id ); - } - debug!( - "📤 Prefill request payload: {}", - serde_json::to_string_pretty(&prefill_request).unwrap_or_default() - ); - - // Start profiling on prefill server - self.start_profiling(&prefill_base_url).await; - - let mut prefill_request_builder = self - .pd_router - .client - .post(&prefill_url) - .header("Content-Type", "application/json") - .header( - "Authorization", - format!( - "Bearer {}", - std::env::var("OPENAI_API_KEY").unwrap_or_default() - ), - ) - .header("X-Request-Id", &request_id); - - // Add X-data-parallel-rank header using shared utilities - prefill_request_builder = - dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); - - let prefill_response = match otel_http::send_client_request( - prefill_request_builder.json(&prefill_request), - headers, - ClientRequestOptions { - method: "POST", - url: &prefill_url, - route: Some(path), - request_phase: Some("prefill"), - }, - ) - .await - { - Ok(resp) => resp, - Err(e) => { - prefill_worker.decrement_load(); - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(&prefill_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - return Err(PDRouterError::NetworkError { - message: format!("Prefill request failed to {}: {}", prefill_url, full_error), - }); + if let Some(rank) = prefill_dp_rank { + debug!("📤 Prefill request headers: Authorization=Bearer [REDACTED], X-Request-Id={}, X-data-parallel-rank={}", request_id, rank); + } else { + debug!( + "📤 Prefill request headers: Authorization=Bearer [REDACTED], X-Request-Id={}", + request_id + ); } - }; + debug!( + "📤 Prefill request payload: {}", + serde_json::to_string_pretty(&prefill_request).unwrap_or_default() + ); - debug!("📥 Prefill response status: {}", prefill_response.status()); - debug!( - "📥 Prefill response headers: {:?}", - prefill_response.headers() - ); + self.start_profiling(&prefill_base_url).await; - // Extract prefill response body to get kv_transfer_params - let prefill_bytes = match prefill_response.bytes().await { - Ok(bytes) => bytes, - Err(e) => { - prefill_worker.decrement_load(); - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(&prefill_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - return Err(PDRouterError::NetworkError { - message: format!( - "Failed to read prefill response from {}: {}", - prefill_url, full_error + let mut prefill_request_builder = self + .pd_router + .client + .post(&prefill_url) + .header("Content-Type", "application/json") + .header( + "Authorization", + format!( + "Bearer {}", + std::env::var("OPENAI_API_KEY").unwrap_or_default() ), - }); - } - }; + ) + .header("X-Request-Id", &request_id); + prefill_request_builder = + dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); - debug!( - "📥 Prefill response body size: {} bytes", - prefill_bytes.len() - ); - if prefill_bytes.len() < 1024 { + let prefill_response = match otel_http::send_client_request( + prefill_request_builder.json(&prefill_request), + headers, + ClientRequestOptions { + method: "POST", + url: &prefill_url, + route: Some(path), + request_phase: Some("prefill"), + }, + ) + .await + { + Ok(resp) => resp, + Err(e) => { + prefill_worker.decrement_load(); + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + return Err(PDRouterError::NetworkError { + message: format!( + "Prefill request failed to {}: {}", + prefill_url, full_error + ), + }); + } + }; + + debug!("📥 Prefill response status: {}", prefill_response.status()); debug!( - "📥 Prefill response body content: {}", - String::from_utf8_lossy(&prefill_bytes) + "📥 Prefill response headers: {:?}", + prefill_response.headers() ); - } - // Parse prefill response to extract kv_transfer_params - let prefill_response_json: Value = match serde_json::from_slice(&prefill_bytes) { - Ok(json) => json, - Err(e) => { - prefill_worker.decrement_load(); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(&prefill_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - return Err(PDRouterError::NetworkError { - message: format!("Failed to parse prefill response as JSON: {}", e), - }); - } - }; - - // Extract kv_transfer_params from prefill response if present - let kv_transfer_params = prefill_response_json.get("kv_transfer_params").cloned(); + let prefill_bytes = match prefill_response.bytes().await { + Ok(bytes) => bytes, + Err(e) => { + prefill_worker.decrement_load(); + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + return Err(PDRouterError::NetworkError { + message: format!( + "Failed to read prefill response from {}: {}", + prefill_url, full_error + ), + }); + } + }; - if let Some(ref params) = kv_transfer_params { debug!( - "Extracted kv_transfer_params from prefill response: {}", - serde_json::to_string_pretty(params).unwrap_or_default() + "📥 Prefill response body size: {} bytes", + prefill_bytes.len() ); - } else { - debug!("No kv_transfer_params found in prefill response, will proceed without them"); - } + if prefill_bytes.len() < 1024 { + debug!( + "📥 Prefill response body content: {}", + String::from_utf8_lossy(&prefill_bytes) + ); + } - // Stop profiling on prefill server after its work is done - self.stop_profiling(&prefill_base_url).await; + let prefill_json: Value = match serde_json::from_slice(&prefill_bytes) { + Ok(json) => json, + Err(e) => { + prefill_worker.decrement_load(); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + return Err(PDRouterError::NetworkError { + message: format!("Failed to parse prefill response as JSON: {}", e), + }); + } + }; + + self.stop_profiling(&prefill_base_url).await; + + prefill_worker.decrement_load(); + decode_worker.increment_load(); - // Prefill phase complete: decrement prefill load, increment decode load - prefill_worker.decrement_load(); - decode_worker.increment_load(); + debug!("✅ vLLM Stage 1 completed, starting Stage 2 - Decode"); - debug!("✅ vLLM Stage 1 completed, starting Stage 2 - Decode"); + Some(prefill_json) + }; // end prefill_response_json - // Stage 2: Prepare decode request with kv_transfer_params + // Stage 2: Prepare decode request let mut decode_request = original_request.clone(); - if matches!(self.kv_connector, KvConnector::Mooncake) { - // Mooncake: set decode params proactively from bootstrap info - if let Some((bootstrap_addr, engine_id)) = self - .get_mooncake_info(&prefill_base_url, prefill_dp_rank) - .await - { - decode_request["kv_transfer_params"] = self - .build_mooncake_decode_kv_transfer_params( - transfer_id.as_deref().unwrap_or(""), - &bootstrap_addr, - &engine_id, - ); - debug!( - "Set Mooncake decode kv_transfer_params with bootstrap_addr={}, engine_id={}", - bootstrap_addr, engine_id - ); - } else { - warn!( - "No Mooncake bootstrap info for prefill {}, decode will proceed without kv_transfer_params", - prefill_base_url - ); - } - } else { - // NIXL and MoRI-IO: extract kv_transfer_params from prefill response - if let Some(mut params) = kv_transfer_params { - if matches!(self.kv_connector, KvConnector::MoriIO) { - // MoRI-IO decode connector needs to know how many prefill DP ranks to handshake with. - params["remote_dp_size"] = json!(self.intra_node_data_parallel_size); - } - decode_request["kv_transfer_params"] = params; - debug!( - "Added kv_transfer_params to decode request for {:?} connector", - self.kv_connector - ); - } + if let Some(params) = self + .build_decode_kv_params( + &prefill_base_url, + prefill_response_json.as_ref(), + transfer_id.as_deref(), + prefill_dp_rank.map(|r| r as u32), + ) + .await + { + decode_request["kv_transfer_params"] = params; + debug!( + "Added kv_transfer_params to decode request for {:?} connector", + self.kv_connector + ); } - // Use endpoint_url() to get the base URL without @rank suffix, - // avoiding IPv6+DP URL corruption (same fix as Router and PDRouter) - let decode_base_url = decode_worker.base_url().to_string(); - let decode_dp_rank = decode_worker.dp_rank(); - let decode_url = decode_worker.endpoint_url(path); - debug!( "🚀 vLLM Stage 2 - Decode: {} with request_id: {}", decode_url, request_id @@ -1238,8 +1293,10 @@ impl VllmPDRouter { })?; // Merge logprobs from prefill into decode response + let empty_json = Value::Null; + let prefill_json_ref = prefill_response_json.as_ref().unwrap_or(&empty_json); let merged = - logprobs_merge::merge_logprobs_in_json(&prefill_response_json, &mut decode_json); + logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); if merged { debug!("Successfully merged logprobs from prefill and decode responses"); } else { @@ -1315,7 +1372,7 @@ impl VllmPDRouter { info!("Starting vLLM service discovery on {}", addr); service_registry - .start_listener(addr) + .start_listener(addr, kv_connector) .await .map_err(|e| format!("Failed to start service discovery: {}", e))?; @@ -2253,4 +2310,72 @@ mod tests { let connector = KvConnector::default(); assert_eq!(connector, KvConnector::Nixl); } + + // --- MoRI-IO WRITE mode parameter tests --- + + fn moriio_write_prefill_params(transfer_id: Option<&str>, dp_size: usize) -> Value { + // Mirror the WRITE mode branch in build_prefill_kv_transfer_params. + json!({ + "remote_dp_size": dp_size, + "remote_tp_size": 1, + "transfer_id": transfer_id.unwrap_or(""), + }) + } + + fn moriio_write_decode_params( + transfer_id: Option<&str>, + dp_size: usize, + prefill_dp_rank: Option, + ) -> Value { + // Mirror build_moriio_write_decode_kv_transfer_params. + let mut params = json!({ + "do_remote_decode": false, + "do_remote_prefill": true, + "remote_engine_id": serde_json::Value::Null, + "remote_block_ids": serde_json::Value::Null, + "transfer_id": transfer_id.unwrap_or(""), + "remote_dp_size": dp_size, + "remote_tp_size": 1, + }); + if dp_size > 1 { + if let Some(rank) = prefill_dp_rank { + params["remote_dp_rank"] = json!(rank); + } + } + params + } + + #[test] + fn test_moriio_write_prefill_params_no_remote_decode_prefill_flags() { + let params = moriio_write_prefill_params(Some("tx-abc"), 1); + assert!(params.get("do_remote_decode").is_none()); + assert!(params.get("do_remote_prefill").is_none()); + assert_eq!(params["remote_tp_size"], 1); + assert_eq!(params["remote_dp_size"], 1); + assert_eq!(params["transfer_id"], "tx-abc"); + } + + #[test] + fn test_moriio_write_decode_params_have_correct_fields() { + let params = moriio_write_decode_params(Some("tx-abc"), 1, None); + assert_eq!(params["do_remote_decode"], false); + assert_eq!(params["do_remote_prefill"], true); + assert!(params["remote_engine_id"].is_null()); + assert!(params["remote_block_ids"].is_null()); + assert_eq!(params["transfer_id"], "tx-abc"); + assert_eq!(params["remote_dp_size"], 1); + assert_eq!(params["remote_tp_size"], 1); + } + + #[test] + fn test_moriio_write_decode_params_no_remote_dp_rank_when_dp_size_is_1() { + let params = moriio_write_decode_params(Some("tx-abc"), 1, Some(0)); + assert!(params.get("remote_dp_rank").is_none()); + } + + #[test] + fn test_moriio_write_decode_params_includes_remote_dp_rank_when_dp_size_gt_1() { + let params = moriio_write_decode_params(Some("tx-abc"), 4, Some(2)); + assert_eq!(params["remote_dp_rank"], 2); + } } diff --git a/src/routers/http/vllm_service_discovery.rs b/src/routers/http/vllm_service_discovery.rs index a655a453..7364898e 100644 --- a/src/routers/http/vllm_service_discovery.rs +++ b/src/routers/http/vllm_service_discovery.rs @@ -1,9 +1,10 @@ // vLLM Service Discovery Implementation // This module implements service discovery for vLLM P2P NCCL coordination +use crate::config::KvConnector; use serde::{Deserialize, Serialize}; use std::collections::HashMap; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, OnceLock}; use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::broadcast; use tracing::{debug, info, warn}; @@ -36,6 +37,32 @@ pub struct ServiceRegistration { pub zmq_address: String, } +/// MoRI-IO transfer mode, determined from the first instance registration. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MoriIOTransferMode { + Read, + Write, +} + +/// MoRI-IO-specific service registration data. The base fields are flattened so +/// the wire format is identical to `ServiceRegistration` plus `transfer_mode`. +#[derive(Debug, Clone, Deserialize)] +pub struct MoriIOServiceRegistration { + #[serde(flatten)] + pub base: ServiceRegistration, + pub transfer_mode: String, // "READ" or "WRITE" +} + +impl MoriIOServiceRegistration { + pub fn parsed_transfer_mode(&self) -> Option { + match self.transfer_mode.as_str() { + "READ" => Some(MoriIOTransferMode::Read), + "WRITE" => Some(MoriIOTransferMode::Write), + _ => None, + } + } +} + /// Service instance with expiration timestamp #[derive(Debug, Clone)] pub struct ServiceInstance { @@ -49,6 +76,62 @@ pub struct ServiceRegistry { prefill_instances: Arc>>, decode_instances: Arc>>, shutdown_tx: Option>, + /// Set on first MoRI-IO registration; subsequent mismatches are rejected. + pub moriio_transfer_mode: Arc>, +} + +/// Deserialize a registration message according to the active KV connector. +/// Returns the base `ServiceRegistration` on success, or `None` if the message +/// should be dropped (parse error, unknown mode, or transfer-mode mismatch). +fn parse_registration( + message_data: &[u8], + remote_address: &[u8], + kv_connector: KvConnector, + moriio_transfer_mode: &Arc>, +) -> Option { + if matches!(kv_connector, KvConnector::MoriIO) { + let reg: MoriIOServiceRegistration = match rmp_serde::from_slice(message_data) { + Ok(r) => r, + Err(e) => { + warn!("Failed to parse MoRI-IO service registration: {}", e); + return None; + } + }; + let mode = match reg.parsed_transfer_mode() { + Some(m) => m, + None => { + warn!( + "Unknown MoRI-IO transfer_mode '{}' from {}; skipping", + reg.transfer_mode, + String::from_utf8_lossy(remote_address) + ); + return None; + } + }; + if moriio_transfer_mode.set(mode).is_err() { + let stored = moriio_transfer_mode.get().copied().unwrap(); + if stored != mode { + warn!( + "MoRI-IO transfer_mode mismatch: expected {:?}, got {:?} from {}; skipping", + stored, + mode, + String::from_utf8_lossy(remote_address) + ); + return None; + } + } else { + info!("MoRI-IO transfer mode set to {:?}", mode); + } + Some(reg.base) + } else { + match rmp_serde::from_slice(message_data) { + Ok(data) => Some(data), + Err(e) => { + warn!("Failed to parse service registration: {}", e); + None + } + } + } } impl Default for ServiceRegistry { @@ -64,11 +147,16 @@ impl ServiceRegistry { prefill_instances: Arc::new(Mutex::new(HashMap::new())), decode_instances: Arc::new(Mutex::new(HashMap::new())), shutdown_tx: None, + moriio_transfer_mode: Arc::new(OnceLock::new()), } } /// Start the ZMQ service discovery listener - pub async fn start_listener(&mut self, bind_address: &str) -> Result<(), String> { + pub async fn start_listener( + &mut self, + bind_address: &str, + kv_connector: KvConnector, + ) -> Result<(), String> { info!( "Starting vLLM service discovery listener on {}", bind_address @@ -79,6 +167,7 @@ impl ServiceRegistry { let prefill_instances = Arc::clone(&self.prefill_instances); let decode_instances = Arc::clone(&self.decode_instances); + let moriio_transfer_mode = Arc::clone(&self.moriio_transfer_mode); let bind_addr = bind_address.to_string(); tokio::spawn(async move { @@ -115,6 +204,8 @@ impl ServiceRegistry { &remote_address, &prefill_instances, &decode_instances, + kv_connector, + &moriio_transfer_mode, ) .await; } @@ -143,14 +234,17 @@ impl ServiceRegistry { remote_address: &[u8], prefill_instances: &Arc>>, decode_instances: &Arc>>, + kv_connector: KvConnector, + moriio_transfer_mode: &Arc>, ) { - // Parse MessagePack data - let data: ServiceRegistration = match rmp_serde::from_slice(message_data) { - Ok(data) => data, - Err(e) => { - warn!("Failed to parse service registration: {}", e); - return; - } + let data = match parse_registration( + message_data, + remote_address, + kv_connector, + moriio_transfer_mode, + ) { + Some(parsed) => parsed, + None => return, }; let current_time = SystemTime::now() @@ -308,7 +402,7 @@ impl ServiceRegistry { .map(|instance| instance.zmq_address.clone()) } - /// Get all available prefill instances + /// Get all available prefill instances as `(http_address, zmq_address)`. pub fn get_prefill_instances(&self) -> Vec<(String, String)> { let guard = self.prefill_instances.lock().unwrap(); guard @@ -317,7 +411,7 @@ impl ServiceRegistry { .collect() } - /// Get all available decode instances + /// Get all available decode instances as `(http_address, zmq_address)`. pub fn get_decode_instances(&self) -> Vec<(String, String)> { let guard = self.decode_instances.lock().unwrap(); guard @@ -346,3 +440,58 @@ impl Drop for ServiceRegistry { self.shutdown(); } } + +#[cfg(test)] +mod tests { + use super::*; + + fn make_moriio_msgpack(service_type: &str, transfer_mode: &str) -> Vec { + let mut map = rmp_serde::to_vec_named(&serde_json::json!({ + "type": service_type, + "http_address": "1.2.3.4:8000", + "zmq_address": "host:1.2.3.4,handshake:6301,notify:61005", + "transfer_mode": transfer_mode, + })) + .unwrap(); + // rmp_serde::to_vec_named produces msgpack — suitable for our deserializer. + map + } + + #[test] + fn test_moriio_service_registration_deserializes_read() { + let bytes = make_moriio_msgpack("P", "READ"); + let reg: MoriIOServiceRegistration = rmp_serde::from_slice(&bytes).unwrap(); + assert_eq!(reg.base.service_type, "P"); + assert_eq!(reg.base.http_address, "1.2.3.4:8000"); + assert_eq!(reg.parsed_transfer_mode(), Some(MoriIOTransferMode::Read)); + } + + #[test] + fn test_moriio_service_registration_deserializes_write() { + let bytes = make_moriio_msgpack("D", "WRITE"); + let reg: MoriIOServiceRegistration = rmp_serde::from_slice(&bytes).unwrap(); + assert_eq!(reg.base.service_type, "D"); + assert_eq!(reg.parsed_transfer_mode(), Some(MoriIOTransferMode::Write)); + } + + #[test] + fn test_moriio_transfer_mode_mismatch_rejected() { + let lock: Arc> = Arc::new(OnceLock::new()); + lock.set(MoriIOTransferMode::Read).unwrap(); + + // Simulate a second registration with WRITE — should detect mismatch. + let stored = lock.get().copied().unwrap(); + let incoming = MoriIOTransferMode::Write; + assert_ne!(stored, incoming, "mismatch should be detected"); + } + + #[test] + fn test_moriio_transfer_mode_consistent_registration_accepted() { + let lock: Arc> = Arc::new(OnceLock::new()); + lock.set(MoriIOTransferMode::Write).unwrap(); + + let stored = lock.get().copied().unwrap(); + let incoming = MoriIOTransferMode::Write; + assert_eq!(stored, incoming); + } +} From 1584560a2c631917274deb4928f55d2e57e1248c Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 27 Apr 2026 14:37:50 +0200 Subject: [PATCH 03/29] fix: add do_remote_* to WRITE path creation of kv transfer params Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 22098252..378a6017 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -177,8 +177,13 @@ impl VllmPDRouter { } KvConnector::MoriIO => { if matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)) { - // WRITE mode: prefill pushes blocks to decode; no do_remote_* flags needed. + // WRITE mode: prefill pushes KV blocks to decode. + // do_remote_decode=true tells the prefill connector to initiate the transfer. json!({ + "do_remote_decode": true, + "do_remote_prefill": false, + "remote_engine_id": serde_json::Value::Null, + "remote_block_ids": serde_json::Value::Null, "remote_dp_size": self.intra_node_data_parallel_size, "remote_tp_size": 1, "transfer_id": transfer_id.unwrap_or(""), @@ -2316,6 +2321,10 @@ mod tests { fn moriio_write_prefill_params(transfer_id: Option<&str>, dp_size: usize) -> Value { // Mirror the WRITE mode branch in build_prefill_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_dp_size": dp_size, "remote_tp_size": 1, "transfer_id": transfer_id.unwrap_or(""), @@ -2346,10 +2355,12 @@ mod tests { } #[test] - fn test_moriio_write_prefill_params_no_remote_decode_prefill_flags() { + fn test_moriio_write_prefill_params_has_do_remote_decode_true() { let params = moriio_write_prefill_params(Some("tx-abc"), 1); - assert!(params.get("do_remote_decode").is_none()); - assert!(params.get("do_remote_prefill").is_none()); + assert_eq!(params["do_remote_decode"], true); + assert_eq!(params["do_remote_prefill"], false); + assert!(params["remote_engine_id"].is_null()); + assert!(params["remote_block_ids"].is_null()); assert_eq!(params["remote_tp_size"], 1); assert_eq!(params["remote_dp_size"], 1); assert_eq!(params["transfer_id"], "tx-abc"); From 19a72b4caf67ab09ece1e98f1cef38c497020959 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 27 Apr 2026 15:20:36 +0200 Subject: [PATCH 04/29] refactor: use tokio::join! rather than spawn to avoid race condition Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 240 ++++++++++++++++++++--------- 1 file changed, 167 insertions(+), 73 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 378a6017..1d0cb887 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -665,25 +665,15 @@ impl VllmPDRouter { .map_err(|e| format!("Failed to serialize prefill request: {}", e))?; // Stage 1: dispatch prefill. - // WRITE mode: fire-and-forget (decode does not need the prefill response). + // WRITE mode: defer the prefill send until Stage 2 so we can run both concurrently + // via tokio::join!, guaranteeing the prefill task always executes. // READ mode: await and parse — decode params come from the prefill response body. - let prefill_response_json: Option = if is_moriio_write { - let http_client = self.http_client.clone(); - let prefill_url = format!("http://{}{}", prefill_base_http, path); - let prefill_request_id = request_id.clone(); - let prefill_dp_rank_copy = prefill_dp_rank; - tokio::spawn(async move { - let mut builder = http_client - .post(&prefill_url) - .header("Content-Type", "application/json") - .header("X-Request-Id", &prefill_request_id); - builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); - match builder.body(prefill_request_str).send().await { - Ok(resp) => debug!("MoRI-IO WRITE prefill completed with status {}", resp.status()), - Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), - } - }); - None + // + // moriio_write_prefill_str holds the serialized prefill body in WRITE mode so it + // can be moved into the join! closure later (prefill_request_str is consumed here). + let (moriio_write_prefill_str, prefill_response_json): (Option, Option) = + if is_moriio_write { + (Some(prefill_request_str), None) } else { debug!( "Stage 1: Sending prefill-only request (max_tokens=1) to prefill server at http://{}", @@ -752,8 +742,8 @@ impl VllmPDRouter { let prefill_json: Value = serde_json::from_str(&prefill_response_text) .map_err(|e| format!("Failed to parse prefill response as JSON: {}", e))?; - Some(prefill_json) - }; // end prefill_response_json + (None, Some(prefill_json)) + }; // end (moriio_write_prefill_str, prefill_response_json) // Prepare decode request let mut decode_request = request_json.clone(); @@ -808,6 +798,78 @@ impl VllmPDRouter { } let decode_request_url = format!("http://{}{}", decode_base_http, path); + + // WRITE mode: run prefill and decode concurrently via tokio::join! so the prefill + // task is always guaranteed to execute (unlike fire-and-forget tokio::spawn, which + // can be silently dropped under load) and both HTTP sends start at the same time. + if let Some(write_prefill_str) = moriio_write_prefill_str { + let http_client = self.http_client.clone(); + let prefill_url = format!("http://{}{}", prefill_base_http, path); + let prefill_request_id = request_id.clone(); + let prefill_dp_rank_copy = prefill_dp_rank; + let prefill_fut = async move { + let mut builder = http_client + .post(&prefill_url) + .header("Content-Type", "application/json") + .header("X-Request-Id", &prefill_request_id); + builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); + match builder.body(write_prefill_str).send().await { + Ok(resp) => debug!("MoRI-IO WRITE prefill completed with status {}", resp.status()), + Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), + } + }; + let decode_fut = otel_http::send_client_request( + decode_request_builder.body(decode_request_str), + headers, + ClientRequestOptions { + method: "POST", + url: &decode_request_url, + route: Some(path), + request_phase: Some("decode"), + }, + ); + let (_, decode_result) = tokio::join!(prefill_fut, decode_fut); + let decode_response = match decode_result { + Ok(resp) => resp, + Err(e) => { + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_decode_error(decode_http); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(prefill_http); + return Err(format!( + "Decode request failed to {}: {}", + decode_http, full_error + )); + } + }; + debug!("Decode server responded with status: {}", decode_response.status()); + // Stop profiling and record metrics then return decode response + self.stop_profiling(&format!("http://{}", decode_base_http)).await; + let duration = start_time.elapsed(); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(prefill_http); + RouterMetrics::record_pd_decode_request(decode_http); + if !decode_response.status().is_success() { + RouterMetrics::record_pd_decode_error(decode_http); + } + let status = decode_response.status(); + let resp_headers = decode_response.headers().clone(); + let body = decode_response + .bytes() + .await + .map_err(|e| format!("Failed to read decode response: {}", e))?; + let mut response_builder = axum::http::Response::builder().status(status); + for (name, value) in resp_headers.iter() { + response_builder = response_builder.header(name, value); + } + return response_builder + .body(axum::body::Body::from(body)) + .map_err(|e| format!("Failed to build response: {}", e)); + } + let decode_response = match otel_http::send_client_request( decode_request_builder.body(decode_request_str), headers, @@ -1001,38 +1063,17 @@ impl VllmPDRouter { && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)); // Stage 1: dispatch prefill. - // WRITE mode: fire-and-forget (decode does not need the prefill response). + // WRITE mode: defer the prefill send until Stage 2 so we can run both concurrently + // via tokio::join!, guaranteeing the prefill task always executes. // READ mode: await and parse — decode params come from the prefill response body. - let prefill_response_json: Option = if is_moriio_write { - debug!("MoRI-IO WRITE mode: dispatching prefill fire-and-forget"); - let http_client = self.pd_router.client.clone(); - let prefill_request_clone = prefill_request.clone(); - let prefill_request_id = request_id.clone(); - let prefill_url_clone = prefill_url.clone(); - let prefill_dp_rank_copy = prefill_dp_rank; - tokio::spawn(async move { - let mut builder = http_client - .post(&prefill_url_clone) - .header("Content-Type", "application/json") - .header( - "Authorization", - format!( - "Bearer {}", - std::env::var("OPENAI_API_KEY").unwrap_or_default() - ), - ) - .header("X-Request-Id", &prefill_request_id); - builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); - match builder.json(&prefill_request_clone).send().await { - Ok(resp) => debug!( - "MoRI-IO WRITE prefill completed with status {}", - resp.status() - ), - Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), - } - }); + let moriio_write_prefill_request: Option = if is_moriio_write { prefill_worker.decrement_load(); decode_worker.increment_load(); + Some(prefill_request.clone()) + } else { + None + }; + let prefill_response_json: Option = if is_moriio_write { None } else { debug!( @@ -1159,6 +1200,7 @@ impl VllmPDRouter { }; // end prefill_response_json // Stage 2: Prepare decode request + // (moriio_write_prefill_request carries the prefill body for WRITE mode join below) let mut decode_request = original_request.clone(); if let Some(params) = self .build_decode_kv_params( @@ -1214,30 +1256,82 @@ impl VllmPDRouter { decode_request_builder = dp_utils::add_dp_rank_header(decode_request_builder, decode_dp_rank); - let decode_response = match otel_http::send_client_request( - decode_request_builder.json(&decode_request), - headers, - ClientRequestOptions { - method: "POST", - url: &decode_url, - route: Some(path), - request_phase: Some("decode"), - }, - ) - .await - { - Ok(resp) => resp, - Err(e) => { - decode_worker.decrement_load(); - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_decode_error(&decode_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - RouterMetrics::record_pd_prefill_request(&prefill_base_url); - return Err(PDRouterError::NetworkError { - message: format!("Decode request failed to {}: {}", decode_url, full_error), - }); + // WRITE mode: run prefill and decode concurrently via tokio::join! + let decode_response = if let Some(write_prefill_req) = moriio_write_prefill_request { + let http_client = self.pd_router.client.clone(); + let prefill_url_clone = prefill_url.clone(); + let prefill_request_id = request_id.clone(); + let prefill_dp_rank_copy = prefill_dp_rank; + let prefill_fut = async move { + let mut builder = http_client + .post(&prefill_url_clone) + .header("Content-Type", "application/json") + .header( + "Authorization", + format!( + "Bearer {}", + std::env::var("OPENAI_API_KEY").unwrap_or_default() + ), + ) + .header("X-Request-Id", &prefill_request_id); + builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); + match builder.json(&write_prefill_req).send().await { + Ok(resp) => debug!("MoRI-IO WRITE prefill completed with status {}", resp.status()), + Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), + } + }; + let decode_fut = otel_http::send_client_request( + decode_request_builder.json(&decode_request), + headers, + ClientRequestOptions { + method: "POST", + url: &decode_url, + route: Some(path), + request_phase: Some("decode"), + }, + ); + let (_, decode_result) = tokio::join!(prefill_fut, decode_fut); + match decode_result { + Ok(resp) => resp, + Err(e) => { + decode_worker.decrement_load(); + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_decode_error(&decode_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(&prefill_base_url); + return Err(PDRouterError::NetworkError { + message: format!("Decode request failed to {}: {}", decode_url, full_error), + }); + } + } + } else { + match otel_http::send_client_request( + decode_request_builder.json(&decode_request), + headers, + ClientRequestOptions { + method: "POST", + url: &decode_url, + route: Some(path), + request_phase: Some("decode"), + }, + ) + .await + { + Ok(resp) => resp, + Err(e) => { + decode_worker.decrement_load(); + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_decode_error(&decode_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(&prefill_base_url); + return Err(PDRouterError::NetworkError { + message: format!("Decode request failed to {}: {}", decode_url, full_error), + }); + } } }; From da3902367f14a6def2f4ab26fd834226b67dd285 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 07:52:11 +0000 Subject: [PATCH 05/29] chore: move docstring to generate_transfer_id Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 2f822b9c..cb0adeb0 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -144,12 +144,12 @@ impl VllmPDRouter { )) } - /// Generate a connector-specific transfer ID for correlating prefill and decode legs. /// Returns the MoRI-IO transfer mode if it has been set by a registration, or `None`. fn moriio_transfer_mode(&self) -> Option { self.service_registry.moriio_transfer_mode.get().copied() } + /// Generate a connector-specific transfer ID for correlating prefill and decode legs. /// Returns `None` for connectors that do not use a transfer_id (e.g. NIXL). fn generate_transfer_id(&self) -> Option { match self.kv_connector { From 7d2f7e2d5138bead7b526e0a651c77be6875f988 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 07:54:36 +0000 Subject: [PATCH 06/29] chroe: rename build_decode_kv_params Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index cb0adeb0..bd309c32 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -215,7 +215,7 @@ impl VllmPDRouter { /// Build decode kv_transfer_params for all connectors. Returns `None` if params cannot be /// determined (e.g. Mooncake bootstrap info missing, or no kv_transfer_params in prefill response). - async fn build_decode_kv_params( + async fn build_decode_kv_transfer_params( &self, prefill_url: &str, prefill_response_json: Option<&Value>, @@ -749,7 +749,7 @@ impl VllmPDRouter { let mut decode_request = request_json.clone(); let prefill_url_key = format!("http://{}", prefill_base_http); if let Some(params) = self - .build_decode_kv_params( + .build_decode_kv_transfer_params( &prefill_url_key, prefill_response_json.as_ref(), transfer_id.as_deref(), @@ -1232,7 +1232,7 @@ impl VllmPDRouter { // (moriio_write_prefill_request carries the prefill body for WRITE mode join below) let mut decode_request = original_request.clone(); if let Some(params) = self - .build_decode_kv_params( + .build_decode_kv_transfer_params( &prefill_base_url, prefill_response_json.as_ref(), transfer_id.as_deref(), From 0846eeddb2c499604d9da88beeeae21e8d722782 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 08:14:33 +0000 Subject: [PATCH 07/29] docs: mention mori in module docstring in vllm_service_discovery Signed-off-by: simondanielsson --- src/routers/http/vllm_service_discovery.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/routers/http/vllm_service_discovery.rs b/src/routers/http/vllm_service_discovery.rs index 7364898e..ec1512b2 100644 --- a/src/routers/http/vllm_service_discovery.rs +++ b/src/routers/http/vllm_service_discovery.rs @@ -1,5 +1,6 @@ // vLLM Service Discovery Implementation -// This module implements service discovery for vLLM P2P NCCL coordination +// This module implements service discovery for vLLM P2P NCCL and MoRI-IO +// connector coordination use crate::config::KvConnector; use serde::{Deserialize, Serialize}; From 79d7efd785e8ff23bff3e96b2b3145a7223b023d Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 08:14:51 +0000 Subject: [PATCH 08/29] chore: cargo fmt Signed-off-by: simondanielsson --- src/config/validation.rs | 4 +- src/routers/http/vllm_pd_router.rs | 153 ++++++++++++++++------------- 2 files changed, 85 insertions(+), 72 deletions(-) diff --git a/src/config/validation.rs b/src/config/validation.rs index acea1c9d..cd822b2d 100644 --- a/src/config/validation.rs +++ b/src/config/validation.rs @@ -763,7 +763,9 @@ mod tests { let result = ConfigValidator::validate(&config); assert!(result.is_err()); if let Err(e) = result { - assert!(e.to_string().contains("MoRI-IO KV connector requires service discovery")); + assert!(e + .to_string() + .contains("MoRI-IO KV connector requires service discovery")); } } diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index bd309c32..f6a2aa23 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -224,8 +224,9 @@ impl VllmPDRouter { ) -> Option { match self.kv_connector { KvConnector::Mooncake => { - let Some((bootstrap_addr, engine_id)) = - self.get_mooncake_info(prefill_url, prefill_dp_rank.map(|r| r as usize)).await + let Some((bootstrap_addr, engine_id)) = self + .get_mooncake_info(prefill_url, prefill_dp_rank.map(|r| r as usize)) + .await else { warn!( "No Mooncake bootstrap info for prefill {}, decode will proceed without kv_transfer_params", @@ -264,9 +265,7 @@ impl VllmPDRouter { Some(params) } } - KvConnector::Nixl => { - Some(prefill_response_json?.get("kv_transfer_params")?.clone()) - } + KvConnector::Nixl => Some(prefill_response_json?.get("kv_transfer_params")?.clone()), } } @@ -672,78 +671,82 @@ impl VllmPDRouter { // moriio_write_prefill_str holds the serialized prefill body in WRITE mode so it // can be moved into the join! closure later (prefill_request_str is consumed here). let (moriio_write_prefill_str, prefill_response_json): (Option, Option) = - if is_moriio_write { - (Some(prefill_request_str), None) - } else { - debug!( + if is_moriio_write { + (Some(prefill_request_str), None) + } else { + debug!( "Stage 1: Sending prefill-only request (max_tokens=1) to prefill server at http://{}", prefill_http ); - self.start_profiling(&format!("http://{}", prefill_base_http)).await; + self.start_profiling(&format!("http://{}", prefill_base_http)) + .await; - let mut prefill_request_builder = self - .http_client - .post(format!("http://{}{}", prefill_base_http, path)) - .header("Content-Type", "application/json") - .header("X-Request-Id", &request_id); - prefill_request_builder = - dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); + let mut prefill_request_builder = self + .http_client + .post(format!("http://{}{}", prefill_base_http, path)) + .header("Content-Type", "application/json") + .header("X-Request-Id", &request_id); + prefill_request_builder = + dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); - let prefill_request_url = format!("http://{}{}", prefill_base_http, path); - let prefill_response = match otel_http::send_client_request( - prefill_request_builder.body(prefill_request_str), - headers, - ClientRequestOptions { - method: "POST", - url: &prefill_request_url, - route: Some(path), - request_phase: Some("prefill"), - }, - ) - .await - { - Ok(resp) => resp, - Err(e) => { - let full_error = error_chain(&e); + let prefill_request_url = format!("http://{}{}", prefill_base_http, path); + let prefill_response = match otel_http::send_client_request( + prefill_request_builder.body(prefill_request_str), + headers, + ClientRequestOptions { + method: "POST", + url: &prefill_request_url, + route: Some(path), + request_phase: Some("prefill"), + }, + ) + .await + { + Ok(resp) => resp, + Err(e) => { + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(prefill_http); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + return Err(format!( + "Prefill request failed to {}: {}", + prefill_http, full_error + )); + } + }; + + let prefill_status = prefill_response.status(); + debug!("Prefill server responded with status: {}", prefill_status); + + if !prefill_status.is_success() { let duration = start_time.elapsed(); RouterMetrics::record_pd_prefill_error(prefill_http); RouterMetrics::record_pd_request(path); RouterMetrics::record_pd_request_duration(path, duration); - return Err(format!("Prefill request failed to {}: {}", prefill_http, full_error)); + let error_body = prefill_response.text().await.unwrap_or_default(); + return Err(format!( + "Prefill server error {}: {}", + prefill_status, error_body + )); } - }; - - let prefill_status = prefill_response.status(); - debug!("Prefill server responded with status: {}", prefill_status); - if !prefill_status.is_success() { - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(prefill_http); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - let error_body = prefill_response.text().await.unwrap_or_default(); - return Err(format!( - "Prefill server error {}: {}", - prefill_status, error_body - )); - } - - // Extract kv_transfer_params from prefill response - let prefill_response_text = prefill_response.text().await.map_err(|e| { - let full_error = error_chain(&e); - format!( - "Failed to read prefill response from {}: {}", - prefill_http, full_error - ) - })?; + // Extract kv_transfer_params from prefill response + let prefill_response_text = prefill_response.text().await.map_err(|e| { + let full_error = error_chain(&e); + format!( + "Failed to read prefill response from {}: {}", + prefill_http, full_error + ) + })?; - debug!("Prefill response body: {}", prefill_response_text); + debug!("Prefill response body: {}", prefill_response_text); - let prefill_json: Value = serde_json::from_str(&prefill_response_text) - .map_err(|e| format!("Failed to parse prefill response as JSON: {}", e))?; + let prefill_json: Value = serde_json::from_str(&prefill_response_text) + .map_err(|e| format!("Failed to parse prefill response as JSON: {}", e))?; - (None, Some(prefill_json)) - }; // end (moriio_write_prefill_str, prefill_response_json) + (None, Some(prefill_json)) + }; // end (moriio_write_prefill_str, prefill_response_json) // Prepare decode request let mut decode_request = request_json.clone(); @@ -814,7 +817,10 @@ impl VllmPDRouter { .header("X-Request-Id", &prefill_request_id); builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); match builder.body(write_prefill_str).send().await { - Ok(resp) => debug!("MoRI-IO WRITE prefill completed with status {}", resp.status()), + Ok(resp) => debug!( + "MoRI-IO WRITE prefill completed with status {}", + resp.status() + ), Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), } }; @@ -844,9 +850,13 @@ impl VllmPDRouter { )); } }; - debug!("Decode server responded with status: {}", decode_response.status()); + debug!( + "Decode server responded with status: {}", + decode_response.status() + ); // Stop profiling and record metrics then return decode response - self.stop_profiling(&format!("http://{}", decode_base_http)).await; + self.stop_profiling(&format!("http://{}", decode_base_http)) + .await; let duration = start_time.elapsed(); RouterMetrics::record_pd_request(path); RouterMetrics::record_pd_request_duration(path, duration); @@ -949,8 +959,7 @@ impl VllmPDRouter { // Merge logprobs from prefill into decode response let empty_json = Value::Null; let prefill_json_ref = prefill_response_json.as_ref().unwrap_or(&empty_json); - let merged = - logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); + let merged = logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); if merged { debug!("Successfully merged logprobs from prefill and decode responses"); } else { @@ -1305,7 +1314,10 @@ impl VllmPDRouter { .header("X-Request-Id", &prefill_request_id); builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); match builder.json(&write_prefill_req).send().await { - Ok(resp) => debug!("MoRI-IO WRITE prefill completed with status {}", resp.status()), + Ok(resp) => debug!( + "MoRI-IO WRITE prefill completed with status {}", + resp.status() + ), Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), } }; @@ -1423,8 +1435,7 @@ impl VllmPDRouter { // Merge logprobs from prefill into decode response let empty_json = Value::Null; let prefill_json_ref = prefill_response_json.as_ref().unwrap_or(&empty_json); - let merged = - logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); + let merged = logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); if merged { debug!("Successfully merged logprobs from prefill and decode responses"); } else { From 12f7fbaabad6a585db5099cac85456939b7e5d3f Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 08:23:04 +0000 Subject: [PATCH 09/29] chore: shorten docstring for MoriIOServiceRegistratoin Signed-off-by: simondanielsson --- src/routers/http/vllm_service_discovery.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/routers/http/vllm_service_discovery.rs b/src/routers/http/vllm_service_discovery.rs index ec1512b2..5f6de5f2 100644 --- a/src/routers/http/vllm_service_discovery.rs +++ b/src/routers/http/vllm_service_discovery.rs @@ -45,8 +45,7 @@ pub enum MoriIOTransferMode { Write, } -/// MoRI-IO-specific service registration data. The base fields are flattened so -/// the wire format is identical to `ServiceRegistration` plus `transfer_mode`. +/// MoRI-IO-specific service registration data. #[derive(Debug, Clone, Deserialize)] pub struct MoriIOServiceRegistration { #[serde(flatten)] From 6ddf9af5c597e4b682df025905844cda982be083 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 08:32:54 +0000 Subject: [PATCH 10/29] chore: add trackking issue in comment for why remote_tp_size is 1 Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index f6a2aa23..440d4e8f 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -185,6 +185,8 @@ impl VllmPDRouter { "remote_engine_id": serde_json::Value::Null, "remote_block_ids": serde_json::Value::Null, "remote_dp_size": self.intra_node_data_parallel_size, + // remote_tp_size is not yet consumed by the vLLM MoRI-IO connector; + // hardcoded to 1 until https://github.com/vllm-project/vllm/issues/41211 is resolved. "remote_tp_size": 1, "transfer_id": transfer_id.unwrap_or(""), }) @@ -250,6 +252,8 @@ impl VllmPDRouter { "remote_block_ids": serde_json::Value::Null, "transfer_id": transfer_id.unwrap_or(""), "remote_dp_size": self.intra_node_data_parallel_size, + // remote_tp_size is not yet consumed by the vLLM MoRI-IO connector; + // hardcoded to 1 until https://github.com/vllm-project/vllm/issues/41211 is resolved. "remote_tp_size": 1, }); if self.intra_node_data_parallel_size > 1 { @@ -2482,6 +2486,8 @@ mod tests { "remote_engine_id": serde_json::Value::Null, "remote_block_ids": serde_json::Value::Null, "remote_dp_size": dp_size, + // remote_tp_size is not yet consumed by the vLLM MoRI-IO connector; + // hardcoded to 1 until https://github.com/vllm-project/vllm/issues/41211 is resolved. "remote_tp_size": 1, "transfer_id": transfer_id.unwrap_or(""), }) @@ -2500,6 +2506,8 @@ mod tests { "remote_block_ids": serde_json::Value::Null, "transfer_id": transfer_id.unwrap_or(""), "remote_dp_size": dp_size, + // remote_tp_size is not yet consumed by the vLLM MoRI-IO connector; + // hardcoded to 1 until https://github.com/vllm-project/vllm/issues/41211 is resolved. "remote_tp_size": 1, }); if dp_size > 1 { From d605ce2268a744fb37bde1e83a4852fb7fb20205 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 08:55:51 +0000 Subject: [PATCH 11/29] feat: error if dispatching before registration Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 64 +++++++++++++++++------------- 1 file changed, 36 insertions(+), 28 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 440d4e8f..5828ee3c 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -165,21 +165,30 @@ impl VllmPDRouter { } } - /// Build kv_transfer_params for the prefill request - fn build_prefill_kv_transfer_params(&self, transfer_id: Option<&str>) -> Value { + /// Build kv_transfer_params for the prefill request. + /// + /// Returns an error for MoRI-IO when no transfer mode has been registered yet, so that + /// requests are not silently dispatched in READ mode when no instances have registered. + fn build_prefill_kv_transfer_params( + &self, + transfer_id: Option<&str>, + ) -> Result { match self.kv_connector { - KvConnector::Mooncake => { - json!({ - "do_remote_decode": true, - "do_remote_prefill": false, - "transfer_id": transfer_id.unwrap_or(""), - }) - } + KvConnector::Mooncake => Ok(json!({ + "do_remote_decode": true, + "do_remote_prefill": false, + "transfer_id": transfer_id.unwrap_or(""), + })), KvConnector::MoriIO => { - if matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)) { + let mode = self.moriio_transfer_mode().ok_or_else(|| { + "No MoRI-IO instances have registered a transfer mode; \ + cannot dispatch request without knowing READ vs WRITE mode" + .to_string() + })?; + if matches!(mode, MoriIOTransferMode::Write) { // WRITE mode: prefill pushes KV blocks to decode. // do_remote_decode=true tells the prefill connector to initiate the transfer. - json!({ + Ok(json!({ "do_remote_decode": true, "do_remote_prefill": false, "remote_engine_id": serde_json::Value::Null, @@ -189,29 +198,27 @@ impl VllmPDRouter { // hardcoded to 1 until https://github.com/vllm-project/vllm/issues/41211 is resolved. "remote_tp_size": 1, "transfer_id": transfer_id.unwrap_or(""), - }) + })) } else { - // READ mode (default): prefill waits for decode to pull blocks. - json!({ + // READ mode: prefill waits for decode to pull blocks. + Ok(json!({ "do_remote_decode": true, "do_remote_prefill": false, "remote_engine_id": serde_json::Value::Null, "remote_block_ids": serde_json::Value::Null, "transfer_id": transfer_id.unwrap_or(""), "remote_dp_size": self.intra_node_data_parallel_size, - }) + })) } } - KvConnector::Nixl => { - 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 - }) - } + KvConnector::Nixl => Ok(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 + })), } } @@ -634,7 +641,7 @@ impl VllmPDRouter { // Add kv_transfer_params for KV connector support at top level prefill_request["kv_transfer_params"] = - self.build_prefill_kv_transfer_params(transfer_id.as_deref()); + self.build_prefill_kv_transfer_params(transfer_id.as_deref())?; debug!( "Added kv_transfer_params to prefill request for {:?} connector", @@ -1083,8 +1090,9 @@ impl VllmPDRouter { let transfer_id = self.generate_transfer_id(); // Add kv_transfer_params for KV connector support at top level - prefill_request["kv_transfer_params"] = - self.build_prefill_kv_transfer_params(transfer_id.as_deref()); + prefill_request["kv_transfer_params"] = self + .build_prefill_kv_transfer_params(transfer_id.as_deref()) + .map_err(|reason| PDRouterError::InvalidConfiguration { reason })?; debug!( "Added kv_transfer_params to prefill request for {:?} connector", From 4528365af8c14933a48a7386d1384e07679b4ebc Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 09:10:40 +0000 Subject: [PATCH 12/29] refactor: make more DRY with extract_base_http_and_dp_ran Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 39 +++++++++++++++--------------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 5828ee3c..45b23952 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -64,6 +64,22 @@ pub struct VllmPDRouter { /// Must match `MoRIIOConstants.TRANSFER_PREFIX` in the vLLM Python connector. const MORIIO_TRANSFER_PREFIX: &str = "tx"; +/// Strip the DP-rank suffix from a worker's HTTP address and return the base address +/// plus the parsed rank. Returns `(original, None)` when DP is disabled. +fn extract_base_http_and_dp_rank( + http: &str, + intra_node_data_parallel_size: usize, +) -> (String, Option) { + if intra_node_data_parallel_size > 1 { + let url = format!("http://{}", http); + let (base, rank) = dp_utils::parse_worker_url(&url); + let base_http = base.replace("http://", "").replace("https://", ""); + (base_http, rank) + } else { + (http.to_string(), None) + } +} + impl VllmPDRouter { /// Query the Mooncake bootstrap server on a prefill node to get engine_id per dp_rank. /// Retries with backoff since the prefill server may not be ready at router startup. @@ -648,25 +664,10 @@ impl VllmPDRouter { self.kv_connector ); - // Extract dp_rank from prefill_http if intra_node_data_parallel_size > 1 - let (prefill_base_http, prefill_dp_rank) = if self.intra_node_data_parallel_size > 1 { - let prefill_url = format!("http://{}", prefill_http); - let (base, rank) = dp_utils::parse_worker_url(&prefill_url); - let base_http = base.replace("http://", "").replace("https://", ""); - (base_http, rank) - } else { - (prefill_http.to_string(), None) - }; - - // Extract dp_rank from decode_http if intra_node_data_parallel_size > 1 - let (decode_base_http, decode_dp_rank) = if self.intra_node_data_parallel_size > 1 { - let decode_url = format!("http://{}", decode_http); - let (base, rank) = dp_utils::parse_worker_url(&decode_url); - let base_http = base.replace("http://", "").replace("https://", ""); - (base_http, rank) - } else { - (decode_http.to_string(), None) - }; + let (prefill_base_http, prefill_dp_rank) = + extract_base_http_and_dp_rank(prefill_http, self.intra_node_data_parallel_size); + let (decode_base_http, decode_dp_rank) = + extract_base_http_and_dp_rank(decode_http, self.intra_node_data_parallel_size); let is_moriio_write = matches!(self.kv_connector, KvConnector::MoriIO) && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)); From 0f1d3fd502747d445632b85c54ac298cd0dfb8a4 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 09:22:08 +0000 Subject: [PATCH 13/29] Move prefill stop profiling into big else block Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 45b23952..22034cb7 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -677,7 +677,6 @@ impl VllmPDRouter { // Stage 1: dispatch prefill. // WRITE mode: defer the prefill send until Stage 2 so we can run both concurrently - // via tokio::join!, guaranteeing the prefill task always executes. // READ mode: await and parse — decode params come from the prefill response body. // // moriio_write_prefill_str holds the serialized prefill body in WRITE mode so it @@ -757,6 +756,10 @@ impl VllmPDRouter { let prefill_json: Value = serde_json::from_str(&prefill_response_text) .map_err(|e| format!("Failed to parse prefill response as JSON: {}", e))?; + // Stop profiling on prefill server once we have its response. + self.stop_profiling(&format!("http://{}", prefill_base_http)) + .await; + (None, Some(prefill_json)) }; // end (moriio_write_prefill_str, prefill_response_json) @@ -782,10 +785,6 @@ impl VllmPDRouter { let decode_request_str = serde_json::to_string(&decode_request) .map_err(|e| format!("Failed to serialize decode request: {}", e))?; - // Stop profiling on prefill server after its work is done - self.stop_profiling(&format!("http://{}", prefill_base_http)) - .await; - // Stage 2: Send to decode server with original request and same P2P coordination header debug!( "Stage 2: Sending original request to decode server at http://{}", From 988d99507d2908b91de8b86c887f9d24f6cc8495 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 09:22:27 +0000 Subject: [PATCH 14/29] chore: cargo fmt Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 22034cb7..94b73b4f 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -185,10 +185,7 @@ impl VllmPDRouter { /// /// Returns an error for MoRI-IO when no transfer mode has been registered yet, so that /// requests are not silently dispatched in READ mode when no instances have registered. - fn build_prefill_kv_transfer_params( - &self, - transfer_id: Option<&str>, - ) -> Result { + fn build_prefill_kv_transfer_params(&self, transfer_id: Option<&str>) -> Result { match self.kv_connector { KvConnector::Mooncake => Ok(json!({ "do_remote_decode": true, From ec42d8a41242697f5ae7dc33098bf42f0ee316f0 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 09:43:14 +0000 Subject: [PATCH 15/29] feat: support prefill profiling in WRITE mode Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 94b73b4f..38311176 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -814,10 +814,16 @@ impl VllmPDRouter { // task is always guaranteed to execute (unlike fire-and-forget tokio::spawn, which // can be silently dropped under load) and both HTTP sends start at the same time. if let Some(write_prefill_str) = moriio_write_prefill_str { + self.start_profiling(&format!("http://{}", prefill_base_http)) + .await; let http_client = self.http_client.clone(); let prefill_url = format!("http://{}{}", prefill_base_http, path); let prefill_request_id = request_id.clone(); let prefill_dp_rank_copy = prefill_dp_rank; + let enable_profiling = self.enable_profiling; + let profiling_tasks = self.profiling_tasks.clone(); + let pd_router = self.pd_router.clone(); + let prefill_base_http_clone = prefill_base_http.clone(); let prefill_fut = async move { let mut builder = http_client .post(&prefill_url) @@ -831,6 +837,14 @@ impl VllmPDRouter { ), Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), } + if enable_profiling { + let worker_url = format!("http://{}", prefill_base_http_clone); + let mut tasks = profiling_tasks.lock().await; + if let Some(handle) = tasks.remove(&worker_url) { + handle.abort(); + } + pd_router.stop_profiling(&worker_url).await; + } }; let decode_fut = otel_http::send_client_request( decode_request_builder.body(decode_request_str), From 7a42f50272c297dd580a7404ddf2185fc28a240a Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 10:49:41 +0000 Subject: [PATCH 16/29] refactor: make more DRY by adding handle_decode_response Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 291 ++++++++++++++--------------- 1 file changed, 140 insertions(+), 151 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 38311176..31d3ee52 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -618,6 +618,112 @@ impl VllmPDRouter { } } + /// Handle the decode response shared by both READ and WRITE dispatch paths. + /// Stops decode profiling, records metrics, then routes to streaming, logprobs-merge, + /// or plain full-body response depending on the original request. + async fn handle_decode_response( + &self, + decode_response: reqwest::Response, + prefill_response_json: Option<&Value>, + path: &str, + prefill_http: &str, + decode_http: &str, + decode_base_http: &str, + start_time: Instant, + is_streaming: bool, + needs_logprobs: bool, + ) -> Result { + debug!( + "Decode server responded with status: {}", + decode_response.status() + ); + + self.stop_profiling(&format!("http://{}", decode_base_http)) + .await; + + let duration = start_time.elapsed(); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(prefill_http); + RouterMetrics::record_pd_decode_request(decode_http); + + if !decode_response.status().is_success() { + RouterMetrics::record_pd_decode_error(decode_http); + } + + if needs_logprobs && !is_streaming { + debug!("Logprobs requested and non-streaming - merging prefill and decode logprobs"); + + let status = decode_response.status(); + let resp_headers = decode_response.headers().clone(); + let decode_body = decode_response + .bytes() + .await + .map_err(|e| format!("Failed to read decode response: {}", e))?; + + let mut decode_json: Value = serde_json::from_slice(&decode_body) + .map_err(|e| format!("Failed to parse decode response as JSON: {}", e))?; + + let empty_json = Value::Null; + let prefill_json_ref = prefill_response_json.unwrap_or(&empty_json); + let merged = logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); + if merged { + debug!("Successfully merged logprobs from prefill and decode responses"); + } else { + warn!("No logprobs were merged (might be expected if logprobs not in response)"); + } + + let merged_body = serde_json::to_vec(&decode_json) + .map_err(|e| format!("Failed to serialize merged response: {}", e))?; + + let mut response_builder = axum::http::Response::builder().status(status); + for (name, value) in resp_headers.iter() { + response_builder = response_builder.header(name, value); + } + return response_builder + .body(axum::body::Body::from(merged_body)) + .map_err(|e| format!("Failed to build response: {}", e)); + } + + debug!( + "No logprobs merging needed (streaming={}, needs_logprobs={})", + is_streaming, needs_logprobs + ); + + let status = decode_response.status(); + + if is_streaming { + let mut response_builder = axum::http::Response::builder().status(status); + let mut decode_headers = + header_utils::preserve_response_headers(decode_response.headers()); + decode_headers.remove(axum::http::header::CONTENT_LENGTH); + for (name, value) in decode_headers.iter() { + response_builder = response_builder.header(name, value); + } + let body = axum::body::Body::from_stream(decode_response.bytes_stream()); + return response_builder.body(body).map_err(|e| { + format!( + "Failed to build streaming response from {}: {}", + decode_http, e + ) + }); + } + + // Non-streaming, no logprobs: read entire body + let decode_headers = decode_response.headers().clone(); + let body = decode_response + .bytes() + .await + .map_err(|e| format!("Failed to read decode response: {}", e))?; + let mut response_builder = axum::http::Response::builder().status(status); + for (name, value) in decode_headers.iter() { + response_builder = response_builder.header(name, value); + } + response_builder + .body(axum::body::Body::from(body)) + .map_err(|e| format!("Failed to build response: {}", e)) + } + /// Two-stage request processing for vLLM disaggregated mode using discovered endpoints async fn process_vllm_two_stage_request_discovered( &self, @@ -669,6 +775,16 @@ impl VllmPDRouter { let is_moriio_write = matches!(self.kv_connector, KvConnector::MoriIO) && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)); + let needs_logprobs = request_json.get("logprobs").is_some() + || request_json + .get("echo") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let is_streaming = request_json + .get("stream") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let prefill_request_str = serde_json::to_string(&prefill_request) .map_err(|e| format!("Failed to serialize prefill request: {}", e))?; @@ -872,34 +988,19 @@ impl VllmPDRouter { )); } }; - debug!( - "Decode server responded with status: {}", - decode_response.status() - ); - // Stop profiling and record metrics then return decode response - self.stop_profiling(&format!("http://{}", decode_base_http)) + return self + .handle_decode_response( + decode_response, + None, // WRITE mode: no prefill response JSON for logprobs + path, + prefill_http, + decode_http, + &decode_base_http, + start_time, + is_streaming, + needs_logprobs, + ) .await; - let duration = start_time.elapsed(); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - RouterMetrics::record_pd_prefill_request(prefill_http); - RouterMetrics::record_pd_decode_request(decode_http); - if !decode_response.status().is_success() { - RouterMetrics::record_pd_decode_error(decode_http); - } - let status = decode_response.status(); - let resp_headers = decode_response.headers().clone(); - let body = decode_response - .bytes() - .await - .map_err(|e| format!("Failed to read decode response: {}", e))?; - let mut response_builder = axum::http::Response::builder().status(status); - for (name, value) in resp_headers.iter() { - response_builder = response_builder.header(name, value); - } - return response_builder - .body(axum::body::Body::from(body)) - .map_err(|e| format!("Failed to build response: {}", e)); } let decode_response = match otel_http::send_client_request( @@ -929,130 +1030,18 @@ impl VllmPDRouter { } }; - debug!( - "Decode server responded with status: {}", - decode_response.status() - ); - - // Stop profiling on decode server after response received - self.stop_profiling(&format!("http://{}", decode_base_http)) - .await; - - // Record PD metrics - let duration = start_time.elapsed(); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - RouterMetrics::record_pd_prefill_request(prefill_http); - RouterMetrics::record_pd_decode_request(decode_http); - - if !decode_response.status().is_success() { - RouterMetrics::record_pd_decode_error(decode_http); - } - - // Check if logprobs merging is needed - let needs_logprobs = request_json.get("logprobs").is_some() - || request_json - .get("echo") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - - // Check if original request was streaming - let is_streaming = request_json - .get("stream") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - // Accept header "text/event-stream" for streaming requests seems unnecessary. - - // If logprobs requested and non-streaming, merge prefill and decode logprobs - if needs_logprobs && !is_streaming { - debug!("Logprobs requested and non-streaming - merging prefill and decode logprobs"); - - let status = decode_response.status(); - let headers = decode_response.headers().clone(); - let decode_body = decode_response - .bytes() - .await - .map_err(|e| format!("Failed to read decode response: {}", e))?; - - // Parse decode response as JSON - let mut decode_json: Value = serde_json::from_slice(&decode_body) - .map_err(|e| format!("Failed to parse decode response as JSON: {}", e))?; - - // Merge logprobs from prefill into decode response - let empty_json = Value::Null; - let prefill_json_ref = prefill_response_json.as_ref().unwrap_or(&empty_json); - let merged = logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); - if merged { - debug!("Successfully merged logprobs from prefill and decode responses"); - } else { - warn!("No logprobs were merged (might be expected if logprobs not in response)"); - } - - // Serialize merged response - let merged_body = serde_json::to_vec(&decode_json) - .map_err(|e| format!("Failed to serialize merged response: {}", e))?; - - let mut response_builder = axum::http::Response::builder().status(status); - for (name, value) in headers.iter() { - response_builder = response_builder.header(name, value); - } - - let response = response_builder - .body(axum::body::Body::from(merged_body)) - .map_err(|e| format!("Failed to build response: {}", e))?; - - Ok(response) - } else { - // No logprobs merging needed - return decode response as-is (streaming or no logprobs) - debug!( - "No logprobs merging needed (streaming={}, needs_logprobs={})", - is_streaming, needs_logprobs - ); - - let status = decode_response.status(); - - // For streaming responses, use stream passthrough; otherwise read entire body - if is_streaming { - let mut response_builder = axum::http::Response::builder().status(status); - - // Preserve headers from decode response (filtering out hop-by-hop headers) - let mut decode_headers = - header_utils::preserve_response_headers(decode_response.headers()); - decode_headers.remove(axum::http::header::CONTENT_LENGTH); - for (name, value) in decode_headers.iter() { - response_builder = response_builder.header(name, value); - } - - // Stream the response body directly - let body = axum::body::Body::from_stream(decode_response.bytes_stream()); - let response = response_builder.body(body).map_err(|e| { - format!( - "Failed to build streaming response from {}: {}", - decode_http, e - ) - })?; - - Ok(response) - } else { - // Non-streaming: read entire body - let decode_headers = decode_response.headers().clone(); - let body = decode_response - .bytes() - .await - .map_err(|e| format!("Failed to read decode response: {}", e))?; - - let mut response_builder = axum::http::Response::builder().status(status); - for (name, value) in decode_headers.iter() { - response_builder = response_builder.header(name, value); - } - - let response = response_builder - .body(axum::body::Body::from(body)) - .map_err(|e| format!("Failed to build response: {}", e))?; - - Ok(response) - } - } + self.handle_decode_response( + decode_response, + prefill_response_json.as_ref(), + path, + prefill_http, + decode_http, + &decode_base_http, + start_time, + is_streaming, + needs_logprobs, + ) + .await } /// Two-stage request processing for vLLM disaggregated mode From e39f17ae25eccdca3ea55c46f835d5b778294e67 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 11:21:25 +0000 Subject: [PATCH 17/29] fix: revert all changes in process_vllm_two_stage_request as mori does not support it anyways Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 415 +++++++++++++---------------- 1 file changed, 188 insertions(+), 227 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 31d3ee52..a42b7bbd 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -1090,9 +1090,8 @@ impl VllmPDRouter { let transfer_id = self.generate_transfer_id(); // Add kv_transfer_params for KV connector support at top level - prefill_request["kv_transfer_params"] = self - .build_prefill_kv_transfer_params(transfer_id.as_deref()) - .map_err(|reason| PDRouterError::InvalidConfiguration { reason })?; + prefill_request["kv_transfer_params"] = + self.build_prefill_kv_transfer_params(transfer_id.as_deref()); debug!( "Added kv_transfer_params to prefill request for {:?} connector", @@ -1105,169 +1104,187 @@ impl VllmPDRouter { let prefill_dp_rank = prefill_worker.dp_rank(); let prefill_url = prefill_worker.endpoint_url(path); - let decode_base_url = decode_worker.base_url().to_string(); - let decode_dp_rank = decode_worker.dp_rank(); - let decode_url = decode_worker.endpoint_url(path); - - let is_moriio_write = matches!(self.kv_connector, KvConnector::MoriIO) - && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)); - - // Stage 1: dispatch prefill. - // WRITE mode: defer the prefill send until Stage 2 so we can run both concurrently - // via tokio::join!, guaranteeing the prefill task always executes. - // READ mode: await and parse — decode params come from the prefill response body. - let moriio_write_prefill_request: Option = if is_moriio_write { - prefill_worker.decrement_load(); - decode_worker.increment_load(); - Some(prefill_request.clone()) - } else { - None - }; - let prefill_response_json: Option = if is_moriio_write { - None + debug!( + "🚀 vLLM Stage 1 - Prefill: {} with request_id: {}", + prefill_url, request_id + ); + if let Some(rank) = prefill_dp_rank { + debug!("📤 Prefill request headers: Authorization=Bearer [REDACTED], X-Request-Id={}, X-data-parallel-rank={}", request_id, rank); } else { debug!( - "🚀 vLLM Stage 1 - Prefill: {} with request_id: {}", - prefill_url, request_id - ); - if let Some(rank) = prefill_dp_rank { - debug!("📤 Prefill request headers: Authorization=Bearer [REDACTED], X-Request-Id={}, X-data-parallel-rank={}", request_id, rank); - } else { - debug!( - "📤 Prefill request headers: Authorization=Bearer [REDACTED], X-Request-Id={}", - request_id - ); - } - debug!( - "📤 Prefill request payload: {}", - serde_json::to_string_pretty(&prefill_request).unwrap_or_default() + "📤 Prefill request headers: Authorization=Bearer [REDACTED], X-Request-Id={}", + request_id ); + } + debug!( + "📤 Prefill request payload: {}", + serde_json::to_string_pretty(&prefill_request).unwrap_or_default() + ); - self.start_profiling(&prefill_base_url).await; + // Start profiling on prefill server + self.start_profiling(&prefill_base_url).await; - let mut prefill_request_builder = self - .pd_router - .client - .post(&prefill_url) - .header("Content-Type", "application/json") - .header( - "Authorization", - format!( - "Bearer {}", - std::env::var("OPENAI_API_KEY").unwrap_or_default() - ), - ) - .header("X-Request-Id", &request_id); - prefill_request_builder = - dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); - - let prefill_response = match otel_http::send_client_request( - prefill_request_builder.json(&prefill_request), - headers, - ClientRequestOptions { - method: "POST", - url: &prefill_url, - route: Some(path), - request_phase: Some("prefill"), - }, + let mut prefill_request_builder = self + .pd_router + .client + .post(&prefill_url) + .header("Content-Type", "application/json") + .header( + "Authorization", + format!( + "Bearer {}", + std::env::var("OPENAI_API_KEY").unwrap_or_default() + ), ) - .await - { - Ok(resp) => resp, - Err(e) => { - prefill_worker.decrement_load(); - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(&prefill_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - return Err(PDRouterError::NetworkError { - message: format!( - "Prefill request failed to {}: {}", - prefill_url, full_error - ), - }); - } - }; + .header("X-Request-Id", &request_id); - debug!("📥 Prefill response status: {}", prefill_response.status()); - debug!( - "📥 Prefill response headers: {:?}", - prefill_response.headers() - ); + // Add X-data-parallel-rank header using shared utilities + prefill_request_builder = + dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); - let prefill_bytes = match prefill_response.bytes().await { - Ok(bytes) => bytes, - Err(e) => { - prefill_worker.decrement_load(); - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(&prefill_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - return Err(PDRouterError::NetworkError { - message: format!( - "Failed to read prefill response from {}: {}", - prefill_url, full_error - ), - }); - } - }; + let prefill_response = match otel_http::send_client_request( + prefill_request_builder.json(&prefill_request), + headers, + ClientRequestOptions { + method: "POST", + url: &prefill_url, + route: Some(path), + request_phase: Some("prefill"), + }, + ) + .await + { + Ok(resp) => resp, + Err(e) => { + prefill_worker.decrement_load(); + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + return Err(PDRouterError::NetworkError { + message: format!("Prefill request failed to {}: {}", prefill_url, full_error), + }); + } + }; + + debug!("📥 Prefill response status: {}", prefill_response.status()); + debug!( + "📥 Prefill response headers: {:?}", + prefill_response.headers() + ); + + // Extract prefill response body to get kv_transfer_params + let prefill_bytes = match prefill_response.bytes().await { + Ok(bytes) => bytes, + Err(e) => { + prefill_worker.decrement_load(); + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + return Err(PDRouterError::NetworkError { + message: format!( + "Failed to read prefill response from {}: {}", + prefill_url, full_error + ), + }); + } + }; + debug!( + "📥 Prefill response body size: {} bytes", + prefill_bytes.len() + ); + if prefill_bytes.len() < 1024 { debug!( - "📥 Prefill response body size: {} bytes", - prefill_bytes.len() + "📥 Prefill response body content: {}", + String::from_utf8_lossy(&prefill_bytes) ); - if prefill_bytes.len() < 1024 { - debug!( - "📥 Prefill response body content: {}", - String::from_utf8_lossy(&prefill_bytes) - ); + } + + // Parse prefill response to extract kv_transfer_params + let prefill_response_json: Value = match serde_json::from_slice(&prefill_bytes) { + Ok(json) => json, + Err(e) => { + prefill_worker.decrement_load(); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + return Err(PDRouterError::NetworkError { + message: format!("Failed to parse prefill response as JSON: {}", e), + }); } + }; - let prefill_json: Value = match serde_json::from_slice(&prefill_bytes) { - Ok(json) => json, - Err(e) => { - prefill_worker.decrement_load(); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(&prefill_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - return Err(PDRouterError::NetworkError { - message: format!("Failed to parse prefill response as JSON: {}", e), - }); - } - }; + // Extract kv_transfer_params from prefill response if present + let kv_transfer_params = prefill_response_json.get("kv_transfer_params").cloned(); - self.stop_profiling(&prefill_base_url).await; + if let Some(ref params) = kv_transfer_params { + debug!( + "Extracted kv_transfer_params from prefill response: {}", + serde_json::to_string_pretty(params).unwrap_or_default() + ); + } else { + debug!("No kv_transfer_params found in prefill response, will proceed without them"); + } - prefill_worker.decrement_load(); - decode_worker.increment_load(); + // Stop profiling on prefill server after its work is done + self.stop_profiling(&prefill_base_url).await; - debug!("✅ vLLM Stage 1 completed, starting Stage 2 - Decode"); + // Prefill phase complete: decrement prefill load, increment decode load + prefill_worker.decrement_load(); + decode_worker.increment_load(); - Some(prefill_json) - }; // end prefill_response_json + debug!("✅ vLLM Stage 1 completed, starting Stage 2 - Decode"); - // Stage 2: Prepare decode request - // (moriio_write_prefill_request carries the prefill body for WRITE mode join below) + // Stage 2: Prepare decode request with kv_transfer_params let mut decode_request = original_request.clone(); - if let Some(params) = self - .build_decode_kv_transfer_params( - &prefill_base_url, - prefill_response_json.as_ref(), - transfer_id.as_deref(), - prefill_dp_rank.map(|r| r as u32), - ) - .await - { - decode_request["kv_transfer_params"] = params; - debug!( - "Added kv_transfer_params to decode request for {:?} connector", - self.kv_connector - ); + if matches!(self.kv_connector, KvConnector::Mooncake) { + // Mooncake: set decode params proactively from bootstrap info + if let Some((bootstrap_addr, engine_id)) = self + .get_mooncake_info(&prefill_base_url, prefill_dp_rank) + .await + { + decode_request["kv_transfer_params"] = self + .build_mooncake_decode_kv_transfer_params( + transfer_id.as_deref().unwrap_or(""), + &bootstrap_addr, + &engine_id, + ); + debug!( + "Set Mooncake decode kv_transfer_params with bootstrap_addr={}, engine_id={}", + bootstrap_addr, engine_id + ); + } else { + warn!( + "No Mooncake bootstrap info for prefill {}, decode will proceed without kv_transfer_params", + prefill_base_url + ); + } + } else { + // NIXL and MoRI-IO: extract kv_transfer_params from prefill response + if let Some(mut params) = kv_transfer_params { + if matches!(self.kv_connector, KvConnector::MoriIO) { + // MoRI-IO decode connector needs to know how many prefill DP ranks to handshake with. + params["remote_dp_size"] = json!(self.intra_node_data_parallel_size); + } + decode_request["kv_transfer_params"] = params; + debug!( + "Added kv_transfer_params to decode request for {:?} connector", + self.kv_connector + ); + } } + // Use endpoint_url() to get the base URL without @rank suffix, + // avoiding IPv6+DP URL corruption (same fix as Router and PdRouterBase) + let decode_base_url = decode_worker.base_url().to_string(); + let decode_dp_rank = decode_worker.dp_rank(); + let decode_url = decode_worker.endpoint_url(path); + debug!( "🚀 vLLM Stage 2 - Decode: {} with request_id: {}", decode_url, request_id @@ -1306,85 +1323,30 @@ impl VllmPDRouter { decode_request_builder = dp_utils::add_dp_rank_header(decode_request_builder, decode_dp_rank); - // WRITE mode: run prefill and decode concurrently via tokio::join! - let decode_response = if let Some(write_prefill_req) = moriio_write_prefill_request { - let http_client = self.pd_router.client.clone(); - let prefill_url_clone = prefill_url.clone(); - let prefill_request_id = request_id.clone(); - let prefill_dp_rank_copy = prefill_dp_rank; - let prefill_fut = async move { - let mut builder = http_client - .post(&prefill_url_clone) - .header("Content-Type", "application/json") - .header( - "Authorization", - format!( - "Bearer {}", - std::env::var("OPENAI_API_KEY").unwrap_or_default() - ), - ) - .header("X-Request-Id", &prefill_request_id); - builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); - match builder.json(&write_prefill_req).send().await { - Ok(resp) => debug!( - "MoRI-IO WRITE prefill completed with status {}", - resp.status() - ), - Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), - } - }; - let decode_fut = otel_http::send_client_request( - decode_request_builder.json(&decode_request), - headers, - ClientRequestOptions { - method: "POST", - url: &decode_url, - route: Some(path), - request_phase: Some("decode"), - }, - ); - let (_, decode_result) = tokio::join!(prefill_fut, decode_fut); - match decode_result { - Ok(resp) => resp, - Err(e) => { - decode_worker.decrement_load(); - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_decode_error(&decode_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - RouterMetrics::record_pd_prefill_request(&prefill_base_url); - return Err(PDRouterError::NetworkError { - message: format!("Decode request failed to {}: {}", decode_url, full_error), - }); - } - } - } else { - match otel_http::send_client_request( - decode_request_builder.json(&decode_request), - headers, - ClientRequestOptions { - method: "POST", - url: &decode_url, - route: Some(path), - request_phase: Some("decode"), - }, - ) - .await - { - Ok(resp) => resp, - Err(e) => { - decode_worker.decrement_load(); - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_decode_error(&decode_base_url); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - RouterMetrics::record_pd_prefill_request(&prefill_base_url); - return Err(PDRouterError::NetworkError { - message: format!("Decode request failed to {}: {}", decode_url, full_error), - }); - } + let decode_response = match otel_http::send_client_request( + decode_request_builder.json(&decode_request), + headers, + ClientRequestOptions { + method: "POST", + url: &decode_url, + route: Some(path), + request_phase: Some("decode"), + }, + ) + .await + { + Ok(resp) => resp, + Err(e) => { + decode_worker.decrement_load(); + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_decode_error(&decode_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(&prefill_base_url); + return Err(PDRouterError::NetworkError { + message: format!("Decode request failed to {}: {}", decode_url, full_error), + }); } }; @@ -1445,9 +1407,8 @@ impl VllmPDRouter { })?; // Merge logprobs from prefill into decode response - let empty_json = Value::Null; - let prefill_json_ref = prefill_response_json.as_ref().unwrap_or(&empty_json); - let merged = logprobs_merge::merge_logprobs_in_json(prefill_json_ref, &mut decode_json); + let merged = + logprobs_merge::merge_logprobs_in_json(&prefill_response_json, &mut decode_json); if merged { debug!("Successfully merged logprobs from prefill and decode responses"); } else { From 2b497bdff90990d1bbf511785dd580b591d28fea Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 11:47:48 +0000 Subject: [PATCH 18/29] refactor: make more DRY! Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 187 +++++++++++++++-------------- 1 file changed, 100 insertions(+), 87 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index a42b7bbd..3f75fad3 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -80,6 +80,20 @@ fn extract_base_http_and_dp_rank( } } +/// Build a prefill reqwest::RequestBuilder with the standard headers and dp-rank header. +fn build_prefill_request_builder( + http_client: &reqwest::Client, + url: &str, + request_id: &str, + dp_rank: Option, +) -> reqwest::RequestBuilder { + let builder = http_client + .post(url) + .header("Content-Type", "application/json") + .header("X-Request-Id", request_id); + dp_utils::add_dp_rank_header(builder, dp_rank) +} + impl VllmPDRouter { /// Query the Mooncake bootstrap server on a prefill node to get engine_id per dp_rank. /// Retries with backoff since the prefill server may not be ready at router startup. @@ -788,93 +802,88 @@ impl VllmPDRouter { let prefill_request_str = serde_json::to_string(&prefill_request) .map_err(|e| format!("Failed to serialize prefill request: {}", e))?; + let prefill_request_url = format!("http://{}{}", prefill_base_http, path); + // Stage 1: dispatch prefill. - // WRITE mode: defer the prefill send until Stage 2 so we can run both concurrently - // READ mode: await and parse — decode params come from the prefill response body. - // - // moriio_write_prefill_str holds the serialized prefill body in WRITE mode so it - // can be moved into the join! closure later (prefill_request_str is consumed here). - let (moriio_write_prefill_str, prefill_response_json): (Option, Option) = - if is_moriio_write { - (Some(prefill_request_str), None) - } else { - debug!( + // READ mode: send now and await — decode kv_transfer_params come from the prefill response. + // WRITE mode: skip; prefill is sent concurrently with decode in Stage 2 via tokio::join!. + let prefill_response_json: Option = if is_moriio_write { + None + } else { + debug!( "Stage 1: Sending prefill-only request (max_tokens=1) to prefill server at http://{}", prefill_http ); - self.start_profiling(&format!("http://{}", prefill_base_http)) - .await; - - let mut prefill_request_builder = self - .http_client - .post(format!("http://{}{}", prefill_base_http, path)) - .header("Content-Type", "application/json") - .header("X-Request-Id", &request_id); - prefill_request_builder = - dp_utils::add_dp_rank_header(prefill_request_builder, prefill_dp_rank); - - let prefill_request_url = format!("http://{}{}", prefill_base_http, path); - let prefill_response = match otel_http::send_client_request( - prefill_request_builder.body(prefill_request_str), - headers, - ClientRequestOptions { - method: "POST", - url: &prefill_request_url, - route: Some(path), - request_phase: Some("prefill"), - }, - ) - .await - { - Ok(resp) => resp, - Err(e) => { - let full_error = error_chain(&e); - let duration = start_time.elapsed(); - RouterMetrics::record_pd_prefill_error(prefill_http); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - return Err(format!( - "Prefill request failed to {}: {}", - prefill_http, full_error - )); - } - }; - - let prefill_status = prefill_response.status(); - debug!("Prefill server responded with status: {}", prefill_status); + self.start_profiling(&format!("http://{}", prefill_base_http)) + .await; - if !prefill_status.is_success() { + let prefill_response = match otel_http::send_client_request( + build_prefill_request_builder( + &self.http_client, + &prefill_request_url, + &request_id, + prefill_dp_rank, + ) + .body(prefill_request_str.clone()), + headers, + ClientRequestOptions { + method: "POST", + url: &prefill_request_url, + route: Some(path), + request_phase: Some("prefill"), + }, + ) + .await + { + Ok(resp) => resp, + Err(e) => { + let full_error = error_chain(&e); let duration = start_time.elapsed(); RouterMetrics::record_pd_prefill_error(prefill_http); RouterMetrics::record_pd_request(path); RouterMetrics::record_pd_request_duration(path, duration); - let error_body = prefill_response.text().await.unwrap_or_default(); return Err(format!( - "Prefill server error {}: {}", - prefill_status, error_body + "Prefill request failed to {}: {}", + prefill_http, full_error )); } + }; - // Extract kv_transfer_params from prefill response - let prefill_response_text = prefill_response.text().await.map_err(|e| { - let full_error = error_chain(&e); - format!( - "Failed to read prefill response from {}: {}", - prefill_http, full_error - ) - })?; + let prefill_status = prefill_response.status(); + debug!("Prefill server responded with status: {}", prefill_status); - debug!("Prefill response body: {}", prefill_response_text); + if !prefill_status.is_success() { + let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(prefill_http); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + let error_body = prefill_response.text().await.unwrap_or_default(); + return Err(format!( + "Prefill server error {}: {}", + prefill_status, error_body + )); + } - let prefill_json: Value = serde_json::from_str(&prefill_response_text) - .map_err(|e| format!("Failed to parse prefill response as JSON: {}", e))?; + // Extract kv_transfer_params from prefill response + let prefill_response_text = prefill_response.text().await.map_err(|e| { + let full_error = error_chain(&e); + format!( + "Failed to read prefill response from {}: {}", + prefill_http, full_error + ) + })?; + + debug!("Prefill response body: {}", prefill_response_text); - // Stop profiling on prefill server once we have its response. - self.stop_profiling(&format!("http://{}", prefill_base_http)) - .await; + let prefill_json: Value = serde_json::from_str(&prefill_response_text) + .map_err(|e| format!("Failed to parse prefill response as JSON: {}", e))?; - (None, Some(prefill_json)) - }; // end (moriio_write_prefill_str, prefill_response_json) + // Stop profiling on prefill server once we have its response. + self.stop_profiling(&format!("http://{}", prefill_base_http)) + .await; + + Some(prefill_json) + }; // Prepare decode request let mut decode_request = request_json.clone(); @@ -929,24 +938,27 @@ impl VllmPDRouter { // WRITE mode: run prefill and decode concurrently via tokio::join! so the prefill // task is always guaranteed to execute (unlike fire-and-forget tokio::spawn, which // can be silently dropped under load) and both HTTP sends start at the same time. - if let Some(write_prefill_str) = moriio_write_prefill_str { + if is_moriio_write { self.start_profiling(&format!("http://{}", prefill_base_http)) .await; - let http_client = self.http_client.clone(); - let prefill_url = format!("http://{}{}", prefill_base_http, path); - let prefill_request_id = request_id.clone(); - let prefill_dp_rank_copy = prefill_dp_rank; + // Capture references rather than clones — tokio::join! polls both futures on the + // same task so they don't need to be Send, and local borrows are valid for the + // duration of the join. + let http_client = &self.http_client; let enable_profiling = self.enable_profiling; - let profiling_tasks = self.profiling_tasks.clone(); - let pd_router = self.pd_router.clone(); - let prefill_base_http_clone = prefill_base_http.clone(); + let profiling_tasks = &self.profiling_tasks; + let pd_router = &self.pd_router; let prefill_fut = async move { - let mut builder = http_client - .post(&prefill_url) - .header("Content-Type", "application/json") - .header("X-Request-Id", &prefill_request_id); - builder = dp_utils::add_dp_rank_header(builder, prefill_dp_rank_copy); - match builder.body(write_prefill_str).send().await { + match build_prefill_request_builder( + http_client, + &prefill_request_url, + &request_id, + prefill_dp_rank, + ) + .body(prefill_request_str) + .send() + .await + { Ok(resp) => debug!( "MoRI-IO WRITE prefill completed with status {}", resp.status() @@ -954,7 +966,7 @@ impl VllmPDRouter { Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), } if enable_profiling { - let worker_url = format!("http://{}", prefill_base_http_clone); + let worker_url = format!("http://{}", prefill_base_http); let mut tasks = profiling_tasks.lock().await; if let Some(handle) = tasks.remove(&worker_url) { handle.abort(); @@ -1090,8 +1102,9 @@ impl VllmPDRouter { let transfer_id = self.generate_transfer_id(); // Add kv_transfer_params for KV connector support at top level - prefill_request["kv_transfer_params"] = - self.build_prefill_kv_transfer_params(transfer_id.as_deref()); + prefill_request["kv_transfer_params"] = self + .build_prefill_kv_transfer_params(transfer_id.as_deref()) + .map_err(|reason| PDRouterError::InvalidConfiguration { reason })?; debug!( "Added kv_transfer_params to prefill request for {:?} connector", From d3d623322a314558d720cad988213d06fb59a667 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 29 Apr 2026 12:01:48 +0000 Subject: [PATCH 19/29] chore: fix clippy Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 1 + src/routers/http/vllm_service_discovery.rs | 7 +++---- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 3f75fad3..5787ebef 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -635,6 +635,7 @@ impl VllmPDRouter { /// Handle the decode response shared by both READ and WRITE dispatch paths. /// Stops decode profiling, records metrics, then routes to streaming, logprobs-merge, /// or plain full-body response depending on the original request. + #[allow(clippy::too_many_arguments)] async fn handle_decode_response( &self, decode_response: reqwest::Response, diff --git a/src/routers/http/vllm_service_discovery.rs b/src/routers/http/vllm_service_discovery.rs index 5f6de5f2..cb5073be 100644 --- a/src/routers/http/vllm_service_discovery.rs +++ b/src/routers/http/vllm_service_discovery.rs @@ -446,15 +446,14 @@ mod tests { use super::*; fn make_moriio_msgpack(service_type: &str, transfer_mode: &str) -> Vec { - let mut map = rmp_serde::to_vec_named(&serde_json::json!({ + // rmp_serde::to_vec_named produces msgpack — suitable for our deserializer. + rmp_serde::to_vec_named(&serde_json::json!({ "type": service_type, "http_address": "1.2.3.4:8000", "zmq_address": "host:1.2.3.4,handshake:6301,notify:61005", "transfer_mode": transfer_mode, })) - .unwrap(); - // rmp_serde::to_vec_named produces msgpack — suitable for our deserializer. - map + .unwrap() } #[test] From 898650b3d1c40362ada7983735d12bd390e5c513 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 30 Apr 2026 10:08:13 +0200 Subject: [PATCH 20/29] fix: propagate prefill errors Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 41 ++++++++++++++++++++++-------- 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 5787ebef..8e38a8f9 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -950,7 +950,7 @@ impl VllmPDRouter { let profiling_tasks = &self.profiling_tasks; let pd_router = &self.pd_router; let prefill_fut = async move { - match build_prefill_request_builder( + let result = build_prefill_request_builder( http_client, &prefill_request_url, &request_id, @@ -958,14 +958,7 @@ impl VllmPDRouter { ) .body(prefill_request_str) .send() - .await - { - Ok(resp) => debug!( - "MoRI-IO WRITE prefill completed with status {}", - resp.status() - ), - Err(e) => warn!("MoRI-IO WRITE prefill request failed: {}", e), - } + .await; if enable_profiling { let worker_url = format!("http://{}", prefill_base_http); let mut tasks = profiling_tasks.lock().await; @@ -974,6 +967,24 @@ impl VllmPDRouter { } pd_router.stop_profiling(&worker_url).await; } + match result { + Ok(resp) if resp.status().is_success() => { + debug!( + "MoRI-IO WRITE prefill completed with status {}", + resp.status() + ); + Ok(()) + } + Ok(resp) => { + let status = resp.status(); + warn!("MoRI-IO WRITE prefill returned non-2xx status: {}", status); + Err(format!("Prefill request failed with status {}", status)) + } + Err(e) => { + warn!("MoRI-IO WRITE prefill request failed: {}", e); + Err(format!("Prefill request failed: {}", e)) + } + } }; let decode_fut = otel_http::send_client_request( decode_request_builder.body(decode_request_str), @@ -985,7 +996,17 @@ impl VllmPDRouter { request_phase: Some("decode"), }, ); - let (_, decode_result) = tokio::join!(prefill_fut, decode_fut); + let (prefill_result, decode_result) = tokio::join!(prefill_fut, decode_fut); + if let Err(prefill_err) = prefill_result { + let duration = start_time.elapsed(); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(prefill_http); + return Err(format!( + "Prefill request failed to {}: {}", + prefill_http, prefill_err + )); + } let decode_response = match decode_result { Ok(resp) => resp, Err(e) => { From 1a820dda1742d832c070a463755a09bd8c6a90c2 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 30 Apr 2026 10:36:03 +0200 Subject: [PATCH 21/29] fix: support logsprobs from prefill also in WRITE mode Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 53 +++++++++++++++++++++--------- 1 file changed, 37 insertions(+), 16 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 8e38a8f9..71156139 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -969,11 +969,29 @@ impl VllmPDRouter { } match result { Ok(resp) if resp.status().is_success() => { - debug!( - "MoRI-IO WRITE prefill completed with status {}", - resp.status() - ); - Ok(()) + let status = resp.status(); + match resp.bytes().await { + Ok(bytes) => match serde_json::from_slice::(&bytes) { + Ok(json) => { + debug!( + "MoRI-IO WRITE prefill completed with status {}", + status + ); + Ok(Some(json)) + } + Err(_) => { + debug!( + "MoRI-IO WRITE prefill completed with status {} (non-JSON body)", + status + ); + Ok(None) + } + }, + Err(e) => { + warn!("MoRI-IO WRITE prefill: failed to read response body: {}", e); + Ok(None) + } + } } Ok(resp) => { let status = resp.status(); @@ -997,16 +1015,19 @@ impl VllmPDRouter { }, ); let (prefill_result, decode_result) = tokio::join!(prefill_fut, decode_fut); - if let Err(prefill_err) = prefill_result { - let duration = start_time.elapsed(); - RouterMetrics::record_pd_request(path); - RouterMetrics::record_pd_request_duration(path, duration); - RouterMetrics::record_pd_prefill_request(prefill_http); - return Err(format!( - "Prefill request failed to {}: {}", - prefill_http, prefill_err - )); - } + let write_prefill_response_json: Option = match prefill_result { + Err(prefill_err) => { + let duration = start_time.elapsed(); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(prefill_http); + return Err(format!( + "Prefill request failed to {}: {}", + prefill_http, prefill_err + )); + } + Ok(json) => json, + }; let decode_response = match decode_result { Ok(resp) => resp, Err(e) => { @@ -1025,7 +1046,7 @@ impl VllmPDRouter { return self .handle_decode_response( decode_response, - None, // WRITE mode: no prefill response JSON for logprobs + write_prefill_response_json.as_ref(), path, prefill_http, decode_http, From b04a049743339e31f8f2e5476dc4b926b863e3db Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 30 Apr 2026 11:09:24 +0200 Subject: [PATCH 22/29] fix: enforce mori runs with zmq svc discovery Signed-off-by: simondanielsson --- src/config/validation.rs | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/config/validation.rs b/src/config/validation.rs index cd822b2d..ae9e7a18 100644 --- a/src/config/validation.rs +++ b/src/config/validation.rs @@ -455,7 +455,7 @@ impl ConfigValidator { // MoRI-IO requires service discovery: ZMQ addresses are obtained via instance // registration and are not available in direct URL mode. - if config.kv_connector == KvConnector::MoriIO && !has_service_discovery { + if config.kv_connector == KvConnector::MoriIO && !has_vllm_discovery { return Err(ConfigError::IncompatibleConfig { reason: "MoRI-IO KV connector requires service discovery to be enabled \ (ZMQ addresses are obtained via instance registration). Please \ @@ -772,17 +772,16 @@ mod tests { #[test] fn test_moriio_with_service_discovery_is_valid() { let mut config = RouterConfig::new( - RoutingMode::Regular { - worker_urls: vec!["http://worker:8000".to_string()], + RoutingMode::VllmPrefillDecode { + prefill_urls: vec![], + decode_urls: vec![], + prefill_policy: None, + decode_policy: None, + discovery_address: Some("0.0.0.0:36367".to_string()), }, PolicyConfig::Random, ); config.kv_connector = KvConnector::MoriIO; - config.discovery = Some(DiscoveryConfig { - enabled: true, - selector: [("app".to_string(), "vllm".to_string())].into(), - ..Default::default() - }); let result = ConfigValidator::validate(&config); assert!(result.is_ok()); From 0d08cda7bbb844a062c700624e9888dcdba80b9b Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 30 Apr 2026 16:32:27 +0200 Subject: [PATCH 23/29] fix: commit xfer mode only after registration validation succeeds Signed-off-by: simondanielsson --- src/routers/http/vllm_service_discovery.rs | 32 +++++++++++++++------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/src/routers/http/vllm_service_discovery.rs b/src/routers/http/vllm_service_discovery.rs index cb5073be..ff4dd51e 100644 --- a/src/routers/http/vllm_service_discovery.rs +++ b/src/routers/http/vllm_service_discovery.rs @@ -81,14 +81,16 @@ pub struct ServiceRegistry { } /// Deserialize a registration message according to the active KV connector. -/// Returns the base `ServiceRegistration` on success, or `None` if the message -/// should be dropped (parse error, unknown mode, or transfer-mode mismatch). +/// Returns the parsed `ServiceRegistration` and, for MoRI-IO, the validated transfer mode. +/// Returns `None` if the message should be dropped (parse error, unknown mode, or +/// transfer-mode mismatch). The caller is responsible for committing the transfer mode +/// to the `OnceLock` only after the service type has been validated. fn parse_registration( message_data: &[u8], remote_address: &[u8], kv_connector: KvConnector, moriio_transfer_mode: &Arc>, -) -> Option { +) -> Option<(ServiceRegistration, Option)> { if matches!(kv_connector, KvConnector::MoriIO) { let reg: MoriIOServiceRegistration = match rmp_serde::from_slice(message_data) { Ok(r) => r, @@ -108,8 +110,8 @@ fn parse_registration( return None; } }; - if moriio_transfer_mode.set(mode).is_err() { - let stored = moriio_transfer_mode.get().copied().unwrap(); + // Check for mismatch against already-committed mode without committing yet. + if let Some(&stored) = moriio_transfer_mode.get() { if stored != mode { warn!( "MoRI-IO transfer_mode mismatch: expected {:?}, got {:?} from {}; skipping", @@ -119,13 +121,11 @@ fn parse_registration( ); return None; } - } else { - info!("MoRI-IO transfer mode set to {:?}", mode); } - Some(reg.base) + Some((reg.base, Some(mode))) } else { match rmp_serde::from_slice(message_data) { - Ok(data) => Some(data), + Ok(data) => Some((data, None)), Err(e) => { warn!("Failed to parse service registration: {}", e); None @@ -237,7 +237,7 @@ impl ServiceRegistry { kv_connector: KvConnector, moriio_transfer_mode: &Arc>, ) { - let data = match parse_registration( + let (data, parsed_mode) = match parse_registration( message_data, remote_address, kv_connector, @@ -259,6 +259,18 @@ impl ServiceRegistry { let remote_addr_str = String::from_utf8_lossy(remote_address); + match data.service_type.as_str() { + "P" | "D" => { + // Commit transfer mode only after service type is confirmed valid. + if let Some(mode) = parsed_mode { + if moriio_transfer_mode.set(mode).is_ok() { + info!("MoRI-IO transfer mode set to {:?}", mode); + } + } + } + _ => {} + } + match data.service_type.as_str() { "P" => { let mut prefill = prefill_instances.lock().unwrap(); From 6657952e041c3b580633bcaeeeefb7cded0d7c7e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 30 Apr 2026 16:45:01 +0200 Subject: [PATCH 24/29] fix: use send_client_request for prefill request in WRITE mode Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 71156139..364e9514 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -950,14 +950,22 @@ impl VllmPDRouter { let profiling_tasks = &self.profiling_tasks; let pd_router = &self.pd_router; let prefill_fut = async move { - let result = build_prefill_request_builder( - http_client, - &prefill_request_url, - &request_id, - prefill_dp_rank, + let result = otel_http::send_client_request( + build_prefill_request_builder( + http_client, + &prefill_request_url, + &request_id, + prefill_dp_rank, + ) + .body(prefill_request_str), + headers, + ClientRequestOptions { + method: "POST", + url: &prefill_request_url, + route: Some(path), + request_phase: Some("prefill"), + }, ) - .body(prefill_request_str) - .send() .await; if enable_profiling { let worker_url = format!("http://{}", prefill_base_http); From 3dbd46e8f8a1d1bbb6e535d4a24efb80bdaaefe9 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 30 Apr 2026 16:45:51 +0200 Subject: [PATCH 25/29] fix: record failed prefills in WRITE mode with record_pd_prefill_error Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 364e9514..9eefb337 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -1026,6 +1026,7 @@ impl VllmPDRouter { let write_prefill_response_json: Option = match prefill_result { Err(prefill_err) => { let duration = start_time.elapsed(); + RouterMetrics::record_pd_prefill_error(prefill_http); RouterMetrics::record_pd_request(path); RouterMetrics::record_pd_request_duration(path, duration); RouterMetrics::record_pd_prefill_request(prefill_http); From 45ee5b4b70cbef4a2f526e47c0404254aa7a5a4f Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 30 Apr 2026 17:05:39 +0200 Subject: [PATCH 26/29] fix: stop profile decode instance upon prefill error in WRITE mode Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 9eefb337..ba4a02f0 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -1025,6 +1025,8 @@ impl VllmPDRouter { let (prefill_result, decode_result) = tokio::join!(prefill_fut, decode_fut); let write_prefill_response_json: Option = match prefill_result { Err(prefill_err) => { + self.stop_profiling(&format!("http://{}", decode_base_http)) + .await; let duration = start_time.elapsed(); RouterMetrics::record_pd_prefill_error(prefill_http); RouterMetrics::record_pd_request(path); From 11dba53f883a1866a5d604c6b594c70ac654fa7b Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 5 May 2026 08:36:00 +0000 Subject: [PATCH 27/29] chore: update comments and rename variables to concurrent dispatch isntead of WRITRE mode Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 34 ++++++++++++++---------------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index ba4a02f0..38632c1a 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -632,7 +632,7 @@ impl VllmPDRouter { } } - /// Handle the decode response shared by both READ and WRITE dispatch paths. + /// Handle the decode response shared by both sequential and concurrent dispatch paths. /// Stops decode profiling, records metrics, then routes to streaming, logprobs-merge, /// or plain full-body response depending on the original request. #[allow(clippy::too_many_arguments)] @@ -787,7 +787,8 @@ impl VllmPDRouter { let (decode_base_http, decode_dp_rank) = extract_base_http_and_dp_rank(decode_http, self.intra_node_data_parallel_size); - let is_moriio_write = matches!(self.kv_connector, KvConnector::MoriIO) + // Concurrent dispatch: e.g. MoRI-IO WRITE mode + let is_concurrent_dispatch = matches!(self.kv_connector, KvConnector::MoriIO) && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)); let needs_logprobs = request_json.get("logprobs").is_some() @@ -806,9 +807,9 @@ impl VllmPDRouter { let prefill_request_url = format!("http://{}{}", prefill_base_http, path); // Stage 1: dispatch prefill. - // READ mode: send now and await — decode kv_transfer_params come from the prefill response. - // WRITE mode: skip; prefill is sent concurrently with decode in Stage 2 via tokio::join!. - let prefill_response_json: Option = if is_moriio_write { + // Sequential mode: send now and await — decode kv_transfer_params come from the prefill response. + // Concurrent mode: skip; prefill is sent concurrently with decode in Stage 2 via tokio::join!. + let prefill_response_json: Option = if is_concurrent_dispatch { None } else { debug!( @@ -936,10 +937,10 @@ impl VllmPDRouter { let decode_request_url = format!("http://{}{}", decode_base_http, path); - // WRITE mode: run prefill and decode concurrently via tokio::join! so the prefill + // Concurrent dispatch: run prefill and decode concurrently via tokio::join! so the prefill // task is always guaranteed to execute (unlike fire-and-forget tokio::spawn, which // can be silently dropped under load) and both HTTP sends start at the same time. - if is_moriio_write { + if is_concurrent_dispatch { self.start_profiling(&format!("http://{}", prefill_base_http)) .await; // Capture references rather than clones — tokio::join! polls both futures on the @@ -981,33 +982,30 @@ impl VllmPDRouter { match resp.bytes().await { Ok(bytes) => match serde_json::from_slice::(&bytes) { Ok(json) => { - debug!( - "MoRI-IO WRITE prefill completed with status {}", - status - ); + debug!("Concurrent prefill completed with status {}", status); Ok(Some(json)) } Err(_) => { debug!( - "MoRI-IO WRITE prefill completed with status {} (non-JSON body)", + "Concurrent prefill completed with status {} (non-JSON body)", status ); Ok(None) } }, Err(e) => { - warn!("MoRI-IO WRITE prefill: failed to read response body: {}", e); + warn!("Concurrent prefill: failed to read response body: {}", e); Ok(None) } } } Ok(resp) => { let status = resp.status(); - warn!("MoRI-IO WRITE prefill returned non-2xx status: {}", status); + warn!("Concurrent prefill returned non-2xx status: {}", status); Err(format!("Prefill request failed with status {}", status)) } Err(e) => { - warn!("MoRI-IO WRITE prefill request failed: {}", e); + warn!("Concurrent prefill request failed: {}", e); Err(format!("Prefill request failed: {}", e)) } } @@ -1023,7 +1021,7 @@ impl VllmPDRouter { }, ); let (prefill_result, decode_result) = tokio::join!(prefill_fut, decode_fut); - let write_prefill_response_json: Option = match prefill_result { + let concurrent_prefill_response_json: Option = match prefill_result { Err(prefill_err) => { self.stop_profiling(&format!("http://{}", decode_base_http)) .await; @@ -1057,7 +1055,7 @@ impl VllmPDRouter { return self .handle_decode_response( decode_response, - write_prefill_response_json.as_ref(), + concurrent_prefill_response_json.as_ref(), path, prefill_http, decode_http, @@ -1332,7 +1330,7 @@ impl VllmPDRouter { ); } } else { - // NIXL and MoRI-IO: extract kv_transfer_params from prefill response + // Sequential dispatch (NIXL, MoRI-IO READ): extract kv_transfer_params from prefill response if let Some(mut params) = kv_transfer_params { if matches!(self.kv_connector, KvConnector::MoriIO) { // MoRI-IO decode connector needs to know how many prefill DP ranks to handshake with. From edfa3e8632363b9acdc69292aa486757b2f1c717 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 5 May 2026 09:03:40 +0000 Subject: [PATCH 28/29] feat: add p2pnccl enum Signed-off-by: simondanielsson --- src/config/types.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/config/types.rs b/src/config/types.rs index 413022f7..8c010c12 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -121,6 +121,10 @@ pub enum KvConnector { #[serde(rename = "moriio")] #[value(name = "moriio")] MoriIO, + /// P2P NCCL KV transfer (PUT, PUT_ASYNC, GET modes) + #[serde(rename = "p2pnccl")] + #[value(name = "p2pnccl")] + P2pNccl, } #[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)] From 2f59324d7806922bf5b76e81dc122d33a9ac516c Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 5 May 2026 09:03:57 +0000 Subject: [PATCH 29/29] feat: run p2pnccl in concurrent dispatch mode Signed-off-by: simondanielsson --- src/routers/http/vllm_pd_router.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 38632c1a..f71e8dc6 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -192,6 +192,8 @@ impl VllmPDRouter { Uuid::new_v4().simple() )), KvConnector::Nixl => None, + // P2pNccl embeds P/D addresses in the request_id; no transfer_id needed. + KvConnector::P2pNccl => None, } } @@ -246,6 +248,8 @@ impl VllmPDRouter { "remote_host": serde_json::Value::Null, "remote_port": serde_json::Value::Null })), + // P2pNccl resolves P/D addresses from the request_id itself; no params needed. + KvConnector::P2pNccl => Ok(json!({})), } } @@ -304,6 +308,8 @@ impl VllmPDRouter { } } KvConnector::Nixl => Some(prefill_response_json?.get("kv_transfer_params")?.clone()), + // P2pNccl resolves P/D addresses from the request_id; decode needs no injected params. + KvConnector::P2pNccl => None, } } @@ -787,9 +793,10 @@ impl VllmPDRouter { let (decode_base_http, decode_dp_rank) = extract_base_http_and_dp_rank(decode_http, self.intra_node_data_parallel_size); - // Concurrent dispatch: e.g. MoRI-IO WRITE mode - let is_concurrent_dispatch = matches!(self.kv_connector, KvConnector::MoriIO) - && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)); + // Concurrent dispatch: e.g. MoRI-IO WRITE mode, P2pNccl + let is_concurrent_dispatch = matches!(self.kv_connector, KvConnector::P2pNccl) + || (matches!(self.kv_connector, KvConnector::MoriIO) + && matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write))); let needs_logprobs = request_json.get("logprobs").is_some() || request_json