@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5bf32605a9

[Rust Frontend] Deduplicate request preprocessing for `/tokenize` (#50448)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Sage Ahrac <sagiahrak@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Sage committed Aug 5, 2026 at 11:27 UTC 5bf32605a950ee50ed05f230243d54e0baf7d12d
5 files changed +170 -50
rust/src/chat/src/lib.rs
+25 -33
@@ -155,17 +155,9 @@ impl ChatRequestProcessor {
155 Ok(Some(features))
156 }
157
158 - /// Prepare one chat request without submitting it to an engine.
159 - pub async fn prepare(
160 - &self,
161 - mut request: ChatRequest,
162 - options: NewChatOutputProcessorOptions<'_>,
163 - ) -> Result<(TextRequest, DynChatOutputProcessor)> {
164 - request.validate()?;
165 -
158 + async fn prepare_text_request(&self, request: ChatRequest) -> Result<TextRequest> {
159 // Stamp before rendering so render and tokenize count toward TTFT/e2e.
160 let arrival_time = vllm_llm::current_unix_timestamp_secs();
168 - let output_processor = self.backend.new_chat_output_processor(&mut request, options)?;
161 let rendered = self.backend.chat_renderer().render(&request)?;
162 let reasoning_parser_kwargs =
163 request
@@ -176,7 +168,7 @@ impl ChatRequestProcessor {
168 chat_template_kwargs: rendered.effective_template_kwargs.clone(),
169 });
170 let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
179 - let text_request = TextRequest {
171 + Ok(TextRequest {
172 request_id: request.request_id,
173 prompt,
174 mm_features,
@@ -191,7 +183,24 @@ impl ChatRequestProcessor {
183 reasoning_parser_kwargs,
184 lora_request: request.lora_request,
185 arrival_time: Some(arrival_time),
194 - };
186 + })
187 + }
188 +
189 + /// Prepare one chat request for tokenization without constructing an output processor.
190 + pub async fn prepare_for_tokenization(&self, request: ChatRequest) -> Result<TextRequest> {
191 + request.validate()?;
192 + self.prepare_text_request(request).await
193 + }
194 +
195 + /// Prepare one chat request without submitting it to an engine.
196 + pub async fn prepare(
197 + &self,
198 + mut request: ChatRequest,
199 + options: NewChatOutputProcessorOptions<'_>,
200 + ) -> Result<(TextRequest, DynChatOutputProcessor)> {
201 + request.validate()?;
202 + let output_processor = self.backend.new_chat_output_processor(&mut request, options)?;
203 + let text_request = self.prepare_text_request(request).await?;
204 Ok((text_request, output_processor))
205 }
206 }
@@ -265,6 +274,11 @@ impl ChatLlm {
274 &self.text
275 }
276
277 + /// Return the chat request processor.
278 + pub fn request_processor(&self) -> &ChatRequestProcessor {
279 + &self.processor
280 + }
281 +
282 /// Return the model ID reported by the underlying text backend.
283 pub fn model_id(&self) -> &str {
284 self.text.model_id()
@@ -332,28 +346,6 @@ impl ChatLlm {
346 Ok(ChatEventStream::new(request_id, structured_stream))
347 }
348
335 - /// Render through the chat template and tokenize, without submitting to the engine.
336 - ///
337 - /// Same render → [`multimodal::finalize_rendered_prompt`] → encode pipeline as
338 - /// [`Self::chat`], but stops after token IDs so `/tokenize` counts match what
339 - /// generation would see. Used by `POST /tokenize` (chat form).
340 - pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
341 - request.validate()?;
342 -
343 - let rendered = self.processor.backend.chat_renderer().render(&request)?;
344 - let (prompt, _mm_features) =
345 - self.processor.finalize_rendered_prompt(&request, rendered).await?;
346 -
347 - let tokenizer = self.text.tokenizer();
348 - let token_ids = match prompt {
349 - // Rendered string from the template (usual chat path).
350 - Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
351 - // Already tokenized (e.g. multimodal path); pass through unchanged.
352 - Prompt::TokenIds(ids) => ids,
353 - };
354 - Ok(token_ids)
355 - }
356 -
349 /// Abort in-flight requests by their external (user-supplied) request ids.
350 pub async fn abort(&self, external_ids: &[String]) -> Result<()> {
351 self.text.abort(external_ids).await?;
rust/src/server/src/routes/tests.rs
+69
@@ -6069,6 +6069,75 @@ async fn tokenize_chat_includes_generation_prompt_in_token_count() {
6069 assert!(with_len > without_len);
6070 }
6071
6072 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6073 +#[serial]
6074 +async fn tokenize_chat_with_tools_does_not_require_output_parser() {
6075 + let mut app = test_app().await;
6076 +
6077 + let (status, json) = post_json(
6078 + &mut app,
6079 + "/tokenize",
6080 + json!({
6081 + "model": "Qwen/Qwen1.5-0.5B-Chat",
6082 + "messages": [{"role": "user", "content": "hi"}],
6083 + "tools": [{
6084 + "type": "function",
6085 + "function": {
6086 + "name": "get_weather",
6087 + "description": "Get weather",
6088 + "parameters": {
6089 + "type": "object",
6090 + "properties": {"city": {"type": "string"}},
6091 + },
6092 + },
6093 + }],
6094 + }),
6095 + )
6096 + .await;
6097 +
6098 + assert_eq!(status, StatusCode::OK);
6099 + assert!(!json["tokens"].as_array().expect("tokens").is_empty());
6100 +}
6101 +
6102 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6103 +#[serial]
6104 +async fn tokenize_allows_prompts_at_or_above_max_model_len() {
6105 + let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse {
6106 + max_model_len: 4,
6107 + ..default_ready_response()
6108 + };
6109 + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;
6110 +
6111 + let (completion_status, completion_json) = post_json(
6112 + &mut app,
6113 + "/tokenize",
6114 + json!({
6115 + "model": "Qwen/Qwen1.5-0.5B-Chat",
6116 + "prompt": "hello",
6117 + "add_special_tokens": false,
6118 + }),
6119 + )
6120 + .await;
6121 + assert_eq!(completion_status, StatusCode::OK);
6122 + assert_eq!(completion_json["count"], 5);
6123 + assert_eq!(completion_json["max_model_len"], 4);
6124 +
6125 + let (chat_status, chat_json) = post_json(
6126 + &mut app,
6127 + "/tokenize",
6128 + json!({
6129 + "model": "Qwen/Qwen1.5-0.5B-Chat",
6130 + "messages": [{"role": "user", "content": "hello"}],
6131 + "add_generation_prompt": false,
6132 + "add_special_tokens": false,
6133 + }),
6134 + )
6135 + .await;
6136 + assert_eq!(chat_status, StatusCode::OK);
6137 + assert!(chat_json["count"].as_u64().expect("count") > 4);
6138 + assert_eq!(chat_json["max_model_len"], 4);
6139 +}
6140 +
6141 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6142 #[serial]
6143 async fn tokenize_chat_conflicting_generation_flags_returns_400() {
rust/src/server/src/routes/tokenize.rs
+20 -11
@@ -60,9 +60,7 @@ pub async fn tokenize(
60 let max_model_len = state.chat.engine_core_client().max_model_len();
61
62 let result = match body {
63 - // Completion form: encode the raw `prompt` string (no chat template).
64 - TokenizeRequest::Completion(req) => tokenize_completion(&state, &tokenizer, req),
65 - // Chat form: render `messages` through the template, then encode (see `tokenize_chat`).
63 + TokenizeRequest::Completion(req) => tokenize_completion(&state, &request_id, req),
64 TokenizeRequest::Chat(req) => tokenize_chat(&state, &request_id, req).await,
65 };
66
@@ -83,20 +81,24 @@ pub async fn tokenize(
81
82 fn tokenize_completion(
83 state: &AppState,
86 - tokenizer: &vllm_text::tokenizer::DynTokenizer,
84 + request_id: &str,
85 req: TokenizeCompletionRequest,
86 ) -> Result<(Vec<u32>, bool), ApiError> {
87 check_model(state, req.model.as_deref())?;
90 - let tokens = tokenizer
91 - .encode(&req.prompt, req.add_special_tokens)
88 + let return_token_strs = req.return_token_strs;
89 + let tokens = state
90 + .chat
91 + .text()
92 + .request_processor()
93 + .tokenize(req.into_text_request(request_id.to_string()))
94 .map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?;
93 - Ok((tokens, req.return_token_strs))
95 + Ok((tokens, return_token_strs))
96 }
97
98 /// HTTP adapter for the chat-shaped `/tokenize` body.
99 ///
98 -/// Not [`vllm_chat::ChatLlm::tokenize_chat`]: this checks the model name and maps
99 -/// errors to [`ApiError`]; the chat-crate method does render → finalize → encode.
100 +/// Checks the model name, prepares without engine submission, and maps errors
101 +/// to [`ApiError`].
102 async fn tokenize_chat(
103 state: &AppState,
104 request_id: &str,
@@ -106,11 +108,18 @@ async fn tokenize_chat(
108 let return_token_strs = req.return_token_strs;
109 // `continue_final_message` / `add_generation_prompt` mutual exclusion is
110 // enforced in `normalize_generation_prompt_mode` inside `into_chat_request`.
109 - let tokens = state
111 + let text_request = state
112 .chat
111 - .tokenize_chat(req.into_chat_request(request_id.to_string())?)
113 + .request_processor()
114 + .prepare_for_tokenization(req.into_chat_request(request_id.to_string())?)
115 .await
116 .map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?;
117 + let tokens = state
118 + .chat
119 + .text()
120 + .request_processor()
121 + .tokenize(text_request)
122 + .map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?;
123 Ok((tokens, return_token_strs))
124 }
125
rust/src/server/src/routes/tokenize/types.rs
+23
@@ -9,6 +9,7 @@ use serde_json::Value;
9 use validator::{Validate, ValidationErrors};
10 use vllm_chat::{ChatOptions, ChatRequest, ChatToolChoice, SamplingParams};
11 use vllm_text::output::TextDecodeOptions;
12 +use vllm_text::{Prompt, TextRequest};
13
14 use crate::error::ApiError;
15 use crate::routes::openai::chat_completions::convert::{
@@ -37,6 +38,28 @@ pub struct TokenizeCompletionRequest {
38 pub return_token_strs: bool,
39 }
40
41 +impl TokenizeCompletionRequest {
42 + /// Lower this tokenize body into a [`TextRequest`] for prompt tokenization.
43 + pub fn into_text_request(self, request_id: String) -> TextRequest {
44 + TextRequest {
45 + request_id,
46 + prompt: Prompt::Text(self.prompt),
47 + mm_features: None,
48 + sampling_params: SamplingParams::default(),
49 + decode_options: TextDecodeOptions::default(),
50 + intermediate: false,
51 + priority: 0,
52 + cache_salt: None,
53 + add_special_tokens: self.add_special_tokens,
54 + data_parallel_rank: None,
55 + session_id: None,
56 + reasoning_parser_kwargs: None,
57 + lora_request: None,
58 + arrival_time: None,
59 + }
60 + }
61 +}
62 +
63 #[derive(Debug, Clone, Deserialize, Validate)]
64 pub struct TokenizeChatRequest {
65 pub model: Option<String>,
rust/src/text/src/lib.rs
+33 -6
@@ -78,6 +78,29 @@ impl TextRequestProcessor {
78 self.max_model_len
79 }
80
81 + fn tokenize_prompt(
82 + tokenizer: &DynTokenizer,
83 + prompt: Prompt,
84 + add_special_tokens: bool,
85 + ) -> Result<Vec<u32>> {
86 + match prompt {
87 + Prompt::Text(text) => tokenizer.encode(&text, add_special_tokens).map_err(Into::into),
88 + // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark
89 + // and infra workloads bypass chat rendering and tokenizer overhead entirely.
90 + Prompt::TokenIds(token_ids) => Ok(token_ids),
91 + }
92 + }
93 +
94 + /// Tokenize one request without generation-specific lowering.
95 + pub fn tokenize(&self, request: TextRequest) -> Result<Vec<u32>> {
96 + request.validate()?;
97 + Self::tokenize_prompt(
98 + &self.backend.tokenizer(),
99 + request.prompt,
100 + request.add_special_tokens,
101 + )
102 + }
103 +
104 /// Tokenize and lower one request without submitting it to an engine.
105 pub fn prepare(&self, mut request: TextRequest) -> Result<PreparedTextRequest> {
106 request.validate()?;
@@ -87,12 +110,11 @@ impl TextRequestProcessor {
110 }
111
112 let tokenizer = self.backend.tokenizer();
90 - let prompt_token_ids = match take(&mut request.prompt) {
91 - Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
92 - // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark
93 - // and infra workloads bypass chat rendering and tokenizer overhead entirely.
94 - Prompt::TokenIds(token_ids) => token_ids,
95 - };
113 + let prompt_token_ids = Self::tokenize_prompt(
114 + &tokenizer,
115 + take(&mut request.prompt),
116 + request.add_special_tokens,
117 + )?;
118 let sampling_hints = self.backend.sampling_hints()?;
119 let sampling_limits = SamplingLimits {
120 max_model_len: self.max_model_len,
@@ -153,6 +175,11 @@ impl TextLlm {
175 self.llm.engine_core_client()
176 }
177
178 + /// Return the text request processor.
179 + pub fn request_processor(&self) -> &TextRequestProcessor {
180 + &self.processor
181 + }
182 +
183 /// Return the tokenizer used by this text backend.
184 pub fn tokenizer(&self) -> DynTokenizer {
185 self.processor.tokenizer()