Skip to content

[Rust Frontend] Truncate prompts against max_model_len like Python #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
richard-for-bugen wants to merge 1 commit into from
Closed

[Rust Frontend] Truncate prompts against max_model_len like Python #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions rust/src/server/src/error.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
8 changes: 4 additions & 4 deletions rust/src/text/src/error.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down Expand Up @@ -68,7 +68,7 @@ impl Error {
| Self::SamplingParams(_)
| Self::MinTokensExceedsMaxTokens { .. }
| Self::InvalidThinkingTokenBudget
| Self::TruncatePromptTokensExceedsBudget { .. }
| Self::TruncatePromptTokensExceedsMaxModelLen { .. }
| Self::InvalidTruncatePromptTokens { .. }
| Self::TruncateUnsupportedWithMultimodal
| Self::InvalidRepetitionDetection { .. }
Expand Down
37 changes: 31 additions & 6 deletions rust/src/text/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,6 @@ impl TextRequestProcessor {
prompt: Prompt,
add_special_tokens: bool,
prompt_truncation: Option<PromptTruncation>,
max_output_tokens: Option<u32>,
) -> Result<Vec<u32>> {
let tokenizer = self.backend.tokenizer();
let mut prompt_token_ids = match prompt {
Expand All @@ -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)
}
Expand All @@ -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 {
Expand All @@ -132,7 +129,6 @@ impl TextRequestProcessor {
request.prompt,
request.add_special_tokens,
request.prompt_truncation,
request.sampling_params.max_tokens,
)
}

Expand All @@ -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()?;
Expand Down Expand Up @@ -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::<Vec<_>>()
);
assert_eq!(generate_request.sampling_params.max_tokens, 1);
}
}
}
35 changes: 24 additions & 11 deletions rust/src/text/src/truncation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -44,14 +48,14 @@ impl PromptTruncation {
Ok(Self { limit, side })
}

pub(crate) fn apply(self, prompt_token_ids: &mut Vec<u32>, input_budget: u32) -> Result<()> {
pub(crate) fn apply(self, prompt_token_ids: &mut Vec<u32>, 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
Expand Down Expand Up @@ -121,19 +125,19 @@ 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,
side: TruncationSide::Left,
}
.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),
Expand All @@ -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]
Expand Down