forked from Karylab-cklius/vllm
[EC Connector] Add EC Transfer Params (#42433)
Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
@@ -202,5 +202,8 @@ pub enum ChatEvent {
|
||||
finish_reason: FinishReason,
|
||||
/// Connector-specific KV transfer parameters for disaggregated serving.
|
||||
kv_transfer_params: Option<serde_json::Value>,
|
||||
/// Connector-specific encoder cache transfer parameters for
|
||||
/// disaggregated serving.
|
||||
ec_transfer_params: Option<serde_json::Value>,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -257,6 +257,7 @@ pub(crate) async fn unified_event_stream(
|
||||
usage: finished.usage,
|
||||
finish_reason: finished.finish_reason,
|
||||
kv_transfer_params: finished.kv_transfer_params,
|
||||
ec_transfer_params: finished.ec_transfer_params,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
@@ -387,6 +388,7 @@ mod tests {
|
||||
usage: vllm_llm::TokenUsage::default(),
|
||||
finish_reason: crate::FinishReason::Stop(None),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -628,6 +630,7 @@ mod tests {
|
||||
usage: vllm_llm::TokenUsage::default(),
|
||||
finish_reason: crate::FinishReason::Stop(None),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
},
|
||||
]
|
||||
);
|
||||
@@ -671,6 +674,7 @@ mod tests {
|
||||
usage: vllm_llm::TokenUsage::default(),
|
||||
finish_reason: crate::FinishReason::Stop(None),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
},
|
||||
]
|
||||
);
|
||||
|
||||
@@ -370,6 +370,7 @@ async fn harmony_assistant_event_stream(
|
||||
usage: finished.usage,
|
||||
finish_reason: finished.finish_reason,
|
||||
kv_transfer_params: finished.kv_transfer_params,
|
||||
ec_transfer_params: finished.ec_transfer_params,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -52,6 +52,7 @@ fn finished() -> Finished {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -115,6 +116,7 @@ fn interrupted_final_message_is_preserved() {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -175,6 +177,7 @@ fn interrupted_analysis_message_is_preserved() {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
@@ -48,6 +48,9 @@ pub(crate) enum AssistantEvent {
|
||||
finish_reason: FinishReason,
|
||||
/// Connector-specific KV transfer parameters for disaggregated serving.
|
||||
kv_transfer_params: Option<serde_json::Value>,
|
||||
/// Connector-specific encoder cache transfer parameters for
|
||||
/// disaggregated serving.
|
||||
ec_transfer_params: Option<serde_json::Value>,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -146,6 +146,7 @@ impl StructuredEventState {
|
||||
usage: vllm_llm::TokenUsage,
|
||||
finish_reason: FinishReason,
|
||||
kv_transfer_params: Option<serde_json::Value>,
|
||||
ec_transfer_params: Option<serde_json::Value>,
|
||||
) -> Result<Vec<ChatEvent>> {
|
||||
let mut events = Vec::new();
|
||||
self.close_open_text_block(&mut events);
|
||||
@@ -155,6 +156,7 @@ impl StructuredEventState {
|
||||
usage,
|
||||
finish_reason,
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
});
|
||||
Ok(events)
|
||||
}
|
||||
@@ -296,8 +298,11 @@ pub(crate) async fn structured_chat_event_stream(
|
||||
usage,
|
||||
finish_reason,
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
} => {
|
||||
for next in state.finish(usage, finish_reason, kv_transfer_params)? {
|
||||
for next in
|
||||
state.finish(usage, finish_reason, kv_transfer_params, ec_transfer_params)?
|
||||
{
|
||||
y.yield_ok(next).await;
|
||||
}
|
||||
}
|
||||
@@ -334,6 +339,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -388,6 +394,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -439,6 +446,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -490,6 +498,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -557,6 +566,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
|
||||
@@ -22,6 +22,9 @@ pub struct CollectedAssistantMessage {
|
||||
pub finish_reason: FinishReason,
|
||||
/// Connector-specific KV transfer parameters for disaggregated serving.
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
/// Connector-specific encoder cache transfer parameters for disaggregated
|
||||
/// serving.
|
||||
pub ec_transfer_params: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Per-request stream of chat events.
|
||||
@@ -77,6 +80,7 @@ impl ChatEventStream {
|
||||
usage,
|
||||
finish_reason,
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
} => {
|
||||
return Ok(CollectedAssistantMessage {
|
||||
message: done,
|
||||
@@ -89,6 +93,7 @@ impl ChatEventStream {
|
||||
usage,
|
||||
finish_reason,
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
});
|
||||
}
|
||||
ChatEvent::ToolCallEnd { call, .. } => {
|
||||
@@ -194,6 +199,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]),
|
||||
);
|
||||
@@ -234,6 +240,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@@ -50,6 +50,7 @@ fn request_output(
|
||||
stop_reason,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -75,6 +76,7 @@ fn request_output_with_logprobs(
|
||||
stop_reason,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
|
||||
@@ -676,6 +676,7 @@ fn decoded_completion_stream(
|
||||
usage: Default::default(),
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}
|
||||
});
|
||||
@@ -686,6 +687,7 @@ fn decoded_completion_stream(
|
||||
usage: Default::default(),
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
});
|
||||
events.push(DecodedTextEvent::TextDelta {
|
||||
delta: chunk.delta,
|
||||
|
||||
@@ -97,6 +97,8 @@ pub struct EngineCoreOutput {
|
||||
#[serde(default)]
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub ec_transfer_params: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<OpaqueValue>,
|
||||
/// Breakdown of the scheduled prefill computation, set on the first output
|
||||
/// of a newly scheduled prefill and elided for subsequent decode outputs.
|
||||
@@ -374,6 +376,7 @@ mod tests {
|
||||
stop_reason: Some(StopReason::Text("stop".to_string())),
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -426,6 +429,7 @@ mod tests {
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
|
||||
@@ -226,6 +226,7 @@ fn request_output(
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -2517,6 +2518,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
|
||||
@@ -91,6 +91,7 @@ class EngineCoreOutput(
|
||||
stop_reason: int | str | None = None
|
||||
events: object | None = None
|
||||
kv_transfer_params: object | None = None
|
||||
ec_transfer_params: object | None = None
|
||||
trace_headers: object | None = None
|
||||
prefill_stats: object | None = None
|
||||
routed_experts: object | None = None
|
||||
|
||||
@@ -38,6 +38,9 @@ pub struct CollectedGenerateOutput {
|
||||
pub usage: TokenUsage,
|
||||
/// Connector-specific KV transfer parameters for disaggregated serving.
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
/// Connector-specific encoder cache transfer parameters for disaggregated
|
||||
/// serving.
|
||||
pub ec_transfer_params: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for
|
||||
@@ -146,6 +149,9 @@ pub struct GenerateOutput {
|
||||
pub cached_token_count: usize,
|
||||
/// Connector-specific KV transfer parameters for disaggregated serving.
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
/// Connector-specific encoder cache transfer parameters for disaggregated
|
||||
/// serving.
|
||||
pub ec_transfer_params: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl GenerateOutput {
|
||||
@@ -192,6 +198,7 @@ impl GenerateOutput {
|
||||
finish_reason,
|
||||
cached_token_count: 0,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -280,6 +287,7 @@ impl Stream for GenerateOutputStream {
|
||||
finish_reason,
|
||||
cached_token_count,
|
||||
kv_transfer_params: raw.kv_transfer_params,
|
||||
ec_transfer_params: raw.ec_transfer_params,
|
||||
};
|
||||
|
||||
Poll::Ready(Some(Ok(output)))
|
||||
@@ -362,6 +370,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
|
||||
cached_token_count,
|
||||
},
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -374,6 +383,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
|
||||
cached_token_count,
|
||||
};
|
||||
collected.kv_transfer_params = output.kv_transfer_params;
|
||||
collected.ec_transfer_params = output.ec_transfer_params;
|
||||
return Ok(collected);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -51,6 +51,7 @@ fn request_output_with_events(
|
||||
stop_reason: None,
|
||||
events,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -75,6 +76,7 @@ fn request_output_with_logprobs(
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -89,6 +91,7 @@ fn request_output_with_logprobs_and_kv(
|
||||
new_logprobs: Option<Logprobs>,
|
||||
prompt_logprobs: Option<Logprobs>,
|
||||
kv_transfer_params: Option<serde_json::Value>,
|
||||
ec_transfer_params: Option<serde_json::Value>,
|
||||
) -> EngineCoreOutput {
|
||||
EngineCoreOutput {
|
||||
request_id: request_id.to_string(),
|
||||
@@ -100,6 +103,7 @@ fn request_output_with_logprobs_and_kv(
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -356,6 +360,7 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
|
||||
Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)),
|
||||
None,
|
||||
Some(serde_json::json!({"connector": "x"})),
|
||||
None,
|
||||
),
|
||||
],
|
||||
..Default::default()
|
||||
|
||||
@@ -69,6 +69,11 @@ pub fn to_text_request(
|
||||
let map = sampling_params.vllm_xargs.get_or_insert_with(Default::default);
|
||||
map.insert("kv_transfer_params".to_string(), kv_json);
|
||||
}
|
||||
if let Some(ec_struct) = kv.ec_transfer_params.as_ref() {
|
||||
let ec_json = proto_struct_to_json(ec_struct);
|
||||
let map = sampling_params.vllm_xargs.get_or_insert_with(Default::default);
|
||||
map.insert("ec_transfer_params".to_string(), ec_json);
|
||||
}
|
||||
if kv.bypass_prefix_cache {
|
||||
sampling_params.skip_reading_prefix_cache = Some(true);
|
||||
}
|
||||
@@ -343,6 +348,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
|
||||
finish_reason,
|
||||
stop_reason,
|
||||
kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
|
||||
ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -586,6 +592,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: reason,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -71,6 +71,7 @@ impl pb::generate_server::Generate for GenerateServiceImpl {
|
||||
usage: collected.usage,
|
||||
finish_reason: collected.finish_reason,
|
||||
kv_transfer_params: collected.kv_transfer_params,
|
||||
ec_transfer_params: collected.ec_transfer_params,
|
||||
};
|
||||
|
||||
let outputs = convert::to_sequence_output(
|
||||
|
||||
@@ -104,6 +104,7 @@ fn request_output(
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
|
||||
@@ -100,6 +100,7 @@ fn request_output(
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
|
||||
@@ -266,6 +266,7 @@ fn collect_generate(
|
||||
}],
|
||||
prompt_logprobs,
|
||||
kv_transfer_params: collected.kv_transfer_params,
|
||||
ec_transfer_params: collected.ec_transfer_params,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -404,6 +405,7 @@ mod tests {
|
||||
finish_reason: None,
|
||||
cached_token_count: 0,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
Ok(GenerateOutput {
|
||||
request_id: String::new(),
|
||||
@@ -416,6 +418,7 @@ mod tests {
|
||||
finish_reason: Some(FinishReason::stop_eos()),
|
||||
cached_token_count: 2,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ use super::types::GenerateRequest;
|
||||
use super::validate;
|
||||
use crate::error::ApiError;
|
||||
use crate::lora::LoraModelResolution;
|
||||
use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
|
||||
use crate::utils::{ResolvedRequestContext, merge_ec_transfer_params, merge_kv_transfer_params};
|
||||
|
||||
/// Lowered generate request plus the response request ID.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
@@ -56,6 +56,10 @@ pub(super) fn prepare_generate_request(
|
||||
sampling_params.vllm_xargs,
|
||||
request.kv_transfer_params.as_ref(),
|
||||
);
|
||||
sampling_params.vllm_xargs = merge_ec_transfer_params(
|
||||
sampling_params.vllm_xargs,
|
||||
request.ec_transfer_params.as_ref(),
|
||||
);
|
||||
|
||||
let text_request = TextRequest {
|
||||
request_id: ctx.request_id.clone(),
|
||||
|
||||
@@ -22,6 +22,7 @@ pub struct GenerateRequest {
|
||||
#[serde(default)]
|
||||
pub priority: i32,
|
||||
pub kv_transfer_params: Option<HashMap<String, Value>>,
|
||||
pub ec_transfer_params: Option<HashMap<String, Value>>,
|
||||
#[serde(flatten)]
|
||||
pub other: Map<String, Value>,
|
||||
}
|
||||
@@ -66,6 +67,7 @@ pub(super) struct GenerateResponse {
|
||||
pub choices: Vec<GenerateResponseChoice>,
|
||||
pub prompt_logprobs: Option<Vec<Option<HashMap<u32, GenerateLogprob>>>>,
|
||||
pub kv_transfer_params: Option<Value>,
|
||||
pub ec_transfer_params: Option<Value>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads.
|
||||
|
||||
@@ -144,6 +144,7 @@ async fn collect_chat_completion(
|
||||
usage,
|
||||
finish_reason,
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
} = collected;
|
||||
let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
|
||||
let saw_tool_calls = message.tool_calls().next().is_some();
|
||||
@@ -224,6 +225,7 @@ async fn collect_chat_completion(
|
||||
prompt_logprobs,
|
||||
prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()),
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -951,6 +953,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -1031,6 +1034,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -1086,6 +1090,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -1167,6 +1172,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -1300,6 +1306,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -1381,6 +1388,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
|
||||
@@ -13,7 +13,9 @@ use crate::routes::openai::utils::structured_outputs::convert_from_response_form
|
||||
use crate::routes::openai::utils::types::{
|
||||
ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue,
|
||||
};
|
||||
use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params};
|
||||
use crate::utils::{
|
||||
ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params,
|
||||
};
|
||||
|
||||
/// Lowered chat request plus the public response metadata carried by every SSE
|
||||
/// chunk.
|
||||
@@ -132,7 +134,7 @@ pub(super) fn prepare_chat_request(
|
||||
structured_outputs,
|
||||
skip_reading_prefix_cache: None,
|
||||
vllm_xargs: merge_kv_transfer_params(
|
||||
request.vllm_xargs,
|
||||
merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()),
|
||||
request.kv_transfer_params.as_ref(),
|
||||
),
|
||||
},
|
||||
|
||||
@@ -232,6 +232,9 @@ pub struct ChatCompletionRequest {
|
||||
/// KV transfer parameters for disaggregated serving
|
||||
pub kv_transfer_params: Option<HashMap<String, Value>>,
|
||||
|
||||
/// Encoder cache transfer parameters for disaggregated serving
|
||||
pub ec_transfer_params: Option<HashMap<String, Value>>,
|
||||
|
||||
/// Additional request parameters with string or numeric values for custom
|
||||
/// extensions
|
||||
pub vllm_xargs: Option<HashMap<String, Value>>,
|
||||
@@ -299,6 +302,7 @@ impl Default for ChatCompletionRequest {
|
||||
return_token_ids: None,
|
||||
cache_salt: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
vllm_xargs: None,
|
||||
repetition_detection: None,
|
||||
}
|
||||
@@ -346,6 +350,7 @@ pub(super) struct ChatCompletionResponse {
|
||||
pub prompt_logprobs: Option<Vec<Option<HashMap<String, f32>>>>,
|
||||
pub prompt_token_ids: Option<Vec<u32>>,
|
||||
pub kv_transfer_params: Option<Value>,
|
||||
pub ec_transfer_params: Option<Value>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `ChatCompletionResponseChoice` class.
|
||||
|
||||
@@ -212,6 +212,7 @@ async fn collect_completion(
|
||||
usage: Some(usage),
|
||||
system_fingerprint: None,
|
||||
kv_transfer_params: collected.kv_transfer_params,
|
||||
ec_transfer_params: collected.ec_transfer_params,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -682,6 +683,7 @@ mod tests {
|
||||
"repetition_detected".to_string(),
|
||||
))),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
@@ -786,6 +788,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::Length,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
@@ -837,6 +840,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::Length,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
@@ -891,6 +895,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::Length,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
@@ -962,6 +967,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::Length,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
@@ -1035,6 +1041,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::Length,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
|
||||
@@ -7,7 +7,9 @@ use crate::error::ApiError;
|
||||
use crate::lora::LoraModelResolution;
|
||||
use crate::routes::openai::completions::validate;
|
||||
use crate::routes::openai::utils::structured_outputs::convert_from_response_format_value;
|
||||
use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params};
|
||||
use crate::utils::{
|
||||
ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params,
|
||||
};
|
||||
|
||||
/// Lowered completion request plus the public response metadata carried by
|
||||
/// every SSE chunk.
|
||||
@@ -127,7 +129,7 @@ pub(super) fn prepare_completion_request(
|
||||
structured_outputs,
|
||||
skip_reading_prefix_cache: None,
|
||||
vllm_xargs: merge_kv_transfer_params(
|
||||
request.vllm_xargs,
|
||||
merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()),
|
||||
request.kv_transfer_params.as_ref(),
|
||||
),
|
||||
},
|
||||
|
||||
@@ -174,6 +174,9 @@ pub struct CompletionRequest {
|
||||
/// KV transfer parameters for disaggregated serving
|
||||
pub kv_transfer_params: Option<HashMap<String, Value>>,
|
||||
|
||||
/// Encoder cache transfer parameters for disaggregated serving
|
||||
pub ec_transfer_params: Option<HashMap<String, Value>>,
|
||||
|
||||
/// Additional request parameters with string or numeric values for custom
|
||||
/// extensions
|
||||
pub vllm_xargs: Option<HashMap<String, Value>>,
|
||||
@@ -209,6 +212,7 @@ pub(super) struct CompletionResponse {
|
||||
pub usage: Option<Usage>,
|
||||
pub system_fingerprint: Option<String>,
|
||||
pub kv_transfer_params: Option<Value>,
|
||||
pub ec_transfer_params: Option<Value>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `CompletionResponseChoice` class.
|
||||
|
||||
@@ -76,6 +76,7 @@ fn request_output_with_stop_reason(
|
||||
stop_reason,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -101,6 +102,7 @@ fn request_output_with_logprobs(
|
||||
stop_reason,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -116,6 +118,7 @@ fn request_output_with_logprobs_and_kv(
|
||||
new_logprobs: Option<Logprobs>,
|
||||
new_prompt_logprobs_tensors: Option<Logprobs>,
|
||||
kv_transfer_params: Option<serde_json::Value>,
|
||||
ec_transfer_params: Option<serde_json::Value>,
|
||||
) -> EngineCoreOutput {
|
||||
EngineCoreOutput {
|
||||
request_id: request_id.to_string(),
|
||||
@@ -127,6 +130,7 @@ fn request_output_with_logprobs_and_kv(
|
||||
stop_reason,
|
||||
events: None,
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
@@ -3634,6 +3638,7 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
|
||||
Some(sample_logprobs_for_token(44, 45)),
|
||||
None,
|
||||
Some(json!({"connector": "x"})),
|
||||
None,
|
||||
),
|
||||
],
|
||||
..Default::default()
|
||||
|
||||
@@ -45,6 +45,24 @@ pub fn merge_kv_transfer_params(
|
||||
xargs
|
||||
}
|
||||
|
||||
/// Merge `ec_transfer_params` into the `vllm_xargs` map, mirroring the Python
|
||||
/// vLLM behavior where `ec_transfer_params` is injected into `extra_args` for
|
||||
/// engine-core consumption.
|
||||
pub fn merge_ec_transfer_params(
|
||||
mut xargs: Option<HashMap<String, Value>>,
|
||||
ec_transfer_params: Option<&HashMap<String, Value>>,
|
||||
) -> Option<HashMap<String, Value>> {
|
||||
if let Some(ec_params) = ec_transfer_params {
|
||||
let map = xargs.get_or_insert_with(HashMap::new);
|
||||
map.insert(
|
||||
"ec_transfer_params".to_string(),
|
||||
// This is safe because we know that `ec_params` is already valid JSON.
|
||||
serde_json::to_value(ec_params).unwrap(),
|
||||
);
|
||||
}
|
||||
xargs
|
||||
}
|
||||
|
||||
/// Convert OpenAI-style `logit_bias` with string token-ID keys into the
|
||||
/// internal `HashMap<u32, f32>` representation, validating that every key
|
||||
/// parses as a `u32`.
|
||||
|
||||
@@ -44,6 +44,9 @@ pub struct Finished {
|
||||
pub finish_reason: FinishReason,
|
||||
/// Connector-specific KV transfer parameters for disaggregated serving.
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
/// Connector-specific encoder cache transfer parameters for disaggregated
|
||||
/// serving.
|
||||
pub ec_transfer_params: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Internal decoded-text event emitted before higher-level assistant
|
||||
@@ -151,6 +154,7 @@ pub async fn decoded_text_event_stream(
|
||||
let decoder = decoder.as_mut().unwrap();
|
||||
|
||||
let kv_transfer_params = output.kv_transfer_params;
|
||||
let ec_transfer_params = output.ec_transfer_params;
|
||||
let mut finish_reason = output.finish_reason;
|
||||
let mut stop_str_matched = false;
|
||||
let suppress_terminal_stop_token = finish_reason.as_ref().is_some_and(|r| r.is_stop())
|
||||
@@ -275,6 +279,7 @@ pub async fn decoded_text_event_stream(
|
||||
},
|
||||
finish_reason: reason,
|
||||
kv_transfer_params,
|
||||
ec_transfer_params,
|
||||
}),
|
||||
})
|
||||
.await;
|
||||
|
||||
@@ -26,6 +26,9 @@ pub struct CollectedTextOutput {
|
||||
pub usage: vllm_llm::TokenUsage,
|
||||
/// Connector-specific KV transfer parameters for disaggregated serving.
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
/// Connector-specific encoder cache transfer parameters for disaggregated
|
||||
/// serving.
|
||||
pub ec_transfer_params: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[allow(clippy::manual_async_fn, reason = "specify `Send` bound")]
|
||||
@@ -77,6 +80,7 @@ impl<T: TextOutputStream> T {
|
||||
finish_reason: FinishReason::Error,
|
||||
usage: vllm_llm::TokenUsage::default(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
})
|
||||
};
|
||||
|
||||
@@ -85,6 +89,7 @@ impl<T: TextOutputStream> T {
|
||||
collected.finish_reason = finished.finish_reason;
|
||||
collected.usage = finished.usage;
|
||||
collected.kv_transfer_params = finished.kv_transfer_params;
|
||||
collected.ec_transfer_params = finished.ec_transfer_params;
|
||||
return Ok(collected);
|
||||
}
|
||||
}
|
||||
@@ -156,6 +161,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
@@ -273,6 +279,7 @@ mod tests {
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
}),
|
||||
}),
|
||||
]);
|
||||
|
||||
Reference in New Issue
Block a user