[Rust Frontend] Extract request preparation from the inference path (#49045)
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
Sage committed
Jul 22, 2026 at 14:31 UTC
0f6cf7f628de8f09f35bc871f28b28a5ecdf563b
2 files changed
+190
-109
rust/src/chat/src/lib.rs
+106
-61
@@ -54,6 +54,7 @@ mod stream;
54
55
use vllm_engine_core_client::EngineCoreClient;
56
use vllm_engine_core_client::protocol::dtype::ModelDtype;
57
+use vllm_engine_core_client::protocol::multimodal::MmFeatures;
58
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
59
use vllm_llm::Llm;
60
use vllm_text::{Prompt, TextLlm, TextRequest};
@@ -88,6 +89,92 @@ pub fn validate_parser_overrides(
89
Ok(())
90
}
91
92
+/// Chat request preparation shared by inference and render-only frontends.
93
+pub struct ChatRequestProcessor {
94
+ backend: DynChatBackend,
95
+ /// Effective model dtype reported by the engine.
96
+ /// Absent for text-only frontends without an engine handshake.
97
+ model_dtype: Option<ModelDtype>,
98
+}
99
+
100
+impl ChatRequestProcessor {
101
+ /// Create a processor with multimodal support using the effective model
102
+ /// dtype reported by the engine.
103
+ fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self {
104
+ Self {
105
+ backend,
106
+ model_dtype: Some(model_dtype),
107
+ }
108
+ }
109
+
110
+ /// Create a render-only processor that rejects multimodal requests.
111
+ pub fn render_only(backend: DynChatBackend) -> Self {
112
+ Self {
113
+ backend,
114
+ model_dtype: None,
115
+ }
116
+ }
117
+
118
+ async fn finalize_rendered_prompt(
119
+ &self,
120
+ request: &ChatRequest,
121
+ rendered: RenderedPrompt,
122
+ ) -> Result<(Prompt, Option<MmFeatures>)> {
123
+ match self.model_dtype {
124
+ Some(model_dtype) => {
125
+ multimodal::finalize_rendered_prompt(
126
+ request,
127
+ rendered,
128
+ self.backend.multimodal_model_info(),
129
+ model_dtype,
130
+ )
131
+ .await
132
+ }
133
+ None if !request.has_multimodal() => Ok((rendered.prompt, None)),
134
+ None => Err(Error::UnsupportedMultimodalRenderer),
135
+ }
136
+ }
137
+
138
+ /// Prepare one chat request without submitting it to an engine.
139
+ pub async fn prepare(
140
+ &self,
141
+ mut request: ChatRequest,
142
+ options: NewChatOutputProcessorOptions<'_>,
143
+ ) -> Result<(TextRequest, DynChatOutputProcessor)> {
144
+ request.validate()?;
145
+
146
+ // Stamp before rendering so render and tokenize count toward TTFT/e2e.
147
+ let arrival_time = vllm_llm::current_unix_timestamp_secs();
148
+ let output_processor = self.backend.new_chat_output_processor(&mut request, options)?;
149
+ let rendered = self.backend.chat_renderer().render(&request)?;
150
+ let reasoning_parser_kwargs =
151
+ request
152
+ .sampling_params
153
+ .structured_outputs
154
+ .is_some()
155
+ .then(|| ReasoningParserKwargs {
156
+ chat_template_kwargs: rendered.effective_template_kwargs.clone(),
157
+ });
158
+ let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
159
+ let text_request = TextRequest {
160
+ request_id: request.request_id,
161
+ prompt,
162
+ mm_features,
163
+ sampling_params: request.sampling_params,
164
+ decode_options: request.decode_options,
165
+ intermediate: request.intermediate,
166
+ priority: request.priority,
167
+ cache_salt: request.cache_salt,
168
+ add_special_tokens: request.add_special_tokens,
169
+ data_parallel_rank: request.data_parallel_rank,
170
+ reasoning_parser_kwargs,
171
+ lora_request: request.lora_request,
172
+ arrival_time: Some(arrival_time),
173
+ };
174
+ Ok((text_request, output_processor))
175
+ }
176
+}
177
+
178
/// Structured chat facade above [`TextLlm`].
179
///
180
/// This layer stays above raw text semantics: it takes care of chat-template
@@ -95,9 +182,7 @@ pub fn validate_parser_overrides(
182
/// request semantics such as tool calls.
183
pub struct ChatLlm {
184
text: TextLlm,
98
- backend: DynChatBackend,
99
- /// Effective model dtype reported by the engine.
100
- model_dtype: ModelDtype,
185
+ processor: ChatRequestProcessor,
186
/// Tool-call parser selection.
187
tool_call_parser: ParserSelection,
188
/// Reasoning parser selection.
@@ -112,8 +197,7 @@ impl ChatLlm {
197
198
Self {
199
text,
115
- backend,
116
- model_dtype,
200
+ processor: ChatRequestProcessor::new(backend, model_dtype),
201
tool_call_parser: ParserSelection::Auto,
202
reasoning_parser: ParserSelection::Auto,
203
}
@@ -140,7 +224,7 @@ impl ChatLlm {
224
225
/// Override the effective model dtype used for multimodal tensor encoding.
226
pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
143
- self.model_dtype = model_dtype;
227
+ self.processor.model_dtype = Some(model_dtype);
228
self
229
}
230
@@ -172,57 +256,23 @@ impl ChatLlm {
256
}
257
258
/// Render, tokenize, and submit one chat request.
175
- pub async fn chat(&self, mut request: ChatRequest) -> Result<ChatEventStream> {
176
- request.validate()?;
177
-
178
- // Stamp before rendering so render and tokenize count toward TTFT/e2e.
179
- let arrival_time = vllm_llm::current_unix_timestamp_secs();
180
-
181
- let output_processor = self.backend.new_chat_output_processor(
182
- &mut request,
183
- NewChatOutputProcessorOptions {
184
- tool_call_parser: &self.tool_call_parser,
185
- reasoning_parser: &self.reasoning_parser,
186
- },
187
- )?;
188
- let rendered = self.backend.chat_renderer().render(&request)?;
189
- let reasoning_parser_kwargs =
190
- request
191
- .sampling_params
192
- .structured_outputs
193
- .is_some()
194
- .then(|| ReasoningParserKwargs {
195
- chat_template_kwargs: rendered.effective_template_kwargs.clone(),
196
- });
197
-
198
- let (prompt, mm_features) = multimodal::finalize_rendered_prompt(
199
- &request,
200
- rendered,
201
- self.backend.multimodal_model_info(),
202
- self.model_dtype,
203
- )
204
- .await?;
205
-
206
- let text_request = TextRequest {
207
- request_id: request.request_id.clone(),
208
- prompt,
209
- mm_features,
210
- sampling_params: request.sampling_params,
211
- decode_options: request.decode_options,
212
- intermediate: request.intermediate,
213
- priority: request.priority,
214
- cache_salt: request.cache_salt,
215
- add_special_tokens: request.add_special_tokens,
216
- data_parallel_rank: request.data_parallel_rank,
217
- reasoning_parser_kwargs,
218
- lora_request: request.lora_request,
219
- arrival_time: Some(arrival_time),
220
- };
259
+ pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
260
+ let (text_request, output_processor) = self
261
+ .processor
262
+ .prepare(
263
+ request,
264
+ NewChatOutputProcessorOptions {
265
+ tool_call_parser: &self.tool_call_parser,
266
+ reasoning_parser: &self.reasoning_parser,
267
+ },
268
+ )
269
+ .await?;
270
+ let request_id = text_request.request_id.clone();
271
let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
272
273
let structured_stream = output_processor.process(decoded_stream)?;
274
225
- Ok(ChatEventStream::new(request.request_id, structured_stream))
275
+ Ok(ChatEventStream::new(request_id, structured_stream))
276
}
277
278
/// Render through the chat template and tokenize, without submitting to the engine.
@@ -233,14 +283,9 @@ impl ChatLlm {
283
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
284
request.validate()?;
285
236
- let rendered = self.backend.chat_renderer().render(&request)?;
237
- let (prompt, _mm_features) = multimodal::finalize_rendered_prompt(
238
- &request,
239
- rendered,
240
- self.backend.multimodal_model_info(),
241
- self.model_dtype,
242
- )
243
- .await?;
286
+ let rendered = self.processor.backend.chat_renderer().render(&request)?;
287
+ let (prompt, _mm_features) =
288
+ self.processor.finalize_rendered_prompt(&request, rendered).await?;
289
290
let tokenizer = self.text.tokenizer();
291
let token_ids = match prompt {
rust/src/text/src/lib.rs
+84
-48
@@ -38,22 +38,90 @@ trait_set! {
38
pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
39
}
40
41
-/// Raw text facade above [`Llm`].
42
-///
43
-/// This layer stays below chat semantics: prompt text or prompt token IDs flow
44
-/// in, decoded text deltas and terminal metadata flow out.
45
-pub struct TextLlm {
46
- /// Generate-only client owned by this text facade.
47
- llm: Llm,
41
+/// Text request preparation shared by inference and render-only frontends.
42
+pub struct TextRequestProcessor {
43
/// Tokenizer/model metadata backend responsible for prompt encode/decode
44
/// and sampling hints.
45
backend: DynTextBackend,
46
/// Runtime context window size reported by the engine startup handshake.
47
+ /// Render-only frontends supply the downstream engine's effective value.
48
max_model_len: u32,
49
/// Maximum number of top log probabilities accepted by this text facade.
50
max_logprobs: i32,
51
}
52
53
+impl TextRequestProcessor {
54
+ /// Create a processor with the effective model context length.
55
+ pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self {
56
+ Self {
57
+ backend,
58
+ max_model_len,
59
+ max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
60
+ }
61
+ }
62
+
63
+ /// Override the maximum accepted logprobs count.
64
+ pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
65
+ if let Some(max_logprobs) = max_logprobs {
66
+ self.max_logprobs = max_logprobs;
67
+ }
68
+ self
69
+ }
70
+
71
+ /// Return the tokenizer used by this processor.
72
+ pub fn tokenizer(&self) -> DynTokenizer {
73
+ self.backend.tokenizer()
74
+ }
75
+
76
+ /// Return the effective model context length.
77
+ pub fn max_model_len(&self) -> u32 {
78
+ self.max_model_len
79
+ }
80
+
81
+ /// Tokenize and lower one request without submitting it to an engine.
82
+ pub fn prepare(&self, mut request: TextRequest) -> Result<PreparedTextRequest> {
83
+ request.validate()?;
84
+
85
+ if request.arrival_time.is_none() {
86
+ request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs());
87
+ }
88
+
89
+ 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
+ };
96
+ let sampling_hints = self.backend.sampling_hints()?;
97
+ let sampling_limits = SamplingLimits {
98
+ max_model_len: self.max_model_len,
99
+ max_logprobs: self.max_logprobs,
100
+ model_vocab_size: self.backend.model_vocab_size(),
101
+ tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
102
+ };
103
+
104
+ lower_text_request(
105
+ request,
106
+ prompt_token_ids,
107
+ sampling_hints,
108
+ sampling_limits,
109
+ tokenizer.as_ref(),
110
+ )
111
+ }
112
+}
113
+
114
+/// Raw text facade above [`Llm`].
115
+///
116
+/// This layer stays below chat semantics: prompt text or prompt token IDs flow
117
+/// in, decoded text deltas and terminal metadata flow out.
118
+pub struct TextLlm {
119
+ /// Generate-only client owned by this text facade.
120
+ llm: Llm,
121
+ /// Shared engine-free request preparation.
122
+ processor: TextRequestProcessor,
123
+}
124
+
125
impl TextLlm {
126
/// Create a new text-generation facade from a shared LLM client plus a text
127
/// backend.
@@ -64,23 +132,19 @@ impl TextLlm {
132
133
Self {
134
llm,
67
- backend,
68
- max_model_len,
69
- max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
135
+ processor: TextRequestProcessor::new(backend, max_model_len),
136
}
137
}
138
139
/// Override the maximum accepted logprobs count.
140
pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
75
- if let Some(max_logprobs) = max_logprobs {
76
- self.max_logprobs = max_logprobs;
77
- }
141
+ self.processor = self.processor.with_max_logprobs(max_logprobs);
142
self
143
}
144
145
/// Return the backend model ID.
146
pub fn model_id(&self) -> &str {
83
- self.backend.model_id()
147
+ self.processor.backend.model_id()
148
}
149
150
/// Expose the underlying engine-core client for low-level utility/admin
@@ -91,19 +155,19 @@ impl TextLlm {
155
156
/// Return the tokenizer used by this text backend.
157
pub fn tokenizer(&self) -> DynTokenizer {
94
- self.backend.tokenizer()
158
+ self.processor.tokenizer()
159
}
160
161
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
162
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
163
pub fn tokenizer_vocab_size(&self) -> usize {
100
- self.backend.tokenizer_vocab_size()
164
+ self.processor.backend.tokenizer_vocab_size()
165
}
166
167
/// Model vocabulary size from the model config, used to bound generated
168
/// token IDs and logits-domain sampling controls.
169
pub fn model_vocab_size(&self) -> usize {
106
- self.backend.model_vocab_size()
170
+ self.processor.backend.model_vocab_size()
171
}
172
173
/// Tokenize if needed, lower to a generate request, and return the raw
@@ -117,7 +181,7 @@ impl TextLlm {
181
/// incrementally decoded text.
182
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
183
let (text_request, raw_stream) = self.generate_inner(request).await?;
120
- let tokenizer = self.backend.tokenizer();
184
+ let tokenizer = self.processor.tokenizer();
185
let decoded_stream = output::decoded_text_event_stream(
186
text_request.request_id,
187
tokenizer,
@@ -131,40 +195,12 @@ impl TextLlm {
195
196
async fn generate_inner(
197
&self,
134
- mut request: TextRequest,
198
+ request: TextRequest,
199
) -> Result<(TextRequest, GenerateOutputStream)> {
136
- request.validate()?;
137
-
138
- if request.arrival_time.is_none() {
139
- request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs());
140
- }
141
-
142
- let tokenizer = self.backend.tokenizer();
143
- let prompt_token_ids = match take(&mut request.prompt) {
144
- Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
145
- // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark
146
- // and infra workloads bypass chat rendering and tokenizer overhead entirely.
147
- Prompt::TokenIds(token_ids) => token_ids,
148
- };
149
-
150
- let sampling_hints = self.backend.sampling_hints()?;
151
- let sampling_limits = SamplingLimits {
152
- max_model_len: self.max_model_len,
153
- max_logprobs: self.max_logprobs,
154
- model_vocab_size: self.backend.model_vocab_size(),
155
- tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
156
- };
157
-
200
let PreparedTextRequest {
201
text_request,
202
generate_request,
161
- } = lower_text_request(
162
- request,
163
- prompt_token_ids,
164
- sampling_hints,
165
- sampling_limits,
166
- &*tokenizer,
167
- )?;
203
+ } = self.processor.prepare(request)?;
204
205
let raw_stream = self.llm.generate(generate_request).await?;
206
Ok((text_request, raw_stream))