A high-throughput and memory-efficient inference and serving engine for LLMs
[Rust Frontend] Lower out-of-vocab validation to `text` layer (#45685)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
Jun 16, 2026 at 11:37 UTC
f99260d2aa43d779fe4fc9d69bd57ad4353a0f3f
14 files changed
+278
-212
rust/src/server/src/error.rs
+12
@@ -102,6 +102,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool {
102
vllm_text::Error::PromptTooLong { .. }
103
| vllm_text::Error::EmptyPromptTokenIds { .. }
104
| vllm_text::Error::Logprobs(_)
105
+ | vllm_text::Error::OutOfVocab(_)
106
// An empty tokenized prompt detected later, at request prepare
107
// time, surfaces through the transparent Llm wrapper.
108
| vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })
@@ -172,6 +173,17 @@ mod tests {
173
assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
174
}
175
176
+ #[test]
177
+ fn out_of_vocab_validation_maps_to_invalid_request() {
178
+ let error = vllm_text::Error::OutOfVocab(vllm_text::OutOfVocabError {
179
+ parameter: "logprob_token_ids",
180
+ token_ids: vec![1000],
181
+ vocab_size: 1000,
182
+ });
183
+ let api_error = text_submit_error("failed to submit completion request", error);
184
+ assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
185
+ }
186
+
187
#[test]
188
fn other_submit_errors_stay_internal() {
189
let error = vllm_text::Error::Tokenizer("backend exploded".to_string());
rust/src/server/src/routes/openai/chat_completions.rs
-7
@@ -53,13 +53,6 @@ pub async fn chat_completions(
53
let request_context = resolve_request_context(&headers, body.request_id.as_deref());
54
let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
55
56
- if let Err(err) = validate::validate_token_id_ranges(
57
- &body,
58
- state.tokenizer_vocab_size(),
59
- state.model_vocab_size(),
60
- ) {
61
- return err.into_response();
62
- }
56
let prepared = match prepare_chat_request(body, &lora_resolution, request_context) {
57
Ok(prepared) => prepared,
58
Err(error) => return error.into_response(),
rust/src/server/src/routes/openai/chat_completions/validate.rs
+1
-43
@@ -1,6 +1,5 @@
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};
3
use crate::routes::openai::utils::types::{ChatMessage, Tool, ToolChoice, ToolChoiceValue};
4
5
/// 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
153
Ok(())
154
}
155
157
-/// Reject out-of-vocab token ids, mirroring the Python input processor:
158
-/// `allowed_token_ids` against the tokenizer vocab, `logit_bias` keys against the
159
-/// model vocab (skipped when the model size is unknown).
160
-pub(super) fn validate_token_id_ranges(
161
- request: &ChatCompletionRequest,
162
- tokenizer_vocab_size: usize,
163
- model_vocab_size: Option<usize>,
164
-) -> Result<(), ApiError> {
165
- validate_allowed_token_ids(request.allowed_token_ids.as_deref(), tokenizer_vocab_size)?;
166
- validate_logit_bias(
167
- request.logit_bias.as_ref(),
168
- model_vocab_size.unwrap_or(usize::MAX),
169
- )
170
-}
171
-
156
#[cfg(test)]
157
mod tests {
158
use std::collections::HashMap;
@@ -176,7 +160,7 @@ mod tests {
160
use serde_json::json;
161
use vllm_chat::ReasoningEffort;
162
179
- use super::{validate_request_compat, validate_token_id_ranges};
163
+ use super::validate_request_compat;
164
use crate::routes::openai::chat_completions::types::ChatCompletionRequest;
165
use crate::routes::openai::utils::structured_outputs::ResponseFormat;
166
use crate::routes::openai::utils::types::{
@@ -188,32 +172,6 @@ mod tests {
172
names.iter().map(|s| s.to_string()).collect()
173
}
174
191
- #[test]
192
- fn validate_token_id_ranges_rejects_oob_and_accepts_in_vocab() {
193
- // allowed_token_ids are bounded by the tokenizer vocab
194
- let mut request = base_request();
195
- request.allowed_token_ids = Some(vec![5, 1_000_000]);
196
- assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
197
- // logit_bias is bounded by the larger model vocab: an id between the two
198
- // vocabs is valid and must not be rejected (the parity regression we fix)
199
- let mut request = base_request();
200
- request.logit_bias = Some(HashMap::from([("150".to_string(), 1.0)]));
201
- assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
202
- // logit_bias beyond the model vocab -> reject
203
- let mut request = base_request();
204
- request.logit_bias = Some(HashMap::from([("1000000".to_string(), 1.0)]));
205
- assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
206
- // all in-vocab -> accept
207
- let mut request = base_request();
208
- request.allowed_token_ids = Some(vec![5, 50]);
209
- request.logit_bias = Some(HashMap::from([("50".to_string(), 1.0)]));
210
- assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
211
- // unknown sizes -> skip
212
- let mut request = base_request();
213
- request.allowed_token_ids = Some(vec![1_000_000]);
214
- assert!(validate_token_id_ranges(&request, usize::MAX, None).is_ok());
215
- }
216
-
175
fn base_request() -> ChatCompletionRequest {
176
ChatCompletionRequest {
177
model: "Qwen/Qwen1.5-0.5B-Chat".to_string(),
rust/src/server/src/routes/openai/completions.rs
-7
@@ -47,13 +47,6 @@ pub async fn completions(
47
let request_context = resolve_request_context(&headers, body.request_id.as_deref());
48
let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
49
50
- if let Err(err) = validate::validate_token_id_ranges(
51
- &body,
52
- state.tokenizer_vocab_size(),
53
- state.model_vocab_size(),
54
- ) {
55
- return err.into_response();
56
- }
50
let prepared = match prepare_completion_request(body, &lora_resolution, request_context) {
51
Ok(prepared) => prepared,
52
Err(error) => return error.into_response(),
rust/src/server/src/routes/openai/completions/validate.rs
+1
-54
@@ -2,9 +2,6 @@ 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
-};
5
6
/// Enforce the minimal compatibility contract for the Rust OpenAI server.
7
pub(super) fn validate_request_compat(
@@ -98,63 +95,13 @@ pub(super) fn validate_request_compat(
95
Ok(())
96
}
97
101
-/// Reject out-of-vocab token ids, mirroring the Python input processor. A token-id
102
-/// prompt may reference ids the engine embeds beyond either vocab alone (Qwen3
103
-/// extra LM tokens, multimodal placeholders), so it is bounded by the union of the
104
-/// tokenizer and model vocabularies; `allowed_token_ids` by the tokenizer vocab;
105
-/// `logit_bias` keys by the model vocab (skipped when the model size is unknown).
106
-pub(super) fn validate_token_id_ranges(
107
- request: &CompletionRequest,
108
- tokenizer_vocab_size: usize,
109
- model_vocab_size: Option<usize>,
110
-) -> Result<(), ApiError> {
111
- let prompt_bound = tokenizer_vocab_size.max(model_vocab_size.unwrap_or(0));
112
- validate_prompt_token_ids(&request.prompt, prompt_bound)?;
113
- validate_allowed_token_ids(request.allowed_token_ids.as_deref(), tokenizer_vocab_size)?;
114
- validate_logit_bias(
115
- request.logit_bias.as_ref(),
116
- model_vocab_size.unwrap_or(usize::MAX),
117
- )
118
-}
119
-
98
#[cfg(test)]
99
mod tests {
100
use serde_json::json;
123
- use vllm_text::Prompt;
101
125
- use super::{validate_request_compat, validate_token_id_ranges};
102
+ use super::validate_request_compat;
103
use crate::routes::openai::completions::types::CompletionRequest;
104
128
- #[test]
129
- fn validate_token_id_ranges_rejects_oob_prompt_and_params() {
130
- // a token-id prompt below both vocabs is accepted (the engine can embed it)
131
- let mut request = base_request();
132
- request.prompt = Prompt::TokenIds(vec![5, 150]);
133
- assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
134
- // an id at or above the union of the two vocabs is rejected
135
- let mut request = base_request();
136
- request.prompt = Prompt::TokenIds(vec![5, 200]);
137
- assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
138
- // an id beyond the model vocab but within the (larger) tokenizer vocab is
139
- // accepted: the engine embeds added/placeholder ids above the model vocab,
140
- // matching the Python input processor's max(tokenizer, model) bound
141
- let mut request = base_request();
142
- request.prompt = Prompt::TokenIds(vec![150]);
143
- assert!(validate_token_id_ranges(&request, 200, Some(100)).is_ok());
144
- // falls back to the tokenizer vocab when the model size is unknown
145
- let mut request = base_request();
146
- request.prompt = Prompt::TokenIds(vec![150]);
147
- assert!(validate_token_id_ranges(&request, 100, None).is_err());
148
- // allowed_token_ids are bounded by the tokenizer vocab -> reject
149
- let mut request = base_request();
150
- request.allowed_token_ids = Some(vec![150]);
151
- assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
152
- // unknown sizes -> skip
153
- let mut request = base_request();
154
- request.prompt = Prompt::TokenIds(vec![1_000_000]);
155
- assert!(validate_token_id_ranges(&request, usize::MAX, None).is_ok());
156
- }
157
-
105
fn base_request() -> CompletionRequest {
106
serde_json::from_value(json!({
107
"model": "Qwen/Qwen1.5-0.5B-Chat",
rust/src/server/src/routes/openai/utils/mod.rs
-1
@@ -1,6 +1,5 @@
1
pub mod logprobs;
2
pub mod structured_outputs;
3
-pub mod token_ids;
3
pub mod types;
4
pub mod usage;
5
pub mod validated_json;
rust/src/server/src/routes/openai/utils/token_ids.rs
deleted
-55
@@ -1,55 +0,0 @@
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,16 +114,6 @@ 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
-
117
/// Return base served model names plus dynamically loaded LoRA adapter
118
/// names.
119
pub async fn served_model_names_with_loras(&self) -> Vec<String> {
rust/src/text/src/backend/mod.rs
+15
-3
@@ -29,10 +29,12 @@ pub struct SamplingLimits {
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.
32
+
33
+ /// Model vocabulary size from the model config, used to bound
34
+ /// `logit_bias` keys when available.
35
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
+ /// Tokenizer vocabulary size, used to bound `allowed_token_ids` and
37
+ /// token-ID prompts.
38
pub tokenizer_vocab_size: usize,
39
}
40
@@ -48,6 +50,16 @@ impl SamplingLimits {
50
pub fn logprobs_vocab_size(&self) -> usize {
51
self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size)
52
}
53
+
54
+ /// Return the vocabulary size used to validate generated stop token IDs.
55
+ pub fn stop_token_vocab_size(&self) -> usize {
56
+ self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size)
57
+ }
58
+
59
+ /// Return the union bound used to validate token-ID prompts.
60
+ pub fn prompt_token_vocab_size(&self) -> usize {
61
+ self.tokenizer_vocab_size.max(self.model_vocab_size.unwrap_or(0))
62
+ }
63
}
64
65
/// Minimal text-processing backend needed by `vllm-text`.
rust/src/text/src/error.rs
+3
@@ -3,6 +3,7 @@ use vllm_engine_core_client::Error as EngineCoreError;
3
use vllm_llm::Error as LlmError;
4
5
pub use crate::lower::logprobs::LogprobsError;
6
+pub use crate::lower::token_ids::OutOfVocabError;
7
8
#[derive(Debug, Error)]
9
pub enum Error {
@@ -17,6 +18,8 @@ pub enum Error {
18
PromptTooLong { max_model_len: u32, prompt_len: u32 },
19
#[error(transparent)]
20
Logprobs(#[from] LogprobsError),
21
+ #[error(transparent)]
22
+ OutOfVocab(#[from] OutOfVocabError),
23
#[error("text request stream `{request_id}` closed before terminal output")]
24
StreamClosedBeforeTerminalOutput { request_id: String },
25
#[error(transparent)]
rust/src/text/src/lib.rs
+1
-1
@@ -7,7 +7,7 @@
7
use std::mem::take;
8
9
pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend};
10
-pub use error::{Error, LogprobsError, Result};
10
+pub use error::{Error, LogprobsError, OutOfVocabError, Result};
11
use futures::Stream;
12
pub use lower::{
13
PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens,
rust/src/text/src/lower.rs
+141
-6
@@ -1,6 +1,7 @@
1
use std::collections::BTreeSet;
2
3
pub(crate) mod logprobs;
4
+pub(crate) mod token_ids;
5
6
use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
7
use vllm_llm::GenerateRequest;
@@ -10,6 +11,7 @@ use crate::backend::{SamplingHints, SamplingLimits};
11
use crate::error::{Error, Result};
12
use crate::request::{SamplingParams, TextRequest};
13
use logprobs::validate_logprobs;
14
+use token_ids::{validate_prompt_token_ids, validate_vocab_range};
15
16
/// One text request after it has been lowered into the raw generate boundary.
17
#[derive(Debug)]
@@ -31,6 +33,8 @@ pub fn lower_text_request(
33
tokenizer: &dyn Tokenizer,
34
) -> Result<PreparedTextRequest> {
35
let prompt_len = prompt_token_ids.len() as u32;
36
+ validate_prompt_token_ids(&prompt_token_ids, &sampling_limits)?;
37
+
38
let generate_request = GenerateRequest {
39
request_id: request.request_id.clone(),
40
prompt_token_ids,
@@ -100,7 +104,6 @@ pub fn lower_sampling_params(
104
vllm_xargs,
105
} = sampling_params;
106
103
- // Validate logprobs-related fields first with runtime sampling limits first.
107
validate_logprobs(
108
logprobs,
109
prompt_logprobs,
@@ -139,7 +142,7 @@ pub fn lower_sampling_params(
142
merge_unique_token_ids(&mut stop_token_ids, extra_eos_token_ids.iter().copied());
143
}
144
142
- Ok(EngineCoreSamplingParams {
145
+ let params = EngineCoreSamplingParams {
146
temperature,
147
top_p,
148
top_k,
@@ -162,7 +165,9 @@ pub fn lower_sampling_params(
165
logprob_token_ids,
166
skip_reading_prefix_cache,
167
extra_args: vllm_xargs,
165
- })
168
+ };
169
+ validate_vocab_range(¶ms, &sampling_limits)?;
170
+ Ok(params)
171
}
172
173
/// Convert bad-word strings into token-ID sequences, following the Python vLLM
@@ -243,14 +248,14 @@ fn merge_unique_token_ids(
248
249
#[cfg(test)]
250
mod tests {
246
- use std::collections::BTreeSet;
251
+ use std::collections::{BTreeSet, HashMap};
252
253
use serial_test::file_serial;
254
255
use super::*;
256
use crate::backend::hf::HfTextBackend;
257
use crate::backend::{SamplingHints, TextBackend as _};
253
- use crate::error::LogprobsError;
258
+ use crate::error::{LogprobsError, OutOfVocabError};
259
use crate::request::{Prompt, TextRequest};
260
261
/// Stub tokenizer that returns empty token IDs — sufficient for tests that
@@ -430,6 +435,57 @@ mod tests {
435
.assert_debug_eq(¶ms);
436
}
437
438
+ #[test]
439
+ fn lower_text_request_uses_union_vocab_for_prompt_token_ids() {
440
+ lower_text_request(
441
+ sample_request(),
442
+ vec![1500],
443
+ sample_sampling_hints(),
444
+ SamplingLimits {
445
+ model_vocab_size: Some(2000),
446
+ tokenizer_vocab_size: 1000,
447
+ ..sample_sampling_limits()
448
+ },
449
+ &stub_tokenizer(),
450
+ )
451
+ .expect("model vocab extends prompt token range");
452
+
453
+ lower_text_request(
454
+ sample_request(),
455
+ vec![1500],
456
+ sample_sampling_hints(),
457
+ SamplingLimits {
458
+ model_vocab_size: Some(1000),
459
+ tokenizer_vocab_size: 2000,
460
+ ..sample_sampling_limits()
461
+ },
462
+ &stub_tokenizer(),
463
+ )
464
+ .expect("tokenizer vocab extends prompt token range");
465
+
466
+ let error = lower_text_request(
467
+ sample_request(),
468
+ vec![2000],
469
+ sample_sampling_hints(),
470
+ SamplingLimits {
471
+ model_vocab_size: Some(1000),
472
+ tokenizer_vocab_size: 2000,
473
+ ..sample_sampling_limits()
474
+ },
475
+ &stub_tokenizer(),
476
+ )
477
+ .unwrap_err();
478
+
479
+ assert!(matches!(
480
+ error,
481
+ Error::OutOfVocab(OutOfVocabError {
482
+ parameter: "prompt",
483
+ token_ids,
484
+ vocab_size: 2000,
485
+ }) if token_ids == vec![2000]
486
+ ));
487
+ }
488
+
489
#[tokio::test]
490
#[file_serial(hf_qwen3)]
491
async fn lower_text_request_uses_real_qwen_generation_defaults() {
@@ -747,13 +803,92 @@ mod tests {
803
804
assert!(matches!(
805
error,
750
- Error::Logprobs(LogprobsError::InvalidTokenIds {
806
+ Error::OutOfVocab(OutOfVocabError {
807
+ parameter: "logprob_token_ids",
808
token_ids,
809
vocab_size: 1000,
810
}) if token_ids == vec![1000]
811
));
812
}
813
814
+ #[test]
815
+ fn lower_sampling_params_rejects_out_of_vocab_stop_token_ids() {
816
+ let error = lower_sampling_params_with_limits(
817
+ SamplingParams {
818
+ stop_token_ids: Some(vec![999, 1000]),
819
+ ..Default::default()
820
+ },
821
+ sample_sampling_limits(),
822
+ )
823
+ .unwrap_err();
824
+
825
+ assert!(matches!(
826
+ error,
827
+ Error::OutOfVocab(OutOfVocabError {
828
+ parameter: "stop_token_ids",
829
+ token_ids,
830
+ vocab_size: 1000,
831
+ }) if token_ids == vec![1000]
832
+ ));
833
+ }
834
+
835
+ #[test]
836
+ fn lower_sampling_params_rejects_out_of_vocab_allowed_token_ids() {
837
+ let error = lower_sampling_params_with_limits(
838
+ SamplingParams {
839
+ allowed_token_ids: Some(vec![1999, 2000]),
840
+ ..Default::default()
841
+ },
842
+ sample_sampling_limits(),
843
+ )
844
+ .unwrap_err();
845
+
846
+ assert!(matches!(
847
+ error,
848
+ Error::OutOfVocab(OutOfVocabError {
849
+ parameter: "allowed_token_ids",
850
+ token_ids,
851
+ vocab_size: 2000,
852
+ }) if token_ids == vec![2000]
853
+ ));
854
+ }
855
+
856
+ #[test]
857
+ fn lower_sampling_params_rejects_out_of_vocab_logit_bias() {
858
+ let error = lower_sampling_params_with_limits(
859
+ SamplingParams {
860
+ logit_bias: Some(HashMap::from([(1000, 1.0)])),
861
+ ..Default::default()
862
+ },
863
+ sample_sampling_limits(),
864
+ )
865
+ .unwrap_err();
866
+
867
+ assert!(matches!(
868
+ error,
869
+ Error::OutOfVocab(OutOfVocabError {
870
+ parameter: "logit_bias",
871
+ token_ids,
872
+ vocab_size: 1000,
873
+ }) if token_ids == vec![1000]
874
+ ));
875
+ }
876
+
877
+ #[test]
878
+ fn lower_sampling_params_skips_logit_bias_range_when_model_vocab_is_unknown() {
879
+ lower_sampling_params_with_limits(
880
+ SamplingParams {
881
+ logit_bias: Some(HashMap::from([(1_000_000, 1.0)])),
882
+ ..Default::default()
883
+ },
884
+ SamplingLimits {
885
+ model_vocab_size: None,
886
+ ..sample_sampling_limits()
887
+ },
888
+ )
889
+ .expect("logit_bias range check is skipped without model vocab size");
890
+ }
891
+
892
#[test]
893
fn lower_sampling_params_uses_generation_defaults_when_user_omits_values() {
894
let params = lower_sampling_params(
rust/src/text/src/lower/logprobs.rs
+3
-25
@@ -26,14 +26,6 @@ pub enum LogprobsError {
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
- },
29
#[error(
30
"when both logprobs and logprob_token_ids are set, logprobs must equal \
31
len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}."
@@ -41,8 +33,7 @@ pub enum LogprobsError {
33
TokenIdsMismatch { logprobs: i32, num_token_ids: usize },
34
}
35
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.
36
+/// Validate logprobs count sampling parameters.
37
pub(super) fn validate_logprobs(
38
logprobs: Option<i32>,
39
prompt_logprobs: Option<i32>,
@@ -55,7 +46,7 @@ pub(super) fn validate_logprobs(
46
47
validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?;
48
validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?;
58
- validate_logprob_token_ids(logprobs, logprob_token_ids, vocab_size)
49
+ validate_logprob_token_ids(logprobs, logprob_token_ids)
50
}
51
52
fn validate_logprobs_count(
@@ -80,10 +71,9 @@ fn validate_logprobs_count(
71
Ok(())
72
}
73
83
-fn validate_logprob_token_ids(
74
+pub(super) fn validate_logprob_token_ids(
75
logprobs: Option<i32>,
76
logprob_token_ids: Option<&[u32]>,
86
- vocab_size: usize,
77
) -> Result<(), LogprobsError> {
78
let Some(logprob_token_ids) = logprob_token_ids else {
79
return Ok(());
@@ -97,18 +87,6 @@ fn validate_logprob_token_ids(
87
});
88
}
89
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
-
90
if let Some(logprobs) = logprobs
91
&& logprobs != n as i32
92
{
rust/src/text/src/lower/token_ids.rs
new
+101
@@ -0,0 +1,101 @@
1
+use std::result::Result;
2
+
3
+use thiserror::Error;
4
+use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
5
+
6
+use crate::SamplingLimits;
7
+
8
+#[derive(Debug, Error)]
9
+#[error(
10
+ "token_id(s) {token_ids:?} in {parameter} contain out-of-vocab token ids. \
11
+ Vocabulary size: {vocab_size}"
12
+)]
13
+pub struct OutOfVocabError {
14
+ pub parameter: &'static str,
15
+ pub token_ids: Vec<u32>,
16
+ pub vocab_size: usize,
17
+}
18
+
19
+fn validate_param(
20
+ parameter: &'static str,
21
+ token_ids: impl IntoIterator<Item = u32>,
22
+ vocab_size: usize,
23
+) -> Result<(), OutOfVocabError> {
24
+ let invalid_token_ids: Vec<_> = token_ids
25
+ .into_iter()
26
+ .filter(|&token_id| token_id as usize >= vocab_size)
27
+ .collect();
28
+ if invalid_token_ids.is_empty() {
29
+ return Ok(());
30
+ }
31
+
32
+ Err(OutOfVocabError {
33
+ parameter,
34
+ token_ids: invalid_token_ids,
35
+ vocab_size,
36
+ })
37
+}
38
+
39
+/// Validate that pre-tokenized prompt IDs are within the engine-visible prompt
40
+/// vocabulary range.
41
+pub(crate) fn validate_prompt_token_ids(
42
+ prompt_token_ids: &[u32],
43
+ limits: &SamplingLimits,
44
+) -> Result<(), OutOfVocabError> {
45
+ validate_param(
46
+ "prompt",
47
+ prompt_token_ids.iter().copied(),
48
+ limits.prompt_token_vocab_size(),
49
+ )
50
+}
51
+
52
+/// Validate that token IDs in text sampling parameters are within their
53
+/// parameter-specific vocabulary ranges.
54
+pub(crate) fn validate_vocab_range(
55
+ params: &EngineCoreSamplingParams,
56
+ limits: &SamplingLimits,
57
+) -> Result<(), OutOfVocabError> {
58
+ validate_param(
59
+ "stop_token_ids",
60
+ params.stop_token_ids.iter().copied(),
61
+ limits.stop_token_vocab_size(),
62
+ )?;
63
+
64
+ if let Some(token_ids) = params.allowed_token_ids.as_deref() {
65
+ validate_param(
66
+ "allowed_token_ids",
67
+ token_ids.iter().copied(),
68
+ limits.tokenizer_vocab_size,
69
+ )?;
70
+ }
71
+
72
+ if let (Some(logit_bias), Some(vocab_size)) =
73
+ (params.logit_bias.as_ref(), limits.model_vocab_size)
74
+ {
75
+ validate_param("logit_bias", logit_bias.keys().copied(), vocab_size)?;
76
+ }
77
+
78
+ if let Some(token_ids) = params.logprob_token_ids.as_deref() {
79
+ validate_param(
80
+ "logprob_token_ids",
81
+ token_ids.iter().copied(),
82
+ limits.logprobs_vocab_size(),
83
+ )?;
84
+ }
85
+
86
+ Ok(())
87
+}
88
+
89
+#[cfg(test)]
90
+mod tests {
91
+ use super::*;
92
+
93
+ #[test]
94
+ fn validate_vocab_range_rejects_out_of_vocab_ids() {
95
+ let error = validate_param("logprob_token_ids", [5_u32, 1000, 1001], 1000).unwrap_err();
96
+
97
+ assert_eq!(error.parameter, "logprob_token_ids");
98
+ assert_eq!(error.token_ids, vec![1000, 1001]);
99
+ assert_eq!(error.vocab_size, 1000);
100
+ }
101
+}