[Bugfix][Rust Frontend] Validate out-of-vocab token ids in request params (#44680)

Signed-off-by: Ting Sun <suntcrick@gmail.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>

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<usize> {
+        self.text.model_vocab_size()
+    }
+
     /// Expose the underlying text facade for raw text-generation routes such as
     /// `/v1/completions`.
     pub fn text(&self) -> &TextLlm {

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(),

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<usize>,
+) -> 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(),

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(),

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<usize>,
+) -> 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",

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;

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<String, f32>>,
+    bound: usize,
+) -> Result<(), ApiError> {
+    if let Some(bias) = logit_bias {
+        for key in bias.keys() {
+            if let Ok(id) = key.parse::<u32>()
+                && id as usize >= bound
+            {
+                bail_invalid_request!(
+                    param = "logit_bias",
+                    "logit_bias contains out-of-vocab token id {id}; vocabulary size is {bound}."
+                );
+            }
+        }
+    }
+    Ok(())
+}

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<usize> {
+        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<String> {

rust/src/text/src/backend/hf/config.rs

@@ -92,6 +92,7 @@ impl HfSpecialTokens {
 pub struct ModelConfig {
     model_type: Option<String>,
     max_position_embeddings: Option<u32>,
+    vocab_size: Option<u32>,
     num_attention_heads: Option<u32>,
     num_experts: Option<OneOrManyExpertCount>,
     moe_num_experts: Option<OneOrManyExpertCount>,
@@ -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<u32> {
+        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.
     ///

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<usize> {
+        self.model_config.vocab_size().map(|v| v as usize)
+    }
+
     fn model_id(&self) -> &str {
         &self.model_id
     }

rust/src/text/src/backend/mod.rs

@@ -41,6 +41,18 @@ pub trait TextBackend: Send + Sync {
     fn sampling_hints(&self) -> Result<SamplingHints> {
         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<usize> {
+        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`].

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<usize> {
+        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<GenerateOutputStream> {

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<String> {
         match &self.backend {
             Backend::Hf(t) => t.id_to_token(id),

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

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)
     }

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