A high-throughput and memory-efficient inference and serving engine for LLMs
[Rust Frontend] Validate tokenized bad_words vocabulary range (#45876)
Signed-off-by: reidliu41 <reid201711@gmail.com>
Reid committed
Jun 18, 2026 at 10:28 UTC
731fb3323d5c42f0a6fe2843084b782a7f7bf035
2 files changed
+61
rust/src/text/src/lower.rs
+53
@@ -288,6 +288,32 @@ mod tests {
288
StubTokenizer
289
}
290
291
+ struct FixedTokenizer {
292
+ token_ids: Vec<u32>,
293
+ }
294
+
295
+ impl Tokenizer for FixedTokenizer {
296
+ fn encode(
297
+ &self,
298
+ _text: &str,
299
+ _add_special_tokens: bool,
300
+ ) -> vllm_tokenizer::Result<Vec<u32>> {
301
+ Ok(self.token_ids.clone())
302
+ }
303
+
304
+ fn decode(
305
+ &self,
306
+ _token_ids: &[u32],
307
+ _skip_special_tokens: bool,
308
+ ) -> vllm_tokenizer::Result<String> {
309
+ Ok(String::new())
310
+ }
311
+
312
+ fn token_to_id(&self, _token: &str) -> Option<u32> {
313
+ None
314
+ }
315
+ }
316
+
317
fn sample_request() -> TextRequest {
318
TextRequest {
319
prompt: Prompt::TokenIds(vec![1, 2, 3]),
@@ -827,6 +853,33 @@ mod tests {
853
));
854
}
855
856
+ #[test]
857
+ fn lower_sampling_params_rejects_out_of_vocab_bad_words() {
858
+ let tokenizer = FixedTokenizer {
859
+ token_ids: vec![1999, 2000],
860
+ };
861
+ let error = lower_sampling_params(
862
+ SamplingParams {
863
+ bad_words: Some(vec!["blocked".to_string()]),
864
+ ..Default::default()
865
+ },
866
+ SamplingHints::default(),
867
+ sample_sampling_limits(),
868
+ 3,
869
+ &tokenizer,
870
+ )
871
+ .unwrap_err();
872
+
873
+ assert!(matches!(
874
+ error,
875
+ Error::OutOfVocab(OutOfVocabError {
876
+ parameter: "bad_words",
877
+ token_ids,
878
+ vocab_size: 2000,
879
+ }) if token_ids == vec![2000]
880
+ ));
881
+ }
882
+
883
#[test]
884
fn lower_sampling_params_rejects_out_of_vocab_logit_bias() {
885
let error = lower_sampling_params_with_limits(
rust/src/text/src/lower/token_ids.rs
+8
@@ -85,6 +85,14 @@ pub(crate) fn validate_vocab_range(
85
)?;
86
}
87
88
+ if let Some(bad_words_token_ids) = params.bad_words_token_ids.as_deref() {
89
+ validate_param(
90
+ "bad_words",
91
+ bad_words_token_ids.iter().flatten().copied(),
92
+ limits.tokenizer_vocab_size,
93
+ )?;
94
+ }
95
+
96
Ok(())
97
}
98