@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1f9dd7900d
Fork 0

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

[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>

Ting SUN committed Jun 11, 2026 at 18:14 UTC 1f9dd7900dcc2deb6714efa922a1e8e2b49a3f81
16 files changed +258 -2
rust/src/chat/src/lib.rs
+10
@@ -140,6 +140,16 @@ impl ChatLlm {
140 self
141 }
142
143 + /// Tokenizer vocabulary size.
144 + pub fn tokenizer_vocab_size(&self) -> usize {
145 + self.text.tokenizer_vocab_size()
146 + }
147 +
148 + /// Model vocabulary size, else `None`.
149 + pub fn model_vocab_size(&self) -> Option<usize> {
150 + self.text.model_vocab_size()
151 + }
152 +
153 /// Expose the underlying text facade for raw text-generation routes such as
154 /// `/v1/completions`.
155 pub fn text(&self) -> &TextLlm {
rust/src/server/src/routes/openai/chat_completions.rs
+7
@@ -52,6 +52,13 @@ pub async fn chat_completions(
52 let request_context = resolve_request_context(&headers, body.request_id.as_deref());
53 let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
54
55 + if let Err(err) = validate::validate_token_id_ranges(
56 + &body,
57 + state.tokenizer_vocab_size(),
58 + state.model_vocab_size(),
59 + ) {
60 + return err.into_response();
61 + }
62 let prepared = match prepare_chat_request(body, &lora_resolution, request_context) {
63 Ok(prepared) => prepared,
64 Err(error) => return error.into_response(),
rust/src/server/src/routes/openai/chat_completions/validate.rs
+43 -1
@@ -1,5 +1,6 @@
1 use super::types::ChatCompletionRequest;
2 use crate::error::{ApiError, bail_invalid_request};
3 +use crate::routes::openai::utils::token_ids::{validate_allowed_token_ids, validate_logit_bias};
4 use crate::routes::openai::utils::types::{ChatMessage, Tool, ToolChoice, ToolChoiceValue};
5
6 /// 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
170 Ok(())
171 }
172
173 +/// Reject out-of-vocab token ids, mirroring the Python input processor:
174 +/// `allowed_token_ids` against the tokenizer vocab, `logit_bias` keys against the
175 +/// model vocab (skipped when the model size is unknown).
176 +pub(super) fn validate_token_id_ranges(
177 + request: &ChatCompletionRequest,
178 + tokenizer_vocab_size: usize,
179 + model_vocab_size: Option<usize>,
180 +) -> Result<(), ApiError> {
181 + validate_allowed_token_ids(request.allowed_token_ids.as_deref(), tokenizer_vocab_size)?;
182 + validate_logit_bias(
183 + request.logit_bias.as_ref(),
184 + model_vocab_size.unwrap_or(usize::MAX),
185 + )
186 +}
187 +
188 #[cfg(test)]
189 mod tests {
190 use std::collections::HashMap;
@@ -176,7 +192,7 @@ mod tests {
192 use serde_json::json;
193 use vllm_chat::ReasoningEffort;
194
179 - use super::validate_request_compat;
195 + use super::{validate_request_compat, validate_token_id_ranges};
196 use crate::routes::openai::chat_completions::types::ChatCompletionRequest;
197 use crate::routes::openai::utils::structured_outputs::ResponseFormat;
198 use crate::routes::openai::utils::types::{
@@ -188,6 +204,32 @@ mod tests {
204 names.iter().map(|s| s.to_string()).collect()
205 }
206
207 + #[test]
208 + fn validate_token_id_ranges_rejects_oob_and_accepts_in_vocab() {
209 + // allowed_token_ids are bounded by the tokenizer vocab
210 + let mut request = base_request();
211 + request.allowed_token_ids = Some(vec![5, 1_000_000]);
212 + assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
213 + // logit_bias is bounded by the larger model vocab: an id between the two
214 + // vocabs is valid and must not be rejected (the parity regression we fix)
215 + let mut request = base_request();
216 + request.logit_bias = Some(HashMap::from([("150".to_string(), 1.0)]));
217 + assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
218 + // logit_bias beyond the model vocab -> reject
219 + let mut request = base_request();
220 + request.logit_bias = Some(HashMap::from([("1000000".to_string(), 1.0)]));
221 + assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
222 + // all in-vocab -> accept
223 + let mut request = base_request();
224 + request.allowed_token_ids = Some(vec![5, 50]);
225 + request.logit_bias = Some(HashMap::from([("50".to_string(), 1.0)]));
226 + assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
227 + // unknown sizes -> skip
228 + let mut request = base_request();
229 + request.allowed_token_ids = Some(vec![1_000_000]);
230 + assert!(validate_token_id_ranges(&request, usize::MAX, None).is_ok());
231 + }
232 +
233 fn base_request() -> ChatCompletionRequest {
234 ChatCompletionRequest {
235 model: "Qwen/Qwen1.5-0.5B-Chat".to_string(),
rust/src/server/src/routes/openai/completions.rs
+7
@@ -46,6 +46,13 @@ pub async fn completions(
46 let request_context = resolve_request_context(&headers, body.request_id.as_deref());
47 let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
48
49 + if let Err(err) = validate::validate_token_id_ranges(
50 + &body,
51 + state.tokenizer_vocab_size(),
52 + state.model_vocab_size(),
53 + ) {
54 + return err.into_response();
55 + }
56 let prepared = match prepare_completion_request(body, &lora_resolution, request_context) {
57 Ok(prepared) => prepared,
58 Err(error) => return error.into_response(),
rust/src/server/src/routes/openai/completions/validate.rs
+54 -1
@@ -2,6 +2,9 @@ use vllm_text::Prompt;
2
3 use super::types::CompletionRequest;
4 use crate::error::{ApiError, bail_invalid_request};
5 +use crate::routes::openai::utils::token_ids::{
6 + validate_allowed_token_ids, validate_logit_bias, validate_prompt_token_ids,
7 +};
8
9 /// Enforce the minimal compatibility contract for the Rust OpenAI server.
10 pub(super) fn validate_request_compat(
@@ -104,13 +107,63 @@ pub(super) fn validate_request_compat(
107 Ok(())
108 }
109
110 +/// Reject out-of-vocab token ids, mirroring the Python input processor. A token-id
111 +/// prompt may reference ids the engine embeds beyond either vocab alone (Qwen3
112 +/// extra LM tokens, multimodal placeholders), so it is bounded by the union of the
113 +/// tokenizer and model vocabularies; `allowed_token_ids` by the tokenizer vocab;
114 +/// `logit_bias` keys by the model vocab (skipped when the model size is unknown).
115 +pub(super) fn validate_token_id_ranges(
116 + request: &CompletionRequest,
117 + tokenizer_vocab_size: usize,
118 + model_vocab_size: Option<usize>,
119 +) -> Result<(), ApiError> {
120 + let prompt_bound = tokenizer_vocab_size.max(model_vocab_size.unwrap_or(0));
121 + validate_prompt_token_ids(&request.prompt, prompt_bound)?;
122 + validate_allowed_token_ids(request.allowed_token_ids.as_deref(), tokenizer_vocab_size)?;
123 + validate_logit_bias(
124 + request.logit_bias.as_ref(),
125 + model_vocab_size.unwrap_or(usize::MAX),
126 + )
127 +}
128 +
129 #[cfg(test)]
130 mod tests {
131 use serde_json::json;
132 + use vllm_text::Prompt;
133
111 - use super::validate_request_compat;
134 + use super::{validate_request_compat, validate_token_id_ranges};
135 use crate::routes::openai::completions::types::CompletionRequest;
136
137 + #[test]
138 + fn validate_token_id_ranges_rejects_oob_prompt_and_params() {
139 + // a token-id prompt below both vocabs is accepted (the engine can embed it)
140 + let mut request = base_request();
141 + request.prompt = Prompt::TokenIds(vec![5, 150]);
142 + assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
143 + // an id at or above the union of the two vocabs is rejected
144 + let mut request = base_request();
145 + request.prompt = Prompt::TokenIds(vec![5, 200]);
146 + assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
147 + // an id beyond the model vocab but within the (larger) tokenizer vocab is
148 + // accepted: the engine embeds added/placeholder ids above the model vocab,
149 + // matching the Python input processor's max(tokenizer, model) bound
150 + let mut request = base_request();
151 + request.prompt = Prompt::TokenIds(vec![150]);
152 + assert!(validate_token_id_ranges(&request, 200, Some(100)).is_ok());
153 + // falls back to the tokenizer vocab when the model size is unknown
154 + let mut request = base_request();
155 + request.prompt = Prompt::TokenIds(vec![150]);
156 + assert!(validate_token_id_ranges(&request, 100, None).is_err());
157 + // allowed_token_ids are bounded by the tokenizer vocab -> reject
158 + let mut request = base_request();
159 + request.allowed_token_ids = Some(vec![150]);
160 + assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
161 + // unknown sizes -> skip
162 + let mut request = base_request();
163 + request.prompt = Prompt::TokenIds(vec![1_000_000]);
164 + assert!(validate_token_id_ranges(&request, usize::MAX, None).is_ok());
165 + }
166 +
167 fn base_request() -> CompletionRequest {
168 serde_json::from_value(json!({
169 "model": "Qwen/Qwen1.5-0.5B-Chat",
rust/src/server/src/routes/openai/utils/mod.rs
+1
@@ -1,4 +1,5 @@
1 pub mod logprobs;
2 pub mod structured_outputs;
3 +pub mod token_ids;
4 pub mod types;
5 pub mod validated_json;
rust/src/server/src/routes/openai/utils/token_ids.rs new
+55
@@ -0,0 +1,55 @@
1 +use std::collections::HashMap;
2 +
3 +use vllm_text::Prompt;
4 +
5 +use crate::error::{ApiError, bail_invalid_request};
6 +
7 +/// Reject token-id prompt entries at or above `bound` (the highest in-vocab id is
8 +/// `bound - 1`).
9 +pub(crate) fn validate_prompt_token_ids(prompt: &Prompt, bound: usize) -> Result<(), ApiError> {
10 + if let Prompt::TokenIds(ids) = prompt
11 + && let Some(&bad) = ids.iter().find(|&&id| id as usize >= bound)
12 + {
13 + bail_invalid_request!(
14 + param = "prompt",
15 + "prompt contains out-of-vocab token id {bad}; vocabulary size is {bound}."
16 + );
17 + }
18 + Ok(())
19 +}
20 +
21 +/// Reject `allowed_token_ids` entries at or above `bound`.
22 +pub(crate) fn validate_allowed_token_ids(
23 + allowed_token_ids: Option<&[u32]>,
24 + bound: usize,
25 +) -> Result<(), ApiError> {
26 + if let Some(ids) = allowed_token_ids
27 + && let Some(&bad) = ids.iter().find(|&&id| id as usize >= bound)
28 + {
29 + bail_invalid_request!(
30 + param = "allowed_token_ids",
31 + "allowed_token_ids contains out-of-vocab token id {bad}; vocabulary size is {bound}."
32 + );
33 + }
34 + Ok(())
35 +}
36 +
37 +/// Reject `logit_bias` keys at or above `bound`.
38 +pub(crate) fn validate_logit_bias(
39 + logit_bias: Option<&HashMap<String, f32>>,
40 + bound: usize,
41 +) -> Result<(), ApiError> {
42 + if let Some(bias) = logit_bias {
43 + for key in bias.keys() {
44 + if let Ok(id) = key.parse::<u32>()
45 + && id as usize >= bound
46 + {
47 + bail_invalid_request!(
48 + param = "logit_bias",
49 + "logit_bias contains out-of-vocab token id {id}; vocabulary size is {bound}."
50 + );
51 + }
52 + }
53 + }
54 + Ok(())
55 +}
rust/src/server/src/state.rs
+10
@@ -114,6 +114,16 @@ impl AppState {
114 &self.served_model_names
115 }
116
117 + /// Tokenizer vocabulary size.
118 + pub fn tokenizer_vocab_size(&self) -> usize {
119 + self.chat.tokenizer_vocab_size()
120 + }
121 +
122 + /// Model vocabulary size, else `None`.
123 + pub fn model_vocab_size(&self) -> Option<usize> {
124 + self.chat.model_vocab_size()
125 + }
126 +
127 /// Return base served model names plus dynamically loaded LoRA adapter
128 /// names.
129 pub async fn served_model_names_with_loras(&self) -> Vec<String> {
rust/src/text/src/backend/hf/config.rs
+8
@@ -92,6 +92,7 @@ impl HfSpecialTokens {
92 pub struct ModelConfig {
93 model_type: Option<String>,
94 max_position_embeddings: Option<u32>,
95 + vocab_size: Option<u32>,
96 num_attention_heads: Option<u32>,
97 num_experts: Option<OneOrManyExpertCount>,
98 moe_num_experts: Option<OneOrManyExpertCount>,
@@ -179,6 +180,13 @@ impl ModelConfig {
180 self.model_type.as_deref().or_else(|| self.text_config.as_deref()?.model_type())
181 }
182
183 + /// Return the effective model vocabulary size, following the same simplified
184 + /// text-config selection as `model_type`: the top-level config wins,
185 + /// otherwise a single nested `text_config` may provide it.
186 + pub fn vocab_size(&self) -> Option<u32> {
187 + self.vocab_size.or_else(|| self.text_config.as_deref()?.vocab_size())
188 + }
189 +
190 /// Reject partially nested `text_config` payloads that are unlikely to be
191 /// valid LLM configs for our current use.
192 ///
rust/src/text/src/backend/hf/mod.rs
+4
@@ -100,6 +100,10 @@ impl TextBackend for HfTextBackend {
100 self.model_config.is_moe()
101 }
102
103 + fn model_vocab_size(&self) -> Option<usize> {
104 + self.model_config.vocab_size().map(|v| v as usize)
105 + }
106 +
107 fn model_id(&self) -> &str {
108 &self.model_id
109 }
rust/src/text/src/backend/mod.rs
+12
@@ -41,6 +41,18 @@ pub trait TextBackend: Send + Sync {
41 fn sampling_hints(&self) -> Result<SamplingHints> {
42 Ok(SamplingHints::default())
43 }
44 +
45 + /// Return the model vocabulary size from the model config, if known. Used to
46 + /// range-check request token ids against the engine embedding table.
47 + fn model_vocab_size(&self) -> Option<usize> {
48 + None
49 + }
50 +
51 + /// Return the full tokenizer vocabulary size (Python `len(tokenizer)`).
52 + /// Used to range-check `allowed_token_ids` and token-id prompts.
53 + fn tokenizer_vocab_size(&self) -> usize {
54 + self.tokenizer().vocab_size()
55 + }
56 }
57
58 /// Shared trait-object form of [`TextBackend`].
rust/src/text/src/lib.rs
+12
@@ -91,6 +91,18 @@ impl TextLlm {
91 self.backend.tokenizer()
92 }
93
94 + /// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
95 + /// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
96 + pub fn tokenizer_vocab_size(&self) -> usize {
97 + self.backend.tokenizer_vocab_size()
98 + }
99 +
100 + /// Model vocabulary size from the model config, used to bound `logit_bias`
101 + /// keys and token-id prompts against the engine embedding table.
102 + pub fn model_vocab_size(&self) -> Option<usize> {
103 + self.backend.model_vocab_size()
104 + }
105 +
106 /// Tokenize if needed, lower to a generate request, and return the raw
107 /// token stream.
108 pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
rust/src/tokenizer/src/hf.rs
+7
@@ -174,6 +174,13 @@ impl Tokenizer for HuggingFaceTokenizer {
174 }
175 }
176
177 + fn vocab_size(&self) -> usize {
178 + match &self.backend {
179 + Backend::Hf(t) => t.get_vocab_size(true),
180 + Backend::Fastokens(t) | Backend::FastokensByteLevel(t) => t.vocab_size(),
181 + }
182 + }
183 +
184 fn id_to_token(&self, id: u32) -> Option<String> {
185 match &self.backend {
186 Backend::Hf(t) => t.id_to_token(id),
rust/src/tokenizer/src/lib.rs
+6
@@ -34,6 +34,12 @@ pub trait Tokenizer: Send + Sync {
34 None
35 }
36
37 + /// Return the vocabulary size. Backends that cannot report it fall back to
38 + /// `usize::MAX`, an effectively unbounded value used only by test stubs.
39 + fn vocab_size(&self) -> usize {
40 + usize::MAX
41 + }
42 +
43 /// Return whether the given token ID is special.
44 fn is_special_id(&self, _token_id: u32) -> bool {
45 false
rust/src/tokenizer/src/tekken.rs
+4
@@ -56,6 +56,10 @@ impl Tokenizer for TekkenTokenizer {
56 self.inner.id_to_piece(id).ok()
57 }
58
59 + fn vocab_size(&self) -> usize {
60 + self.inner.vocab_size()
61 + }
62 +
63 fn is_special_id(&self, token_id: u32) -> bool {
64 self.inner.is_special_token(token_id)
65 }
rust/src/tokenizer/src/tiktoken.rs
+18
@@ -503,6 +503,13 @@ impl Tokenizer for TiktokenTokenizer {
503 fn is_special_id(&self, token_id: u32) -> bool {
504 self.metadata.is_special_id(token_id)
505 }
506 +
507 + fn vocab_size(&self) -> usize {
508 + // Exclusive upper bound on token ids the tokenizer can decode (BPE base
509 + // tokens plus the registered special/reserved slots), used to range-check
510 + // `allowed_token_ids` so tiktoken models are not exempt from validation.
511 + self.metadata.vocab_upper_bound as usize
512 + }
513 }
514
515 /// Select the BPE regex pattern for a tiktoken model based on `config.json`.
@@ -614,6 +621,17 @@ mod tests {
621 }
622 }
623
624 + #[test]
625 + fn tiktoken_vocab_size_reports_upper_bound() {
626 + // The synthetic BPE file has 256 base tokens (bytes 0..=255) and ships no
627 + // sibling config, so the constructor uses the 256-slot reserved fallback,
628 + // giving a vocab upper bound of 512.
629 + let (backends, _dir) = tiktoken_backends();
630 + for backend in backends {
631 + assert_eq!(backend.vocab_size(), 512);
632 + }
633 + }
634 +
635 /// When `config.json` exposes a `vocab_size`, the reserved-token range must
636 /// be sized to it rather than to the 256-slot fallback. This is the
637 /// general (non-Kimi-specific) path: any tiktoken model whose own