From cf30ef60cfe3397c937867f03463dbc2c16998bb Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 22 Jun 2026 19:30:10 +0800 Subject: [PATCH] use newtype for logprobs count (-1 for all) Signed-off-by: Bugen Zhao --- rust/src/chat/tests/chat.rs | 7 +- rust/src/cmd/src/cli.rs | 11 +- rust/src/cmd/src/cli/tests.rs | 7 +- .../examples/external_engine_logprobs.rs | 14 +- .../src/protocol/logprobs_count.rs | 135 ++++++++++++++++++ .../engine-core-client/src/protocol/mod.rs | 10 +- rust/src/server/src/config.rs | 14 +- rust/src/server/src/grpc/convert.rs | 22 +-- .../src/routes/inference/generate/convert.rs | 8 +- .../src/routes/inference/generate/validate.rs | 10 -- .../routes/openai/chat_completions/convert.rs | 20 ++- .../routes/openai/chat_completions/types.rs | 11 +- .../openai/chat_completions/validate.rs | 36 ++--- .../src/routes/openai/completions/convert.rs | 38 ++--- .../src/routes/openai/completions/types.rs | 5 +- .../src/routes/openai/completions/validate.rs | 29 ++-- rust/src/text/src/backend/mod.rs | 7 +- rust/src/text/src/lib.rs | 5 +- rust/src/text/src/lower.rs | 21 +-- rust/src/text/src/lower/logprobs.rs | 37 ++--- rust/src/text/src/request.rs | 10 +- 21 files changed, 293 insertions(+), 164 deletions(-) create mode 100644 rust/src/engine-core-client/src/protocol/logprobs_count.rs diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index 07aa304af00..8020ec1cfbd 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -16,7 +16,8 @@ use vllm_engine_core_client::protocol::logprobs::{ Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob, }; use vllm_engine_core_client::protocol::{ - EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason, + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, LogprobsCount, + StopReason, }; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig}; @@ -1387,8 +1388,8 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() { .await; let mut request = sample_request("chat-logprobs"); - request.sampling_params.logprobs = Some(1); - request.sampling_params.prompt_logprobs = Some(1); + request.sampling_params.logprobs = Some(LogprobsCount::Top(1)); + request.sampling_params.prompt_logprobs = Some(LogprobsCount::Top(1)); let mut stream = chat.chat(request.clone()).await.unwrap(); match next_semantic(&mut stream).await.unwrap().unwrap() { diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 003d96fa92b..78d1ea32dd4 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -20,6 +20,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as}; use thiserror_ext::AsReport as _; use uuid::Uuid; use vllm_engine_core_client::TransportMode; +use vllm_engine_core_client::protocol::LogprobsCount; use vllm_managed_engine::ManagedEngineConfig; use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args}; use vllm_server::{ @@ -136,9 +137,9 @@ pub struct SharedRuntimeArgs { pub max_model_len: Option, /// Maximum number of log probabilities to return when `logprobs` is /// specified in sampling parameters. `-1` means no cap. - #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)] + #[arg(long, allow_negative_numbers = true)] #[serde(default)] - pub max_logprobs: Option, + pub max_logprobs: Option, /// TCP port for the gRPC Generate service. When not set, no gRPC server is /// started. #[arg(long)] @@ -529,7 +530,7 @@ impl ServeArgs { self.managed_engine.clone().into_config( self.runtime.model.clone(), self.runtime.max_model_len, - self.runtime.max_logprobs, + self.runtime.max_logprobs.map(managed_max_logprobs_to_i32), self.runtime.language_model_only, self.runtime.disable_log_stats, self.runtime.shutdown_timeout, @@ -555,5 +556,9 @@ fn frontend_ipc_addresses() -> (String, String) { ) } +fn managed_max_logprobs_to_i32(count: LogprobsCount) -> i32 { + i32::try_from(count).expect("max_logprobs is parsed through i32") +} + #[cfg(test)] mod tests; diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index c57c23e017c..8bb801d9d9a 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -1,5 +1,6 @@ use expect_test::expect; use vllm_engine_core_client::TransportMode; +use vllm_engine_core_client::protocol::LogprobsCount; use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection}; use super::{Cli, Command}; @@ -165,10 +166,10 @@ fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() { let Command::Serve(args) = cli.command else { panic!("expected serve args"); }; - assert_eq!(args.runtime.max_logprobs, Some(-1)); + assert_eq!(args.runtime.max_logprobs, Some(LogprobsCount::All)); let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string()); - assert_eq!(frontend_config.max_logprobs, Some(-1)); + assert_eq!(frontend_config.max_logprobs, Some(LogprobsCount::All)); let engine_config = args.to_managed_engine_config(5555); assert_eq!(engine_config.python_args, vec!["--max-logprobs", "-1"]); @@ -529,7 +530,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() { assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32); assert!(args.runtime.language_model_only); assert_eq!(args.runtime.max_model_len, Some(8192)); - assert_eq!(args.runtime.max_logprobs, Some(-1)); + assert_eq!(args.runtime.max_logprobs, Some(LogprobsCount::All)); assert_eq!(args.runtime.shutdown_timeout, 3); } diff --git a/rust/src/engine-core-client/examples/external_engine_logprobs.rs b/rust/src/engine-core-client/examples/external_engine_logprobs.rs index 08290c69bf5..8b2ce57725c 100644 --- a/rust/src/engine-core-client/examples/external_engine_logprobs.rs +++ b/rust/src/engine-core-client/examples/external_engine_logprobs.rs @@ -6,7 +6,7 @@ use futures::StreamExt as _; use tokio::time::timeout; use tracing_subscriber::EnvFilter; use vllm_engine_core_client::protocol::{ - EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams, + EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams, LogprobsCount, }; use vllm_engine_core_client::{ EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode, @@ -33,10 +33,10 @@ struct Args { output_timeout_secs: u64, #[arg(long, default_value_t = 1)] max_tokens: u32, - #[arg(long, default_value_t = 2)] - logprobs: i32, - #[arg(long, default_value_t = 1)] - prompt_logprobs: i32, + #[arg(long, default_value_t = LogprobsCount::Top(2), allow_negative_numbers = true)] + logprobs: LogprobsCount, + #[arg(long, default_value_t = LogprobsCount::Top(1), allow_negative_numbers = true)] + prompt_logprobs: LogprobsCount, #[arg(long, default_value_t = 96)] prompt_repeats: usize, } @@ -64,8 +64,8 @@ fn build_request( request_id: String, prompt_token_ids: Vec, max_tokens: u32, - logprobs: i32, - prompt_logprobs: i32, + logprobs: LogprobsCount, + prompt_logprobs: LogprobsCount, client_index: u32, ) -> EngineCoreRequest { EngineCoreRequest { diff --git a/rust/src/engine-core-client/src/protocol/logprobs_count.rs b/rust/src/engine-core-client/src/protocol/logprobs_count.rs new file mode 100644 index 00000000000..df53450cbac --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/logprobs_count.rs @@ -0,0 +1,135 @@ +use std::fmt; +use std::str::FromStr; + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// Number of log probabilities requested for a token position. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LogprobsCount { + /// Return the full model vocabulary. + All, + /// Return the top-N tokens by probability. + Top(u32), +} + +impl LogprobsCount { + /// Expands the count to the actual number of logprobs to return, given the vocabulary size. + pub fn expanded(self, vocab_size: usize) -> usize { + match self { + Self::All => vocab_size, + Self::Top(count) => count as usize, + } + } +} + +impl TryFrom for LogprobsCount { + type Error = String; + + fn try_from(value: i32) -> Result { + match value { + -1 => Ok(Self::All), + value if value < -1 => Err(format!("must be non-negative or -1, got {value}")), + value => Ok(Self::Top(value as u32)), + } + } +} + +impl TryFrom for i32 { + type Error = String; + + fn try_from(value: LogprobsCount) -> Result { + match value { + LogprobsCount::All => Ok(-1), + LogprobsCount::Top(count) => { + i32::try_from(count).map_err(|_| format!("must fit within i32, got {count}")) + } + } + } +} + +impl FromStr for LogprobsCount { + type Err = String; + + fn from_str(s: &str) -> Result { + let value = s + .parse::() + .map_err(|e| format!("must be an i32 integer, got {s:?}: {e}"))?; + Self::try_from(value) + } +} + +impl fmt::Display for LogprobsCount { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::All => (-1).fmt(f), + Self::Top(count) => count.fmt(f), + } + } +} + +impl Serialize for LogprobsCount { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let value: i32 = (*self).try_into().map_err(serde::ser::Error::custom)?; + value.serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for LogprobsCount { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = i32::deserialize(deserializer)?; + Self::try_from(value).map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use rmpv::Value; + + use super::*; + use crate::protocol::{decode_msgpack, encode_msgpack}; + + #[test] + fn logprobs_count_serializes_as_wire_integer() { + assert_eq!(serde_json::to_value(LogprobsCount::All).unwrap(), -1); + assert_eq!(serde_json::to_value(LogprobsCount::Top(3)).unwrap(), 3); + } + + #[test] + fn logprobs_count_deserializes_wire_integer() { + assert_eq!( + serde_json::from_value::(serde_json::json!(-1)).unwrap(), + LogprobsCount::All + ); + assert_eq!( + serde_json::from_value::(serde_json::json!(3)).unwrap(), + LogprobsCount::Top(3) + ); + assert!(serde_json::from_value::(serde_json::json!(-2)).is_err()); + assert!( + serde_json::from_value::(serde_json::json!(i64::from(i32::MAX) + 1)) + .is_err() + ); + } + + #[test] + fn logprobs_count_decodes_msgpack_signed_and_unsigned() { + let mut encoded = Vec::new(); + rmpv::encode::write_value(&mut encoded, &Value::from(-1)).unwrap(); + assert_eq!( + decode_msgpack::(&encoded).unwrap(), + LogprobsCount::All + ); + + let encoded = encode_msgpack(&LogprobsCount::Top(7)).unwrap(); + assert_eq!( + decode_msgpack::(&encoded).unwrap(), + LogprobsCount::Top(7) + ); + } +} diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs index 8862b077c9d..2d9b5948f9b 100644 --- a/rust/src/engine-core-client/src/protocol/mod.rs +++ b/rust/src/engine-core-client/src/protocol/mod.rs @@ -56,6 +56,7 @@ mod classified_outputs; pub mod dtype; pub mod handshake; pub mod logprobs; +mod logprobs_count; pub mod lora; pub mod multimodal; pub mod stats; @@ -66,6 +67,7 @@ pub use classified_outputs::{ }; pub use dtype::ModelDtype; pub use logprobs::decode_engine_core_outputs; +pub use logprobs_count::LogprobsCount; /// Request types are encoded as single-byte protocol constants so they can be /// sent over the ZMQ socket without an extra encoding step. @@ -285,12 +287,12 @@ pub struct EngineCoreSamplingParams { pub thinking_token_budget: Option, /// Number of log probabilities to return per generated token. /// - /// `None` disables sample logprobs. `-1` requests the full vocabulary. - pub logprobs: Option, + /// `None` disables sample logprobs. + pub logprobs: Option, /// Number of log probabilities to return per prompt token. /// - /// `None` disables prompt logprobs. `-1` requests the full vocabulary. - pub prompt_logprobs: Option, + /// `None` disables prompt logprobs. + pub prompt_logprobs: Option, /// Minimum probability threshold for token sampling. pub min_p: f32, /// Frequency penalty applied by the sampler. diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index c601bbfb634..7a259554191 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -2,12 +2,13 @@ use std::collections::HashMap; use std::fmt; use std::time::Duration; -use anyhow::{Result, bail}; +use anyhow::Result; use axum::http::{HeaderName, HeaderValue, Method}; use educe::Educe; use serde::Serialize; use serde_json::Value; use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection}; +use vllm_engine_core_client::protocol::LogprobsCount; use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; /// How the HTTP server obtains its listening socket. @@ -133,7 +134,7 @@ pub struct Config { pub chat_template_content_format: ChatTemplateContentFormatOption, /// Optional maximum number of top log probabilities accepted by the /// frontend. `None` delegates to the text layer default. - pub max_logprobs: Option, + pub max_logprobs: Option, /// HTTP/API-server behavior switches. pub api_server_options: ApiServerOptions, /// CORS settings applied to every HTTP response. @@ -158,15 +159,6 @@ impl Config { pub fn validate(&self) -> Result<()> { vllm_chat::validate_parser_overrides(&self.tool_call_parser, &self.reasoning_parser)?; self.cors.validate()?; - if let Some(max_logprobs) = self.max_logprobs - && max_logprobs < -1 - { - bail!( - "max_logprobs must be non-negative or -1, got {}", - max_logprobs - ); - } - Ok(()) } diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 0bfe7a63beb..e429ad86704 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -3,7 +3,7 @@ use tonic::Status; use uuid::Uuid; -use vllm_engine_core_client::protocol::{StopReason, StructuredOutputsParams}; +use vllm_engine_core_client::protocol::{LogprobsCount, StopReason, StructuredOutputsParams}; use vllm_text::{ DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams, TextDecodeOptions, TextRequest, @@ -202,18 +202,22 @@ fn build_sampling_params( /// Map the proto `CandidateTokens` selector to a `(logprobs_count, /// logprob_token_ids)` pair. /// -/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability -/// - `all` → `(-1, None)` — return the full vocabulary +/// - `top_n(k)` → `(Top(k), None)` — return top-k candidates by probability +/// - `all` → `(All, None)` — return the full vocabulary /// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the /// count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs /// are carried via `logprob_token_ids` on `SamplingParams`) -/// - absent → `(1, None)` — just the sampled/scored token -fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option>) { +/// - absent → `(Top(1), None)` — just the sampled/scored token +fn candidate_logprob_spec( + candidates: Option<&pb::CandidateTokens>, +) -> (LogprobsCount, Option>) { match candidates.and_then(|c| c.select.as_ref()) { - Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None), - Some(pb::candidate_tokens::Select::All(true)) => (-1, None), - Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())), - _ => (1, None), + Some(pb::candidate_tokens::Select::TopN(n)) => (LogprobsCount::Top(*n), None), + Some(pb::candidate_tokens::Select::All(true)) => (LogprobsCount::All, None), + Some(pb::candidate_tokens::Select::TokenIds(ids)) => { + (LogprobsCount::Top(1), Some(ids.ids.clone())) + } + _ => (LogprobsCount::Top(1), None), } } diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index df3ba337357..3b04ef0a676 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -87,6 +87,7 @@ pub(super) fn prepare_generate_request( #[cfg(test)] mod tests { use serde_json::json; + use vllm_engine_core_client::protocol::LogprobsCount; use vllm_text::Prompt; use super::prepare_generate_request; @@ -132,10 +133,13 @@ mod tests { Prompt::TokenIds(vec![11, 22, 33]) ); assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7)); - assert_eq!(prepared.text_request.sampling_params.logprobs, Some(2)); + assert_eq!( + prepared.text_request.sampling_params.logprobs, + Some(LogprobsCount::Top(2)) + ); assert_eq!( prepared.text_request.sampling_params.prompt_logprobs, - Some(1) + Some(LogprobsCount::Top(1)) ); assert!(prepared.text_request.sampling_params.ignore_eos); assert_eq!(prepared.text_request.priority, -3); diff --git a/rust/src/server/src/routes/inference/generate/validate.rs b/rust/src/server/src/routes/inference/generate/validate.rs index 43347c60b57..c31010b01b8 100644 --- a/rust/src/server/src/routes/inference/generate/validate.rs +++ b/rust/src/server/src/routes/inference/generate/validate.rs @@ -34,16 +34,6 @@ pub(super) fn validate_request_compat( ); } - if let Some(prompt_logprobs) = request.sampling_params.prompt_logprobs - && prompt_logprobs < 0 - && prompt_logprobs != -1 - { - bail_invalid_request!( - param = "sampling_params", - "`prompt_logprobs` must be a non-negative value or -1." - ); - } - Ok(()) } diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 0462294b0c2..245fbace210 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -4,6 +4,7 @@ use vllm_chat::{ ChatMessage as VllmChatMessage, ChatOptions, ChatRequest, ChatTool, ChatToolChoice, GenerationPromptMode, SamplingParams, }; +use vllm_engine_core_client::protocol::LogprobsCount; use super::types::ChatCompletionRequest; use super::validate; @@ -94,7 +95,7 @@ pub(super) fn prepare_chat_request( // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's // behavior. - let top_logprobs = request.top_logprobs.unwrap_or(0); + let top_logprobs = request.top_logprobs.unwrap_or(LogprobsCount::Top(0)); let prompt_logprobs = request .prompt_logprobs .or((request.echo && !request.stream).then_some(top_logprobs)); @@ -378,6 +379,7 @@ mod tests { ChatTool as VllmChatTool, ChatToolChoice, GenerationPromptMode, SamplingParams as VllmSamplingParams, }; + use vllm_engine_core_client::protocol::LogprobsCount; use vllm_text::output::TextDecodeOptions; use super::prepare_chat_request; @@ -967,7 +969,7 @@ mod tests { let request = ChatCompletionRequest { stream: false, logprobs: true, - prompt_logprobs: Some(2), + prompt_logprobs: Some(LogprobsCount::Top(2)), ..base_request() }; @@ -980,10 +982,13 @@ mod tests { assert!(prepared.options.requested_logprobs); assert!(prepared.options.include_prompt_logprobs); - assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(0)); + assert_eq!( + prepared.chat_request.sampling_params.logprobs, + Some(LogprobsCount::Top(0)) + ); assert_eq!( prepared.chat_request.sampling_params.prompt_logprobs, - Some(2) + Some(LogprobsCount::Top(2)) ); } @@ -991,7 +996,7 @@ mod tests { fn prepare_chat_request_keeps_prompt_logprobs_independent_from_echo() { let request = ChatCompletionRequest { logprobs: true, - top_logprobs: Some(3), + top_logprobs: Some(LogprobsCount::Top(3)), echo: true, ..base_request() }; @@ -1003,7 +1008,10 @@ mod tests { ) .expect("request is valid"); - assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(3)); + assert_eq!( + prepared.chat_request.sampling_params.logprobs, + Some(LogprobsCount::Top(3)) + ); assert_eq!(prepared.chat_request.sampling_params.prompt_logprobs, None); assert!(!prepared.options.include_prompt_logprobs); } diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index a76a52b8ad0..74398f91361 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -6,6 +6,7 @@ use serde_json::Value; use serde_with::SerializeDisplay; use validator::Validate; use vllm_chat::ReasoningEffort; +use vllm_engine_core_client::protocol::LogprobsCount; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ @@ -44,10 +45,8 @@ pub struct ChatCompletionRequest { #[serde(default)] pub logprobs: bool, - /// An integer specifying the number of most likely tokens to return - /// -1 means return all - #[validate(range(min = -1))] - pub top_logprobs: Option, + /// Number of most likely tokens to return. `-1` means return full vocab. + pub top_logprobs: Option, /// Deprecated: Replaced by max_completion_tokens #[deprecated(note = "Use max_completion_tokens instead")] @@ -155,8 +154,8 @@ pub struct ChatCompletionRequest { /// Truncate prompt tokens to this length pub truncate_prompt_tokens: Option, - /// Number of prompt logprobs to return - pub prompt_logprobs: Option, + /// Number of prompt logprobs to return. `-1` means return full vocab. + pub prompt_logprobs: Option, /// Restrict output to these token IDs only pub allowed_token_ids: Option>, diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index 2e789573d21..e7b21e8adbc 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -1,6 +1,7 @@ use super::types::ChatCompletionRequest; use crate::error::{ApiError, bail_invalid_request}; use crate::routes::openai::utils::types::{ChatMessage, Tool, ToolChoice, ToolChoiceValue}; +use vllm_engine_core_client::protocol::LogprobsCount; /// Enforce the minimal compatibility contract for the Rust OpenAI server. pub(super) fn validate_request_compat( @@ -30,14 +31,12 @@ pub(super) fn validate_request_compat( } if let Some(prompt_logprobs) = request.prompt_logprobs { - if prompt_logprobs < 0 && prompt_logprobs != -1 { - bail_invalid_request!( - param = "prompt_logprobs", - "prompt_logprobs must be a non-negative value or -1." - ); - } - - if request.stream && (prompt_logprobs > 0 || prompt_logprobs == -1) { + if request.stream + && matches!( + prompt_logprobs, + LogprobsCount::All | LogprobsCount::Top(1..) + ) + { bail_invalid_request!( param = "prompt_logprobs", "prompt_logprobs are not available when stream=true." @@ -154,6 +153,7 @@ mod tests { use serde_json::json; use vllm_chat::ReasoningEffort; + use vllm_engine_core_client::protocol::LogprobsCount; use super::validate_request_compat; use crate::routes::openai::chat_completions::types::ChatCompletionRequest; @@ -299,7 +299,7 @@ mod tests { #[test] fn validate_request_compat_rejects_top_logprobs_without_logprobs() { let request = ChatCompletionRequest { - top_logprobs: Some(0), + top_logprobs: Some(LogprobsCount::Top(0)), ..base_request() }; assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()); @@ -308,26 +308,26 @@ mod tests { #[test] fn validate_request_compat_rejects_streaming_prompt_logprobs_requests() { let request = ChatCompletionRequest { - prompt_logprobs: Some(1), + prompt_logprobs: Some(LogprobsCount::Top(1)), ..base_request() }; assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()); let request = ChatCompletionRequest { - prompt_logprobs: Some(-1), + prompt_logprobs: Some(LogprobsCount::All), ..base_request() }; assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()); } #[test] - fn validate_request_compat_rejects_invalid_prompt_logprobs_value() { - let request = ChatCompletionRequest { - stream: false, - prompt_logprobs: Some(-2), - ..base_request() - }; - assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()); + fn chat_request_deserialization_rejects_invalid_prompt_logprobs_value() { + let result = serde_json::from_value::(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [{"role": "user", "content": "hello"}], + "prompt_logprobs": -2 + })); + assert!(result.is_err()); } #[test] diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 517bb93b0d8..0c95cf912f7 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -1,3 +1,4 @@ +use vllm_engine_core_client::protocol::LogprobsCount; use vllm_text::{SamplingParams, TextDecodeOptions, TextRequest}; use super::types::CompletionRequest; @@ -61,15 +62,7 @@ pub(super) fn prepare_completion_request( .map(|request| request.lora_name.clone()) .unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default()); - let logprobs = match request.logprobs { - Some(logprobs) => Some(i32::try_from(logprobs).map_err(|_| { - ApiError::invalid_request( - "`logprobs` must fit within a signed 32-bit integer.".to_string(), - Some("logprobs"), - ) - })?), - None => None, - }; + let logprobs = request.logprobs.map(LogprobsCount::Top); let prompt_only = request.echo && request.max_tokens == Some(0); let prompt_logprobs = request.prompt_logprobs.or(if request.echo && (!request.stream || prompt_only) { @@ -163,6 +156,7 @@ pub(super) fn prepare_completion_request( mod tests { use axum::http::HeaderMap; use serde_json::json; + use vllm_engine_core_client::protocol::LogprobsCount; use vllm_text::Prompt; use super::prepare_completion_request; @@ -247,7 +241,10 @@ mod tests { Prompt::TokenIds(vec![11, 22, 33]) ); assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7)); - assert_eq!(prepared.text_request.sampling_params.logprobs, Some(2)); + assert_eq!( + prepared.text_request.sampling_params.logprobs, + Some(LogprobsCount::Top(2)) + ); assert_eq!(prepared.text_request.sampling_params.top_p, Some(0.9)); assert_eq!(prepared.text_request.sampling_params.top_k, Some(42)); assert_eq!(prepared.text_request.sampling_params.min_p, Some(0.1)); @@ -410,10 +407,13 @@ mod tests { .expect("prepare"); assert!(prepared.options.prompt_only); - assert_eq!(prepared.text_request.sampling_params.logprobs, Some(3)); + assert_eq!( + prepared.text_request.sampling_params.logprobs, + Some(LogprobsCount::Top(3)) + ); assert_eq!( prepared.text_request.sampling_params.prompt_logprobs, - Some(3) + Some(LogprobsCount::Top(3)) ); } @@ -435,10 +435,13 @@ mod tests { ) .expect("prepare"); - assert_eq!(prepared.text_request.sampling_params.logprobs, Some(3)); + assert_eq!( + prepared.text_request.sampling_params.logprobs, + Some(LogprobsCount::Top(3)) + ); assert_eq!( prepared.text_request.sampling_params.prompt_logprobs, - Some(3) + Some(LogprobsCount::Top(3)) ); } @@ -479,10 +482,13 @@ mod tests { ResolvedRequestContext::default(), ) .expect("prepare"); - assert_eq!(prepared.text_request.sampling_params.logprobs, Some(1)); + assert_eq!( + prepared.text_request.sampling_params.logprobs, + Some(LogprobsCount::Top(1)) + ); assert_eq!( prepared.text_request.sampling_params.prompt_logprobs, - Some(2) + Some(LogprobsCount::Top(2)) ); } diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 93005cd8b83..1978c58c7f5 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; +use vllm_engine_core_client::protocol::LogprobsCount; use vllm_text::Prompt; use crate::routes::openai::utils::types::{ @@ -131,8 +132,8 @@ pub struct CompletionRequest { /// Restrict output to these token IDs only pub allowed_token_ids: Option>, - /// Number of prompt logprobs to return - pub prompt_logprobs: Option, + /// Number of prompt logprobs to return. `-1` means return full vocab. + pub prompt_logprobs: Option, // -------- Extra vLLM Parameters -------- /// Whether to add special tokens (e.g. BOS) to the prompt diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs index cbb040b90d0..ff14ab27352 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -1,3 +1,4 @@ +use vllm_engine_core_client::protocol::LogprobsCount; use vllm_text::Prompt; use super::types::CompletionRequest; @@ -44,29 +45,18 @@ pub(super) fn validate_request_compat( bail_invalid_request!(param = "suffix", "suffix is not supported."); } - if let Some(logprobs) = request.logprobs - && logprobs > i32::MAX as u32 - { - bail_invalid_request!( - param = "logprobs", - "`logprobs` must fit within a signed 32-bit integer." - ); - } - if let Some(prompt_logprobs) = request.prompt_logprobs { - if request.stream && (prompt_logprobs > 0 || prompt_logprobs == -1) { + if request.stream + && matches!( + prompt_logprobs, + LogprobsCount::All | LogprobsCount::Top(1..) + ) + { bail_invalid_request!( param = "prompt_logprobs", "`prompt_logprobs` are not available when `stream=true`." ); } - - if prompt_logprobs < 0 && prompt_logprobs != -1 { - bail_invalid_request!( - param = "prompt_logprobs", - "`prompt_logprobs` must be a non-negative value or -1." - ); - } } if request.use_beam_search { @@ -101,6 +91,7 @@ pub(super) fn validate_request_compat( #[cfg(test)] mod tests { use serde_json::json; + use vllm_engine_core_client::protocol::LogprobsCount; use super::validate_request_compat; use crate::routes::openai::completions::types::CompletionRequest; @@ -150,7 +141,7 @@ mod tests { #[test] fn validate_request_compat_rejects_streaming_prompt_logprobs() { let request = CompletionRequest { - prompt_logprobs: Some(1), + prompt_logprobs: Some(LogprobsCount::Top(1)), ..base_request() }; assert!( @@ -162,7 +153,7 @@ mod tests { fn validate_request_compat_accepts_non_stream_prompt_logprobs() { let request = CompletionRequest { stream: false, - prompt_logprobs: Some(-1), + prompt_logprobs: Some(LogprobsCount::All), ..base_request() }; assert!( diff --git a/rust/src/text/src/backend/mod.rs b/rust/src/text/src/backend/mod.rs index 8bc834aeae2..61dec0cbf1b 100644 --- a/rust/src/text/src/backend/mod.rs +++ b/rust/src/text/src/backend/mod.rs @@ -2,6 +2,7 @@ pub mod hf; use std::sync::Arc; +use vllm_engine_core_client::protocol::LogprobsCount; use vllm_tokenizer::DynTokenizer; use crate::error::Result; @@ -26,9 +27,7 @@ pub struct SamplingLimits { /// Runtime context window size reported by the engine startup handshake. pub max_model_len: u32, /// Maximum number of top log probabilities accepted by this frontend. - /// - /// `-1` means allowing requests up to the model vocabulary size. - pub max_logprobs: i32, + pub max_logprobs: LogprobsCount, /// Model vocabulary size from the model config, used to bound generated /// token IDs and logits-domain sampling controls. @@ -41,7 +40,7 @@ pub struct SamplingLimits { impl SamplingLimits { /// Original Python definition: /// - pub const DEFAULT_MAX_LOGPROBS: i32 = 20; + pub const DEFAULT_MAX_LOGPROBS: LogprobsCount = LogprobsCount::Top(20); /// Original Python definition: /// pub const MAX_LOGPROB_TOKEN_IDS: usize = 128; diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index a4fb86d19c5..d2b6db3109a 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -19,6 +19,7 @@ pub use output::{ pub use request::{Prompt, SamplingParams, TextRequest}; use trait_set::trait_set; use vllm_engine_core_client::EngineCoreClient; +use vllm_engine_core_client::protocol::LogprobsCount; pub use vllm_llm::FinishReason; use vllm_llm::{GenerateOutputStream, Llm}; use vllm_tokenizer::DynTokenizer; @@ -48,7 +49,7 @@ pub struct TextLlm { /// Runtime context window size reported by the engine startup handshake. max_model_len: u32, /// Maximum number of top log probabilities accepted by this text facade. - max_logprobs: i32, + max_logprobs: LogprobsCount, } impl TextLlm { @@ -68,7 +69,7 @@ impl TextLlm { } /// Override the maximum accepted logprobs count. - pub fn with_max_logprobs(mut self, max_logprobs: Option) -> Self { + pub fn with_max_logprobs(mut self, max_logprobs: Option) -> Self { if let Some(max_logprobs) = max_logprobs { self.max_logprobs = max_logprobs; } diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index c5f1e675687..e57ba7efca6 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -269,6 +269,7 @@ mod tests { use std::collections::{BTreeSet, HashMap}; use serial_test::file_serial; + use vllm_engine_core_client::protocol::LogprobsCount; use super::*; use crate::backend::hf::HfTextBackend; @@ -774,8 +775,8 @@ mod tests { #[test] fn lower_sampling_params_passes_logprobs_fields_through() { let sampling_params = SamplingParams { - logprobs: Some(3), - prompt_logprobs: Some(-1), + logprobs: Some(LogprobsCount::Top(3)), + prompt_logprobs: Some(LogprobsCount::All), ..Default::default() }; @@ -792,7 +793,7 @@ mod tests { default_max_tokens: None, }, SamplingLimits { - max_logprobs: -1, + max_logprobs: LogprobsCount::All, ..sample_sampling_limits() }, 3, @@ -800,15 +801,15 @@ mod tests { ) .unwrap(); - assert_eq!(params.logprobs, Some(3)); - assert_eq!(params.prompt_logprobs, Some(-1)); + assert_eq!(params.logprobs, Some(LogprobsCount::Top(3))); + assert_eq!(params.prompt_logprobs, Some(LogprobsCount::All)); } #[test] fn lower_sampling_params_rejects_full_vocab_logprobs_over_default_cap() { let error = lower_sampling_params_with_limits( SamplingParams { - logprobs: Some(-1), + logprobs: Some(LogprobsCount::All), ..Default::default() }, sample_sampling_limits(), @@ -829,24 +830,24 @@ mod tests { fn lower_sampling_params_expands_full_vocab_logprobs_from_model_vocab() { let params = lower_sampling_params_with_limits( SamplingParams { - logprobs: Some(-1), + logprobs: Some(LogprobsCount::All), ..Default::default() }, SamplingLimits { - max_logprobs: 1500, + max_logprobs: LogprobsCount::Top(1500), ..sample_sampling_limits() }, ) .unwrap(); - assert_eq!(params.logprobs, Some(-1)); + assert_eq!(params.logprobs, Some(LogprobsCount::All)); } #[test] fn lower_sampling_params_rejects_invalid_logprob_token_ids() { let error = lower_sampling_params_with_limits( SamplingParams { - logprobs: Some(1), + logprobs: Some(LogprobsCount::Top(1)), logprob_token_ids: Some(vec![1000]), ..Default::default() }, diff --git a/rust/src/text/src/lower/logprobs.rs b/rust/src/text/src/lower/logprobs.rs index 3c90f339107..803af8937b5 100644 --- a/rust/src/text/src/lower/logprobs.rs +++ b/rust/src/text/src/lower/logprobs.rs @@ -1,15 +1,14 @@ //! Python-compatible validation for logprobs sampling params. //! -//! `-1` is expanded only for bounds checks. The original request values are +//! `All` is expanded only for bounds checks. The original request values are //! passed through to engine-core. use crate::backend::SamplingLimits; use thiserror::Error; +use vllm_engine_core_client::protocol::LogprobsCount; #[derive(Debug, Error)] pub enum LogprobsError { - #[error("{parameter} must be non-negative or -1, got {value}")] - InvalidCount { parameter: &'static str, value: i32 }, #[error( "requested {parameter} of {requested}, which is greater than max allowed: {max_allowed}" )] @@ -30,19 +29,21 @@ pub enum LogprobsError { "when both logprobs and logprob_token_ids are set, logprobs must equal \ len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}." )] - TokenIdsMismatch { logprobs: i32, num_token_ids: usize }, + TokenIdsMismatch { + logprobs: LogprobsCount, + num_token_ids: usize, + }, } /// Validate logprobs count sampling parameters. pub(super) fn validate_logprobs( - logprobs: Option, - prompt_logprobs: Option, + logprobs: Option, + prompt_logprobs: Option, logprob_token_ids: Option<&[u32]>, sampling_limits: SamplingLimits, ) -> Result<(), LogprobsError> { let vocab_size = sampling_limits.model_vocab_size; - let max_logprobs = - normalize_logprobs_count(sampling_limits.max_logprobs, vocab_size, "max_logprobs")?; + let max_logprobs = sampling_limits.max_logprobs.expanded(vocab_size); validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?; validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?; @@ -50,7 +51,7 @@ pub(super) fn validate_logprobs( } fn validate_logprobs_count( - requested: Option, + requested: Option, max_logprobs: usize, vocab_size: usize, parameter: &'static str, @@ -59,7 +60,7 @@ fn validate_logprobs_count( return Ok(()); }; - let requested = normalize_logprobs_count(requested, vocab_size, parameter)?; + let requested = requested.expanded(vocab_size); if requested > max_logprobs { return Err(LogprobsError::TooManyCount { parameter, @@ -72,7 +73,7 @@ fn validate_logprobs_count( } pub(super) fn validate_logprob_token_ids( - logprobs: Option, + logprobs: Option, logprob_token_ids: Option<&[u32]>, ) -> Result<(), LogprobsError> { let Some(logprob_token_ids) = logprob_token_ids else { @@ -88,7 +89,7 @@ pub(super) fn validate_logprob_token_ids( } if let Some(logprobs) = logprobs - && logprobs != n as i32 + && logprobs != LogprobsCount::Top(n as u32) { return Err(LogprobsError::TokenIdsMismatch { logprobs, @@ -98,15 +99,3 @@ pub(super) fn validate_logprob_token_ids( Ok(()) } - -fn normalize_logprobs_count( - value: i32, - vocab_size: usize, - parameter: &'static str, -) -> Result { - match value { - -1 => Ok(vocab_size), - value if value < 0 => Err(LogprobsError::InvalidCount { parameter, value }), - value => Ok(value as usize), - } -} diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 682d85390d3..778ad347c72 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -3,9 +3,9 @@ use std::collections::HashMap; use enum_as_inner::EnumAsInner; use serde::{Deserialize, Serialize}; use serde_json::Value; -use vllm_engine_core_client::protocol::StructuredOutputsParams; use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::multimodal::MmFeatures; +use vllm_engine_core_client::protocol::{LogprobsCount, StructuredOutputsParams}; use crate::error::{Error, Result}; use crate::output::TextDecodeOptions; @@ -64,12 +64,12 @@ pub struct SamplingParams { pub thinking_token_budget: Option, /// Number of log probabilities to return per generated token. /// - /// `None` disables sample logprobs. `-1` requests the full vocabulary. - pub logprobs: Option, + /// `None` disables sample logprobs. + pub logprobs: Option, /// Number of log probabilities to return per prompt token. /// - /// `None` disables prompt logprobs. `-1` requests the full vocabulary. - pub prompt_logprobs: Option, + /// `None` disables prompt logprobs. + pub prompt_logprobs: Option, /// Minimum probability threshold for token sampling. `None` means no /// explicit user override. pub min_p: Option,