karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Rust Frontend] Extract request preparation from the inference path (#49045)

Browse Source 
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
This commit is contained in:
Sage
2026-07-22 11:31:36 +00:00
committed by GitHub
parent c79ad3ae21
commit 0f6cf7f628
2 changed files with 214 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/chat/src/lib.rs
+106 -61
View File
@@ -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<ModelDtype>,
}
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<MmFeatures>)> {
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<ChatEventStream> {
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<ChatEventStream> {
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<Vec<u32>> {
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 {
rust/src/text/src/lib.rs
+108 -72
View File
@@ -38,32 +38,22 @@ trait_set! {
pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + 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<GenerateOutputStream> {
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<impl TextOutputStream> {
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<PreparedTextRequest> {
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<i32>) -> 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<GenerateOutputStream> {
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<impl TextOutputStream> {
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))