From 0f6cf7f628de8f09f35bc871f28b28a5ecdf563b Mon Sep 17 00:00:00 2001 From: Sage <80211083+sagearc@users.noreply.github.com> Date: Wed, 22 Jul 2026 14:31:36 +0300 Subject: [PATCH] [Rust Frontend] Extract request preparation from the inference path (#49045) Signed-off-by: Sage Ahrac --- rust/src/chat/src/lib.rs | 167 +++++++++++++++++++++++------------- rust/src/text/src/lib.rs | 180 +++++++++++++++++++++++---------------- 2 files changed, 214 insertions(+), 133 deletions(-) diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 3188657b909..e584257eaa4 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -54,6 +54,7 @@ mod stream; use vllm_engine_core_client::EngineCoreClient; use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::multimodal::MmFeatures; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_llm::Llm; use vllm_text::{Prompt, TextLlm, TextRequest}; @@ -88,6 +89,92 @@ pub fn validate_parser_overrides( Ok(()) } +/// Chat request preparation shared by inference and render-only frontends. +pub struct ChatRequestProcessor { + backend: DynChatBackend, + /// Effective model dtype reported by the engine. + /// Absent for text-only frontends without an engine handshake. + model_dtype: Option, +} + +impl ChatRequestProcessor { + /// Create a processor with multimodal support using the effective model + /// dtype reported by the engine. + fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self { + Self { + backend, + model_dtype: Some(model_dtype), + } + } + + /// Create a render-only processor that rejects multimodal requests. + pub fn render_only(backend: DynChatBackend) -> Self { + Self { + backend, + model_dtype: None, + } + } + + async fn finalize_rendered_prompt( + &self, + request: &ChatRequest, + rendered: RenderedPrompt, + ) -> Result<(Prompt, Option)> { + match self.model_dtype { + Some(model_dtype) => { + multimodal::finalize_rendered_prompt( + request, + rendered, + self.backend.multimodal_model_info(), + model_dtype, + ) + .await + } + None if !request.has_multimodal() => Ok((rendered.prompt, None)), + None => Err(Error::UnsupportedMultimodalRenderer), + } + } + + /// Prepare one chat request without submitting it to an engine. + pub async fn prepare( + &self, + mut request: ChatRequest, + options: NewChatOutputProcessorOptions<'_>, + ) -> Result<(TextRequest, DynChatOutputProcessor)> { + request.validate()?; + + // Stamp before rendering so render and tokenize count toward TTFT/e2e. + let arrival_time = vllm_llm::current_unix_timestamp_secs(); + let output_processor = self.backend.new_chat_output_processor(&mut request, options)?; + let rendered = self.backend.chat_renderer().render(&request)?; + let reasoning_parser_kwargs = + request + .sampling_params + .structured_outputs + .is_some() + .then(|| ReasoningParserKwargs { + chat_template_kwargs: rendered.effective_template_kwargs.clone(), + }); + let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?; + let text_request = TextRequest { + request_id: request.request_id, + prompt, + mm_features, + sampling_params: request.sampling_params, + decode_options: request.decode_options, + intermediate: request.intermediate, + priority: request.priority, + cache_salt: request.cache_salt, + add_special_tokens: request.add_special_tokens, + data_parallel_rank: request.data_parallel_rank, + reasoning_parser_kwargs, + lora_request: request.lora_request, + arrival_time: Some(arrival_time), + }; + Ok((text_request, output_processor)) + } +} + /// Structured chat facade above [`TextLlm`]. /// /// This layer stays above raw text semantics: it takes care of chat-template @@ -95,9 +182,7 @@ pub fn validate_parser_overrides( /// request semantics such as tool calls. pub struct ChatLlm { text: TextLlm, - backend: DynChatBackend, - /// Effective model dtype reported by the engine. - model_dtype: ModelDtype, + processor: ChatRequestProcessor, /// Tool-call parser selection. tool_call_parser: ParserSelection, /// Reasoning parser selection. @@ -112,8 +197,7 @@ impl ChatLlm { Self { text, - backend, - model_dtype, + processor: ChatRequestProcessor::new(backend, model_dtype), tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, } @@ -140,7 +224,7 @@ impl ChatLlm { /// Override the effective model dtype used for multimodal tensor encoding. pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self { - self.model_dtype = model_dtype; + self.processor.model_dtype = Some(model_dtype); self } @@ -172,57 +256,23 @@ impl ChatLlm { } /// Render, tokenize, and submit one chat request. - pub async fn chat(&self, mut request: ChatRequest) -> Result { - request.validate()?; - - // Stamp before rendering so render and tokenize count toward TTFT/e2e. - let arrival_time = vllm_llm::current_unix_timestamp_secs(); - - let output_processor = self.backend.new_chat_output_processor( - &mut request, - NewChatOutputProcessorOptions { - tool_call_parser: &self.tool_call_parser, - reasoning_parser: &self.reasoning_parser, - }, - )?; - let rendered = self.backend.chat_renderer().render(&request)?; - let reasoning_parser_kwargs = - request - .sampling_params - .structured_outputs - .is_some() - .then(|| ReasoningParserKwargs { - chat_template_kwargs: rendered.effective_template_kwargs.clone(), - }); - - let (prompt, mm_features) = multimodal::finalize_rendered_prompt( - &request, - rendered, - self.backend.multimodal_model_info(), - self.model_dtype, - ) - .await?; - - let text_request = TextRequest { - request_id: request.request_id.clone(), - prompt, - mm_features, - sampling_params: request.sampling_params, - decode_options: request.decode_options, - intermediate: request.intermediate, - priority: request.priority, - cache_salt: request.cache_salt, - add_special_tokens: request.add_special_tokens, - data_parallel_rank: request.data_parallel_rank, - reasoning_parser_kwargs, - lora_request: request.lora_request, - arrival_time: Some(arrival_time), - }; + pub async fn chat(&self, request: ChatRequest) -> Result { + let (text_request, output_processor) = self + .processor + .prepare( + request, + NewChatOutputProcessorOptions { + tool_call_parser: &self.tool_call_parser, + reasoning_parser: &self.reasoning_parser, + }, + ) + .await?; + let request_id = text_request.request_id.clone(); let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed(); let structured_stream = output_processor.process(decoded_stream)?; - Ok(ChatEventStream::new(request.request_id, structured_stream)) + Ok(ChatEventStream::new(request_id, structured_stream)) } /// Render through the chat template and tokenize, without submitting to the engine. @@ -233,14 +283,9 @@ impl ChatLlm { pub async fn tokenize_chat(&self, request: ChatRequest) -> Result> { request.validate()?; - let rendered = self.backend.chat_renderer().render(&request)?; - let (prompt, _mm_features) = multimodal::finalize_rendered_prompt( - &request, - rendered, - self.backend.multimodal_model_info(), - self.model_dtype, - ) - .await?; + let rendered = self.processor.backend.chat_renderer().render(&request)?; + let (prompt, _mm_features) = + self.processor.finalize_rendered_prompt(&request, rendered).await?; let tokenizer = self.text.tokenizer(); let token_ids = match prompt { diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index 1ccd53127ac..b00155999e8 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -38,32 +38,22 @@ trait_set! { pub trait TextOutputStream = Stream> + Send + 'static; } -/// Raw text facade above [`Llm`]. -/// -/// This layer stays below chat semantics: prompt text or prompt token IDs flow -/// in, decoded text deltas and terminal metadata flow out. -pub struct TextLlm { - /// Generate-only client owned by this text facade. - llm: Llm, +/// Text request preparation shared by inference and render-only frontends. +pub struct TextRequestProcessor { /// Tokenizer/model metadata backend responsible for prompt encode/decode /// and sampling hints. backend: DynTextBackend, /// Runtime context window size reported by the engine startup handshake. + /// Render-only frontends supply the downstream engine's effective value. max_model_len: u32, /// Maximum number of top log probabilities accepted by this text facade. max_logprobs: i32, } -impl TextLlm { - /// Create a new text-generation facade from a shared LLM client plus a text - /// backend. - pub fn new(llm: Llm, backend: DynTextBackend) -> Self { - // The engine-reported value reflects the post-profiling, auto-fitted - // KV cache limit used at runtime. - let max_model_len = llm.engine_core_client().max_model_len(); - +impl TextRequestProcessor { + /// Create a processor with the effective model context length. + pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self { Self { - llm, backend, max_model_len, max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS, @@ -78,61 +68,18 @@ impl TextLlm { self } - /// Return the backend model ID. - pub fn model_id(&self) -> &str { - self.backend.model_id() - } - - /// Expose the underlying engine-core client for low-level utility/admin - /// calls. - pub fn engine_core_client(&self) -> &EngineCoreClient { - self.llm.engine_core_client() - } - - /// Return the tokenizer used by this text backend. + /// Return the tokenizer used by this processor. pub fn tokenizer(&self) -> DynTokenizer { self.backend.tokenizer() } - /// Tokenizer vocabulary size (the number of tokens the tokenizer knows), - /// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`. - pub fn tokenizer_vocab_size(&self) -> usize { - self.backend.tokenizer_vocab_size() + /// Return the effective model context length. + pub fn max_model_len(&self) -> u32 { + self.max_model_len } - /// Model vocabulary size from the model config, used to bound generated - /// token IDs and logits-domain sampling controls. - pub fn model_vocab_size(&self) -> usize { - self.backend.model_vocab_size() - } - - /// Tokenize if needed, lower to a generate request, and return the raw - /// token stream. - pub async fn generate_raw(&self, request: TextRequest) -> Result { - let (_, raw_stream) = self.generate_inner(request).await?; - Ok(raw_stream) - } - - /// Tokenize if needed, lower to a generate request, and stream - /// incrementally decoded text. - pub async fn generate(&self, request: TextRequest) -> Result { - let (text_request, raw_stream) = self.generate_inner(request).await?; - let tokenizer = self.backend.tokenizer(); - let decoded_stream = output::decoded_text_event_stream( - text_request.request_id, - tokenizer, - raw_stream, - text_request.decode_options, - text_request.intermediate, - ); - - Ok(decoded_stream) - } - - async fn generate_inner( - &self, - mut request: TextRequest, - ) -> Result<(TextRequest, GenerateOutputStream)> { + /// Tokenize and lower one request without submitting it to an engine. + pub fn prepare(&self, mut request: TextRequest) -> Result { request.validate()?; if request.arrival_time.is_none() { @@ -146,7 +93,6 @@ impl TextLlm { // and infra workloads bypass chat rendering and tokenizer overhead entirely. Prompt::TokenIds(token_ids) => token_ids, }; - let sampling_hints = self.backend.sampling_hints()?; let sampling_limits = SamplingLimits { max_model_len: self.max_model_len, @@ -155,16 +101,106 @@ impl TextLlm { tokenizer_vocab_size: self.backend.tokenizer_vocab_size(), }; - let PreparedTextRequest { - text_request, - generate_request, - } = lower_text_request( + lower_text_request( request, prompt_token_ids, sampling_hints, sampling_limits, - &*tokenizer, - )?; + tokenizer.as_ref(), + ) + } +} + +/// Raw text facade above [`Llm`]. +/// +/// This layer stays below chat semantics: prompt text or prompt token IDs flow +/// in, decoded text deltas and terminal metadata flow out. +pub struct TextLlm { + /// Generate-only client owned by this text facade. + llm: Llm, + /// Shared engine-free request preparation. + processor: TextRequestProcessor, +} + +impl TextLlm { + /// Create a new text-generation facade from a shared LLM client plus a text + /// backend. + pub fn new(llm: Llm, backend: DynTextBackend) -> Self { + // The engine-reported value reflects the post-profiling, auto-fitted + // KV cache limit used at runtime. + let max_model_len = llm.engine_core_client().max_model_len(); + + Self { + llm, + processor: TextRequestProcessor::new(backend, max_model_len), + } + } + + /// Override the maximum accepted logprobs count. + pub fn with_max_logprobs(mut self, max_logprobs: Option) -> Self { + self.processor = self.processor.with_max_logprobs(max_logprobs); + self + } + + /// Return the backend model ID. + pub fn model_id(&self) -> &str { + self.processor.backend.model_id() + } + + /// Expose the underlying engine-core client for low-level utility/admin + /// calls. + pub fn engine_core_client(&self) -> &EngineCoreClient { + self.llm.engine_core_client() + } + + /// Return the tokenizer used by this text backend. + pub fn tokenizer(&self) -> DynTokenizer { + self.processor.tokenizer() + } + + /// Tokenizer vocabulary size (the number of tokens the tokenizer knows), + /// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`. + pub fn tokenizer_vocab_size(&self) -> usize { + self.processor.backend.tokenizer_vocab_size() + } + + /// Model vocabulary size from the model config, used to bound generated + /// token IDs and logits-domain sampling controls. + pub fn model_vocab_size(&self) -> usize { + self.processor.backend.model_vocab_size() + } + + /// Tokenize if needed, lower to a generate request, and return the raw + /// token stream. + pub async fn generate_raw(&self, request: TextRequest) -> Result { + let (_, raw_stream) = self.generate_inner(request).await?; + Ok(raw_stream) + } + + /// Tokenize if needed, lower to a generate request, and stream + /// incrementally decoded text. + pub async fn generate(&self, request: TextRequest) -> Result { + let (text_request, raw_stream) = self.generate_inner(request).await?; + let tokenizer = self.processor.tokenizer(); + let decoded_stream = output::decoded_text_event_stream( + text_request.request_id, + tokenizer, + raw_stream, + text_request.decode_options, + text_request.intermediate, + ); + + Ok(decoded_stream) + } + + async fn generate_inner( + &self, + request: TextRequest, + ) -> Result<(TextRequest, GenerateOutputStream)> { + let PreparedTextRequest { + text_request, + generate_request, + } = self.processor.prepare(request)?; let raw_stream = self.llm.generate(generate_request).await?; Ok((text_request, raw_stream))