diff --git a/e2e_test/router/test_pd_mmlu.py b/e2e_test/router/test_pd_mmlu.py index 496c8016a..e6ed583da 100644 --- a/e2e_test/router/test_pd_mmlu.py +++ b/e2e_test/router/test_pd_mmlu.py @@ -4,7 +4,8 @@ workers for improved throughput and resource utilization. Backends: -- "pd_http": HTTP mode (SGLang only - vLLM does not support HTTP) +- "pd_http": HTTP mode (SGLang parallel bootstrap dispatch; vLLM sequential + prefill-then-decode with kv_transfer_params relay) - "pd_grpc": gRPC mode (both SGLang and vLLM) Requirements: @@ -16,7 +17,7 @@ # SGLang (runs both HTTP and gRPC) pytest e2e_test/router/test_pd_mmlu.py -v - # vLLM (runs gRPC only, HTTP skipped) + # vLLM (runs both HTTP and gRPC) E2E_RUNTIME=vllm pytest e2e_test/router/test_pd_mmlu.py -v """ @@ -31,11 +32,10 @@ logger = logging.getLogger(__name__) -@pytest.mark.engine("sglang") +@pytest.mark.engine("sglang", "vllm") @pytest.mark.gpu(2) @pytest.mark.model("meta-llama/Llama-3.1-8B-Instruct") @pytest.mark.e2e -@pytest.mark.skip_for_runtime("vllm", reason="vLLM does not support HTTP mode") @pytest.mark.parametrize("setup_backend", ["pd_http"], indirect=True) class TestPDMMLUHttp: """MMLU evaluation tests using PD disaggregation (HTTP mode).""" diff --git a/model_gateway/src/routers/common/kv_transfer.rs b/model_gateway/src/routers/common/kv_transfer.rs new file mode 100644 index 000000000..6aab0c438 --- /dev/null +++ b/model_gateway/src/routers/common/kv_transfer.rs @@ -0,0 +1,259 @@ +//! vLLM PD KV-transfer connector handling, shared by both transports. +//! +//! vLLM disaggregation is sequential: the prefill leg is tagged with +//! connector params, the engine returns (or the router mints) the handoff +//! params, and the decode leg carries them. The gRPC pipeline and the HTTP +//! PD router both implement that flow; the connector vocabulary lives here +//! so the two stay in lockstep. + +use tracing::warn; + +use crate::{ + observability::metrics::metrics_labels, + worker::{Worker, DEFAULT_BOOTSTRAP_PORT, MOONCAKE_CONNECTOR, NIXL_CONNECTOR}, +}; + +/// KV-transfer params tagged onto the NIXL prefill leg so the engine pins its +/// KV blocks and returns the handoff params for the decode worker. +pub(crate) const NIXL_PREFILL_KV_PARAMS: &str = + r#"{"do_remote_decode":true,"do_remote_prefill":false}"#; + +/// PD KV-transfer behavior derived from prefill worker metadata. +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum KvConnectorMode { + /// MooncakeConnector: mint a transfer_id, tag both legs, synthesize decode + /// params from worker metadata; legacy host/port injection when the + /// servicer predates kv_engine_id reporting (or DP runs without a pinned rank). + Mooncake { + host: String, + port: u32, + engine_id: Option, + }, + /// NixlConnector: tag prefill with do_remote_decode, relay returned params to decode. + Nixl, + /// Unknown/absent connector: relay returned params opportunistically. + Passthrough, +} + +impl KvConnectorMode { + pub(crate) fn metrics_label(&self) -> &'static str { + match self { + Self::Mooncake { .. } => metrics_labels::KV_CONNECTOR_MOONCAKE, + Self::Nixl => metrics_labels::KV_CONNECTOR_NIXL, + Self::Passthrough => metrics_labels::KV_CONNECTOR_PASSTHROUGH, + } + } +} + +pub(crate) fn kv_connector_mode( + kv_connector: Option<&str>, + bootstrap_host: &str, + bootstrap_port: Option, + kv_engine_id: Option<&str>, +) -> KvConnectorMode { + match kv_connector { + Some(MOONCAKE_CONNECTOR) => KvConnectorMode::Mooncake { + host: bootstrap_host.to_string(), + port: u32::from(bootstrap_port.unwrap_or(DEFAULT_BOOTSTRAP_PORT)), + // Empty means unknown (forces the legacy fallback) + engine_id: kv_engine_id.filter(|s| !s.is_empty()).map(str::to_string), + }, + Some(NIXL_CONNECTOR) => KvConnectorMode::Nixl, + _ => KvConnectorMode::Passthrough, + } +} + +/// Connector id of the engine core serving the prefill leg. With DP the cores +/// suffix the configured id as `{base}_dp{rank}`, so minting needs a pinned +/// rank; unpinned DP>1 yields None (no mint — decode recomputes locally). +pub(crate) fn effective_kv_engine_id( + base: Option<&str>, + dp_size: Option, + dp_rank: Option, +) -> Option { + let base = base.filter(|s| !s.is_empty())?; + if dp_size.unwrap_or(1) > 1 { + dp_rank.map(|rank| format!("{base}_dp{rank}")) + } else { + Some(base.to_string()) + } +} + +/// The connector mode for a PD pair, read off the prefill worker's metadata. +/// Discovered dp_size matters even without `--dp-aware` expansion: a DP>1 +/// engine behind an unexpanded worker must not be minted for. +pub(crate) fn connector_mode_for_worker(worker: &dyn Worker) -> KvConnectorMode { + let meta = worker.metadata(); + let dp_label = meta.spec.labels.get("dp_size"); + let label_dp = dp_label.and_then(|s| s.parse::().ok().filter(|v| *v > 0)); + let engine_id = if worker.dp_size().is_none() && dp_label.is_some() && label_dp.is_none() { + // A dp_size label that does not parse as a positive integer means the + // DP topology is unknown, not absent. Minting an unsuffixed engine id + // could target the wrong engine core, so fail closed: no mint, decode + // recomputes the prompt. + warn!( + worker = %worker.url(), + dp_size = ?dp_label, + "invalid dp_size label; treating DP topology as unknown and \ + skipping KV engine-id minting" + ); + None + } else { + let dp_size = worker.dp_size().or(label_dp); + effective_kv_engine_id(meta.spec.kv_engine_id.as_deref(), dp_size, worker.dp_rank()) + }; + kv_connector_mode( + meta.spec.kv_connector.as_deref(), + &meta.spec.bootstrap_host, + meta.spec.bootstrap_port, + engine_id.as_deref(), + ) +} + +/// Prefill-leg params for Mooncake: the engine pins blocks under the minted id. +pub(crate) fn mooncake_prefill_params(transfer_id: &str) -> String { + serde_json::json!({ + "do_remote_decode": true, + "do_remote_prefill": false, + "transfer_id": transfer_id, + }) + .to_string() +} + +/// Decode-leg params for Mooncake, synthesized from prefill worker metadata +/// (the engine returns nothing to relay; the connector is push-based). +pub(crate) fn mooncake_decode_params( + transfer_id: &str, + engine_id: &str, + host: &str, + port: u32, +) -> String { + serde_json::json!({ + "do_remote_decode": false, + "do_remote_prefill": true, + "transfer_id": transfer_id, + "remote_engine_id": engine_id, + "remote_bootstrap_addr": format!("http://{host}:{port}"), + }) + .to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn kv_connector_mode_mooncake_uses_bootstrap_metadata() { + let mode = kv_connector_mode( + Some(MOONCAKE_CONNECTOR), + "prefill-host", + Some(9090), + Some("engine-1"), + ); + assert_eq!( + mode, + KvConnectorMode::Mooncake { + host: "prefill-host".to_string(), + port: 9090, + engine_id: Some("engine-1".to_string()), + } + ); + } + + #[test] + fn kv_connector_mode_mooncake_defaults_port_and_tolerates_missing_engine_id() { + let mode = kv_connector_mode(Some(MOONCAKE_CONNECTOR), "prefill-host", None, None); + assert_eq!( + mode, + KvConnectorMode::Mooncake { + host: "prefill-host".to_string(), + port: u32::from(DEFAULT_BOOTSTRAP_PORT), + engine_id: None, + } + ); + } + + #[test] + fn kv_connector_mode_mooncake_empty_engine_id_means_legacy() { + let mode = kv_connector_mode(Some(MOONCAKE_CONNECTOR), "host", Some(9090), Some("")); + assert_eq!( + mode, + KvConnectorMode::Mooncake { + host: "host".to_string(), + port: 9090, + engine_id: None, + } + ); + } + + #[test] + fn kv_connector_mode_nixl() { + assert_eq!( + kv_connector_mode(Some(NIXL_CONNECTOR), "ignored", Some(9090), None), + KvConnectorMode::Nixl + ); + } + + #[test] + fn kv_connector_mode_unknown_or_missing_is_passthrough() { + assert_eq!( + kv_connector_mode(Some("LMCacheConnector"), "host", None, None), + KvConnectorMode::Passthrough + ); + assert_eq!( + kv_connector_mode(None, "host", None, None), + KvConnectorMode::Passthrough + ); + } + + #[test] + fn invalid_dp_size_label_fails_closed_on_minting() { + use crate::worker::{BasicWorkerBuilder, WorkerType}; + + let worker = BasicWorkerBuilder::new("http://prefill:8000") + .worker_type(WorkerType::Prefill) + .kv_connector(MOONCAKE_CONNECTOR) + .kv_engine_id("eng") + .label("dp_size", "not-a-number") + .build(); + let mode = connector_mode_for_worker(&worker); + // Unknown DP topology must not mint an unsuffixed engine id. + assert!(matches!( + mode, + KvConnectorMode::Mooncake { + engine_id: None, + .. + } + )); + + let worker = BasicWorkerBuilder::new("http://prefill:8000") + .worker_type(WorkerType::Prefill) + .kv_connector(MOONCAKE_CONNECTOR) + .kv_engine_id("eng") + .label("dp_size", "1") + .build(); + let mode = connector_mode_for_worker(&worker); + assert!(matches!( + mode, + KvConnectorMode::Mooncake { + engine_id: Some(ref id), + .. + } if id == "eng" + )); + } + + #[test] + fn effective_engine_id_requires_pinned_rank_under_dp() { + assert_eq!( + effective_kv_engine_id(Some("eng"), Some(2), Some(1)), + Some("eng_dp1".to_string()) + ); + assert_eq!(effective_kv_engine_id(Some("eng"), Some(2), None), None); + assert_eq!( + effective_kv_engine_id(Some("eng"), None, None), + Some("eng".to_string()) + ); + assert_eq!(effective_kv_engine_id(Some(""), None, None), None); + assert_eq!(effective_kv_engine_id(None, Some(2), Some(0)), None); + } +} diff --git a/model_gateway/src/routers/common/mod.rs b/model_gateway/src/routers/common/mod.rs index d72c86fc1..da5e8940e 100644 --- a/model_gateway/src/routers/common/mod.rs +++ b/model_gateway/src/routers/common/mod.rs @@ -28,6 +28,7 @@ //! responses to clients and parsing upstream SSE byte streams pub mod header_utils; +pub(crate) mod kv_transfer; pub mod mcp_utils; pub mod openai_bridge; pub mod overload; diff --git a/model_gateway/src/routers/grpc/common/stages/request_execution.rs b/model_gateway/src/routers/grpc/common/stages/request_execution.rs index 5bd33302c..32303863c 100644 --- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs +++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs @@ -11,6 +11,10 @@ use super::PipelineStage; use crate::{ observability::metrics::{metrics_labels, Metrics}, routers::{ + common::kv_transfer::{ + connector_mode_for_worker, mooncake_decode_params, mooncake_prefill_params, + KvConnectorMode, NIXL_PREFILL_KV_PARAMS, + }, error, grpc::{ common::stages::encode::EncodeDispatchPlan, @@ -25,9 +29,7 @@ use crate::{ utils::tonic_ext::{TonicResultExt, TonicStatusExt}, }, }, - worker::{ - ConnectionModeExt, RuntimeType, DEFAULT_BOOTSTRAP_PORT, MOONCAKE_CONNECTOR, NIXL_CONNECTOR, - }, + worker::{ConnectionModeExt, RuntimeType}, }; type StreamResult = Result; @@ -48,94 +50,6 @@ fn pd_leg_labels(workers: &WorkerSelection) -> (&'static str, &'static str) { } } -/// KV-transfer params tagged onto the NIXL prefill leg so the engine pins its -/// KV blocks and returns the handoff params for the decode worker. -const NIXL_PREFILL_KV_PARAMS: &str = r#"{"do_remote_decode":true,"do_remote_prefill":false}"#; - -/// PD KV-transfer behavior derived from prefill worker metadata. -#[derive(Debug, Clone, PartialEq)] -enum KvConnectorMode { - /// MooncakeConnector: mint a transfer_id, tag both legs, synthesize decode - /// params from worker metadata; legacy host/port injection when the - /// servicer predates kv_engine_id reporting (or DP runs without a pinned rank). - Mooncake { - host: String, - port: u32, - engine_id: Option, - }, - /// NixlConnector: tag prefill with do_remote_decode, relay returned params to decode. - Nixl, - /// Unknown/absent connector: relay returned params opportunistically. - Passthrough, -} - -impl KvConnectorMode { - fn metrics_label(&self) -> &'static str { - match self { - Self::Mooncake { .. } => metrics_labels::KV_CONNECTOR_MOONCAKE, - Self::Nixl => metrics_labels::KV_CONNECTOR_NIXL, - Self::Passthrough => metrics_labels::KV_CONNECTOR_PASSTHROUGH, - } - } -} - -fn kv_connector_mode( - kv_connector: Option<&str>, - bootstrap_host: &str, - bootstrap_port: Option, - kv_engine_id: Option<&str>, -) -> KvConnectorMode { - match kv_connector { - Some(MOONCAKE_CONNECTOR) => KvConnectorMode::Mooncake { - host: bootstrap_host.to_string(), - port: u32::from(bootstrap_port.unwrap_or(DEFAULT_BOOTSTRAP_PORT)), - // Empty means unknown (forces the legacy fallback) - engine_id: kv_engine_id.filter(|s| !s.is_empty()).map(str::to_string), - }, - Some(NIXL_CONNECTOR) => KvConnectorMode::Nixl, - _ => KvConnectorMode::Passthrough, - } -} - -/// Connector id of the engine core serving the prefill leg. With DP the cores -/// suffix the configured id as `{base}_dp{rank}`, so minting needs a pinned -/// rank; unpinned DP>1 yields None (no mint — decode recomputes locally). -fn effective_kv_engine_id( - base: Option<&str>, - dp_size: Option, - dp_rank: Option, -) -> Option { - let base = base.filter(|s| !s.is_empty())?; - if dp_size.unwrap_or(1) > 1 { - dp_rank.map(|rank| format!("{base}_dp{rank}")) - } else { - Some(base.to_string()) - } -} - -/// Prefill-leg params for Mooncake: the engine pins blocks under the minted id. -fn mooncake_prefill_params(transfer_id: &str) -> String { - serde_json::json!({ - "do_remote_decode": true, - "do_remote_prefill": false, - "transfer_id": transfer_id, - }) - .to_string() -} - -/// Decode-leg params for Mooncake, synthesized from prefill worker metadata -/// (the engine returns nothing to relay; the connector is push-based). -fn mooncake_decode_params(transfer_id: &str, engine_id: &str, host: &str, port: u32) -> String { - serde_json::json!({ - "do_remote_decode": false, - "do_remote_prefill": true, - "transfer_id": transfer_id, - "remote_engine_id": engine_id, - "remote_bootstrap_addr": format!("http://{host}:{port}"), - }) - .to_string() -} - /// Request execution stage: execute the plan produced by request building. pub(crate) struct RequestExecutionStage; @@ -588,22 +502,7 @@ impl RequestExecutionStage { let mode = workers .prefill_worker() - .map(|w| { - let meta = w.metadata(); - // Discovered dp_size matters even without --dp-aware expansion: - // a DP>1 engine behind an unexpanded worker must not be minted for - let dp_size = w - .dp_size() - .or_else(|| meta.spec.labels.get("dp_size").and_then(|s| s.parse().ok())); - let engine_id = - effective_kv_engine_id(meta.spec.kv_engine_id.as_deref(), dp_size, w.dp_rank()); - kv_connector_mode( - meta.spec.kv_connector.as_deref(), - &meta.spec.bootstrap_host, - meta.spec.bootstrap_port, - engine_id.as_deref(), - ) - }) + .map(|w| connector_mode_for_worker(w.as_ref())) .unwrap_or(KvConnectorMode::Passthrough); // Recorded on the success path (after decode established) so failed @@ -878,70 +777,6 @@ mod tests { ); } - #[test] - fn kv_connector_mode_mooncake_uses_bootstrap_metadata() { - let mode = kv_connector_mode( - Some(MOONCAKE_CONNECTOR), - "prefill-host", - Some(9090), - Some("engine-1"), - ); - assert_eq!( - mode, - KvConnectorMode::Mooncake { - host: "prefill-host".to_string(), - port: 9090, - engine_id: Some("engine-1".to_string()), - } - ); - } - - #[test] - fn kv_connector_mode_mooncake_defaults_port_and_tolerates_missing_engine_id() { - let mode = kv_connector_mode(Some(MOONCAKE_CONNECTOR), "prefill-host", None, None); - assert_eq!( - mode, - KvConnectorMode::Mooncake { - host: "prefill-host".to_string(), - port: u32::from(DEFAULT_BOOTSTRAP_PORT), - engine_id: None, - } - ); - } - - #[test] - fn kv_connector_mode_mooncake_empty_engine_id_means_legacy() { - let mode = kv_connector_mode(Some(MOONCAKE_CONNECTOR), "host", Some(9090), Some("")); - assert_eq!( - mode, - KvConnectorMode::Mooncake { - host: "host".to_string(), - port: 9090, - engine_id: None, - } - ); - } - - #[test] - fn kv_connector_mode_nixl() { - assert_eq!( - kv_connector_mode(Some(NIXL_CONNECTOR), "ignored", Some(9090), None), - KvConnectorMode::Nixl - ); - } - - #[test] - fn kv_connector_mode_unknown_or_missing_is_passthrough() { - assert_eq!( - kv_connector_mode(Some("LMCacheConnector"), "host", None, None), - KvConnectorMode::Passthrough - ); - assert_eq!( - kv_connector_mode(None, "host", None, None), - KvConnectorMode::Passthrough - ); - } - #[test] fn mooncake_prefill_params_carry_transfer_id() { let value: serde_json::Value = @@ -1117,39 +952,4 @@ mod tests { let unset = ProtoGenerateComplete::Vllm(vllm::GenerateComplete::default()); assert_eq!(unset.kv_transfer_params_json(), None); } - - #[test] - fn effective_engine_id_passthrough_when_no_dp() { - assert_eq!( - effective_kv_engine_id(Some("eng"), None, None).as_deref(), - Some("eng") - ); - assert_eq!( - effective_kv_engine_id(Some("eng"), Some(1), None).as_deref(), - Some("eng") - ); - } - - #[test] - fn effective_engine_id_suffixes_pinned_dp_rank() { - assert_eq!( - effective_kv_engine_id(Some("eng"), Some(2), Some(1)).as_deref(), - Some("eng_dp1") - ); - assert_eq!( - effective_kv_engine_id(Some("eng"), Some(2), Some(0)).as_deref(), - Some("eng_dp0") - ); - } - - #[test] - fn effective_engine_id_none_for_unpinned_dp() { - assert_eq!(effective_kv_engine_id(Some("eng"), Some(2), None), None); - } - - #[test] - fn effective_engine_id_none_for_missing_or_empty_base() { - assert_eq!(effective_kv_engine_id(None, Some(2), Some(0)), None); - assert_eq!(effective_kv_engine_id(Some(""), None, None), None); - } } diff --git a/model_gateway/src/routers/http/pd_router.rs b/model_gateway/src/routers/http/pd_router.rs index 324f19be5..5147fd0cf 100644 --- a/model_gateway/src/routers/http/pd_router.rs +++ b/model_gateway/src/routers/http/pd_router.rs @@ -37,7 +37,12 @@ use crate::{ policies::{LoadBalancingPolicy, PolicyRegistry, SelectWorkerInfo}, routers::{ common::{ - attach_sized_body, header_utils, overload, + attach_sized_body, header_utils, + kv_transfer::{ + connector_mode_for_worker, mooncake_decode_params, mooncake_prefill_params, + KvConnectorMode, NIXL_PREFILL_KV_PARAMS, + }, + overload, request_lease::{ReleasePoint, RequestLease, RoutingDerivatives}, retry::{is_retryable_response, RetryExecutor}, sse::{SseEncoder, SSE_CHANNEL_BUFFER}, @@ -47,7 +52,10 @@ use crate::{ http::router::send_with_stale_conn_retry, RouterTrait, }, - worker::{HashRing, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType, UNKNOWN_MODEL_ID}, + worker::{ + HashRing, RuntimeType, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType, + UNKNOWN_MODEL_ID, + }, }; /// Why PD pair selection produced nothing. @@ -492,6 +500,73 @@ impl PDRouter { ] }); + if prefill.metadata().spec.runtime_type == RuntimeType::Vllm { + // vLLM PD is sequential: prefill first with connector params, then + // decode carrying the KV handoff — no bootstrap rendezvous exists. + let mode = connector_mode_for_worker(prefill.as_ref()); + let transfer_id = match &mode { + KvConnectorMode::Mooncake { + engine_id: Some(_), .. + } => Some(format!("xfer-{}", uuid::Uuid::now_v7())), + _ => None, + }; + let legs = + lease.serialize_legs_with(|view| -> Result<(Vec, Vec), Box> { + let mut json_request = serde_json::to_value(view.request) + .map_err(|e| Box::new(Self::handle_serialization_error(e)))?; + super::set_request_model(&mut json_request, context.model_id); + + // The KV handoff is single-consumer: with n>1 each fan-out + // child on decode would pull, and the first completion + // frees the prefill blocks under its siblings. + let relay = Self::sampling_n(&json_request) <= 1; + let mut prefill_json = json_request.clone(); + Self::sanitize_prefill_for_kv_handoff(&mut prefill_json, context.route); + if relay { + let params = match (&mode, &transfer_id) { + (KvConnectorMode::Nixl, _) => { + serde_json::from_str::(NIXL_PREFILL_KV_PARAMS).ok() + } + (KvConnectorMode::Mooncake { .. }, Some(id)) => { + serde_json::from_str::(&mooncake_prefill_params(id)).ok() + } + _ => None, + }; + if let (Some(obj), Some(params)) = (prefill_json.as_object_mut(), params) { + obj.insert("kv_transfer_params".to_string(), params); + } + } + + Ok(( + serde_json::to_vec(&prefill_json) + .map_err(|e| Box::new(Self::handle_serialization_error(e)))?, + serde_json::to_vec(&json_request) + .map_err(|e| Box::new(Self::handle_serialization_error(e)))?, + )) + }); + let (prefill_body, decode_body) = match legs { + Ok(pair) => pair, + Err(response) => return *response, + }; + lease.release_dispatch(); + + // Outcome accounting happens per-leg inside the sequential path: + // a prefill-only failure must not feed the decode worker's + // circuit breaker (the decode leg was never contacted). + return self + .execute_sequential_dispatch_internal( + headers, + (prefill_body, decode_body), + mode, + transfer_id, + context, + Arc::clone(&prefill), + Arc::clone(&decode), + load_guards, + ) + .await; + } + let legs = lease.serialize_legs_with(|view| -> Result<(Vec, Vec), Box> { let mut json_request = serde_json::to_value(view.request) .map_err(|e| Box::new(Self::handle_serialization_error(e)))?; @@ -848,6 +923,318 @@ impl PDRouter { prefill_head_elapsed + prefill_drain_start.elapsed(), ); + self.forward_decode_body( + decode_response, + status, + &context, + decode, + load_guards, + prefill_body, + ) + .await + } + + /// The request's fan-out factor, read from the serialized body. + fn sampling_n(json: &Value) -> u64 { + json.get("n").and_then(Value::as_u64).unwrap_or(1) + } + + /// Sanitize the prefill leg for a KV handoff: the prefill engine computes + /// KV for the prompt and must produce (at most) one token, unstreamed. + /// The output-cap key is per-endpoint. + fn sanitize_prefill_for_kv_handoff(json: &mut Value, route: &str) { + let Some(obj) = json.as_object_mut() else { + return; + }; + obj.insert("stream".to_string(), Value::Bool(false)); + // stream_options without stream=true is rejected by strict engines. + obj.remove("stream_options"); + if obj.contains_key("n") { + obj.insert("n".to_string(), Value::from(1)); + } + // A min_tokens floor would force decode-phase work onto the prefill leg. + obj.remove("min_tokens"); + match route { + // The engine rejects requests carrying both cap spellings, so + // overwrite the one the client used. + "/v1/chat/completions" if obj.contains_key("max_completion_tokens") => { + obj.insert("max_completion_tokens".to_string(), Value::from(1)); + obj.remove("max_tokens"); + } + "/v1/responses" => { + obj.insert("max_output_tokens".to_string(), Value::from(1)); + } + _ => { + obj.insert("max_tokens".to_string(), Value::from(1)); + } + } + } + + /// vLLM PD over HTTP: send the tagged prefill leg, wait for it, harvest + /// (or synthesize) the KV handoff params, then dispatch the decode leg + /// carrying them. Mirrors the gRPC pipeline's `execute_sequential_pd`. + #[expect( + clippy::too_many_arguments, + reason = "mirrors execute_dual_dispatch_internal's dispatch surface" + )] + async fn execute_sequential_dispatch_internal( + &self, + headers: Option<&HeaderMap>, + leg_bodies: (Bytes, Bytes), + mode: KvConnectorMode, + transfer_id: Option, + context: PDRequestContext<'_>, + prefill: Arc, + decode: Arc, + load_guards: Vec, + ) -> Response { + let (prefill_body, decode_body) = leg_bodies; + + let mut headers_with_trace = headers.cloned().unwrap_or_default(); + inject_trace_context_http(&mut headers_with_trace); + let headers = Some(&headers_with_trace); + let runtime = prefill.metadata().spec.runtime_type.as_str(); + + let mut decode_json = match serde_json::from_slice::(&decode_body) { + Ok(json) => json, + Err(e) => return Self::handle_serialization_error(e), + }; + // The KV handoff is single-consumer, so an n>1 fan-out cannot use it — + // and without the handoff a prefill leg is pure wasted GPU work plus + // serial latency. Skip prefill entirely and let decode own the prompt. + let relay = Self::sampling_n(&decode_json) <= 1; + Metrics::record_pd_kv_connector_mode(mode.metrics_label()); + + let dispatch_start = Instant::now(); + let harvested = if relay { + events::RequestPDSentEvent { + prefill_url: prefill.url(), + decode_url: decode.url(), + } + .emit(); + + let prefill_request = self.build_post_with_headers( + &self.client, + prefill.as_ref(), + context.route, + prefill_body, + headers, + false, + ); + let prefill_response = match send_with_stale_conn_retry(prefill_request).await { + Ok(response) => response, + Err(e) => { + error!("PD prefill transport error: {e}"); + Self::record_sequential_leg( + prefill.as_ref(), + metrics_labels::WORKER_PREFILL, + StatusCode::BAD_GATEWAY, + ); + return error::bad_gateway( + "prefill_request_failed", + format!("Prefill transport error: {e}"), + ); + } + }; + let prefill_status = StatusCode::from_u16(prefill_response.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + if !prefill_status.is_success() { + error!( + "Prefill server returned error status prefill_url={} status={}", + prefill.url(), + prefill_status + ); + Self::record_sequential_leg( + prefill.as_ref(), + metrics_labels::WORKER_PREFILL, + prefill_status, + ); + return Self::prefill_error_response(prefill_status, prefill_response).await; + } + let prefill_bytes = match prefill_response.bytes().await { + Ok(bytes) => bytes, + Err(e) => { + error!("Failed to read prefill response: {e}"); + Self::record_sequential_leg( + prefill.as_ref(), + metrics_labels::WORKER_PREFILL, + StatusCode::BAD_GATEWAY, + ); + return error::bad_gateway( + "prefill_read_failed", + format!("Failed to read prefill response: {e}"), + ); + } + }; + Self::record_sequential_leg( + prefill.as_ref(), + metrics_labels::WORKER_PREFILL, + prefill_status, + ); + Metrics::record_pd_prefill_duration( + metrics_labels::BACKEND_PD, + context.model_id, + runtime, + dispatch_start.elapsed(), + ); + + // Harvest the handoff params the prefill engine returned (NIXL and + // opportunistic passthrough; Mooncake returns nothing and is minted). + serde_json::from_slice::(&prefill_bytes) + .ok() + .and_then(|json| json.get("kv_transfer_params").cloned()) + .filter(|params| !params.is_null()) + } else { + debug!( + "vLLM PD over HTTP: n>1 fan-out cannot consume a KV handoff; \ + dispatching to decode only" + ); + None + }; + let decode_params = match (&mode, harvested) { + // Modern Mooncake: synthesized params under the minted transfer_id + ( + KvConnectorMode::Mooncake { + host, + port, + engine_id: Some(engine_id), + }, + _, + ) if relay && transfer_id.is_some() => { + let id = transfer_id.as_deref().unwrap_or_default(); + serde_json::from_str::(&mooncake_decode_params(id, engine_id, host, *port)) + .ok() + } + (KvConnectorMode::Mooncake { .. }, _) => { + // Legacy typed host/port injection is a sidecar-proto shape + // with no HTTP equivalent; decode recomputes the prompt. + warn!( + "vLLM PD over HTTP: Mooncake without a discovered kv_engine_id \ + cannot be minted for; decode recomputes the prompt locally" + ); + None + } + (KvConnectorMode::Nixl | KvConnectorMode::Passthrough, Some(params)) if relay => { + Some(params) + } + (KvConnectorMode::Nixl, None) if relay => { + Metrics::record_pd_kv_transfer_failure(); + warn!( + "vLLM PD (NIXL) over HTTP: prefill returned no kv_transfer_params; \ + decode recomputes the prompt locally" + ); + None + } + _ => None, + }; + if let (Some(obj), Some(params)) = (decode_json.as_object_mut(), decode_params) { + obj.insert("kv_transfer_params".to_string(), params); + } + let decode_body = match serde_json::to_vec(&decode_json) { + Ok(body) => Bytes::from(body), + Err(e) => return Self::handle_serialization_error(e), + }; + + let decode_request = self.build_post_with_headers( + &self.client, + decode.as_ref(), + context.route, + decode_body, + headers, + false, + ); + let decode_response = match send_with_stale_conn_retry(decode_request).await { + Ok(response) => response, + Err(e) => { + error!("PD decode transport error: {e}"); + Self::record_sequential_leg( + decode.as_ref(), + metrics_labels::WORKER_DECODE, + StatusCode::BAD_GATEWAY, + ); + return error::bad_gateway( + "decode_request_failed", + format!("Decode transport error: {e}"), + ); + } + }; + + events::RequestReceivedEvent {}.emit(); + + let status = StatusCode::from_u16(decode_response.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + Self::record_sequential_leg(decode.as_ref(), metrics_labels::WORKER_DECODE, status); + if !status.is_success() { + error!( + "Decode server returned error status decode_url={} status={}", + decode.url(), + status + ); + return self + .handle_decode_error_response(decode_response, &context, decode, load_guards) + .await; + } + + // Honest sequential TTFT: prefill dispatch to the decode response + // head — the prefill wait is part of what the client experiences. + Metrics::record_pd_ttft( + metrics_labels::BACKEND_PD, + context.model_id, + runtime, + dispatch_start.elapsed(), + ); + + self.forward_decode_body(decode_response, status, &context, decode, load_guards, None) + .await + } + + /// Map a failed prefill response to a client-facing error. The exact + /// upstream status is preserved (not classified into a fixed set) so + /// retryability and capacity-pushback handling see what the worker sent. + async fn prefill_error_response(status: StatusCode, response: reqwest::Response) -> Response { + let message = match response.bytes().await { + Ok(body) => { + if let Ok(json) = serde_json::from_slice::(&body) { + json.get("error") + .and_then(|e| e.get("message")) + .and_then(Value::as_str) + .or_else(|| json.get("message").and_then(Value::as_str)) + .map(str::to_string) + .unwrap_or_else(|| String::from_utf8_lossy(&body).to_string()) + } else { + String::from_utf8_lossy(&body).to_string() + } + } + Err(e) => format!("Prefill server error: {e}"), + }; + error::create_error(status, "prefill_upstream_error", message) + } + + /// Per-leg outcome accounting for the sequential path: a leg that was + /// never contacted must not feed the other worker's circuit breaker. + fn record_sequential_leg(worker: &dyn Worker, role: &'static str, status: StatusCode) { + worker.record_outcome(status.as_u16()); + if status.is_server_error() { + Metrics::record_worker_error( + role, + metrics_labels::CONNECTION_HTTP, + error_type_from_status(status), + ); + } + } + + /// Forward a successful decode response to the client, streaming or not, + /// merging prefill logprobs when requested. Shared by the parallel and + /// sequential PD dispatch paths. + async fn forward_decode_body( + &self, + decode_response: reqwest::Response, + status: StatusCode, + context: &PDRequestContext<'_>, + decode: Arc, + load_guards: Vec, + prefill_body: Option, + ) -> Response { if context.is_stream { // Streaming response let prefill_logprobs = if context.return_logprob { @@ -1999,7 +2386,9 @@ mod tests { clippy::disallowed_methods, reason = "test stub server lives for the duration of the test process" )] - async fn spawn_recording_stub() -> (String, Arc>>) { + async fn spawn_recording_stub( + reply: &'static str, + ) -> (String, Arc>>) { let seen: Arc>> = Arc::new(std::sync::Mutex::new(Vec::new())); let log = Arc::clone(&seen); @@ -2012,7 +2401,7 @@ mod tests { .unwrap_or_default(); let json = serde_json::from_slice(&bytes).unwrap_or(Value::Null); log.lock().unwrap().push((path, json)); - ([(CONTENT_TYPE, "application/json")], "{}") + ([(CONTENT_TYPE, "application/json")], reply) } })); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -2025,8 +2414,8 @@ mod tests { #[tokio::test] async fn messages_and_responses_dispatch_to_their_worker_routes() { - let (prefill_url, prefill_seen) = spawn_recording_stub().await; - let (decode_url, decode_seen) = spawn_recording_stub().await; + let (prefill_url, prefill_seen) = spawn_recording_stub("{}").await; + let (decode_url, decode_seen) = spawn_recording_stub("{}").await; let router = create_test_pd_router(); router @@ -2085,6 +2474,150 @@ mod tests { } } + #[tokio::test] + async fn vllm_pd_dispatches_sequentially_with_nixl_relay() { + let (prefill_url, prefill_seen) = spawn_recording_stub( + r#"{"kv_transfer_params":{"remote_engine_id":"eng0","remote_block_ids":[1,2]}}"#, + ) + .await; + let (decode_url, decode_seen) = + spawn_recording_stub(r#"{"object":"chat.completion"}"#).await; + + let router = create_test_pd_router(); + let prefill = BasicWorkerBuilder::new(prefill_url) + .worker_type(WorkerType::Prefill) + .runtime_type(RuntimeType::Vllm) + .kv_connector("NixlConnector") + .build(); + prefill.set_status(openai_protocol::worker::WorkerStatus::Ready); + let decode = BasicWorkerBuilder::new(decode_url) + .worker_type(WorkerType::Decode) + .runtime_type(RuntimeType::Vllm) + .build(); + decode.set_status(openai_protocol::worker::WorkerStatus::Ready); + router + .worker_registry + .register_or_replace(Arc::new(prefill)); + router.worker_registry.register_or_replace(Arc::new(decode)); + let tenant = TenantRequestMeta::new(TenantKey::new("test-tenant")); + + let chat: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "m", + "max_tokens": 50, + "stream_options": {"include_usage": true}, + "messages": [{"role": "user", "content": "hi"}], + })) + .expect("valid chat request"); + let response = router + .route_chat(None, &tenant, chat, UNKNOWN_MODEL_ID) + .await; + assert_eq!(response.status(), StatusCode::OK); + + // Prefill leg: sanitized to a one-token unstreamed probe, tagged with + // the NIXL handoff params. + let prefill_seen = prefill_seen + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert_eq!(prefill_seen.len(), 1, "exactly one prefill request"); + let (path, body) = &prefill_seen[0]; + assert_eq!(path, "/v1/chat/completions"); + assert_eq!(body.get("max_tokens"), Some(&Value::from(1))); + assert_eq!(body.get("stream"), Some(&Value::Bool(false))); + assert_eq!(body.get("stream_options"), None); + assert_eq!( + body.pointer("/kv_transfer_params/do_remote_decode"), + Some(&Value::Bool(true)) + ); + + // Decode leg: original sampling, carrying the params the prefill + // engine returned — proof the legs ran sequentially. + let decode_seen = decode_seen + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert_eq!(decode_seen.len(), 1, "exactly one decode request"); + let (path, body) = &decode_seen[0]; + assert_eq!(path, "/v1/chat/completions"); + assert_eq!(body.get("max_tokens"), Some(&Value::from(50))); + assert_eq!( + body.pointer("/kv_transfer_params/remote_engine_id"), + Some(&Value::from("eng0")) + ); + assert_eq!(body.get("bootstrap_room"), None, "no bootstrap on vLLM PD"); + } + + #[tokio::test] + async fn vllm_pd_fanout_skips_the_prefill_leg() { + let (prefill_url, prefill_seen) = spawn_recording_stub("{}").await; + let (decode_url, decode_seen) = + spawn_recording_stub(r#"{"object":"chat.completion"}"#).await; + + let router = create_test_pd_router(); + let prefill = BasicWorkerBuilder::new(prefill_url) + .worker_type(WorkerType::Prefill) + .runtime_type(RuntimeType::Vllm) + .kv_connector("NixlConnector") + .build(); + prefill.set_status(openai_protocol::worker::WorkerStatus::Ready); + let decode = BasicWorkerBuilder::new(decode_url) + .worker_type(WorkerType::Decode) + .runtime_type(RuntimeType::Vllm) + .build(); + decode.set_status(openai_protocol::worker::WorkerStatus::Ready); + router + .worker_registry + .register_or_replace(Arc::new(prefill)); + router.worker_registry.register_or_replace(Arc::new(decode)); + let tenant = TenantRequestMeta::new(TenantKey::new("test-tenant")); + + let chat: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "m", + "n": 2, + "max_tokens": 50, + "messages": [{"role": "user", "content": "hi"}], + })) + .expect("valid chat request"); + let response = router + .route_chat(None, &tenant, chat, UNKNOWN_MODEL_ID) + .await; + assert_eq!(response.status(), StatusCode::OK); + + // The single-consumer KV handoff cannot serve an n>1 fan-out, so the + // prefill leg is skipped entirely rather than burned for nothing. + let prefill_seen = prefill_seen + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert!(prefill_seen.is_empty(), "prefill leg must not be contacted"); + + let decode_seen = decode_seen + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert_eq!(decode_seen.len(), 1, "exactly one decode request"); + let (_, body) = &decode_seen[0]; + assert_eq!(body.get("n"), Some(&Value::from(2))); + assert_eq!(body.get("kv_transfer_params"), None); + } + + #[test] + fn prefill_sanitization_is_route_aware() { + let mut chat = json!({"max_completion_tokens": 99, "max_tokens": 88, "n": 3, + "stream": true, "stream_options": {"include_usage": true}, "min_tokens": 5}); + PDRouter::sanitize_prefill_for_kv_handoff(&mut chat, "/v1/chat/completions"); + assert_eq!(chat.get("max_completion_tokens"), Some(&Value::from(1))); + assert_eq!(chat.get("max_tokens"), None, "both caps would be rejected"); + assert_eq!(chat.get("n"), Some(&Value::from(1))); + assert_eq!(chat.get("stream"), Some(&Value::Bool(false))); + assert_eq!(chat.get("stream_options"), None); + assert_eq!(chat.get("min_tokens"), None); + + let mut responses = json!({"max_output_tokens": 99, "stream": true}); + PDRouter::sanitize_prefill_for_kv_handoff(&mut responses, "/v1/responses"); + assert_eq!(responses.get("max_output_tokens"), Some(&Value::from(1))); + + let mut completion = json!({"prompt": "hi"}); + PDRouter::sanitize_prefill_for_kv_handoff(&mut completion, "/v1/completions"); + assert_eq!(completion.get("max_tokens"), Some(&Value::from(1))); + } + #[tokio::test] async fn messages_and_responses_endpoints_dispatch_through_pd() { let router = create_test_pd_router();