diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 4add5f0f9b..63b4cbdbf4 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -140,6 +140,16 @@ impl ChatLlm { self } + /// Tokenizer vocabulary size. + pub fn tokenizer_vocab_size(&self) -> usize { + self.text.tokenizer_vocab_size() + } + + /// Model vocabulary size, else `None`. + pub fn model_vocab_size(&self) -> Option { + self.text.model_vocab_size() + } + /// Expose the underlying text facade for raw text-generation routes such as /// `/v1/completions`. pub fn text(&self) -> &TextLlm { diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index 8a2df9b25b..e93c049b2d 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -52,6 +52,13 @@ 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 379f2c2d39..fb64428e4b 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -1,5 +1,6 @@ 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. @@ -169,6 +170,21 @@ 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 +192,7 @@ mod tests { use serde_json::json; use vllm_chat::ReasoningEffort; - use super::validate_request_compat; + use super::{validate_request_compat, validate_token_id_ranges}; use crate::routes::openai::chat_completions::types::ChatCompletionRequest; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ @@ -188,6 +204,32 @@ 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 fb0e7bdd87..b6e4383c7d 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -46,6 +46,13 @@ 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 a53609234b..2af8c8add1 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -2,6 +2,9 @@ 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( @@ -104,13 +107,63 @@ 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; + use super::{validate_request_compat, validate_token_id_ranges}; 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 57b1d99690..039df87f9d 100644 --- a/rust/src/server/src/routes/openai/utils/mod.rs +++ b/rust/src/server/src/routes/openai/utils/mod.rs @@ -1,4 +1,5 @@ pub mod logprobs; pub mod structured_outputs; +pub mod token_ids; pub mod types; 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 new file mode 100644 index 0000000000..ffa945ef94 --- /dev/null +++ b/rust/src/server/src/routes/openai/utils/token_ids.rs @@ -0,0 +1,55 @@ +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 2fee91d457..55959b60d9 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -114,6 +114,16 @@ 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/hf/config.rs b/rust/src/text/src/backend/hf/config.rs index 5f2ecf8ba6..1efb31618d 100644 --- a/rust/src/text/src/backend/hf/config.rs +++ b/rust/src/text/src/backend/hf/config.rs @@ -92,6 +92,7 @@ impl HfSpecialTokens { pub struct ModelConfig { model_type: Option, max_position_embeddings: Option, + vocab_size: Option, num_attention_heads: Option, num_experts: Option, moe_num_experts: Option, @@ -179,6 +180,13 @@ impl ModelConfig { self.model_type.as_deref().or_else(|| self.text_config.as_deref()?.model_type()) } + /// Return the effective model vocabulary size, following the same simplified + /// text-config selection as `model_type`: the top-level config wins, + /// otherwise a single nested `text_config` may provide it. + pub fn vocab_size(&self) -> Option { + self.vocab_size.or_else(|| self.text_config.as_deref()?.vocab_size()) + } + /// Reject partially nested `text_config` payloads that are unlikely to be /// valid LLM configs for our current use. /// diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index a5d07dd8fc..b6b79d9914 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -100,6 +100,10 @@ impl TextBackend for HfTextBackend { self.model_config.is_moe() } + fn model_vocab_size(&self) -> Option { + self.model_config.vocab_size().map(|v| v as usize) + } + fn model_id(&self) -> &str { &self.model_id } diff --git a/rust/src/text/src/backend/mod.rs b/rust/src/text/src/backend/mod.rs index 4f2d7093a7..680d454da9 100644 --- a/rust/src/text/src/backend/mod.rs +++ b/rust/src/text/src/backend/mod.rs @@ -41,6 +41,18 @@ pub trait TextBackend: Send + Sync { fn sampling_hints(&self) -> Result { Ok(SamplingHints::default()) } + + /// Return the model vocabulary size from the model config, if known. Used to + /// range-check request token ids against the engine embedding table. + fn model_vocab_size(&self) -> Option { + None + } + + /// Return the full tokenizer vocabulary size (Python `len(tokenizer)`). + /// Used to range-check `allowed_token_ids` and token-id prompts. + fn tokenizer_vocab_size(&self) -> usize { + self.tokenizer().vocab_size() + } } /// Shared trait-object form of [`TextBackend`]. diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index 48828045a2..a550a8afc5 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -91,6 +91,18 @@ impl TextLlm { self.backend.tokenizer() } + /// Tokenizer vocabulary size (the number of tokens the tokenizer knows), + /// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`. + pub fn tokenizer_vocab_size(&self) -> usize { + self.backend.tokenizer_vocab_size() + } + + /// Model vocabulary size from the model config, used to bound `logit_bias` + /// keys and token-id prompts against the engine embedding table. + pub fn model_vocab_size(&self) -> Option { + self.backend.model_vocab_size() + } + /// Tokenize if needed, lower to a generate request, and return the raw /// token stream. pub async fn generate_raw(&self, request: TextRequest) -> Result { diff --git a/rust/src/tokenizer/src/hf.rs b/rust/src/tokenizer/src/hf.rs index bd8052e6fa..2982f8c4aa 100644 --- a/rust/src/tokenizer/src/hf.rs +++ b/rust/src/tokenizer/src/hf.rs @@ -174,6 +174,13 @@ impl Tokenizer for HuggingFaceTokenizer { } } + fn vocab_size(&self) -> usize { + match &self.backend { + Backend::Hf(t) => t.get_vocab_size(true), + Backend::Fastokens(t) | Backend::FastokensByteLevel(t) => t.vocab_size(), + } + } + fn id_to_token(&self, id: u32) -> Option { match &self.backend { Backend::Hf(t) => t.id_to_token(id), diff --git a/rust/src/tokenizer/src/lib.rs b/rust/src/tokenizer/src/lib.rs index 6a512a5a62..6f315bc01b 100644 --- a/rust/src/tokenizer/src/lib.rs +++ b/rust/src/tokenizer/src/lib.rs @@ -34,6 +34,12 @@ pub trait Tokenizer: Send + Sync { None } + /// Return the vocabulary size. Backends that cannot report it fall back to + /// `usize::MAX`, an effectively unbounded value used only by test stubs. + fn vocab_size(&self) -> usize { + usize::MAX + } + /// Return whether the given token ID is special. fn is_special_id(&self, _token_id: u32) -> bool { false diff --git a/rust/src/tokenizer/src/tekken.rs b/rust/src/tokenizer/src/tekken.rs index e8560c65a3..50981efdde 100644 --- a/rust/src/tokenizer/src/tekken.rs +++ b/rust/src/tokenizer/src/tekken.rs @@ -56,6 +56,10 @@ impl Tokenizer for TekkenTokenizer { self.inner.id_to_piece(id).ok() } + fn vocab_size(&self) -> usize { + self.inner.vocab_size() + } + fn is_special_id(&self, token_id: u32) -> bool { self.inner.is_special_token(token_id) } diff --git a/rust/src/tokenizer/src/tiktoken.rs b/rust/src/tokenizer/src/tiktoken.rs index 0c57ff5f6b..9b4c17a855 100644 --- a/rust/src/tokenizer/src/tiktoken.rs +++ b/rust/src/tokenizer/src/tiktoken.rs @@ -503,6 +503,13 @@ impl Tokenizer for TiktokenTokenizer { fn is_special_id(&self, token_id: u32) -> bool { self.metadata.is_special_id(token_id) } + + fn vocab_size(&self) -> usize { + // Exclusive upper bound on token ids the tokenizer can decode (BPE base + // tokens plus the registered special/reserved slots), used to range-check + // `allowed_token_ids` so tiktoken models are not exempt from validation. + self.metadata.vocab_upper_bound as usize + } } /// Select the BPE regex pattern for a tiktoken model based on `config.json`. @@ -614,6 +621,17 @@ mod tests { } } + #[test] + fn tiktoken_vocab_size_reports_upper_bound() { + // The synthetic BPE file has 256 base tokens (bytes 0..=255) and ships no + // sibling config, so the constructor uses the 256-slot reserved fallback, + // giving a vocab upper bound of 512. + let (backends, _dir) = tiktoken_backends(); + for backend in backends { + assert_eq!(backend.vocab_size(), 512); + } + } + /// When `config.json` exposes a `vocab_size`, the reserved-token range must /// be sized to it rather than to the 256-slot fallback. This is the /// general (non-Kimi-specific) path: any tiktoken model whose own