@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3f65e21e32
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Rust Frontend] Support `max_logprobs` validation (#45674)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 16, 2026 at 10:57 UTC 3f65e21e3200038e1f4524144b739ae207c8560d
15 files changed +484 -98
rust/src/cmd/src/cli.rs
+8
@@ -127,6 +127,11 @@ pub struct SharedRuntimeArgs {
127 /// `config.json`.
128 #[arg(long)]
129 pub max_model_len: Option<u32>,
130 + /// Maximum number of log probabilities to return when `logprobs` is
131 + /// specified in sampling parameters. `-1` means no cap.
132 + #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)]
133 + #[serde(default)]
134 + pub max_logprobs: Option<i32>,
135 /// TCP port for the gRPC Generate service. When not set, no gRPC server is
136 /// started.
137 #[arg(long)]
@@ -281,6 +286,7 @@ impl SharedRuntimeArgs {
286 chat_template: self.chat_template,
287 default_chat_template_kwargs: self.default_chat_template_kwargs,
288 chat_template_content_format: self.chat_template_content_format,
289 + max_logprobs: self.max_logprobs,
290 api_server_options,
291 api_keys: self.api_key,
292 disable_log_stats: self.disable_log_stats,
@@ -324,6 +330,7 @@ impl SharedRuntimeArgs {
330 chat_template: self.chat_template,
331 default_chat_template_kwargs: self.default_chat_template_kwargs,
332 chat_template_content_format: self.chat_template_content_format,
333 + max_logprobs: self.max_logprobs,
334 api_server_options,
335 api_keys: self.api_key,
336 disable_log_stats: self.disable_log_stats,
@@ -467,6 +474,7 @@ impl ServeArgs {
474 self.managed_engine.clone().into_config(
475 self.runtime.model.clone(),
476 self.runtime.max_model_len,
477 + self.runtime.max_logprobs,
478 self.runtime.language_model_only,
479 self.runtime.disable_log_stats,
480 self.runtime.shutdown_timeout,
rust/src/cmd/src/cli/tests.rs
+32 -1
@@ -38,6 +38,7 @@ fn serve_args_forward_python_flags_with_separator() {
38 max_model_len: Some(
39 512,
40 ),
41 + max_logprobs: None,
42 grpc_port: None,
43 shutdown_timeout: 0,
44 chat_template: None,
@@ -134,6 +135,29 @@ fn serve_args_forward_disable_log_stats_to_managed_engine() {
135 assert_eq!(config.python_args, vec!["--disable-log-stats"]);
136 }
137
138 +#[test]
139 +fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() {
140 + let cli = Cli::try_parse_from([
141 + "vllm-rs",
142 + "serve",
143 + "Qwen/Qwen3-0.6B",
144 + "--max-logprobs",
145 + "-1",
146 + ])
147 + .unwrap();
148 +
149 + let Command::Serve(args) = cli.command else {
150 + panic!("expected serve args");
151 + };
152 + assert_eq!(args.runtime.max_logprobs, Some(-1));
153 +
154 + let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
155 + assert_eq!(frontend_config.max_logprobs, Some(-1));
156 +
157 + let engine_config = args.to_managed_engine_config(5555);
158 + assert_eq!(engine_config.python_args, vec!["--max-logprobs", "-1"]);
159 +}
160 +
161 #[test]
162 fn serve_args_auto_forward_python_multi_char_alias_without_separator() {
163 let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "-tp", "2"]).unwrap();
@@ -388,6 +412,7 @@ fn frontend_args_accept_json() {
412 renderer: Auto,
413 language_model_only: false,
414 max_model_len: None,
415 + max_logprobs: None,
416 grpc_port: None,
417 shutdown_timeout: 0,
418 chat_template: None,
@@ -431,6 +456,7 @@ fn frontend_args_json_applies_defaults() {
456 assert_eq!(args.runtime.reasoning_parser, ParserSelection::Auto);
457 assert_eq!(args.runtime.renderer, RendererSelection::Auto);
458 assert_eq!(args.runtime.max_model_len, None);
459 + assert_eq!(args.runtime.max_logprobs, None);
460 assert_eq!(args.runtime.shutdown_timeout, 0);
461 }
462
@@ -446,7 +472,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
472 "--output-address",
473 "ipc:///tmp/output.sock",
474 "--args-json",
449 - r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_model_len":8192,"shutdown_timeout":3}"#,
475 + r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_model_len":8192,"max_logprobs":-1,"shutdown_timeout":3}"#,
476 ])
477 .unwrap();
478
@@ -465,6 +491,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
491 assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
492 assert!(args.runtime.language_model_only);
493 assert_eq!(args.runtime.max_model_len, Some(8192));
494 + assert_eq!(args.runtime.max_logprobs, Some(-1));
495 assert_eq!(args.runtime.shutdown_timeout, 3);
496 }
497
@@ -792,6 +819,7 @@ fn serve_args_accept_handshake_aliases() {
819 renderer: Auto,
820 language_model_only: false,
821 max_model_len: None,
822 + max_logprobs: None,
823 grpc_port: None,
824 shutdown_timeout: 0,
825 chat_template: None,
@@ -917,6 +945,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
945 chat_template: None,
946 default_chat_template_kwargs: None,
947 chat_template_content_format: Auto,
948 + max_logprobs: None,
949 api_server_options: ApiServerOptions {
950 enable_log_requests: false,
951 enable_prompt_tokens_details: false,
@@ -985,6 +1014,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
1014 chat_template: None,
1015 default_chat_template_kwargs: None,
1016 chat_template_content_format: Auto,
1017 + max_logprobs: None,
1018 api_server_options: ApiServerOptions {
1019 enable_log_requests: false,
1020 enable_prompt_tokens_details: false,
@@ -1068,6 +1098,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
1098 chat_template: None,
1099 default_chat_template_kwargs: None,
1100 chat_template_content_format: Auto,
1101 + max_logprobs: None,
1102 api_server_options: ApiServerOptions {
1103 enable_log_requests: false,
1104 enable_prompt_tokens_details: false,
rust/src/cmd/src/cli/unsupported.rs
-8
@@ -202,14 +202,6 @@ pub struct EngineUnsupportedArgs {
202 #[arg(long)]
203 pub tokenizer_revision: Option<Unsupported>,
204
205 - /// Maximum number of log probabilities to return when `logprobs` is
206 - /// specified in `SamplingParams`. The default value comes the default for
207 - /// the OpenAI Chat Completions API. -1 means no cap, i.e. all
208 - /// (output_length * vocab_size) logprobs are allowed to be returned and
209 - /// it may cause OOM.
210 - #[arg(long)]
211 - pub max_logprobs: Option<Unsupported>,
212 -
205 /// Skip initialization of tokenizer and detokenizer. Expects valid
206 /// `prompt_token_ids` and `None` for prompt from the input. The generated
207 /// output will contain token ids.
rust/src/managed-engine/src/cli.rs
+5
@@ -71,6 +71,7 @@ impl ManagedEngineArgs {
71 self,
72 model: String,
73 max_model_len: Option<u32>,
74 + max_logprobs: Option<i32>,
75 language_model_only: bool,
76 disable_log_stats: bool,
77 shutdown_timeout: u64,
@@ -82,6 +83,10 @@ impl ManagedEngineArgs {
83 python_args.push("--max-model-len".to_string());
84 python_args.push(max_model_len.to_string());
85 }
86 + if let Some(max_logprobs) = max_logprobs {
87 + python_args.push("--max-logprobs".to_string());
88 + python_args.push(max_logprobs.to_string());
89 + }
90 if language_model_only {
91 python_args.push("--language-model-only".to_string());
92 }
rust/src/server/examples/external_engine_openai_qwen.rs
+1
@@ -68,6 +68,7 @@ async fn main() -> Result<()> {
68 chat_template: None,
69 default_chat_template_kwargs: None,
70 chat_template_content_format: ChatTemplateContentFormatOption::Auto,
71 + max_logprobs: None,
72 api_server_options: ApiServerOptions::default(),
73 api_keys: Vec::new(),
74 disable_log_stats: false,
rust/src/server/src/config.rs
+12 -1
@@ -2,7 +2,7 @@ use std::collections::HashMap;
2 use std::fmt;
3 use std::time::Duration;
4
5 -use anyhow::Result;
5 +use anyhow::{Result, bail};
6 use educe::Educe;
7 use serde::Serialize;
8 use serde_json::Value;
@@ -77,6 +77,9 @@ pub struct Config {
77 pub default_chat_template_kwargs: Option<HashMap<String, Value>>,
78 /// How to serialize `message.content` for chat-template rendering.
79 pub chat_template_content_format: ChatTemplateContentFormatOption,
80 + /// Optional maximum number of top log probabilities accepted by the
81 + /// frontend. `None` delegates to the text layer default.
82 + pub max_logprobs: Option<i32>,
83 /// HTTP/API-server behavior switches.
84 pub api_server_options: ApiServerOptions,
85 /// API keys accepted as bearer tokens for guarded routes.
@@ -98,6 +101,14 @@ impl Config {
101 /// startup.
102 pub fn validate(&self) -> Result<()> {
103 vllm_chat::validate_parser_overrides(&self.tool_call_parser, &self.reasoning_parser)?;
104 + if let Some(max_logprobs) = self.max_logprobs
105 + && max_logprobs < -1
106 + {
107 + bail!(
108 + "max_logprobs must be non-negative or -1, got {}",
109 + max_logprobs
110 + );
111 + }
112
113 Ok(())
114 }
rust/src/server/src/error.rs
+34 -7
@@ -74,12 +74,11 @@ impl IntoResponse for ApiError {
74 }
75 }
76
77 -/// Classify a text-pipeline submit failure: tokenized-prompt validation
78 -/// failures (the prompt is too long for the model, or empty after
79 -/// tokenization) are the client's fault and map to HTTP 400, mirroring the
80 -/// Python frontend. Everything else stays an internal 500.
77 +/// Classify a text-pipeline submit failure: request validation failures are
78 +/// the client's fault and map to HTTP 400, mirroring the Python frontend.
79 +/// Everything else stays an internal 500.
80 pub fn text_submit_error(context: &'static str, error: vllm_text::Error) -> ApiError {
82 - if is_prompt_validation_error(&error) {
81 + if is_request_validation_error(&error) {
82 return invalid_request!("{error}");
83 }
84 server_error!("{}: {}", context, error.to_report_string())
@@ -90,18 +89,19 @@ pub fn text_submit_error(context: &'static str, error: vllm_text::Error) -> ApiE
89 pub fn chat_submit_error(context: &'static str, error: vllm_chat::Error) -> ApiError {
90 match &error {
91 vllm_chat::Error::PromptTooLong { .. } => invalid_request!("{error}"),
93 - vllm_chat::Error::Text(text_error) if is_prompt_validation_error(text_error) => {
92 + vllm_chat::Error::Text(text_error) if is_request_validation_error(text_error) => {
93 invalid_request!("{error}")
94 }
95 _ => server_error!("{}: {}", context, error.to_report_string()),
96 }
97 }
98
100 -fn is_prompt_validation_error(error: &vllm_text::Error) -> bool {
99 +fn is_request_validation_error(error: &vllm_text::Error) -> bool {
100 matches!(
101 error,
102 vllm_text::Error::PromptTooLong { .. }
103 | vllm_text::Error::EmptyPromptTokenIds { .. }
104 + | vllm_text::Error::Logprobs(_)
105 // An empty tokenized prompt detected later, at request prepare
106 // time, surfaces through the transparent Llm wrapper.
107 | vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })
@@ -145,6 +145,33 @@ mod tests {
145 assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
146 }
147
148 + #[test]
149 + fn logprobs_validation_maps_to_invalid_request() {
150 + let error = vllm_text::Error::Logprobs(vllm_text::LogprobsError::TooManyCount {
151 + parameter: "logprobs",
152 + requested: 1000,
153 + max_allowed: 20,
154 + });
155 + let api_error = text_submit_error("failed to submit completion request", error);
156 + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
157 + let response = api_error.to_error_response();
158 + assert_eq!(response.error.error_type, "invalid_request_error");
159 + assert!(response.error.message.contains("logprobs"));
160 + }
161 +
162 + #[test]
163 + fn chat_wrapped_logprobs_validation_maps_to_invalid_request() {
164 + let error = vllm_chat::Error::Text(vllm_text::Error::Logprobs(
165 + vllm_text::LogprobsError::TooManyCount {
166 + parameter: "prompt_logprobs",
167 + requested: 1000,
168 + max_allowed: 20,
169 + },
170 + ));
171 + let api_error = chat_submit_error("failed to submit chat request", error);
172 + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
173 + }
174 +
175 #[test]
176 fn other_submit_errors_stay_internal() {
177 let error = vllm_text::Error::Tokenizer("backend exploded".to_string());
rust/src/server/src/lib.rs
+1 -1
@@ -90,7 +90,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
90 .context("failed to connect to engine core")?;
91
92 let llm = Llm::new(client).with_log_stats(!config.disable_log_stats);
93 - let text = TextLlm::new(llm, text_backend);
93 + let text = TextLlm::new(llm, text_backend).with_max_logprobs(config.max_logprobs);
94
95 let chat = ChatLlm::new(text, chat_backend)
96 .with_tool_call_parser(config.tool_call_parser.clone())
rust/src/text/src/backend/hf/config.rs
+8 -6
@@ -245,10 +245,6 @@ impl ModelConfig {
245 pub(super) fn is_moe(&self) -> bool {
246 self.num_experts() > 0
247 }
248 -
249 - pub(super) fn max_position_embeddings(&self) -> Option<u32> {
250 - self.effective_text_config().max_position_embeddings
251 - }
248 }
249
250 /// Load the tokenizer-side EOS metadata if a config file is present.
@@ -356,7 +352,10 @@ mod tests {
352
353 assert_eq!(config.num_experts(), 8);
354 assert_eq!(config.model_type(), Some("top_level"));
359 - assert_eq!(config.max_position_embeddings(), Some(4096));
355 + assert_eq!(
356 + config.effective_text_config().max_position_embeddings,
357 + Some(4096)
358 + );
359 assert!(config.is_moe());
360 }
361
@@ -367,7 +366,10 @@ mod tests {
366
367 assert_eq!(config.num_experts(), 0);
368 assert!(!config.is_moe());
370 - assert_eq!(config.max_position_embeddings(), Some(4096));
369 + assert_eq!(
370 + config.effective_text_config().max_position_embeddings,
371 + Some(4096)
372 + );
373 }
374
375 #[test]
rust/src/text/src/backend/hf/mod.rs
-1
@@ -118,7 +118,6 @@ impl TextBackend for HfTextBackend {
118 default_min_p: self.generation_config.min_p,
119 default_repetition_penalty: self.generation_config.repetition_penalty,
120 default_max_tokens: self.generation_config.max_new_tokens,
121 - max_model_len: self.model_config.max_position_embeddings(),
121 })
122 }
123 }
rust/src/text/src/backend/mod.rs
+32 -5
@@ -6,8 +6,8 @@ use vllm_tokenizer::DynTokenizer;
6
7 use crate::error::Result;
8
9 -/// Tokenizer/model-derived hints used to enrich text-generation requests before
10 -/// they are lowered into engine-core.
9 +/// Tokenizer/model-derived defaults used to enrich text-generation requests
10 +/// before they are lowered into engine-core.
11 #[derive(Debug, Clone, Default, PartialEq)]
12 pub struct SamplingHints {
13 pub primary_eos_token_id: Option<u32>,
@@ -18,9 +18,36 @@ pub struct SamplingHints {
18 pub default_min_p: Option<f32>,
19 pub default_repetition_penalty: Option<f32>,
20 pub default_max_tokens: Option<u32>,
21 - /// Model context window size (`max_position_embeddings` from
22 - /// `config.json`).
23 - pub max_model_len: Option<u32>,
21 +}
22 +
23 +/// Effective bounds used to validate and lower sampling requests.
24 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25 +pub struct SamplingLimits {
26 + /// Runtime context window size reported by the engine startup handshake.
27 + pub max_model_len: u32,
28 + /// Maximum number of top log probabilities accepted by this frontend.
29 + ///
30 + /// `-1` means allowing requests up to the model vocabulary size.
31 + pub max_logprobs: i32,
32 + /// Model vocabulary size from the model config.
33 + pub model_vocab_size: Option<usize>,
34 + /// Tokenizer vocabulary size, used as a fallback when the model config does
35 + /// not expose a vocabulary size.
36 + pub tokenizer_vocab_size: usize,
37 +}
38 +
39 +impl SamplingLimits {
40 + /// Original Python definition:
41 + /// <https://github.com/vllm-project/vllm/blob/b5adb027ad03c29b46181752ba3b1cb84eff1dd4/vllm/config/model.py#L216-L220>
42 + pub const DEFAULT_MAX_LOGPROBS: i32 = 20;
43 + /// Original Python definition:
44 + /// <https://github.com/vllm-project/vllm/blob/b5adb027ad03c29b46181752ba3b1cb84eff1dd4/vllm/sampling_params.py#L30-L32>
45 + pub const MAX_LOGPROB_TOKEN_IDS: usize = 128;
46 +
47 + /// Return the vocabulary size used to expand `logprobs=-1`.
48 + pub fn logprobs_vocab_size(&self) -> usize {
49 + self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size)
50 + }
51 }
52
53 /// Minimal text-processing backend needed by `vllm-text`.
rust/src/text/src/error.rs
+4
@@ -2,6 +2,8 @@ use thiserror::Error;
2 use vllm_engine_core_client::Error as EngineCoreError;
3 use vllm_llm::Error as LlmError;
4
5 +pub use crate::lower::logprobs::LogprobsError;
6 +
7 #[derive(Debug, Error)]
8 pub enum Error {
9 #[error("tokenizer error: {0}")]
@@ -13,6 +15,8 @@ pub enum Error {
15 but the prompt contains {prompt_len} input tokens"
16 )]
17 PromptTooLong { max_model_len: u32, prompt_len: u32 },
18 + #[error(transparent)]
19 + Logprobs(#[from] LogprobsError),
20 #[error("text request stream `{request_id}` closed before terminal output")]
21 StreamClosedBeforeTerminalOutput { request_id: String },
22 #[error(transparent)]
rust/src/text/src/lib.rs
+28 -16
@@ -6,8 +6,8 @@
6
7 use std::mem::take;
8
9 -pub use backend::{DynTextBackend, SamplingHints, TextBackend};
10 -pub use error::{Error, Result};
9 +pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend};
10 +pub use error::{Error, LogprobsError, Result};
11 use futures::Stream;
12 pub use lower::{
13 PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens,
@@ -45,33 +45,33 @@ pub struct TextLlm {
45 /// Tokenizer/model metadata backend responsible for prompt encode/decode
46 /// and sampling hints.
47 backend: DynTextBackend,
48 - /// Context window size reported by the engine startup handshake, with
49 - /// optional override from config.
48 + /// Runtime context window size reported by the engine startup handshake.
49 max_model_len: u32,
50 + /// Maximum number of top log probabilities accepted by this text facade.
51 + max_logprobs: i32,
52 }
53
54 impl TextLlm {
55 /// Create a new text-generation facade from a shared LLM client plus a text
56 /// backend.
57 pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
57 - // Prefer the engine-reported max_model_len because it reflects the
58 - // post-profiling, auto-fitted KV cache limit rather than static
59 - // frontend metadata.
58 + // The engine-reported value reflects the post-profiling, auto-fitted
59 + // KV cache limit used at runtime.
60 let max_model_len = llm.engine_core_client().max_model_len();
61
62 Self {
63 llm,
64 backend,
65 max_model_len,
66 + max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
67 }
68 }
69
69 - /// Override the maximum model context length explicitly.
70 - ///
71 - /// This takes priority over both the engine-reported default and any
72 - /// tokenizer/model metadata exposed by the backend.
73 - pub fn with_max_model_len(mut self, max_model_len: u32) -> Self {
74 - self.max_model_len = max_model_len;
70 + /// Override the maximum accepted logprobs count.
71 + pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
72 + if let Some(max_logprobs) = max_logprobs {
73 + self.max_logprobs = max_logprobs;
74 + }
75 self
76 }
77
@@ -140,12 +140,24 @@ impl TextLlm {
140 Prompt::TokenIds(token_ids) => token_ids,
141 };
142
143 - let mut sampling_hints = self.backend.sampling_hints()?;
144 - sampling_hints.max_model_len = Some(self.max_model_len);
143 + let sampling_hints = self.backend.sampling_hints()?;
144 + let sampling_limits = SamplingLimits {
145 + max_model_len: self.max_model_len,
146 + max_logprobs: self.max_logprobs,
147 + model_vocab_size: self.backend.model_vocab_size(),
148 + tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
149 + };
150 +
151 let PreparedTextRequest {
152 text_request,
153 generate_request,
148 - } = lower_text_request(request, prompt_token_ids, sampling_hints, &*tokenizer)?;
154 + } = lower_text_request(
155 + request,
156 + prompt_token_ids,
157 + sampling_hints,
158 + sampling_limits,
159 + &*tokenizer,
160 + )?;
161
162 let raw_stream = self.llm.generate(generate_request).await?;
163 Ok((text_request, raw_stream))
rust/src/text/src/lower.rs
+185 -52
@@ -1,12 +1,15 @@
1 use std::collections::BTreeSet;
2
3 +pub(crate) mod logprobs;
4 +
5 use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
6 use vllm_llm::GenerateRequest;
7 use vllm_tokenizer::Tokenizer;
8
7 -use crate::backend::SamplingHints;
9 +use crate::backend::{SamplingHints, SamplingLimits};
10 use crate::error::{Error, Result};
11 use crate::request::{SamplingParams, TextRequest};
12 +use logprobs::validate_logprobs;
13
14 /// One text request after it has been lowered into the raw generate boundary.
15 #[derive(Debug)]
@@ -24,6 +27,7 @@ pub fn lower_text_request(
27 request: TextRequest,
28 prompt_token_ids: Vec<u32>,
29 sampling_hints: SamplingHints,
30 + sampling_limits: SamplingLimits,
31 tokenizer: &dyn Tokenizer,
32 ) -> Result<PreparedTextRequest> {
33 let prompt_len = prompt_token_ids.len() as u32;
@@ -34,6 +38,7 @@ pub fn lower_text_request(
38 sampling_params: lower_sampling_params(
39 request.sampling_params.clone(),
40 sampling_hints,
41 + sampling_limits,
42 prompt_len,
43 tokenizer,
44 )?,
@@ -66,8 +71,8 @@ pub fn lower_sampling_params(
71 default_min_p,
72 default_repetition_penalty,
73 default_max_tokens,
69 - max_model_len,
74 }: SamplingHints,
75 + sampling_limits: SamplingLimits,
76 prompt_len: u32,
77 tokenizer: &dyn Tokenizer,
78 ) -> Result<EngineCoreSamplingParams> {
@@ -95,6 +100,14 @@ pub fn lower_sampling_params(
100 vllm_xargs,
101 } = sampling_params;
102
103 + // Validate logprobs-related fields first with runtime sampling limits first.
104 + validate_logprobs(
105 + logprobs,
106 + prompt_logprobs,
107 + logprob_token_ids.as_deref(),
108 + sampling_limits,
109 + )?;
110 +
111 // Mirrors the model-generation-config inheritance used by vLLM's OpenAI chat
112 // path: https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/openai/chat_completion/protocol.py#L424-L450
113 // If neither the caller nor the model provides a value, fall back to 1.0 — the
@@ -105,7 +118,12 @@ pub fn lower_sampling_params(
118 let top_k = top_k.or(default_top_k).unwrap_or(0);
119 let min_p = min_p.or(default_min_p).unwrap_or(0.0);
120 let repetition_penalty = repetition_penalty.or(default_repetition_penalty).unwrap_or(1.0);
108 - let max_tokens = resolve_max_tokens(max_tokens, default_max_tokens, max_model_len, prompt_len)?;
121 + let max_tokens = resolve_max_tokens(
122 + max_tokens,
123 + default_max_tokens,
124 + sampling_limits.max_model_len,
125 + prompt_len,
126 + )?;
127 let min_tokens = min_tokens.unwrap_or(0);
128 let frequency_penalty = frequency_penalty.unwrap_or(0.0);
129 let presence_penalty = presence_penalty.unwrap_or(0.0);
@@ -189,33 +207,25 @@ fn tokenize_bad_words(
207 /// Resolve the effective `max_tokens` for generation, mirroring vLLM Python's
208 /// `get_max_tokens()` in `vllm/entrypoints/utils.py`.
209 ///
192 -/// Takes the minimum of all available limits (user-specified, generation-config
193 -/// default, and `max_model_len - prompt_len`). When nothing is known, falls
194 -/// back to `u32::MAX` so the engine-core can apply its own context-window
195 -/// limit.
210 +/// Takes the minimum of all available limits: user-specified, generation-config
211 +/// default, and `max_model_len - prompt_len`.
212 pub fn resolve_max_tokens(
213 user_max_tokens: Option<u32>,
214 default_max_tokens: Option<u32>,
199 - max_model_len: Option<u32>,
215 + max_model_len: u32,
216 prompt_len: u32,
217 ) -> Result<u32> {
202 - let model_max_tokens = match max_model_len {
203 - Some(max_model_len) if prompt_len >= max_model_len => {
204 - return Err(Error::PromptTooLong {
205 - max_model_len,
206 - prompt_len,
207 - });
208 - }
209 - Some(max_model_len) => Some(max_model_len - prompt_len),
210 - None => None,
218 + let model_max_tokens = if prompt_len >= max_model_len {
219 + return Err(Error::PromptTooLong {
220 + max_model_len,
221 + prompt_len,
222 + });
223 + } else {
224 + max_model_len - prompt_len
225 };
226
213 - let fallback_max_tokens = user_max_tokens.or(default_max_tokens);
214 - Ok([fallback_max_tokens, model_max_tokens]
215 - .into_iter()
216 - .flatten()
217 - .min()
218 - .unwrap_or(u32::MAX /* TODO: a reasonable fallback? */))
227 + let request_max_tokens = user_max_tokens.or(default_max_tokens);
228 + Ok(request_max_tokens.map_or(model_max_tokens, |n| n.min(model_max_tokens)))
229 }
230
231 fn merge_unique_token_ids(
@@ -240,6 +250,7 @@ mod tests {
250 use super::*;
251 use crate::backend::hf::HfTextBackend;
252 use crate::backend::{SamplingHints, TextBackend as _};
253 + use crate::error::LogprobsError;
254 use crate::request::{Prompt, TextRequest};
255
256 /// Stub tokenizer that returns empty token IDs — sufficient for tests that
@@ -290,16 +301,47 @@ mod tests {
301 default_min_p: None,
302 default_repetition_penalty: None,
303 default_max_tokens: None,
293 - max_model_len: None,
304 }
305 }
306
307 + fn sample_sampling_limits() -> SamplingLimits {
308 + SamplingLimits {
309 + max_model_len: 1_000_000,
310 + max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
311 + model_vocab_size: Some(1000),
312 + tokenizer_vocab_size: 2000,
313 + }
314 + }
315 +
316 + fn lower_sampling_params_with_limits(
317 + sampling_params: SamplingParams,
318 + sampling_limits: SamplingLimits,
319 + ) -> Result<EngineCoreSamplingParams> {
320 + lower_sampling_params(
321 + sampling_params,
322 + SamplingHints {
323 + primary_eos_token_id: None,
324 + extra_eos_token_ids: BTreeSet::new(),
325 + default_temperature: None,
326 + default_top_p: None,
327 + default_top_k: None,
328 + default_min_p: None,
329 + default_repetition_penalty: None,
330 + default_max_tokens: None,
331 + },
332 + sampling_limits,
333 + 3,
334 + &stub_tokenizer(),
335 + )
336 + }
337 +
338 #[test]
339 fn lower_text_request_applies_python_style_eos_hints() {
340 let prepared = lower_text_request(
341 sample_request(),
342 vec![1, 2, 3],
343 sample_sampling_hints(),
344 + sample_sampling_limits(),
345 &stub_tokenizer(),
346 )
347 .unwrap();
@@ -311,7 +353,7 @@ mod tests {
353 top_p: 1.0,
354 top_k: 0,
355 seed: None,
314 - max_tokens: 4294967295,
356 + max_tokens: 999997,
357 min_tokens: 0,
358 logprobs: None,
359 prompt_logprobs: None,
@@ -350,6 +392,7 @@ mod tests {
392 request,
393 vec![1, 2, 3],
394 sample_sampling_hints(),
395 + sample_sampling_limits(),
396 &stub_tokenizer(),
397 )
398 .unwrap();
@@ -361,7 +404,7 @@ mod tests {
404 top_p: 1.0,
405 top_k: 0,
406 seed: None,
364 - max_tokens: 4294967295,
407 + max_tokens: 999997,
408 min_tokens: 0,
409 logprobs: None,
410 prompt_logprobs: None,
@@ -415,16 +458,23 @@ mod tests {
458 default_min_p: None,
459 default_repetition_penalty: None,
460 default_max_tokens: None,
418 - max_model_len: Some(
419 - 40960,
420 - ),
461 }
462 "#]]
463 .assert_debug_eq(&hints);
464
425 - let prepared =
426 - lower_text_request(sample_request(), vec![1, 2, 3], hints, &stub_tokenizer())
427 - .expect("lower request");
465 + let prepared = lower_text_request(
466 + sample_request(),
467 + vec![1, 2, 3],
468 + hints,
469 + SamplingLimits {
470 + max_model_len: 40960,
471 + max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
472 + model_vocab_size: backend.model_vocab_size(),
473 + tokenizer_vocab_size: backend.tokenizer_vocab_size(),
474 + },
475 + &stub_tokenizer(),
476 + )
477 + .expect("lower request");
478 let params = prepared.generate_request.sampling_params;
479
480 expect_test::expect![[r#"
@@ -481,8 +531,8 @@ mod tests {
531 default_min_p: None,
532 default_repetition_penalty: None,
533 default_max_tokens: None,
484 - max_model_len: None,
534 },
535 + sample_sampling_limits(),
536 3,
537 &stub_tokenizer(),
538 )
@@ -494,7 +544,7 @@ mod tests {
544 top_p: 1.0,
545 top_k: 0,
546 seed: None,
497 - max_tokens: 4294967295,
547 + max_tokens: 999997,
548 min_tokens: 0,
549 logprobs: None,
550 prompt_logprobs: None,
@@ -550,8 +600,8 @@ mod tests {
600 default_min_p: Some(0.1),
601 default_repetition_penalty: Some(1.2),
602 default_max_tokens: Some(128),
553 - max_model_len: None,
603 },
604 + sample_sampling_limits(),
605 3,
606 &stub_tokenizer(),
607 )
@@ -605,7 +655,10 @@ mod tests {
655 default_min_p: None,
656 default_repetition_penalty: None,
657 default_max_tokens: None,
608 - max_model_len: None,
658 + },
659 + SamplingLimits {
660 + max_logprobs: -1,
661 + ..sample_sampling_limits()
662 },
663 3,
664 &stub_tokenizer(),
@@ -616,6 +669,91 @@ mod tests {
669 assert_eq!(params.prompt_logprobs, Some(-1));
670 }
671
672 + #[test]
673 + fn lower_sampling_params_rejects_full_vocab_logprobs_over_default_cap() {
674 + let error = lower_sampling_params_with_limits(
675 + SamplingParams {
676 + logprobs: Some(-1),
677 + ..Default::default()
678 + },
679 + sample_sampling_limits(),
680 + )
681 + .unwrap_err();
682 +
683 + assert!(matches!(
684 + error,
685 + Error::Logprobs(LogprobsError::TooManyCount {
686 + parameter: "logprobs",
687 + requested: 1000,
688 + max_allowed: 20,
689 + })
690 + ));
691 + }
692 +
693 + #[test]
694 + fn lower_sampling_params_expands_full_vocab_logprobs_from_model_vocab() {
695 + let params = lower_sampling_params_with_limits(
696 + SamplingParams {
697 + logprobs: Some(-1),
698 + ..Default::default()
699 + },
700 + SamplingLimits {
701 + max_logprobs: 1500,
702 + ..sample_sampling_limits()
703 + },
704 + )
705 + .unwrap();
706 +
707 + assert_eq!(params.logprobs, Some(-1));
708 + }
709 +
710 + #[test]
711 + fn lower_sampling_params_uses_tokenizer_vocab_when_model_vocab_is_unknown() {
712 + let error = lower_sampling_params_with_limits(
713 + SamplingParams {
714 + logprobs: Some(-1),
715 + ..Default::default()
716 + },
717 + SamplingLimits {
718 + max_logprobs: 1500,
719 + model_vocab_size: None,
720 + tokenizer_vocab_size: 2000,
721 + ..sample_sampling_limits()
722 + },
723 + )
724 + .unwrap_err();
725 +
726 + assert!(matches!(
727 + error,
728 + Error::Logprobs(LogprobsError::TooManyCount {
729 + parameter: "logprobs",
730 + requested: 2000,
731 + max_allowed: 1500,
732 + })
733 + ));
734 + }
735 +
736 + #[test]
737 + fn lower_sampling_params_rejects_invalid_logprob_token_ids() {
738 + let error = lower_sampling_params_with_limits(
739 + SamplingParams {
740 + logprobs: Some(1),
741 + logprob_token_ids: Some(vec![1000]),
742 + ..Default::default()
743 + },
744 + sample_sampling_limits(),
745 + )
746 + .unwrap_err();
747 +
748 + assert!(matches!(
749 + error,
750 + Error::Logprobs(LogprobsError::InvalidTokenIds {
751 + token_ids,
752 + vocab_size: 1000,
753 + }) if token_ids == vec![1000]
754 + ));
755 + }
756 +
757 #[test]
758 fn lower_sampling_params_uses_generation_defaults_when_user_omits_values() {
759 let params = lower_sampling_params(
@@ -629,8 +767,8 @@ mod tests {
767 default_min_p: Some(0.1),
768 default_repetition_penalty: Some(1.2),
769 default_max_tokens: Some(128),
632 - max_model_len: None,
770 },
771 + sample_sampling_limits(),
772 3,
773 &stub_tokenizer(),
774 )
@@ -667,7 +805,7 @@ mod tests {
805
806 #[test]
807 fn resolve_max_tokens_caps_by_model_len() {
670 - let result = resolve_max_tokens(Some(150), None, Some(200), 100);
808 + let result = resolve_max_tokens(Some(150), None, 200, 100);
809 assert_eq!(result.unwrap(), 100);
810 }
811
@@ -680,6 +818,7 @@ mod tests {
818 request,
819 vec![1, 2, 3],
820 sample_sampling_hints(),
821 + sample_sampling_limits(),
822 &stub_tokenizer(),
823 )
824 .unwrap();
@@ -690,37 +829,31 @@ mod tests {
829
830 #[test]
831 fn resolve_max_tokens_user_smaller_than_model_limit() {
693 - let result = resolve_max_tokens(Some(50), None, Some(200), 100);
832 + let result = resolve_max_tokens(Some(50), None, 200, 100);
833 assert_eq!(result.unwrap(), 50);
834 }
835
836 #[test]
837 fn resolve_max_tokens_uses_default_when_user_omits() {
699 - let result = resolve_max_tokens(None, Some(64), Some(200), 100);
838 + let result = resolve_max_tokens(None, Some(64), 200, 100);
839 assert_eq!(result.unwrap(), 64);
840 }
841
842 #[test]
843 fn resolve_max_tokens_default_capped_by_model_len() {
705 - let result = resolve_max_tokens(None, Some(256), Some(200), 100);
844 + let result = resolve_max_tokens(None, Some(256), 200, 100);
845 assert_eq!(result.unwrap(), 100);
846 }
847
848 #[test]
710 - fn resolve_max_tokens_no_model_len_falls_back() {
711 - let result = resolve_max_tokens(Some(9999), None, None, 100);
712 - assert_eq!(result.unwrap(), 9999);
713 - }
714 -
715 - #[test]
716 - fn resolve_max_tokens_no_limits_known_falls_back_to_u32_max() {
717 - let result = resolve_max_tokens(None, None, None, 100);
718 - assert_eq!(result.unwrap(), u32::MAX);
849 + fn resolve_max_tokens_uses_model_limit_when_user_omits() {
850 + let result = resolve_max_tokens(None, None, 200, 100);
851 + assert_eq!(result.unwrap(), 100);
852 }
853
854 #[test]
855 fn resolve_max_tokens_prompt_too_long() {
723 - let result = resolve_max_tokens(Some(10), None, Some(100), 100);
856 + let result = resolve_max_tokens(Some(10), None, 100, 100);
857 assert!(matches!(
858 result,
859 Err(Error::PromptTooLong {
@@ -732,7 +865,7 @@ mod tests {
865
866 #[test]
867 fn resolve_max_tokens_prompt_exceeds_model_len() {
735 - let result = resolve_max_tokens(Some(10), None, Some(100), 200);
868 + let result = resolve_max_tokens(Some(10), None, 100, 200);
869 assert!(matches!(
870 result,
871 Err(Error::PromptTooLong {
rust/src/text/src/lower/logprobs.rs new
+134
@@ -0,0 +1,134 @@
1 +//! Python-compatible validation for logprobs sampling params.
2 +//!
3 +//! `-1` is expanded only for bounds checks. The original request values are
4 +//! passed through to engine-core.
5 +
6 +use crate::backend::SamplingLimits;
7 +use thiserror::Error;
8 +
9 +#[derive(Debug, Error)]
10 +pub enum LogprobsError {
11 + #[error("{parameter} must be non-negative or -1, got {value}")]
12 + InvalidCount { parameter: &'static str, value: i32 },
13 + #[error(
14 + "requested {parameter} of {requested}, which is greater than max allowed: {max_allowed}"
15 + )]
16 + TooManyCount {
17 + parameter: &'static str,
18 + requested: usize,
19 + max_allowed: usize,
20 + },
21 + #[error(
22 + "requested logprob_token_ids of length {requested}, \
23 + which is greater than max allowed: {max_allowed}"
24 + )]
25 + TooManyTokenIds {
26 + requested: usize,
27 + max_allowed: usize,
28 + },
29 + #[error(
30 + "token_id(s) {token_ids:?} in logprob_token_ids contain out-of-vocab token ids. \
31 + Vocabulary size: {vocab_size}"
32 + )]
33 + InvalidTokenIds {
34 + token_ids: Vec<u32>,
35 + vocab_size: usize,
36 + },
37 + #[error(
38 + "when both logprobs and logprob_token_ids are set, logprobs must equal \
39 + len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}."
40 + )]
41 + TokenIdsMismatch { logprobs: i32, num_token_ids: usize },
42 +}
43 +
44 +/// Validate logprobs-related sampling parameters, returning an error if any
45 +/// parameter is out of bounds or if the combination of parameters is invalid.
46 +pub(super) fn validate_logprobs(
47 + logprobs: Option<i32>,
48 + prompt_logprobs: Option<i32>,
49 + logprob_token_ids: Option<&[u32]>,
50 + sampling_limits: SamplingLimits,
51 +) -> Result<(), LogprobsError> {
52 + let vocab_size = sampling_limits.logprobs_vocab_size();
53 + let max_logprobs =
54 + normalize_logprobs_count(sampling_limits.max_logprobs, vocab_size, "max_logprobs")?;
55 +
56 + validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?;
57 + validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?;
58 + validate_logprob_token_ids(logprobs, logprob_token_ids, vocab_size)
59 +}
60 +
61 +fn validate_logprobs_count(
62 + requested: Option<i32>,
63 + max_logprobs: usize,
64 + vocab_size: usize,
65 + parameter: &'static str,
66 +) -> Result<(), LogprobsError> {
67 + let Some(requested) = requested else {
68 + return Ok(());
69 + };
70 +
71 + let requested = normalize_logprobs_count(requested, vocab_size, parameter)?;
72 + if requested > max_logprobs {
73 + return Err(LogprobsError::TooManyCount {
74 + parameter,
75 + requested,
76 + max_allowed: max_logprobs,
77 + });
78 + }
79 +
80 + Ok(())
81 +}
82 +
83 +fn validate_logprob_token_ids(
84 + logprobs: Option<i32>,
85 + logprob_token_ids: Option<&[u32]>,
86 + vocab_size: usize,
87 +) -> Result<(), LogprobsError> {
88 + let Some(logprob_token_ids) = logprob_token_ids else {
89 + return Ok(());
90 + };
91 +
92 + let n = logprob_token_ids.len();
93 + if n > SamplingLimits::MAX_LOGPROB_TOKEN_IDS {
94 + return Err(LogprobsError::TooManyTokenIds {
95 + requested: n,
96 + max_allowed: SamplingLimits::MAX_LOGPROB_TOKEN_IDS,
97 + });
98 + }
99 +
100 + let invalid_token_ids: Vec<_> = logprob_token_ids
101 + .iter()
102 + .copied()
103 + .filter(|&token_id| token_id as usize >= vocab_size)
104 + .collect();
105 + if !invalid_token_ids.is_empty() {
106 + return Err(LogprobsError::InvalidTokenIds {
107 + token_ids: invalid_token_ids,
108 + vocab_size,
109 + });
110 + }
111 +
112 + if let Some(logprobs) = logprobs
113 + && logprobs != n as i32
114 + {
115 + return Err(LogprobsError::TokenIdsMismatch {
116 + logprobs,
117 + num_token_ids: n,
118 + });
119 + }
120 +
121 + Ok(())
122 +}
123 +
124 +fn normalize_logprobs_count(
125 + value: i32,
126 + vocab_size: usize,
127 + parameter: &'static str,
128 +) -> Result<usize, LogprobsError> {
129 + match value {
130 + -1 => Ok(vocab_size),
131 + value if value < 0 => Err(LogprobsError::InvalidCount { parameter, value }),
132 + value => Ok(value as usize),
133 + }
134 +}