From f99260d2aa43d779fe4fc9d69bd57ad4353a0f3f Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 16 Jun 2026 11:37:58 +0800 Subject: [PATCH] [Rust Frontend] Lower out-of-vocab validation to `text` layer (#45685) Signed-off-by: Bugen Zhao --- rust/src/server/src/error.rs | 12 ++ .../src/routes/openai/chat_completions.rs | 7 - .../openai/chat_completions/validate.rs | 44 +----- .../server/src/routes/openai/completions.rs | 7 - .../src/routes/openai/completions/validate.rs | 55 +------ .../src/server/src/routes/openai/utils/mod.rs | 1 - .../src/routes/openai/utils/token_ids.rs | 55 ------- rust/src/server/src/state.rs | 10 -- rust/src/text/src/backend/mod.rs | 18 ++- rust/src/text/src/error.rs | 3 + rust/src/text/src/lib.rs | 2 +- rust/src/text/src/lower.rs | 147 +++++++++++++++++- rust/src/text/src/lower/logprobs.rs | 28 +--- rust/src/text/src/lower/token_ids.rs | 101 ++++++++++++ 14 files changed, 278 insertions(+), 212 deletions(-) delete mode 100644 rust/src/server/src/routes/openai/utils/token_ids.rs create mode 100644 rust/src/text/src/lower/token_ids.rs diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index 2096f4876d..e5a5c1a40d 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -102,6 +102,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool { vllm_text::Error::PromptTooLong { .. } | vllm_text::Error::EmptyPromptTokenIds { .. } | vllm_text::Error::Logprobs(_) + | vllm_text::Error::OutOfVocab(_) // An empty tokenized prompt detected later, at request prepare // time, surfaces through the transparent Llm wrapper. | vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. }) @@ -172,6 +173,17 @@ mod tests { assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); } + #[test] + fn out_of_vocab_validation_maps_to_invalid_request() { + let error = vllm_text::Error::OutOfVocab(vllm_text::OutOfVocabError { + parameter: "logprob_token_ids", + token_ids: vec![1000], + vocab_size: 1000, + }); + let api_error = text_submit_error("failed to submit completion request", error); + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); + } + #[test] fn other_submit_errors_stay_internal() { let error = vllm_text::Error::Tokenizer("backend exploded".to_string()); diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index a8c70d273d..60cd14f9a8 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -53,13 +53,6 @@ pub async fn chat_completions( let request_context = resolve_request_context(&headers, body.request_id.as_deref()); let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await; - if let Err(err) = validate::validate_token_id_ranges( - &body, - state.tokenizer_vocab_size(), - state.model_vocab_size(), - ) { - return err.into_response(); - } let prepared = match prepare_chat_request(body, &lora_resolution, request_context) { Ok(prepared) => prepared, Err(error) => return error.into_response(), 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 b83d9035a0..bbf32c6950 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,5 @@ use super::types::ChatCompletionRequest; use crate::error::{ApiError, bail_invalid_request}; -use crate::routes::openai::utils::token_ids::{validate_allowed_token_ids, validate_logit_bias}; use crate::routes::openai::utils::types::{ChatMessage, Tool, ToolChoice, ToolChoiceValue}; /// Enforce the minimal compatibility contract for the Rust OpenAI server. @@ -154,21 +153,6 @@ fn validate_function_tools(tools: &[Tool], param: &'static str) -> Result<(), Ap Ok(()) } -/// Reject out-of-vocab token ids, mirroring the Python input processor: -/// `allowed_token_ids` against the tokenizer vocab, `logit_bias` keys against the -/// model vocab (skipped when the model size is unknown). -pub(super) fn validate_token_id_ranges( - request: &ChatCompletionRequest, - tokenizer_vocab_size: usize, - model_vocab_size: Option, -) -> Result<(), ApiError> { - validate_allowed_token_ids(request.allowed_token_ids.as_deref(), tokenizer_vocab_size)?; - validate_logit_bias( - request.logit_bias.as_ref(), - model_vocab_size.unwrap_or(usize::MAX), - ) -} - #[cfg(test)] mod tests { use std::collections::HashMap; @@ -176,7 +160,7 @@ mod tests { use serde_json::json; use vllm_chat::ReasoningEffort; - use super::{validate_request_compat, validate_token_id_ranges}; + use super::validate_request_compat; use crate::routes::openai::chat_completions::types::ChatCompletionRequest; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ @@ -188,32 +172,6 @@ mod tests { names.iter().map(|s| s.to_string()).collect() } - #[test] - fn validate_token_id_ranges_rejects_oob_and_accepts_in_vocab() { - // allowed_token_ids are bounded by the tokenizer vocab - let mut request = base_request(); - request.allowed_token_ids = Some(vec![5, 1_000_000]); - assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err()); - // logit_bias is bounded by the larger model vocab: an id between the two - // vocabs is valid and must not be rejected (the parity regression we fix) - let mut request = base_request(); - request.logit_bias = Some(HashMap::from([("150".to_string(), 1.0)])); - assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok()); - // logit_bias beyond the model vocab -> reject - let mut request = base_request(); - request.logit_bias = Some(HashMap::from([("1000000".to_string(), 1.0)])); - assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err()); - // all in-vocab -> accept - let mut request = base_request(); - request.allowed_token_ids = Some(vec![5, 50]); - request.logit_bias = Some(HashMap::from([("50".to_string(), 1.0)])); - assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok()); - // unknown sizes -> skip - let mut request = base_request(); - request.allowed_token_ids = Some(vec![1_000_000]); - assert!(validate_token_id_ranges(&request, usize::MAX, None).is_ok()); - } - fn base_request() -> ChatCompletionRequest { ChatCompletionRequest { model: "Qwen/Qwen1.5-0.5B-Chat".to_string(), diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 3dc3bbff6f..de21dc3a1c 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -47,13 +47,6 @@ pub async fn completions( let request_context = resolve_request_context(&headers, body.request_id.as_deref()); let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await; - if let Err(err) = validate::validate_token_id_ranges( - &body, - state.tokenizer_vocab_size(), - state.model_vocab_size(), - ) { - return err.into_response(); - } let prepared = match prepare_completion_request(body, &lora_resolution, request_context) { Ok(prepared) => prepared, Err(error) => return error.into_response(), diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs index 2af41877bf..f19defe5e4 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -2,9 +2,6 @@ use vllm_text::Prompt; use super::types::CompletionRequest; use crate::error::{ApiError, bail_invalid_request}; -use crate::routes::openai::utils::token_ids::{ - validate_allowed_token_ids, validate_logit_bias, validate_prompt_token_ids, -}; /// Enforce the minimal compatibility contract for the Rust OpenAI server. pub(super) fn validate_request_compat( @@ -98,63 +95,13 @@ pub(super) fn validate_request_compat( Ok(()) } -/// Reject out-of-vocab token ids, mirroring the Python input processor. A token-id -/// prompt may reference ids the engine embeds beyond either vocab alone (Qwen3 -/// extra LM tokens, multimodal placeholders), so it is bounded by the union of the -/// tokenizer and model vocabularies; `allowed_token_ids` by the tokenizer vocab; -/// `logit_bias` keys by the model vocab (skipped when the model size is unknown). -pub(super) fn validate_token_id_ranges( - request: &CompletionRequest, - tokenizer_vocab_size: usize, - model_vocab_size: Option, -) -> Result<(), ApiError> { - let prompt_bound = tokenizer_vocab_size.max(model_vocab_size.unwrap_or(0)); - validate_prompt_token_ids(&request.prompt, prompt_bound)?; - validate_allowed_token_ids(request.allowed_token_ids.as_deref(), tokenizer_vocab_size)?; - validate_logit_bias( - request.logit_bias.as_ref(), - model_vocab_size.unwrap_or(usize::MAX), - ) -} - #[cfg(test)] mod tests { use serde_json::json; - use vllm_text::Prompt; - use super::{validate_request_compat, validate_token_id_ranges}; + use super::validate_request_compat; use crate::routes::openai::completions::types::CompletionRequest; - #[test] - fn validate_token_id_ranges_rejects_oob_prompt_and_params() { - // a token-id prompt below both vocabs is accepted (the engine can embed it) - let mut request = base_request(); - request.prompt = Prompt::TokenIds(vec![5, 150]); - assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok()); - // an id at or above the union of the two vocabs is rejected - let mut request = base_request(); - request.prompt = Prompt::TokenIds(vec![5, 200]); - assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err()); - // an id beyond the model vocab but within the (larger) tokenizer vocab is - // accepted: the engine embeds added/placeholder ids above the model vocab, - // matching the Python input processor's max(tokenizer, model) bound - let mut request = base_request(); - request.prompt = Prompt::TokenIds(vec![150]); - assert!(validate_token_id_ranges(&request, 200, Some(100)).is_ok()); - // falls back to the tokenizer vocab when the model size is unknown - let mut request = base_request(); - request.prompt = Prompt::TokenIds(vec![150]); - assert!(validate_token_id_ranges(&request, 100, None).is_err()); - // allowed_token_ids are bounded by the tokenizer vocab -> reject - let mut request = base_request(); - request.allowed_token_ids = Some(vec![150]); - assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err()); - // unknown sizes -> skip - let mut request = base_request(); - request.prompt = Prompt::TokenIds(vec![1_000_000]); - assert!(validate_token_id_ranges(&request, usize::MAX, None).is_ok()); - } - fn base_request() -> CompletionRequest { serde_json::from_value(json!({ "model": "Qwen/Qwen1.5-0.5B-Chat", diff --git a/rust/src/server/src/routes/openai/utils/mod.rs b/rust/src/server/src/routes/openai/utils/mod.rs index 7ec1251ddf..70e9d1466d 100644 --- a/rust/src/server/src/routes/openai/utils/mod.rs +++ b/rust/src/server/src/routes/openai/utils/mod.rs @@ -1,6 +1,5 @@ pub mod logprobs; pub mod structured_outputs; -pub mod token_ids; pub mod types; pub mod usage; pub mod validated_json; diff --git a/rust/src/server/src/routes/openai/utils/token_ids.rs b/rust/src/server/src/routes/openai/utils/token_ids.rs deleted file mode 100644 index ffa945ef94..0000000000 --- a/rust/src/server/src/routes/openai/utils/token_ids.rs +++ /dev/null @@ -1,55 +0,0 @@ -use std::collections::HashMap; - -use vllm_text::Prompt; - -use crate::error::{ApiError, bail_invalid_request}; - -/// Reject token-id prompt entries at or above `bound` (the highest in-vocab id is -/// `bound - 1`). -pub(crate) fn validate_prompt_token_ids(prompt: &Prompt, bound: usize) -> Result<(), ApiError> { - if let Prompt::TokenIds(ids) = prompt - && let Some(&bad) = ids.iter().find(|&&id| id as usize >= bound) - { - bail_invalid_request!( - param = "prompt", - "prompt contains out-of-vocab token id {bad}; vocabulary size is {bound}." - ); - } - Ok(()) -} - -/// Reject `allowed_token_ids` entries at or above `bound`. -pub(crate) fn validate_allowed_token_ids( - allowed_token_ids: Option<&[u32]>, - bound: usize, -) -> Result<(), ApiError> { - if let Some(ids) = allowed_token_ids - && let Some(&bad) = ids.iter().find(|&&id| id as usize >= bound) - { - bail_invalid_request!( - param = "allowed_token_ids", - "allowed_token_ids contains out-of-vocab token id {bad}; vocabulary size is {bound}." - ); - } - Ok(()) -} - -/// Reject `logit_bias` keys at or above `bound`. -pub(crate) fn validate_logit_bias( - logit_bias: Option<&HashMap>, - bound: usize, -) -> Result<(), ApiError> { - if let Some(bias) = logit_bias { - for key in bias.keys() { - if let Ok(id) = key.parse::() - && id as usize >= bound - { - bail_invalid_request!( - param = "logit_bias", - "logit_bias contains out-of-vocab token id {id}; vocabulary size is {bound}." - ); - } - } - } - Ok(()) -} diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index 55959b60d9..2fee91d457 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -114,16 +114,6 @@ impl AppState { &self.served_model_names } - /// Tokenizer vocabulary size. - pub fn tokenizer_vocab_size(&self) -> usize { - self.chat.tokenizer_vocab_size() - } - - /// Model vocabulary size, else `None`. - pub fn model_vocab_size(&self) -> Option { - self.chat.model_vocab_size() - } - /// Return base served model names plus dynamically loaded LoRA adapter /// names. pub async fn served_model_names_with_loras(&self) -> Vec { diff --git a/rust/src/text/src/backend/mod.rs b/rust/src/text/src/backend/mod.rs index 2b11d81d96..06be187414 100644 --- a/rust/src/text/src/backend/mod.rs +++ b/rust/src/text/src/backend/mod.rs @@ -29,10 +29,12 @@ pub struct SamplingLimits { /// /// `-1` means allowing requests up to the model vocabulary size. pub max_logprobs: i32, - /// Model vocabulary size from the model config. + + /// Model vocabulary size from the model config, used to bound + /// `logit_bias` keys when available. pub model_vocab_size: Option, - /// Tokenizer vocabulary size, used as a fallback when the model config does - /// not expose a vocabulary size. + /// Tokenizer vocabulary size, used to bound `allowed_token_ids` and + /// token-ID prompts. pub tokenizer_vocab_size: usize, } @@ -48,6 +50,16 @@ impl SamplingLimits { pub fn logprobs_vocab_size(&self) -> usize { self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size) } + + /// Return the vocabulary size used to validate generated stop token IDs. + pub fn stop_token_vocab_size(&self) -> usize { + self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size) + } + + /// Return the union bound used to validate token-ID prompts. + pub fn prompt_token_vocab_size(&self) -> usize { + self.tokenizer_vocab_size.max(self.model_vocab_size.unwrap_or(0)) + } } /// Minimal text-processing backend needed by `vllm-text`. diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index a6b2fe5af1..f686e56d52 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -3,6 +3,7 @@ use vllm_engine_core_client::Error as EngineCoreError; use vllm_llm::Error as LlmError; pub use crate::lower::logprobs::LogprobsError; +pub use crate::lower::token_ids::OutOfVocabError; #[derive(Debug, Error)] pub enum Error { @@ -17,6 +18,8 @@ pub enum Error { PromptTooLong { max_model_len: u32, prompt_len: u32 }, #[error(transparent)] Logprobs(#[from] LogprobsError), + #[error(transparent)] + OutOfVocab(#[from] OutOfVocabError), #[error("text request stream `{request_id}` closed before terminal output")] StreamClosedBeforeTerminalOutput { request_id: String }, #[error(transparent)] diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index b23e33135c..4085f90478 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -7,7 +7,7 @@ use std::mem::take; pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend}; -pub use error::{Error, LogprobsError, Result}; +pub use error::{Error, LogprobsError, OutOfVocabError, Result}; use futures::Stream; pub use lower::{ PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index d482f8cbf1..082809fe5c 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -1,6 +1,7 @@ use std::collections::BTreeSet; pub(crate) mod logprobs; +pub(crate) mod token_ids; use vllm_engine_core_client::protocol::EngineCoreSamplingParams; use vllm_llm::GenerateRequest; @@ -10,6 +11,7 @@ use crate::backend::{SamplingHints, SamplingLimits}; use crate::error::{Error, Result}; use crate::request::{SamplingParams, TextRequest}; use logprobs::validate_logprobs; +use token_ids::{validate_prompt_token_ids, validate_vocab_range}; /// One text request after it has been lowered into the raw generate boundary. #[derive(Debug)] @@ -31,6 +33,8 @@ pub fn lower_text_request( tokenizer: &dyn Tokenizer, ) -> Result { let prompt_len = prompt_token_ids.len() as u32; + validate_prompt_token_ids(&prompt_token_ids, &sampling_limits)?; + let generate_request = GenerateRequest { request_id: request.request_id.clone(), prompt_token_ids, @@ -100,7 +104,6 @@ pub fn lower_sampling_params( vllm_xargs, } = sampling_params; - // Validate logprobs-related fields first with runtime sampling limits first. validate_logprobs( logprobs, prompt_logprobs, @@ -139,7 +142,7 @@ pub fn lower_sampling_params( merge_unique_token_ids(&mut stop_token_ids, extra_eos_token_ids.iter().copied()); } - Ok(EngineCoreSamplingParams { + let params = EngineCoreSamplingParams { temperature, top_p, top_k, @@ -162,7 +165,9 @@ pub fn lower_sampling_params( logprob_token_ids, skip_reading_prefix_cache, extra_args: vllm_xargs, - }) + }; + validate_vocab_range(¶ms, &sampling_limits)?; + Ok(params) } /// Convert bad-word strings into token-ID sequences, following the Python vLLM @@ -243,14 +248,14 @@ fn merge_unique_token_ids( #[cfg(test)] mod tests { - use std::collections::BTreeSet; + use std::collections::{BTreeSet, HashMap}; use serial_test::file_serial; use super::*; use crate::backend::hf::HfTextBackend; use crate::backend::{SamplingHints, TextBackend as _}; - use crate::error::LogprobsError; + use crate::error::{LogprobsError, OutOfVocabError}; use crate::request::{Prompt, TextRequest}; /// Stub tokenizer that returns empty token IDs — sufficient for tests that @@ -430,6 +435,57 @@ mod tests { .assert_debug_eq(¶ms); } + #[test] + fn lower_text_request_uses_union_vocab_for_prompt_token_ids() { + lower_text_request( + sample_request(), + vec![1500], + sample_sampling_hints(), + SamplingLimits { + model_vocab_size: Some(2000), + tokenizer_vocab_size: 1000, + ..sample_sampling_limits() + }, + &stub_tokenizer(), + ) + .expect("model vocab extends prompt token range"); + + lower_text_request( + sample_request(), + vec![1500], + sample_sampling_hints(), + SamplingLimits { + model_vocab_size: Some(1000), + tokenizer_vocab_size: 2000, + ..sample_sampling_limits() + }, + &stub_tokenizer(), + ) + .expect("tokenizer vocab extends prompt token range"); + + let error = lower_text_request( + sample_request(), + vec![2000], + sample_sampling_hints(), + SamplingLimits { + model_vocab_size: Some(1000), + tokenizer_vocab_size: 2000, + ..sample_sampling_limits() + }, + &stub_tokenizer(), + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::OutOfVocab(OutOfVocabError { + parameter: "prompt", + token_ids, + vocab_size: 2000, + }) if token_ids == vec![2000] + )); + } + #[tokio::test] #[file_serial(hf_qwen3)] async fn lower_text_request_uses_real_qwen_generation_defaults() { @@ -747,13 +803,92 @@ mod tests { assert!(matches!( error, - Error::Logprobs(LogprobsError::InvalidTokenIds { + Error::OutOfVocab(OutOfVocabError { + parameter: "logprob_token_ids", token_ids, vocab_size: 1000, }) if token_ids == vec![1000] )); } + #[test] + fn lower_sampling_params_rejects_out_of_vocab_stop_token_ids() { + let error = lower_sampling_params_with_limits( + SamplingParams { + stop_token_ids: Some(vec![999, 1000]), + ..Default::default() + }, + sample_sampling_limits(), + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::OutOfVocab(OutOfVocabError { + parameter: "stop_token_ids", + token_ids, + vocab_size: 1000, + }) if token_ids == vec![1000] + )); + } + + #[test] + fn lower_sampling_params_rejects_out_of_vocab_allowed_token_ids() { + let error = lower_sampling_params_with_limits( + SamplingParams { + allowed_token_ids: Some(vec![1999, 2000]), + ..Default::default() + }, + sample_sampling_limits(), + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::OutOfVocab(OutOfVocabError { + parameter: "allowed_token_ids", + token_ids, + vocab_size: 2000, + }) if token_ids == vec![2000] + )); + } + + #[test] + fn lower_sampling_params_rejects_out_of_vocab_logit_bias() { + let error = lower_sampling_params_with_limits( + SamplingParams { + logit_bias: Some(HashMap::from([(1000, 1.0)])), + ..Default::default() + }, + sample_sampling_limits(), + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::OutOfVocab(OutOfVocabError { + parameter: "logit_bias", + token_ids, + vocab_size: 1000, + }) if token_ids == vec![1000] + )); + } + + #[test] + fn lower_sampling_params_skips_logit_bias_range_when_model_vocab_is_unknown() { + lower_sampling_params_with_limits( + SamplingParams { + logit_bias: Some(HashMap::from([(1_000_000, 1.0)])), + ..Default::default() + }, + SamplingLimits { + model_vocab_size: None, + ..sample_sampling_limits() + }, + ) + .expect("logit_bias range check is skipped without model vocab size"); + } + #[test] fn lower_sampling_params_uses_generation_defaults_when_user_omits_values() { let params = lower_sampling_params( diff --git a/rust/src/text/src/lower/logprobs.rs b/rust/src/text/src/lower/logprobs.rs index 0372b20c9b..087f4dce2d 100644 --- a/rust/src/text/src/lower/logprobs.rs +++ b/rust/src/text/src/lower/logprobs.rs @@ -26,14 +26,6 @@ pub enum LogprobsError { requested: usize, max_allowed: usize, }, - #[error( - "token_id(s) {token_ids:?} in logprob_token_ids contain out-of-vocab token ids. \ - Vocabulary size: {vocab_size}" - )] - InvalidTokenIds { - token_ids: Vec, - vocab_size: usize, - }, #[error( "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}." @@ -41,8 +33,7 @@ pub enum LogprobsError { TokenIdsMismatch { logprobs: i32, num_token_ids: usize }, } -/// Validate logprobs-related sampling parameters, returning an error if any -/// parameter is out of bounds or if the combination of parameters is invalid. +/// Validate logprobs count sampling parameters. pub(super) fn validate_logprobs( logprobs: Option, prompt_logprobs: Option, @@ -55,7 +46,7 @@ pub(super) fn validate_logprobs( validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?; validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?; - validate_logprob_token_ids(logprobs, logprob_token_ids, vocab_size) + validate_logprob_token_ids(logprobs, logprob_token_ids) } fn validate_logprobs_count( @@ -80,10 +71,9 @@ fn validate_logprobs_count( Ok(()) } -fn validate_logprob_token_ids( +pub(super) fn validate_logprob_token_ids( logprobs: Option, logprob_token_ids: Option<&[u32]>, - vocab_size: usize, ) -> Result<(), LogprobsError> { let Some(logprob_token_ids) = logprob_token_ids else { return Ok(()); @@ -97,18 +87,6 @@ fn validate_logprob_token_ids( }); } - let invalid_token_ids: Vec<_> = logprob_token_ids - .iter() - .copied() - .filter(|&token_id| token_id as usize >= vocab_size) - .collect(); - if !invalid_token_ids.is_empty() { - return Err(LogprobsError::InvalidTokenIds { - token_ids: invalid_token_ids, - vocab_size, - }); - } - if let Some(logprobs) = logprobs && logprobs != n as i32 { diff --git a/rust/src/text/src/lower/token_ids.rs b/rust/src/text/src/lower/token_ids.rs new file mode 100644 index 0000000000..0e8dc8ff87 --- /dev/null +++ b/rust/src/text/src/lower/token_ids.rs @@ -0,0 +1,101 @@ +use std::result::Result; + +use thiserror::Error; +use vllm_engine_core_client::protocol::EngineCoreSamplingParams; + +use crate::SamplingLimits; + +#[derive(Debug, Error)] +#[error( + "token_id(s) {token_ids:?} in {parameter} contain out-of-vocab token ids. \ + Vocabulary size: {vocab_size}" +)] +pub struct OutOfVocabError { + pub parameter: &'static str, + pub token_ids: Vec, + pub vocab_size: usize, +} + +fn validate_param( + parameter: &'static str, + token_ids: impl IntoIterator, + vocab_size: usize, +) -> Result<(), OutOfVocabError> { + let invalid_token_ids: Vec<_> = token_ids + .into_iter() + .filter(|&token_id| token_id as usize >= vocab_size) + .collect(); + if invalid_token_ids.is_empty() { + return Ok(()); + } + + Err(OutOfVocabError { + parameter, + token_ids: invalid_token_ids, + vocab_size, + }) +} + +/// Validate that pre-tokenized prompt IDs are within the engine-visible prompt +/// vocabulary range. +pub(crate) fn validate_prompt_token_ids( + prompt_token_ids: &[u32], + limits: &SamplingLimits, +) -> Result<(), OutOfVocabError> { + validate_param( + "prompt", + prompt_token_ids.iter().copied(), + limits.prompt_token_vocab_size(), + ) +} + +/// Validate that token IDs in text sampling parameters are within their +/// parameter-specific vocabulary ranges. +pub(crate) fn validate_vocab_range( + params: &EngineCoreSamplingParams, + limits: &SamplingLimits, +) -> Result<(), OutOfVocabError> { + validate_param( + "stop_token_ids", + params.stop_token_ids.iter().copied(), + limits.stop_token_vocab_size(), + )?; + + if let Some(token_ids) = params.allowed_token_ids.as_deref() { + validate_param( + "allowed_token_ids", + token_ids.iter().copied(), + limits.tokenizer_vocab_size, + )?; + } + + if let (Some(logit_bias), Some(vocab_size)) = + (params.logit_bias.as_ref(), limits.model_vocab_size) + { + validate_param("logit_bias", logit_bias.keys().copied(), vocab_size)?; + } + + if let Some(token_ids) = params.logprob_token_ids.as_deref() { + validate_param( + "logprob_token_ids", + token_ids.iter().copied(), + limits.logprobs_vocab_size(), + )?; + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn validate_vocab_range_rejects_out_of_vocab_ids() { + let error = validate_param("logprob_token_ids", [5_u32, 1000, 1001], 1000).unwrap_err(); + + assert_eq!(error.parameter, "logprob_token_ids"); + assert_eq!(error.token_ids, vec![1000, 1001]); + assert_eq!(error.vocab_size, 1000); + } +}