diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index 40a0bd44acb6..4ae24f0a3826 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -303,17 +303,17 @@ mod tests { } #[test] - fn truncate_exceeds_budget_maps_to_invalid_request() { - let error = vllm_text::Error::TruncatePromptTokensExceedsBudget { + fn truncate_exceeds_max_model_len_maps_to_invalid_request() { + let error = vllm_text::Error::TruncatePromptTokensExceedsMaxModelLen { value: 8000, - budget: 4000, + max_model_len: 4000, }; let api_error = text_submit_error("failed to submit completion request", error); assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); let response = api_error.to_error_response(); assert_eq!(response.error.error_type, "invalid_request_error"); assert!(response.error.message.contains("truncate_prompt_tokens=8000")); - assert!(response.error.message.contains("max_tokens = 4000")); + assert!(response.error.message.contains("max_model_len = 4000")); } #[test] diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index c3e8ccc11d5d..2d8d4fb0ae08 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -36,10 +36,10 @@ pub enum Error { #[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")] InvalidThinkingTokenBudget, #[error( - "truncate_prompt_tokens={value} cannot be greater than max_model_len - \ - max_tokens = {budget}. Please request a smaller truncation size." + "truncate_prompt_tokens={value} cannot be greater than max_model_len = \ + {max_model_len}. Please request a smaller truncation size." )] - TruncatePromptTokensExceedsBudget { value: u64, budget: u32 }, + TruncatePromptTokensExceedsMaxModelLen { value: u64, max_model_len: u32 }, #[error("invalid truncate_prompt_tokens={value}; must be >= -1")] InvalidTruncatePromptTokens { value: i64 }, #[error("truncate_prompt_tokens is not supported for multimodal requests")] @@ -68,7 +68,7 @@ impl Error { | Self::SamplingParams(_) | Self::MinTokensExceedsMaxTokens { .. } | Self::InvalidThinkingTokenBudget - | Self::TruncatePromptTokensExceedsBudget { .. } + | Self::TruncatePromptTokensExceedsMaxModelLen { .. } | Self::InvalidTruncatePromptTokens { .. } | Self::TruncateUnsupportedWithMultimodal | Self::InvalidRepetitionDetection { .. } diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index a6aed592b795..d20d7bcd24b9 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -88,7 +88,6 @@ impl TextRequestProcessor { prompt: Prompt, add_special_tokens: bool, prompt_truncation: Option, - max_output_tokens: Option, ) -> Result> { let tokenizer = self.backend.tokenizer(); let mut prompt_token_ids = match prompt { @@ -98,8 +97,7 @@ impl TextRequestProcessor { Prompt::TokenIds(token_ids) => token_ids, }; if let Some(prompt_truncation) = prompt_truncation { - let input_budget = self.max_model_len.saturating_sub(max_output_tokens.unwrap_or(0)); - prompt_truncation.apply(&mut prompt_token_ids, input_budget)?; + prompt_truncation.apply(&mut prompt_token_ids, self.max_model_len)?; } Ok(prompt_token_ids) } @@ -116,7 +114,6 @@ impl TextRequestProcessor { take(&mut request.prompt), request.add_special_tokens, request.prompt_truncation.take(), - request.sampling_params.max_tokens, )?; request.prompt = Prompt::TokenIds(prompt_token_ids); match &request.prompt { @@ -132,7 +129,6 @@ impl TextRequestProcessor { request.prompt, request.add_special_tokens, request.prompt_truncation, - request.sampling_params.max_tokens, ) } @@ -148,7 +144,6 @@ impl TextRequestProcessor { take(&mut request.prompt), request.add_special_tokens, request.prompt_truncation, - request.sampling_params.max_tokens, )?; let tokenizer = self.backend.tokenizer(); let sampling_hints = self.backend.sampling_hints()?; @@ -321,4 +316,34 @@ mod tests { vec![3, 4, 5] ); } + + /// `truncate_prompt_tokens=-1` keeps `max_model_len - 1` prompt tokens + /// whether or not `max_tokens` is set, and generation gets the one token + /// left. Without `max_tokens`, truncating to the full context used to fail + /// the request as too long. + #[test] + fn auto_truncation_leaves_one_output_token() { + let processor = TextRequestProcessor::new(Arc::new(TestTextBackend), 10); + for max_tokens in [None, Some(5)] { + let request = TextRequest { + prompt: Prompt::TokenIds((0..15).collect()), + prompt_truncation: Some(PromptTruncation { + limit: PromptTruncationLimit::InputBudget, + side: TruncationSide::Left, + }), + sampling_params: SamplingParams { + max_tokens, + ..SamplingParams::default() + }, + ..TextRequest::for_test() + }; + + let generate_request = processor.prepare(request).unwrap().generate_request; + assert_eq!( + generate_request.prompt_token_ids, + (6..15).collect::>() + ); + assert_eq!(generate_request.sampling_params.max_tokens, 1); + } + } } diff --git a/rust/src/text/src/truncation.rs b/rust/src/text/src/truncation.rs index 91a40ae2234c..c238a277043e 100644 --- a/rust/src/text/src/truncation.rs +++ b/rust/src/text/src/truncation.rs @@ -18,7 +18,11 @@ pub enum TruncationSide { /// Maximum number of prompt tokens to retain. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum PromptTruncationLimit { - /// Use the input budget remaining after reserving output tokens. + /// Retain `max_model_len - 1` tokens, holding one token of context back + /// so generation always has room for at least one output token. + /// + /// Python's `TokenizeParams.max_truncation_tokens`; `max_tokens` is an + /// upper bound that is clamped after tokenization, not a reservation. InputBudget, /// Retain at most the given number of prompt tokens. Fixed(u64), @@ -44,14 +48,14 @@ impl PromptTruncation { Ok(Self { limit, side }) } - pub(crate) fn apply(self, prompt_token_ids: &mut Vec, input_budget: u32) -> Result<()> { + pub(crate) fn apply(self, prompt_token_ids: &mut Vec, max_model_len: u32) -> Result<()> { let max_input_tokens = match self.limit { - PromptTruncationLimit::InputBudget => input_budget, + PromptTruncationLimit::InputBudget => max_model_len.saturating_sub(1), PromptTruncationLimit::Fixed(limit) => { - if limit > u64::from(input_budget) { - return Err(Error::TruncatePromptTokensExceedsBudget { + if limit > u64::from(max_model_len) { + return Err(Error::TruncatePromptTokensExceedsMaxModelLen { value: limit, - budget: input_budget, + max_model_len, }); } limit as u32 @@ -121,7 +125,7 @@ mod tests { } #[test] - fn input_budget_reserves_output_tokens() { + fn input_budget_holds_back_one_token_for_output() { let mut prompt = vec![0; 100]; PromptTruncation { limit: PromptTruncationLimit::InputBudget, @@ -129,11 +133,11 @@ mod tests { } .apply(&mut prompt, 70) .unwrap(); - assert_eq!(prompt.len(), 70); + assert_eq!(prompt.len(), 69); } #[test] - fn fixed_limit_cannot_exceed_input_budget() { + fn fixed_limit_cannot_exceed_max_model_len() { let mut prompt = vec![1, 2, 3]; let error = PromptTruncation { limit: PromptTruncationLimit::Fixed(80), @@ -143,11 +147,20 @@ mod tests { .unwrap_err(); assert!(matches!( error, - Error::TruncatePromptTokensExceedsBudget { + Error::TruncatePromptTokensExceedsMaxModelLen { value: 80, - budget: 70 + max_model_len: 70 } )); + + let mut prompt = vec![1, 2, 3]; + PromptTruncation { + limit: PromptTruncationLimit::Fixed(70), + side: TruncationSide::Left, + } + .apply(&mut prompt, 70) + .unwrap(); + assert_eq!(prompt, vec![1, 2, 3]); } #[test]