diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 9fa8c78f4f..319c995369 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -1338,6 +1338,47 @@ mod tests { )); } + #[test] + fn lower_sampling_params_enforces_bad_words_limits() { + let tokenizer = TestTokenizer::new() + .with_regular_token("blocked", 300) + .with_regular_token(" blocked", 301); + let lower = |bad_words: Option>, bad_words_token_ids: Vec>| { + lower_sampling_params( + SamplingParams { + bad_words, + bad_words_token_ids: Some(bad_words_token_ids), + ..Default::default() + }, + SamplingHints::default(), + sample_sampling_limits(), + 3, + &tokenizer, + ) + }; + + lower(None, vec![vec![7; 8]; 128]).unwrap(); + for (bad_words, raw, expected) in [ + ( + None, + vec![vec![7]; 129], + "Too many bad words: 129. The max number is 128.", + ), + ( + Some(vec!["blocked".to_string()]), + vec![vec![7]; 127], + "Too many bad words: 129. The max number is 128.", + ), + ( + None, + vec![vec![7; 1025]], + "Too many total bad word tokens: 1025. The max is 1024.", + ), + ] { + assert_eq!(lower(bad_words, raw).unwrap_err().to_string(), expected); + } + } + #[test] fn lower_sampling_params_rejects_out_of_vocab_logit_bias() { let error = lower_sampling_params_with_limits( diff --git a/rust/src/text/src/lower/token_ids.rs b/rust/src/text/src/lower/token_ids.rs index 51329f9290..b7f160b919 100644 --- a/rust/src/text/src/lower/token_ids.rs +++ b/rust/src/text/src/lower/token_ids.rs @@ -8,12 +8,21 @@ use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use crate::SamplingLimits; +const MAX_NUM_BAD_WORDS_ENV: &str = "VLLM_MAX_NUM_BAD_WORDS"; +const DEFAULT_MAX_NUM_BAD_WORDS: usize = 128; +const MAX_BAD_WORDS_TOTAL_TOKENS_ENV: &str = "VLLM_MAX_BAD_WORDS_TOTAL_TOKENS"; +const DEFAULT_MAX_BAD_WORDS_TOTAL_TOKENS: usize = 1024; + #[derive(Debug, Error)] pub enum TokenIdsError { #[error("allowed_token_ids should not be empty")] EmptyAllowedTokenIds, #[error("bad_words_token_ids must not contain empty sequences")] EmptyBadWordSequence, + #[error("Too many bad words: {count}. The max number is {max}.")] + TooManyBadWords { count: usize, max: usize }, + #[error("Too many total bad word tokens: {count}. The max is {max}.")] + TooManyBadWordTokens { count: usize, max: usize }, #[error( "token_id(s) {token_ids:?} in {parameter} are out of vocabulary. \ Vocabulary size: {vocab_size}" @@ -45,6 +54,31 @@ fn validate_param( }) } +fn env_limit(name: &str, default: usize) -> usize { + std::env::var(name).ok().and_then(|value| value.parse().ok()).unwrap_or(default) +} + +/// Mirrors Model Runner V2's per-request bad-words buffer limits: +/// +fn validate_bad_words_limits(bad_words_token_ids: &[Vec]) -> Result<(), TokenIdsError> { + let count = bad_words_token_ids.len(); + let max = env_limit(MAX_NUM_BAD_WORDS_ENV, DEFAULT_MAX_NUM_BAD_WORDS); + if count > max { + return Err(TokenIdsError::TooManyBadWords { count, max }); + } + + let count = bad_words_token_ids.iter().map(Vec::len).sum(); + let max = env_limit( + MAX_BAD_WORDS_TOTAL_TOKENS_ENV, + DEFAULT_MAX_BAD_WORDS_TOTAL_TOKENS, + ); + if count > max { + return Err(TokenIdsError::TooManyBadWordTokens { count, max }); + } + + Ok(()) +} + /// Validate that pre-tokenized prompt IDs are within the engine-visible prompt /// vocabulary range. pub(crate) fn validate_prompt_token_ids( @@ -101,6 +135,7 @@ pub(crate) fn validate_vocab_range( if bad_words_token_ids.iter().any(Vec::is_empty) { return Err(TokenIdsError::EmptyBadWordSequence); } + validate_bad_words_limits(bad_words_token_ids)?; validate_param( "bad_words", bad_words_token_ids.iter().flatten().copied(),