diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 3188657b9096..e584257eaa43 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 1ccd53127acd..b00155999e8b 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -38,22 +38,90 @@ 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 TextRequestProcessor { + /// Create a processor with the effective model context length. + pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self { + Self { + backend, + max_model_len, + max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS, + } + } + + /// Override the maximum accepted logprobs count. + pub fn with_max_logprobs(mut self, max_logprobs: Option) -> Self { + if let Some(max_logprobs) = max_logprobs { + self.max_logprobs = max_logprobs; + } + self + } + + /// Return the tokenizer used by this processor. + pub fn tokenizer(&self) -> DynTokenizer { + self.backend.tokenizer() + } + + /// Return the effective model context length. + pub fn max_model_len(&self) -> u32 { + self.max_model_len + } + + /// 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() { + request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs()); + } + + let tokenizer = self.backend.tokenizer(); + let prompt_token_ids = match take(&mut request.prompt) { + Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?, + // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark + // 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, + max_logprobs: self.max_logprobs, + model_vocab_size: self.backend.model_vocab_size(), + tokenizer_vocab_size: self.backend.tokenizer_vocab_size(), + }; + + lower_text_request( + request, + prompt_token_ids, + sampling_hints, + sampling_limits, + 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. @@ -64,23 +132,19 @@ impl TextLlm { Self { llm, - backend, - max_model_len, - max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS, + processor: TextRequestProcessor::new(backend, max_model_len), } } /// Override the maximum accepted logprobs count. pub fn with_max_logprobs(mut self, max_logprobs: Option) -> Self { - if let Some(max_logprobs) = max_logprobs { - self.max_logprobs = max_logprobs; - } + self.processor = self.processor.with_max_logprobs(max_logprobs); self } /// Return the backend model ID. pub fn model_id(&self) -> &str { - self.backend.model_id() + self.processor.backend.model_id() } /// Expose the underlying engine-core client for low-level utility/admin @@ -91,19 +155,19 @@ impl TextLlm { /// Return the tokenizer used by this text backend. pub fn tokenizer(&self) -> DynTokenizer { - self.backend.tokenizer() + 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.backend.tokenizer_vocab_size() + 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.backend.model_vocab_size() + self.processor.backend.model_vocab_size() } /// Tokenize if needed, lower to a generate request, and return the raw @@ -117,7 +181,7 @@ impl TextLlm { /// 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 tokenizer = self.processor.tokenizer(); let decoded_stream = output::decoded_text_event_stream( text_request.request_id, tokenizer, @@ -131,40 +195,12 @@ impl TextLlm { async fn generate_inner( &self, - mut request: TextRequest, + request: TextRequest, ) -> Result<(TextRequest, GenerateOutputStream)> { - request.validate()?; - - if request.arrival_time.is_none() { - request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs()); - } - - let tokenizer = self.backend.tokenizer(); - let prompt_token_ids = match take(&mut request.prompt) { - Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?, - // Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark - // 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, - max_logprobs: self.max_logprobs, - model_vocab_size: self.backend.model_vocab_size(), - tokenizer_vocab_size: self.backend.tokenizer_vocab_size(), - }; - let PreparedTextRequest { text_request, generate_request, - } = lower_text_request( - request, - prompt_token_ids, - sampling_hints, - sampling_limits, - &*tokenizer, - )?; + } = self.processor.prepare(request)?; let raw_stream = self.llm.generate(generate_request).await?; Ok((text_request, raw_stream))