@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 085e9bb07e

[Rust Frontend] Add support for `truncate_prompt_tokens` and `truncation_side` (#48584)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: pranavthakur0-0 <pranavthakurdew@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Pranav Thakur committed Aug 29, 2026 at 05:46 UTC 085e9bb07eb3a1ede9689f4110b6ab35eda3e08f

17 files changed +424 -35

rust/src/chat/src/lib.rs

@@ -195,6 +195,7 @@ impl ChatRequestProcessor {

195 sampling_params: request.sampling_params,

196 decode_options: request.decode_options,

197 intermediate: request.intermediate,

198 + prompt_truncation: request.prompt_truncation,

199 priority: request.priority,

200 cache_salt: request.cache_salt,

201 add_special_tokens: request.add_special_tokens,

rust/src/chat/src/request.rs

+9 -1

@@ -8,8 +8,8 @@ use serde::{Deserialize, Serialize};

8 use serde_json::Value;

9 use vllm_engine_core_client::protocol::lora::LoraRequest;

10 pub use vllm_parser::tool::Tool as ChatTool;

11 -pub use vllm_text::SamplingParams;

11 use vllm_text::TextDecodeOptions;

12 +pub use vllm_text::{PromptTruncation, SamplingParams};

14 use crate::AssistantMessageExt;

15 use crate::error::{Error, Result};

@@ -570,6 +570,8 @@ pub struct ChatRequest {

570 /// output. If `true`, callers may receive zero or more incremental

571 /// content events before the final terminal one.

572 pub intermediate: bool,

573 + /// Prompt-truncation policy resolved by the caller.

574 + pub prompt_truncation: Option<PromptTruncation>,

575 /// Request scheduling priority (lower means earlier handling; default 0).

576 pub priority: i32,

577 /// Documents for RAG (retrieval-augmented generation), passed to the chat

@@ -601,6 +603,7 @@ impl ChatRequest {

603 tool_context: ResolvedToolContext::default(),

604 decode_options: TextDecodeOptions::default(),

605 intermediate: true,

606 + prompt_truncation: None,

607 priority: 0,

608 documents: None,

609 cache_salt: None,

@@ -627,6 +630,11 @@ impl ChatRequest {

630 (GenerationPromptMode::NoGenerationPrompt, _)

631 | (GenerationPromptMode::StartNewAssistant, _) => {}

632 }

633 + if self.has_multimodal() && self.prompt_truncation.is_some() {

634 + return Err(Error::Text(

635 + vllm_text::Error::TruncateUnsupportedWithMultimodal,

636 + ));

637 + }

638 Ok(())

639 }

640

rust/src/server/src/error.rs

+25

@@ -224,6 +224,31 @@ mod tests {

224 assert!(response.error.message.contains("allowed_token_ids"));

225 }

226

227 + #[test]

228 + fn truncate_exceeds_budget_maps_to_invalid_request() {

229 + let error = vllm_text::Error::TruncatePromptTokensExceedsBudget {

230 + value: 8000,

231 + budget: 4000,

232 + };

233 + let api_error = text_submit_error("failed to submit completion request", error);

234 + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);

235 + let response = api_error.to_error_response();

236 + assert_eq!(response.error.error_type, "invalid_request_error");

237 + assert!(response.error.message.contains("truncate_prompt_tokens=8000"));

238 + assert!(response.error.message.contains("max_tokens = 4000"));

239 + }

240 +

241 + #[test]

242 + fn invalid_truncate_prompt_tokens_maps_to_invalid_request() {

243 + let error = vllm_text::Error::InvalidTruncatePromptTokens { value: -2 };

244 + let api_error = text_submit_error("failed to submit completion request", error);

245 + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);

246 + let response = api_error.to_error_response();

247 + assert_eq!(response.error.error_type, "invalid_request_error");

248 + assert!(response.error.message.contains("-2"));

249 + assert!(response.error.message.contains("must be >= -1"));

250 + }

251 +

252 #[test]

253 fn other_submit_errors_stay_internal() {

254 let error = vllm_text::Error::Tokenizer("backend exploded".to_string());

rust/src/server/src/grpc/convert.rs

+9 -7

@@ -11,8 +11,8 @@ use vllm_chat::MediaContentPart;

11 use vllm_engine_core_client::protocol::output::StopReason;

12 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;

13 use vllm_text::{

14 - DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams,

15 - TextDecodeOptions, TextRequest,

14 + DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, PromptTruncation,

15 + PromptTruncationLimit, SamplingParams, TextDecodeOptions, TextRequest, TruncationSide,

16 };

18 use super::pb;

@@ -110,11 +110,12 @@ pub fn to_text_request(

110 )));

111 }

112

113 - if req.truncate_prompt_tokens != 0 {

114 - return Err(Status::invalid_argument(

115 - "truncate_prompt_tokens is not supported",

116 - ));

117 - }

113 + // Proto3 uses zero as unset; positive values select fixed left truncation.

114 + // The -1 input-budget sentinel is outside the uint32 field domain.

115 + let prompt_truncation = (req.truncate_prompt_tokens != 0).then_some(PromptTruncation {

116 + limit: PromptTruncationLimit::Fixed(u64::from(req.truncate_prompt_tokens)),

117 + side: TruncationSide::Left,

118 + });

119

120 let prompt = match req.prompt {

121 Some(pb::generate_request::Prompt::Text(text)) => Prompt::Text(text),

@@ -173,6 +174,7 @@ pub fn to_text_request(

174 sampling_params,

175 decode_options,

176 intermediate: stream,

177 + prompt_truncation,

178 priority: req.priority,

179 cache_salt: kv.map(|k| &k.cache_salt).filter(|s| !s.is_empty()).cloned(),

180 add_special_tokens: true,

rust/src/server/src/routes/inference/generate/convert.rs

@@ -73,6 +73,7 @@ pub(super) fn prepare_generate_request(

73 sampling_params,

74 decode_options: TextDecodeOptions::default(),

75 intermediate: false,

76 + prompt_truncation: None,

77 priority: request.priority,

78 cache_salt: request.cache_salt,

79 add_special_tokens: false,

rust/src/server/src/routes/openai/chat_completions/convert.rs

+33 -2

@@ -10,8 +10,9 @@ use vllm_chat::{

11 use super::types::ChatCompletionRequest;

12 use super::validate;

13 -use crate::error::{ApiError, bail_invalid_request, chat_submit_error};

13 +use crate::error::{ApiError, bail_invalid_request, chat_submit_error, text_submit_error};

14 use crate::lora::LoraModelResolution;

15 +use crate::routes::openai::utils::resolve_generation_prompt_truncation;

16 use crate::routes::openai::utils::structured_outputs::convert_from_response_format;

17 use crate::routes::openai::utils::types::{

18 ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue, ToolReference,

@@ -71,6 +72,12 @@ pub(super) fn prepare_chat_request(

72 ) -> Result<PreparedRequest, ApiError> {

73 validate::validate_request_compat(&request, &lora_resolution.model_names)?;

75 + let prompt_truncation = resolve_generation_prompt_truncation(

76 + request.truncate_prompt_tokens,

77 + request.truncation_side,

78 + )

79 + .map_err(|error| text_submit_error("invalid prompt truncation", error))?;

80 +

81 let request_id = format!("chatcmpl-{}", ctx.request_id);

82 let response_model = lora_resolution

83 .lora_request

@@ -184,6 +191,7 @@ pub(super) fn prepare_chat_request(

191 min_tokens: request.min_tokens.unwrap_or(0),

192 },

193 intermediate: request.stream,

194 + prompt_truncation,

195 priority: ctx.priority.or(request.priority).unwrap_or(0),

196 documents: request.documents,

197 cache_salt: request.cache_salt,

@@ -439,7 +447,9 @@ mod tests {

447 ChatRenderer, ChatTool as VllmChatTool, ChatToolChoice, GenerationPromptMode,

448 KimiK3ChatRenderer, SamplingParams as VllmSamplingParams,

449 };

442 - use vllm_text::{Prompt, output::TextDecodeOptions};

450 + use vllm_text::{

451 + Prompt, PromptTruncation, PromptTruncationLimit, TruncationSide, output::TextDecodeOptions,

452 + };

453 use vllm_tokenizer::{Tokenizer, test_utils::TestTokenizer};

454

455 use super::prepare_chat_request;

@@ -530,6 +540,27 @@ mod tests {

540 }

541 }

542

543 + #[test]

544 + fn prepare_chat_request_preserves_explicit_truncation_side() {

545 + let mut request = base_request();

546 + request.truncate_prompt_tokens = Some(2);

547 + request.truncation_side = Some(TruncationSide::Right);

548 + let prepared = prepare_chat_request(

549 + request,

550 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

551 + ResolvedRequestContext::default(),

552 + )

553 + .expect("prepare truncation");

554 +

555 + assert_eq!(

556 + prepared.chat_request.prompt_truncation,

557 + Some(PromptTruncation {

558 + limit: PromptTruncationLimit::Fixed(2),

559 + side: TruncationSide::Right,

560 + })

561 + );

562 + }

563 +

564 #[test]

565 fn prepare_chat_request_accepts_zero_min_tokens() {

566 let request: ChatCompletionRequest = serde_json::from_value(json!({

rust/src/server/src/routes/openai/chat_completions/types.rs

@@ -11,6 +11,7 @@ use serde_with::SerializeDisplay;

11 use validator::Validate;

12 use vllm_chat::ReasoningEffort;

13 use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams;

14 +use vllm_text::TruncationSide;

16 use crate::routes::openai::utils::structured_outputs::ResponseFormat;

17 use crate::routes::openai::utils::types::{

@@ -159,6 +160,9 @@ pub struct ChatCompletionRequest {

160 /// Truncate prompt tokens to this length

161 pub truncate_prompt_tokens: Option<i64>,

162

163 + /// Which side to truncate from when truncate_prompt_tokens is active

164 + pub truncation_side: Option<TruncationSide>,

165 +

166 /// Number of prompt logprobs to return

167 pub prompt_logprobs: Option<i32>,

168

@@ -290,6 +294,7 @@ impl Default for ChatCompletionRequest {

294 skip_special_tokens: true,

295 spaces_between_special_tokens: true,

296 truncate_prompt_tokens: None,

297 + truncation_side: None,

298 prompt_logprobs: None,

299 allowed_token_ids: None,

300 bad_words: None,

rust/src/server/src/routes/openai/chat_completions/validate.rs

+2 -5

@@ -4,6 +4,7 @@

4 use super::types::ChatCompletionRequest;

5 use crate::error::{ApiError, bail_invalid_request};

6 use crate::routes::openai::utils::types::{ChatMessage, Tool};

7 +use crate::routes::openai::utils::validate_generation_prompt_truncation;

9 /// Enforce the minimal compatibility contract for the Rust OpenAI server.

10 pub(super) fn validate_request_compat(

@@ -84,11 +85,7 @@ pub(super) fn validate_request_compat(

85 "spaces_between_special_tokens is not supported."

86 );

87 }

87 - reject_non_default(

88 - request.truncate_prompt_tokens.as_ref(),

89 - "truncate_prompt_tokens",

90 - "truncate_prompt_tokens is not supported.",

91 - )?;

88 + validate_generation_prompt_truncation(request.truncate_prompt_tokens, request.echo)?;

89 reject_non_default(

90 request.media_io_kwargs.as_ref(),

91 "media_io_kwargs",

rust/src/server/src/routes/openai/completions/convert.rs

+35 -2

@@ -6,9 +6,10 @@ use vllm_text::tokenizer::Tokenizer;

6 use vllm_text::{Prompt, SamplingParams, TextDecodeOptions, TextRequest};

8 use super::types::CompletionRequest;

9 -use crate::error::ApiError;

9 +use crate::error::{ApiError, text_submit_error};

10 use crate::lora::LoraModelResolution;

11 use crate::routes::openai::completions::validate;

12 +use crate::routes::openai::utils::resolve_generation_prompt_truncation;

13 use crate::routes::openai::utils::structured_outputs::convert_from_response_format_value;

14 use crate::utils::{

15 ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params,

@@ -62,6 +63,12 @@ pub(super) fn prepare_completion_request(

63 ) -> Result<PreparedRequest, ApiError> {

64 validate::validate_request_compat(&request, &lora_resolution.model_names)?;

66 + let prompt_truncation = resolve_generation_prompt_truncation(

67 + request.truncate_prompt_tokens,

68 + request.truncation_side,

69 + )

70 + .map_err(|error| text_submit_error("invalid prompt truncation", error))?;

71 +

72 let request_id = format!("cmpl-{}", ctx.request_id);

73 let response_model = lora_resolution

74 .lora_request

@@ -141,6 +148,7 @@ pub(super) fn prepare_completion_request(

148 min_tokens: request.min_tokens.unwrap_or(0),

149 },

150 intermediate: request.stream,

151 + prompt_truncation,

152 priority: ctx.priority.or(request.priority).unwrap_or(0),

153 cache_salt: request.cache_salt,

154 add_special_tokens: request.add_special_tokens,

@@ -200,7 +208,7 @@ mod tests {

208 use axum::http::HeaderMap;

209 use serde_json::json;

210 use validator::Validate;

203 - use vllm_text::Prompt;

211 + use vllm_text::{Prompt, PromptTruncation, PromptTruncationLimit, TruncationSide};

212 use vllm_tokenizer::test_utils::TestTokenizer;

213

214 use super::prepare_completion_request;

@@ -323,6 +331,31 @@ mod tests {

331 }

332 }

333

334 + #[test]

335 + fn prepare_completion_request_uses_left_when_side_is_omitted() {

336 + let mut value = base_request_json();

337 + value

338 + .as_object_mut()

339 + .expect("request object")

340 + .insert("truncate_prompt_tokens".to_string(), json!(-1));

341 + let request = serde_json::from_value(value).expect("parse truncation");

342 + let prepared = prepare_completion_request(

343 + request,

344 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

345 + ResolvedRequestContext::default(),

346 + &test_tokenizer(),

347 + )

348 + .expect("prepare truncation");

349 +

350 + assert_eq!(

351 + prepared.text_request.prompt_truncation,

352 + Some(PromptTruncation {

353 + limit: PromptTruncationLimit::InputBudget,

354 + side: TruncationSide::Left,

355 + })

356 + );

357 + }

358 +

359 #[test]

360 fn normalize_coerces_null_max_tokens_to_default() {

361 // An absent `max_tokens` already gets the serde default.

rust/src/server/src/routes/openai/completions/types.rs

+4 -1

@@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize};

8 use serde_json::{Map, Value};

9 use validator::Validate;

10 use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams;

11 -use vllm_text::Prompt;

11 +use vllm_text::{Prompt, TruncationSide};

13 use crate::routes::openai::utils::types::{

14 LogProbs, Normalizable, PromptLogprobs, StreamOptions, StreamResponseEnvelope, StringOrArray,

@@ -138,6 +138,9 @@ pub struct CompletionRequest {

138 /// Truncate prompt tokens to this length

139 pub truncate_prompt_tokens: Option<i64>,

140

141 + /// Which side to truncate from when truncate_prompt_tokens is active

142 + pub truncation_side: Option<TruncationSide>,

143 +

144 /// Restrict output to these token IDs only

145 pub allowed_token_ids: Option<Vec<u32>>,

146

rust/src/server/src/routes/openai/completions/validate.rs

+2 -6

@@ -3,6 +3,7 @@

4 use super::types::CompletionRequest;

5 use crate::error::{ApiError, bail_invalid_request};

6 +use crate::routes::openai::utils::validate_generation_prompt_truncation;

8 /// Enforce the minimal compatibility contract for the Rust OpenAI server.

9 pub(super) fn validate_request_compat(

@@ -85,12 +86,7 @@ pub(super) fn validate_request_compat(

86 "spaces_between_special_tokens is not supported."

87 );

88 }

88 - if request.truncate_prompt_tokens.is_some() {

89 - bail_invalid_request!(

90 - param = "truncate_prompt_tokens",

91 - "truncate_prompt_tokens is not supported."

92 - );

93 - }

89 + validate_generation_prompt_truncation(request.truncate_prompt_tokens, request.echo)?;

91 Ok(())

92 }

rust/src/server/src/routes/openai/utils/mod.rs

+34

@@ -6,3 +6,37 @@ pub mod structured_outputs;

6 pub mod types;

7 pub mod usage;

8 pub mod validated_json;

9 +

10 +use vllm_text::{PromptTruncation, TruncationSide};

11 +

12 +use crate::error::{ApiError, bail_invalid_request};

13 +

14 +pub(crate) fn validate_generation_prompt_truncation(

15 + limit: Option<i64>,

16 + echo: bool,

17 +) -> Result<(), ApiError> {

18 + if let Some(limit) = limit

19 + && limit < -1

20 + {

21 + bail_invalid_request!(

22 + param = "truncate_prompt_tokens",

23 + "truncate_prompt_tokens must be >= -1."

24 + );

25 + }

26 + if echo && limit.is_some() {

27 + bail_invalid_request!(

28 + param = "echo",

29 + "`echo=true` is not supported with `truncate_prompt_tokens`."

30 + );

31 + }

32 + Ok(())

33 +}

34 +

35 +pub(crate) fn resolve_generation_prompt_truncation(

36 + limit: Option<i64>,

37 + side: Option<TruncationSide>,

38 +) -> vllm_text::Result<Option<PromptTruncation>> {

39 + limit

40 + .map(|limit| PromptTruncation::from_wire(limit, side.unwrap_or(TruncationSide::Left)))

41 + .transpose()

42 +}

rust/src/server/src/routes/tokenize/types.rs

@@ -50,6 +50,7 @@ impl TokenizeCompletionRequest {

50 intermediate: false,

51 priority: 0,

52 cache_salt: None,

53 + prompt_truncation: None,

54 add_special_tokens: self.add_special_tokens,

55 data_parallel_rank: None,

56 session_id: None,

@@ -116,6 +117,7 @@ impl TokenizeChatRequest {

117 tool_context,

118 decode_options: TextDecodeOptions::default(),

119 intermediate: false,

120 + prompt_truncation: None,

121 priority: 0,

122 documents: None,

123 cache_salt: None,

rust/src/text/src/error.rs

+12

@@ -35,6 +35,15 @@ pub enum Error {

35 MinTokensExceedsMaxTokens { min_tokens: u32, max_tokens: u32 },

36 #[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")]

37 InvalidThinkingTokenBudget,

38 + #[error(

39 + "truncate_prompt_tokens={value} cannot be greater than max_model_len - \

40 + max_tokens = {budget}. Please request a smaller truncation size."

41 + )]

42 + TruncatePromptTokensExceedsBudget { value: u64, budget: u32 },

43 + #[error("invalid truncate_prompt_tokens={value}; must be >= -1")]

44 + InvalidTruncatePromptTokens { value: i64 },

45 + #[error("truncate_prompt_tokens is not supported for multimodal requests")]

46 + TruncateUnsupportedWithMultimodal,

47 #[error("invalid repetition detection params: {message}")]

48 InvalidRepetitionDetection { message: String },

49 #[error("text request stream `{request_id}` closed before terminal output")]

@@ -59,6 +68,9 @@ impl Error {

68 | Self::SamplingParams(_)

69 | Self::MinTokensExceedsMaxTokens { .. }

70 | Self::InvalidThinkingTokenBudget

71 + | Self::TruncatePromptTokensExceedsBudget { .. }

72 + | Self::InvalidTruncatePromptTokens { .. }

73 + | Self::TruncateUnsupportedWithMultimodal

74 | Self::InvalidRepetitionDetection { .. }

75 // An empty tokenized prompt detected later, at request prepare

76 // time, surfaces through the transparent Llm wrapper.

rust/src/text/src/lib.rs

+64 -11

@@ -33,6 +33,8 @@ mod error;

33 mod lower;

34 pub mod output;

35 mod request;

36 +mod truncation;

37 +pub use truncation::{PromptTruncation, PromptTruncationLimit, TruncationSide};

38 pub use vllm_tokenizer as tokenizer;

40 trait_set! {

@@ -80,26 +82,35 @@ impl TextRequestProcessor {

82 self.max_model_len

83 }

83 - fn tokenize_prompt(

84 - tokenizer: &DynTokenizer,

85 + fn prepare_prompt_tokens(

86 + &self,

87 prompt: Prompt,

88 add_special_tokens: bool,

89 + prompt_truncation: Option<PromptTruncation>,

90 + max_output_tokens: Option<u32>,

91 ) -> Result<Vec<u32>> {

88 - match prompt {

89 - Prompt::Text(text) => tokenizer.encode(&text, add_special_tokens).map_err(Into::into),

92 + let tokenizer = self.backend.tokenizer();

93 + let mut prompt_token_ids = match prompt {

94 + Prompt::Text(text) => tokenizer.encode(&text, add_special_tokens)?,

95 // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark

96 // and infra workloads bypass chat rendering and tokenizer overhead entirely.

92 - Prompt::TokenIds(token_ids) => Ok(token_ids),

97 + Prompt::TokenIds(token_ids) => token_ids,

98 + };

99 + if let Some(prompt_truncation) = prompt_truncation {

100 + let input_budget = self.max_model_len.saturating_sub(max_output_tokens.unwrap_or(0));

101 + prompt_truncation.apply(&mut prompt_token_ids, input_budget)?;

102 }

103 + Ok(prompt_token_ids)

104 }

105

96 - /// Tokenize one request without generation-specific lowering.

106 + /// Prepare one request's prompt tokens without generation-specific lowering.

107 pub fn tokenize(&self, request: TextRequest) -> Result<Vec<u32>> {

108 request.validate()?;

99 - Self::tokenize_prompt(

100 - &self.backend.tokenizer(),

109 + self.prepare_prompt_tokens(

110 request.prompt,

111 request.add_special_tokens,

112 + request.prompt_truncation,

113 + request.sampling_params.max_tokens,

114 )

115 }

116

@@ -111,12 +122,13 @@ impl TextRequestProcessor {

122 request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs());

123 }

124

114 - let tokenizer = self.backend.tokenizer();

115 - let prompt_token_ids = Self::tokenize_prompt(

116 - &tokenizer,

125 + let prompt_token_ids = self.prepare_prompt_tokens(

126 take(&mut request.prompt),

127 request.add_special_tokens,

128 + request.prompt_truncation,

129 + request.sampling_params.max_tokens,

130 )?;

131 + let tokenizer = self.backend.tokenizer();

132 let sampling_hints = self.backend.sampling_hints()?;

133 let sampling_limits = SamplingLimits {

134 max_model_len: self.max_model_len,

@@ -247,3 +259,44 @@ impl TextLlm {

259 Ok(())

260 }

261 }

262 +

263 +#[cfg(test)]

264 +mod tests {

265 + use std::sync::Arc;

266 +

267 + use vllm_tokenizer::test_utils::TestTokenizer;

268 +

269 + use super::*;

270 +

271 + #[derive(Debug)]

272 + struct TestTextBackend;

273 +

274 + impl TextBackend for TestTextBackend {

275 + fn tokenizer(&self) -> DynTokenizer {

276 + Arc::new(TestTokenizer::new())

277 + }

278 +

279 + fn model_id(&self) -> &str {

280 + "test-model"

281 + }

282 + }

283 +

284 + #[test]

285 + fn tokenize_and_prepare_apply_the_same_prompt_truncation() {

286 + let processor = TextRequestProcessor::new(Arc::new(TestTextBackend), 10);

287 + let request = TextRequest {

288 + prompt: Prompt::TokenIds(vec![1, 2, 3, 4, 5]),

289 + prompt_truncation: Some(PromptTruncation {

290 + limit: PromptTruncationLimit::Fixed(3),

291 + side: TruncationSide::Left,

292 + }),

293 + ..TextRequest::for_test()

294 + };

295 +

296 + assert_eq!(processor.tokenize(request.clone()).unwrap(), vec![3, 4, 5]);

297 + assert_eq!(

298 + processor.prepare(request).unwrap().generate_request.prompt_token_ids,

299 + vec![3, 4, 5]

300 + );

301 + }

302 +}

rust/src/text/src/request.rs

@@ -14,6 +14,7 @@ use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsPara

15 use crate::error::{Error, Result};

16 use crate::output::TextDecodeOptions;

17 +use crate::truncation::PromptTruncation;

19 /// Normalize vLLM's signed `top_k` domain into [`SamplingParams::top_k`].

20 ///

@@ -196,6 +197,9 @@ pub struct TextRequest {

197 pub priority: i32,

198 /// Salt for prefix cache isolation in multi-user environments.

199 pub cache_salt: Option<String>,

200 + /// Prompt-truncation policy resolved by the caller.

201 + #[serde(default)]

202 + pub prompt_truncation: Option<PromptTruncation>,

203 /// Whether to add special tokens (e.g. BOS) during prompt tokenization.

204 pub add_special_tokens: bool,

205 /// Override data parallel rank.

@@ -231,6 +235,7 @@ impl TextRequest {

235 intermediate: true,

236 priority: 0,

237 cache_salt: None,

238 + prompt_truncation: None,

239 add_special_tokens: false,

240 data_parallel_rank: None,

241 session_id: None,

@@ -257,6 +262,9 @@ impl TextRequest {

262 request_id: self.request_id.clone(),

263 });

264 }

265 + if self.mm_features.is_some() && self.prompt_truncation.is_some() {

266 + return Err(Error::TruncateUnsupportedWithMultimodal);

267 + }

268 Ok(())

269 }

270 }

rust/src/text/src/truncation.rs new

+178

@@ -0,0 +1,178 @@

1 +// SPDX-License-Identifier: Apache-2.0

2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +use serde::{Deserialize, Serialize};

5 +

6 +use crate::error::{Error, Result};

7 +

8 +/// Which side of the prompt to discard when truncation is active.

9 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]

10 +#[serde(rename_all = "lowercase")]

11 +pub enum TruncationSide {

12 + /// Keep the last N tokens, discarding the prompt prefix.

13 + Left,

14 + /// Keep the first N tokens, discarding the prompt suffix.

15 + Right,

16 +}

17 +

18 +/// Maximum number of prompt tokens to retain.

19 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]

20 +pub enum PromptTruncationLimit {

21 + /// Use the input budget remaining after reserving output tokens.

22 + InputBudget,

23 + /// Retain at most the given number of prompt tokens.

24 + Fixed(u64),

25 +}

26 +

27 +/// Typed prompt-truncation policy used inside the text frontend.

28 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]

29 +pub struct PromptTruncation {

30 + /// Maximum number of prompt tokens to retain.

31 + pub limit: PromptTruncationLimit,

32 + /// Side from which excess prompt tokens are discarded.

33 + pub side: TruncationSide,

34 +}

35 +

36 +impl PromptTruncation {

37 + /// Convert the OpenAI-compatible integer representation into a typed policy.

38 + pub fn from_wire(limit: i64, side: TruncationSide) -> Result<Self> {

39 + let limit = match limit {

40 + -1 => PromptTruncationLimit::InputBudget,

41 + 0.. => PromptTruncationLimit::Fixed(limit as u64),

42 + _ => return Err(Error::InvalidTruncatePromptTokens { value: limit }),

43 + };

44 + Ok(Self { limit, side })

45 + }

46 +

47 + pub(crate) fn apply(self, prompt_token_ids: &mut Vec<u32>, input_budget: u32) -> Result<()> {

48 + let max_input_tokens = match self.limit {

49 + PromptTruncationLimit::InputBudget => input_budget,

50 + PromptTruncationLimit::Fixed(limit) => {

51 + if limit > u64::from(input_budget) {

52 + return Err(Error::TruncatePromptTokensExceedsBudget {

53 + value: limit,

54 + budget: input_budget,

55 + });

56 + }

57 + limit as u32

58 + }

59 + } as usize;

60 +

61 + if prompt_token_ids.len() <= max_input_tokens {

62 + return Ok(());

63 + }

64 +

65 + match self.side {

66 + TruncationSide::Left => {

67 + let start = prompt_token_ids.len() - max_input_tokens;

68 + prompt_token_ids.drain(0..start);

69 + }

70 + TruncationSide::Right => prompt_token_ids.truncate(max_input_tokens),

71 + }

72 + Ok(())

73 + }

74 +}

75 +

76 +#[cfg(test)]

77 +mod tests {

78 + use super::*;

79 +

80 + #[test]

81 + fn parses_wire_limits() {

82 + assert_eq!(

83 + PromptTruncation::from_wire(-1, TruncationSide::Left).unwrap(),

84 + PromptTruncation {

85 + limit: PromptTruncationLimit::InputBudget,

86 + side: TruncationSide::Left,

87 + }

88 + );

89 + assert_eq!(

90 + PromptTruncation::from_wire(0, TruncationSide::Right).unwrap(),

91 + PromptTruncation {

92 + limit: PromptTruncationLimit::Fixed(0),

93 + side: TruncationSide::Right,

94 + }

95 + );

96 + assert!(matches!(

97 + PromptTruncation::from_wire(-2, TruncationSide::Left),

98 + Err(Error::InvalidTruncatePromptTokens { value: -2 })

99 + ));

100 + }

101 +

102 + #[test]

103 + fn applies_fixed_limit_from_either_side() {

104 + let mut left = vec![1, 2, 3, 4, 5];

105 + PromptTruncation {

106 + limit: PromptTruncationLimit::Fixed(3),

107 + side: TruncationSide::Left,

108 + }

109 + .apply(&mut left, 90)

110 + .unwrap();

111 + assert_eq!(left, vec![3, 4, 5]);

112 +

113 + let mut right = vec![1, 2, 3, 4, 5];

114 + PromptTruncation {

115 + limit: PromptTruncationLimit::Fixed(3),

116 + side: TruncationSide::Right,

117 + }

118 + .apply(&mut right, 90)

119 + .unwrap();

120 + assert_eq!(right, vec![1, 2, 3]);

121 + }

122 +

123 + #[test]

124 + fn input_budget_reserves_output_tokens() {

125 + let mut prompt = vec![0; 100];

126 + PromptTruncation {

127 + limit: PromptTruncationLimit::InputBudget,

128 + side: TruncationSide::Left,

129 + }

130 + .apply(&mut prompt, 70)

131 + .unwrap();

132 + assert_eq!(prompt.len(), 70);

133 + }

134 +

135 + #[test]

136 + fn fixed_limit_cannot_exceed_input_budget() {

137 + let mut prompt = vec![1, 2, 3];

138 + let error = PromptTruncation {

139 + limit: PromptTruncationLimit::Fixed(80),

140 + side: TruncationSide::Left,

141 + }

142 + .apply(&mut prompt, 70)

143 + .unwrap_err();

144 + assert!(matches!(

145 + error,

146 + Error::TruncatePromptTokensExceedsBudget {

147 + value: 80,

148 + budget: 70

149 + }

150 + ));

151 + }

152 +

153 + #[test]

154 + fn zero_limit_clears_prompt() {

155 + let mut prompt = vec![1, 2, 3];

156 + PromptTruncation {

157 + limit: PromptTruncationLimit::Fixed(0),

158 + side: TruncationSide::Left,

159 + }

160 + .apply(&mut prompt, 90)

161 + .unwrap();

162 + assert!(prompt.is_empty());

163 + }

164 +

165 + #[test]

166 + fn limit_at_or_above_prompt_length_is_a_noop() {

167 + for limit in [3, 10] {

168 + let mut prompt = vec![1, 2, 3];

169 + PromptTruncation {

170 + limit: PromptTruncationLimit::Fixed(limit),

171 + side: TruncationSide::Left,

172 + }

173 + .apply(&mut prompt, 90)

174 + .unwrap();

175 + assert_eq!(prompt, vec![1, 2, 3]);

176 + }

177 + }

178 +}