[Rust] Split and rename embedded server components (#37220)
This commit is contained in:
Lianmin Zheng
@@ -22,7 +22,7 @@ features = ["python", "parallel"]
|
||||
# which renames the built artifact.
|
||||
name = "sglang_mm_core"
|
||||
# cdylib: the PyO3 module (`sglang.srt.rust_extensions._multimodal`).
|
||||
# rlib: pure-Rust core linked by sglang-server's native MM path.
|
||||
# rlib: pure-Rust core linked by sglang-server's Rust MM path.
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[features]
|
||||
|
||||
@@ -65,7 +65,7 @@ pub struct QwenVlProcessor {
|
||||
lut: [[f32; 256]; 3],
|
||||
}
|
||||
|
||||
/// `1 / rescale_factor`; `resolve_native_spec` rejects any other factor.
|
||||
/// `1 / rescale_factor`; `resolve_spec` rejects any other factor.
|
||||
const INV_RESCALE: f32 = 255.0;
|
||||
|
||||
/// u8 → normalized f32, rounded as the mirrored processor rounds. The slow one
|
||||
@@ -406,7 +406,7 @@ mod python {
|
||||
|
||||
/// `(pixel_values flat f32, (t, h, w))` for one preprocessed image.
|
||||
type PyProcessedImage<'py> = (Bound<'py, PyArray1<f32>>, (u32, u32, u32));
|
||||
/// Full native pipeline output at the scheduler boundary:
|
||||
/// Full Rust pipeline output at the scheduler boundary:
|
||||
/// `(input_ids, features, grids, hashes, offsets, mrope, mrope_delta)`.
|
||||
type PyNativeOutput<'py> = (
|
||||
Vec<i32>,
|
||||
@@ -486,7 +486,7 @@ mod python {
|
||||
/// `sglang-server` (whose message layer owns the wire-payload parsing).
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (input_ids, images, spec_json))]
|
||||
fn process_native_mm<'py>(
|
||||
fn process_mm<'py>(
|
||||
py: Python<'py>,
|
||||
input_ids: Option<Vec<i32>>,
|
||||
images: Vec<PyImageSource>,
|
||||
@@ -529,7 +529,7 @@ mod python {
|
||||
m.add_function(wrap_pyfunction!(preprocess, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(smart_resize_py, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(mrope_image_only_py, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(process_native_mm, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(process_mm, &m)?)?;
|
||||
parent.add_submodule(&m)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -50,7 +50,7 @@ hf-hub = { version = "0.4", default-features = false }
|
||||
itertools = "0.14"
|
||||
# POSIX shm for the MM feature fan-out (`mm::ShmSegment`).
|
||||
libc = "0.2"
|
||||
# Same major as the workspace pyo3: the zero-copy MM drain (`take_mm`) moves
|
||||
# Same major as the workspace pyo3: the zero-copy MM drain (`take_mm_result`) moves
|
||||
# Rust vectors into numpy arrays.
|
||||
numpy = "0.29.0"
|
||||
# Pinned EXACTLY: this crate's accepted grammar defines the
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
# sglang-server
|
||||
|
||||
`sglang-server` is SGLang's Rust HTTP frontend and request-processing pipeline. It exchanges typed requests and responses with the Python scheduler while keeping latency-sensitive work outside Python.
|
||||
|
||||
## Code review principles
|
||||
|
||||
1. **Use strongly typed boundaries.** Model every supported protocol shape with structs, enums, and validated newtypes; avoid opaque values such as `serde_json::Value` and `rmpv::Value` in production paths.
|
||||
2. **Keep one canonical schema.** Rust and Python must derive their wire contracts from one source of truth, aligned with `io_struct.py`, instead of independently duplicating field names, order, defaults, or validation.
|
||||
3. **Make protocol declarations minimal and declarative.** A reviewer should be able to understand the wire format from its type declarations alone, without tracing fillers, conversion code, macros, or repeated field lists.
|
||||
4. **Use one representation per semantic stage.** Separate external input, normalized domain data, and wire data, and convert between them once at explicit boundaries; do not keep multiple overlapping representations of the same state.
|
||||
5. **Design compatibility and safety explicitly.** Version protocols, reject unsupported or malformed inputs clearly, validate lengths and resource bounds before allocation, preserve invariants in types, and test compatibility across the real Rust and Python codecs.
|
||||
@@ -14,6 +14,9 @@ mod completions;
|
||||
mod models;
|
||||
mod reasoning;
|
||||
mod template;
|
||||
mod template_builtins;
|
||||
mod template_legacy;
|
||||
mod template_loader;
|
||||
mod tools;
|
||||
|
||||
pub(super) use template::ChatFormatter;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,284 @@
|
||||
//! Built-in legacy SGLang conversation templates.
|
||||
|
||||
use crate::message::types::OneOrMany;
|
||||
|
||||
use super::template_legacy::LegacySpec;
|
||||
|
||||
pub(super) fn builtin_template(name: &str) -> Option<LegacySpec> {
|
||||
let spec = match name {
|
||||
"llama-2" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n".into(),
|
||||
roles: ("[INST]".into(), "[/INST]".into()),
|
||||
style: "LLAMA2".into(),
|
||||
sep: " ".into(),
|
||||
sep2: Some(" </s><s>".into()),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"[INST]".into(),
|
||||
"[/INST]".into(),
|
||||
"<<SYS>>".into(),
|
||||
"<</SYS>>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"mistral" | "devstral" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "[SYSTEM_PROMPT]\n{system_message}\n[/SYSTEM_PROMPT]\n\n".into(),
|
||||
roles: ("[INST]".into(), "[/INST]".into()),
|
||||
style: "LLAMA2".into(),
|
||||
sep: " ".into(),
|
||||
sep2: Some(" </s><s>".into()),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"[INST]".into(),
|
||||
"[/INST]".into(),
|
||||
"[SYSTEM_PROMPT]".into(),
|
||||
"[/SYSTEM_PROMPT]".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"llama-4" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|header_start|>system<|header_end|>\n\n{system_message}<|eot|>"
|
||||
.into(),
|
||||
roles: ("user".into(), "assistant".into()),
|
||||
style: "LLAMA4".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|end_of_text|>".into(),
|
||||
"<|eot|>".into(),
|
||||
"<|eom|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"phi-4-mm" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}".into(),
|
||||
roles: ("<|user|>".into(), "<|assistant|>".into()),
|
||||
style: "NO_COLON_SINGLE".into(),
|
||||
sep: "<|end|>".into(),
|
||||
stop_str: Some(OneOrMany::One("<|end|>".into())),
|
||||
image_token: "<|endoftext10|>".into(),
|
||||
audio_token: "<|endoftext11|>".into(),
|
||||
..Default::default()
|
||||
},
|
||||
"chatml" | "chatml-llava" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "CHATML".into(),
|
||||
sep: "<|im_end|>".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|endoftext|>".into(),
|
||||
"<|im_end|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"vicuna_v1.1" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}".into(),
|
||||
system_message: "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.".into(),
|
||||
roles: ("USER".into(), "ASSISTANT".into()),
|
||||
style: "ADD_COLON_TWO".into(),
|
||||
sep: " ".into(),
|
||||
sep2: Some("</s>".into()),
|
||||
..Default::default()
|
||||
},
|
||||
"llama_3_vision" | "llava_llama_3" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>"
|
||||
.into(),
|
||||
system_message: "You are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.".into(),
|
||||
roles: ("user".into(), "assistant".into()),
|
||||
style: "LLAMA3".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|end_of_text|>".into(),
|
||||
"<|eot_id|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"internlm2-chat" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_COLON_SINGLE".into(),
|
||||
sep: "\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|im_end|>".into(),
|
||||
"<|action_end|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"internvl-2-5" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。".into(),
|
||||
roles: ("<|im_start|>user\n".into(), "<|im_start|>assistant\n".into()),
|
||||
style: "MPT".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|im_end|>".into(),
|
||||
"<|action_end|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"qwen2-vl" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|im_end|>".into()])),
|
||||
..Default::default()
|
||||
},
|
||||
"deepseek-ocr" => LegacySpec {
|
||||
name: name.into(),
|
||||
style: "NO_COLON_SINGLE".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<｜end▁of▁sentence｜>".into()])),
|
||||
..Default::default()
|
||||
},
|
||||
"unlimited-ocr" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}".into(),
|
||||
style: "UNLIMITED_OCR".into(),
|
||||
sep2: Some(String::new()),
|
||||
..Default::default()
|
||||
},
|
||||
"paddle-ocr" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|begin_of_sentence|>{system_message}".into(),
|
||||
roles: ("User".into(), "Assistant".into()),
|
||||
style: "PADDLE_OCR".into(),
|
||||
sep: "<|end_of_sentence|>".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|end_of_sentence|>".into()])),
|
||||
image_token: "<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>".into(),
|
||||
..Default::default()
|
||||
},
|
||||
"deepseek-vl2" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}".into(),
|
||||
roles: ("<|User|>".into(), "<|Assistant|>".into()),
|
||||
style: "DeepSeekVL2".into(),
|
||||
sep: "\n\n".into(),
|
||||
sep2: Some("<｜end▁of▁sentence｜>".into()),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"User:".into(),
|
||||
"<｜end▁of▁sentence｜>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"gemma-it" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<start_of_turn>user\n{system_message}\n\n".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<start_of_turn>user\n".into(), "<start_of_turn>model\n".into()),
|
||||
style: "GEMMA3".into(),
|
||||
sep: "<end_of_turn>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<end_of_turn>".into()])),
|
||||
image_token: "<start_of_image>".into(),
|
||||
audio_token: "<start_of_audio>".into(),
|
||||
..Default::default()
|
||||
},
|
||||
"gme-qwen2-vl" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "QWEN2_VL_EMBED".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::One("<|endoftext|>".into())),
|
||||
..Default::default()
|
||||
},
|
||||
"minicpmv" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}.".into(),
|
||||
system_message: "You are a helpful assistant".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|im_end|>".into(),
|
||||
"<|endoftext|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"janus-pro" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}.".into(),
|
||||
system_message: "You are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language".into(),
|
||||
roles: ("User".into(), "Assistant".into()),
|
||||
style: "ADD_COLON_TWO".into(),
|
||||
sep: "\n\n".into(),
|
||||
sep2: Some("<｜end▁of▁sentence｜>".into()),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|User|>".into(),
|
||||
"<｜end▁of▁sentence｜>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"minicpmo" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
|
||||
.into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|im_end|>".into(),
|
||||
"<|endoftext|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"kimi-vl" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_system|>system<|im_middle|>{system_message}".into(),
|
||||
system_message: "You are a helpful assistant".into(),
|
||||
roles: (
|
||||
"<|im_user|>user<|im_middle|>".into(),
|
||||
"<|im_assistant|>assistant<|im_middle|>".into(),
|
||||
),
|
||||
style: "NO_COLON_SINGLE".into(),
|
||||
sep: "<|im_end|>".into(),
|
||||
stop_str: Some(OneOrMany::One("<|im_end|>".into())),
|
||||
..Default::default()
|
||||
},
|
||||
"qwen2-audio" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "QWEN2_AUDIO".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|im_end|>".into()])),
|
||||
audio_token: "Audio {idx}: <|audio_bos|><|AUDIO|><|audio_eos|>\n".into(),
|
||||
..Default::default()
|
||||
},
|
||||
"moss-vl" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|im_end|>".into()])),
|
||||
..Default::default()
|
||||
},
|
||||
"points-v15-chat" => LegacySpec {
|
||||
name: name.into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|im_end|>".into()])),
|
||||
..Default::default()
|
||||
},
|
||||
"whisper" => LegacySpec {
|
||||
name: name.into(),
|
||||
style: "NO_COLON_SINGLE".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|endoftext|>".into()])),
|
||||
audio_token: String::new(),
|
||||
..Default::default()
|
||||
},
|
||||
_ => return None,
|
||||
};
|
||||
Some(spec)
|
||||
}
|
||||
@@ -0,0 +1,605 @@
|
||||
//! Legacy SGLang conversation-template rendering.
|
||||
|
||||
use dynamo_protocols::types::{
|
||||
ChatCompletionRequestAssistantMessageContent, ChatCompletionRequestAssistantMessageContentPart,
|
||||
ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageContent,
|
||||
ChatCompletionRequestSystemMessageContentPart, ChatCompletionRequestUserMessageContent,
|
||||
ChatCompletionRequestUserMessageContentPart, CreateChatCompletionRequest,
|
||||
};
|
||||
|
||||
use crate::message::types::OneOrMany;
|
||||
|
||||
use super::template::TemplateError;
|
||||
|
||||
/// A legacy conversation template, mirroring Python's `Conversation` fields.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LegacySpec {
|
||||
/// Python `Conversation.name` — drives the CHATGLM round-offset quirk.
|
||||
pub(super) name: String,
|
||||
pub(super) system_template: String,
|
||||
pub(super) system_message: String,
|
||||
/// `(user_role, assistant_role)` — Python `Conversation.roles`.
|
||||
pub(super) roles: (String, String),
|
||||
pub(super) style: String,
|
||||
pub(super) sep: String,
|
||||
/// `None` = Python's `Conversation.sep2` default. Styles that alternate
|
||||
/// seps (`seps[i % 2]`) need it set; Python crashes on `None` there and we
|
||||
/// error deliberately.
|
||||
pub(super) sep2: Option<String>,
|
||||
/// Python `Conversation.stop_str` (`str | list[str] | None`).
|
||||
pub(super) stop_str: Option<OneOrMany<String>>,
|
||||
pub(super) image_token: String,
|
||||
pub(super) audio_token: String,
|
||||
}
|
||||
|
||||
impl Default for LegacySpec {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
name: String::new(),
|
||||
system_template: String::new(),
|
||||
system_message: String::new(),
|
||||
roles: (String::new(), String::new()),
|
||||
style: String::new(),
|
||||
sep: String::new(),
|
||||
sep2: None,
|
||||
stop_str: None,
|
||||
image_token: "<image>".into(),
|
||||
audio_token: "<audio>".into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Native port of Python `generate_chat_conv` + `Conversation.get_prompt()`:
|
||||
/// fold system messages into the system prompt, keep user/assistant messages in
|
||||
/// order, always append the assistant opening, then render per `sep_style`.
|
||||
#[derive(Clone)]
|
||||
pub struct LegacyFormatter {
|
||||
pub(super) spec: LegacySpec,
|
||||
}
|
||||
|
||||
impl LegacyFormatter {
|
||||
pub(super) fn render(
|
||||
&self,
|
||||
request: &CreateChatCompletionRequest,
|
||||
) -> Result<String, TemplateError> {
|
||||
let mut system_message = self.spec.system_message.clone();
|
||||
let mut messages: Vec<(String, String)> = Vec::new();
|
||||
for message in &request.messages {
|
||||
match message {
|
||||
ChatCompletionRequestMessage::System(message) => {
|
||||
system_message = extract_system_text(&message.content)?;
|
||||
}
|
||||
ChatCompletionRequestMessage::User(message) => {
|
||||
let content = match &message.content {
|
||||
ChatCompletionRequestUserMessageContent::Text(text) => text.clone(),
|
||||
ChatCompletionRequestUserMessageContent::Array(parts) => {
|
||||
let mut text = String::new();
|
||||
for part in parts {
|
||||
match part {
|
||||
ChatCompletionRequestUserMessageContentPart::Text(part) => {
|
||||
text.push_str(&part.text);
|
||||
}
|
||||
// Python would splice media tokens in here;
|
||||
// the OpenAI adapter rejects media content
|
||||
// upstream, so this is unreachable — error
|
||||
// rather than silently drop.
|
||||
_ => {
|
||||
return Err(TemplateError::MediaContent { role: "user" });
|
||||
}
|
||||
}
|
||||
}
|
||||
text
|
||||
}
|
||||
};
|
||||
messages.push((self.spec.roles.0.clone(), content));
|
||||
}
|
||||
ChatCompletionRequestMessage::Assistant(message) => {
|
||||
let content = message
|
||||
.content
|
||||
.as_ref()
|
||||
.map(extract_assistant_text)
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
messages.push((self.spec.roles.1.clone(), content));
|
||||
}
|
||||
other => {
|
||||
return Err(TemplateError::UnsupportedRole {
|
||||
role: match other {
|
||||
ChatCompletionRequestMessage::Developer(_) => "developer",
|
||||
ChatCompletionRequestMessage::Tool(_) => "tool",
|
||||
ChatCompletionRequestMessage::Function(_) => "function",
|
||||
_ => unreachable!(),
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
// Python's `generate_chat_conv` appends the assistant opening.
|
||||
messages.push((self.spec.roles.1.clone(), String::new()));
|
||||
self.render_prompt(&system_message, &messages)
|
||||
}
|
||||
|
||||
fn render_prompt(
|
||||
&self,
|
||||
system_message: &str,
|
||||
messages: &[(String, String)],
|
||||
) -> Result<String, TemplateError> {
|
||||
let spec = &self.spec;
|
||||
// Python: `self.system_template.format(system_message=self.system_message)`.
|
||||
let system_prompt = spec
|
||||
.system_template
|
||||
.replace("{system_message}", system_message);
|
||||
let user_role = &spec.roles.0;
|
||||
let assistant_role = &spec.roles.1;
|
||||
let mut ret = String::new();
|
||||
|
||||
match spec.style.as_str() {
|
||||
"ADD_COLON_SINGLE" => {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"ADD_COLON_TWO" => {
|
||||
let sep2 = sep2(spec, "ADD_COLON_TWO")?;
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"ADD_COLON_SPACE_SINGLE" => {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}: "));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"ADD_NEW_LINE_SINGLE" => {
|
||||
if !system_message.is_empty() && !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}\n{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"QWEN2_VL_EMBED" => {
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}\n{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
match &spec.stop_str {
|
||||
Some(OneOrMany::One(stop)) => ret.push_str(stop),
|
||||
// Python `ret += self.stop_str` raises TypeError for
|
||||
// `None` / list; error deliberately instead.
|
||||
_ => {
|
||||
return Err(TemplateError::InvalidStopString {
|
||||
style: "QWEN2_VL_EMBED".into(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
"NO_COLON_SINGLE" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"NO_COLON_TWO" => {
|
||||
let sep2 = sep2(spec, "NO_COLON_TWO")?;
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"RWKV" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"{role}: {}",
|
||||
content.replace("\r\n", "\n").replace("\n\n", "\n")
|
||||
));
|
||||
ret.push_str("\n\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
"LLAMA4" => {
|
||||
if !system_message.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("<|header_start|>{role}<|header_end|>\n\n"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"<|header_start|>{role}<|header_end|>\n\n{}<|eot|>",
|
||||
content.trim()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
"LLAMA3" => {
|
||||
if !system_message.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("<|start_header_id|>{role}<|end_header_id|>\n\n"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"<|start_header_id|>{role}<|end_header_id|>\n\n{}<|eot_id|>",
|
||||
content.trim()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
"LLAMA2" => {
|
||||
let sep2 = sep2(spec, "LLAMA2")?;
|
||||
if system_message.is_empty() {
|
||||
ret.push_str("[INST] ");
|
||||
} else {
|
||||
ret.push_str(&system_prompt);
|
||||
}
|
||||
for (i, (_, content)) in messages.iter().enumerate() {
|
||||
// Python: `tag = self.roles[i % 2]` — parity, not the
|
||||
// stored role, and the first message has no tag.
|
||||
let tag = if i % 2 == 0 {
|
||||
user_role
|
||||
} else {
|
||||
assistant_role
|
||||
};
|
||||
if content.is_empty() {
|
||||
ret.push_str(tag);
|
||||
} else if i == 0 {
|
||||
ret.push_str(&format!("{content} "));
|
||||
} else {
|
||||
ret.push_str(&format!("{tag} {content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"CHATGLM" => {
|
||||
// Python: `round_add_n = 1 if self.name == "chatglm2" else 0`.
|
||||
let round_add_n = if spec.name == "chatglm2" { 1 } else { 0 };
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if i % 2 == 0 {
|
||||
ret.push_str(&format!("[Round {}]{}", i / 2 + round_add_n, spec.sep));
|
||||
}
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}："));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}：{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"CHATML" => {
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
ret.push('\n');
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}\n{content}{}\n", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"CHATGLM3" => {
|
||||
if !system_message.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}\n{content}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
"CHATINTERN" => {
|
||||
let sep2 = sep2(spec, "CHATINTERN")?;
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if i % 2 == 0 {
|
||||
ret.push_str("<s>");
|
||||
}
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"{role}:{content}{}\n",
|
||||
sep_even_odd(spec, sep2, i)
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
"DOLLY" => {
|
||||
let sep2 = sep2(spec, "DOLLY")?;
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:\n"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"{role}:\n{content}{}",
|
||||
sep_even_odd(spec, sep2, i)
|
||||
));
|
||||
if i % 2 == 1 {
|
||||
ret.push_str("\n\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"PHOENIX" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}: <s>"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: <s>{content}</s>"));
|
||||
}
|
||||
}
|
||||
}
|
||||
"ROBIN" => {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}:\n{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"FALCON_CHAT" => {
|
||||
if !system_message.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"METAMATH" => {
|
||||
let sep2 = sep2(spec, "METAMATH")?;
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
// Python: sep2 prefixes odd messages; sep ends even ones.
|
||||
if content.is_empty() {
|
||||
if i % 2 == 0 {
|
||||
ret.push_str(&format!("{role}:\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {sep2}"));
|
||||
}
|
||||
} else if i % 2 == 0 {
|
||||
ret.push_str(&format!("{role}:\n{content}{}", spec.sep));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {sep2}{content}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
"DEEPSEEK_CHAT" => {
|
||||
let sep2 = sep2(spec, "DEEPSEEK_CHAT")?;
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"DeepSeekVL2" => {
|
||||
let sep2 = sep2(spec, "DeepSeekVL2")?;
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"GEMMA3" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else if i == 0 {
|
||||
ret.push_str(&format!("{content}{}", spec.sep));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"MPT" => {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"QWEN2_AUDIO" => {
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
let mut counter = 1usize;
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}\n"));
|
||||
} else {
|
||||
let mut message = content.clone();
|
||||
while message.contains(&spec.audio_token) {
|
||||
// Python: `audio_token.format(idx=counter)`. A
|
||||
// token without `{idx}` makes the replace a no-op
|
||||
// and Python's loop infinite; bail out instead of
|
||||
// hanging the server.
|
||||
let indexed = spec.audio_token.replace("{idx}", &counter.to_string());
|
||||
if indexed == spec.audio_token {
|
||||
break;
|
||||
}
|
||||
message = message.replacen(&spec.audio_token, &indexed, 1);
|
||||
counter += 1;
|
||||
}
|
||||
ret.push_str(&format!("{role}\n{message}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"PADDLE_OCR" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}: "));
|
||||
} else if role == user_role {
|
||||
ret.push_str(&format!("{role}: "));
|
||||
if content.contains(&spec.image_token) {
|
||||
ret.push_str(
|
||||
&content
|
||||
.replace(&format!("{}\n", spec.image_token), &spec.image_token),
|
||||
);
|
||||
} else {
|
||||
ret.push_str(content);
|
||||
}
|
||||
ret.push('\n');
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"UNLIMITED_OCR" => {
|
||||
let sep2 = sep2(spec, "UNLIMITED_OCR")?;
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
other => {
|
||||
return Err(TemplateError::InvalidStyle {
|
||||
style: other.to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(ret)
|
||||
}
|
||||
}
|
||||
|
||||
/// Python `seps = [self.sep, self.sep2]` indexed by message parity.
|
||||
fn sep_even_odd<'a>(spec: &'a LegacySpec, sep2: &'a str, index: usize) -> &'a str {
|
||||
if index.is_multiple_of(2) {
|
||||
&spec.sep
|
||||
} else {
|
||||
sep2
|
||||
}
|
||||
}
|
||||
|
||||
fn sep2<'a>(spec: &'a LegacySpec, style: &str) -> Result<&'a str, TemplateError> {
|
||||
spec.sep2
|
||||
.as_deref()
|
||||
.ok_or_else(|| TemplateError::MissingSep2 {
|
||||
style: style.to_owned(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Python `generate_chat_conv` system extraction: a plain string, or an array
|
||||
/// with exactly one `text` part.
|
||||
fn extract_system_text(
|
||||
content: &ChatCompletionRequestSystemMessageContent,
|
||||
) -> Result<String, TemplateError> {
|
||||
match content {
|
||||
ChatCompletionRequestSystemMessageContent::Text(text) => Ok(text.clone()),
|
||||
ChatCompletionRequestSystemMessageContent::Array(parts) => {
|
||||
let mut texts = parts.iter().map(|part| match part {
|
||||
ChatCompletionRequestSystemMessageContentPart::Text(part) => part.text.as_str(),
|
||||
});
|
||||
match (texts.next(), texts.next()) {
|
||||
(Some(text), None) => Ok(text.to_owned()),
|
||||
_ => Err(TemplateError::NonTextContent { role: "system" }),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Python `generate_chat_conv` assistant extraction: a plain string, or an
|
||||
/// array with exactly one `text` part (`refusal` parts are rejected).
|
||||
fn extract_assistant_text(
|
||||
content: &ChatCompletionRequestAssistantMessageContent,
|
||||
) -> Result<String, TemplateError> {
|
||||
match content {
|
||||
ChatCompletionRequestAssistantMessageContent::Text(text) => Ok(text.clone()),
|
||||
ChatCompletionRequestAssistantMessageContent::Array(parts) => {
|
||||
let mut texts = parts.iter().filter_map(|part| match part {
|
||||
ChatCompletionRequestAssistantMessageContentPart::Text(part) => {
|
||||
Some(part.text.as_str())
|
||||
}
|
||||
ChatCompletionRequestAssistantMessageContentPart::Refusal(_) => None,
|
||||
});
|
||||
match (texts.next(), texts.next()) {
|
||||
(Some(text), None) => Ok(text.to_owned()),
|
||||
_ => Err(TemplateError::NonTextContent { role: "assistant" }),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,351 @@
|
||||
//! Chat-template loading and model-path inference.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use dynamo_renderer::{ChatTemplate, ContextMixins, PromptContextMixin, PromptFormatter};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::message::types::OneOrMany;
|
||||
|
||||
use super::template::{ChatFormatter, TemplateError};
|
||||
use super::template_builtins::builtin_template;
|
||||
use super::template_legacy::{LegacyFormatter, LegacySpec};
|
||||
|
||||
const SUPPORTED_STYLES: &[&str] = &[
|
||||
"ADD_COLON_SINGLE",
|
||||
"ADD_COLON_TWO",
|
||||
"ADD_COLON_SPACE_SINGLE",
|
||||
"NO_COLON_SINGLE",
|
||||
"NO_COLON_TWO",
|
||||
"ADD_NEW_LINE_SINGLE",
|
||||
"LLAMA2",
|
||||
"LLAMA3",
|
||||
"LLAMA4",
|
||||
"CHATGLM",
|
||||
"CHATML",
|
||||
"CHATINTERN",
|
||||
"DOLLY",
|
||||
"RWKV",
|
||||
"PHOENIX",
|
||||
"ROBIN",
|
||||
"FALCON_CHAT",
|
||||
"CHATGLM3",
|
||||
"DEEPSEEK_CHAT",
|
||||
"METAMATH",
|
||||
"DeepSeekVL2",
|
||||
"QWEN2_VL_EMBED",
|
||||
"QWEN2_AUDIO",
|
||||
"GEMMA3",
|
||||
"MPT",
|
||||
"PADDLE_OCR",
|
||||
"UNLIMITED_OCR",
|
||||
];
|
||||
|
||||
pub(super) fn load_chat_formatter(
|
||||
config_file: Option<&str>,
|
||||
model_path: Option<&str>,
|
||||
chat_template_arg: Option<&str>,
|
||||
) -> Result<ChatFormatter, TemplateError> {
|
||||
// Python resolves registry names before looking at the filesystem — and
|
||||
// before touching the tokenizer config, so a built-in name works even when
|
||||
// `tokenizer_config.json` is absent.
|
||||
if let Some(argument) = chat_template_arg
|
||||
&& let Some(spec) = builtin_template(argument)
|
||||
{
|
||||
return Ok(ChatFormatter::Legacy(Box::new(LegacyFormatter { spec })));
|
||||
}
|
||||
|
||||
// Python `load_chat_template` (no `--chat-template`): infer a legacy
|
||||
// template from the model path before falling back to the HF template, so
|
||||
// a legacy model whose config has no `chat_template` still gets one.
|
||||
if chat_template_arg.is_none()
|
||||
&& let Some(model_path) = model_path
|
||||
&& let Some(spec) = infer_legacy_template_from_model_path(model_path)
|
||||
{
|
||||
tracing::info!(%model_path, "inferred legacy chat template from model path");
|
||||
return Ok(ChatFormatter::Legacy(Box::new(LegacyFormatter { spec })));
|
||||
}
|
||||
|
||||
// Every remaining source builds the HF renderer around the tokenizer
|
||||
// config (the template itself, or the argument injected into it).
|
||||
let Some(config_file) = config_file else {
|
||||
return Err(TemplateError::MissingConfig);
|
||||
};
|
||||
|
||||
let config_path = Path::new(config_file);
|
||||
let config_text = read_to_string(config_path, "tokenizer config")?;
|
||||
let mut config = parse_json(&config_text, config_path, "tokenizer config")?;
|
||||
|
||||
let Some(argument) = chat_template_arg else {
|
||||
return formatter_from_config(&config);
|
||||
};
|
||||
|
||||
let path = Path::new(argument);
|
||||
if !path.exists() {
|
||||
return Err(TemplateError::NotFound {
|
||||
path: path.to_path_buf(),
|
||||
});
|
||||
}
|
||||
if !path.is_file() {
|
||||
return Err(TemplateError::NotFile {
|
||||
path: path.to_path_buf(),
|
||||
});
|
||||
}
|
||||
|
||||
if path.extension().and_then(|extension| extension.to_str()) == Some("jinja") {
|
||||
let template = read_to_string(path, "chat template")?;
|
||||
set_chat_template(
|
||||
&mut config,
|
||||
Value::String(template.trim_matches('\n').replace("\\n", "\n")),
|
||||
)?;
|
||||
return formatter_from_config(&config);
|
||||
}
|
||||
|
||||
let template_text = read_to_string(path, "chat template")?;
|
||||
let template = parse_json(&template_text, path, "chat template")?;
|
||||
|
||||
// HF-style JSON files may carry chat_template directly. Legacy SGLang
|
||||
// files carry Conversation fields and are translated below.
|
||||
if let Some(chat_template) = template.get("chat_template") {
|
||||
set_chat_template(&mut config, chat_template.clone())?;
|
||||
formatter_from_config(&config)
|
||||
} else {
|
||||
Ok(ChatFormatter::Legacy(Box::new(LegacyFormatter {
|
||||
spec: parse_legacy_template(&template, path)?,
|
||||
})))
|
||||
}
|
||||
}
|
||||
|
||||
/// Port of Python `get_conv_template_by_model_path` (conversation.py
|
||||
/// `matching_function_registry`, run in registration order): infer a legacy
|
||||
/// built-in template from the model path, optionally consulting the model's
|
||||
/// `config.json` `model_type`. `None` when nothing matches — the HF template
|
||||
/// is the fallback then, as in Python.
|
||||
pub(super) fn infer_legacy_template_from_model_path(model_path: &str) -> Option<LegacySpec> {
|
||||
let lower = model_path.to_lowercase();
|
||||
// Regexes without regex: every Python pattern here is a plain substring or
|
||||
// a `prefix.*suffix` pair, both on a lowercased path.
|
||||
let contains = |needle: &str| lower.contains(needle);
|
||||
let precedes = |prefix: &str, suffix: &str| {
|
||||
lower
|
||||
.find(prefix)
|
||||
.is_some_and(|start| lower[start + prefix.len()..].contains(suffix))
|
||||
};
|
||||
|
||||
if lower
|
||||
.split(|c: char| !c.is_alphanumeric())
|
||||
.any(|word| word == "points")
|
||||
{
|
||||
return builtin_template("points-v15-chat");
|
||||
}
|
||||
if precedes("moss", "vl") {
|
||||
return builtin_template("moss-vl");
|
||||
}
|
||||
if contains("internvl") {
|
||||
return builtin_template("internvl-2-5");
|
||||
}
|
||||
if contains("janus") {
|
||||
return builtin_template("janus-pro");
|
||||
}
|
||||
if contains("vicuna") || contains("llava-v1.5") || contains("llava-next-video-7b") {
|
||||
return builtin_template("vicuna_v1.1");
|
||||
}
|
||||
if precedes("deepseek", "vl2") {
|
||||
return builtin_template("deepseek-vl2");
|
||||
}
|
||||
if contains("llava-v1.6-34b")
|
||||
|| contains("llava-v1.6-yi-34b")
|
||||
|| contains("llava-next-video-34b")
|
||||
|| contains("llava-onevision-qwen2")
|
||||
{
|
||||
return builtin_template("chatml-llava");
|
||||
}
|
||||
// MiniCPM: 4.6+ uses its own template and must not fall back to the
|
||||
// legacy conv template.
|
||||
if contains("minicpm-v-4.6")
|
||||
|| contains("minicpm-v-4_6")
|
||||
|| contains("minicpm-o-4.6")
|
||||
|| contains("minicpm-o-4_6")
|
||||
{
|
||||
return None;
|
||||
}
|
||||
if contains("minicpm-v") {
|
||||
return builtin_template("minicpmv");
|
||||
}
|
||||
if contains("minicpm-o") {
|
||||
return builtin_template("minicpmo");
|
||||
}
|
||||
if contains("phi-4-multimodal") {
|
||||
return builtin_template("phi-4-mm");
|
||||
}
|
||||
if contains("deepseek-ocr") {
|
||||
return builtin_template("deepseek-ocr");
|
||||
}
|
||||
if contains("unlimited") {
|
||||
return builtin_template("unlimited-ocr");
|
||||
}
|
||||
if contains("paddleocr") {
|
||||
return builtin_template("paddle-ocr");
|
||||
}
|
||||
if contains("whisper") {
|
||||
return builtin_template("whisper");
|
||||
}
|
||||
|
||||
// Model-type matchers read `<model_path>/config.json` (local dirs only —
|
||||
// Python's `get_model_type` cannot resolve HF repo ids either).
|
||||
let model_type = read_model_type(model_path)?;
|
||||
// Python `MODEL_TYPE_TO_TEMPLATE`; minicpmv4_6 is deliberately absent.
|
||||
let name = match model_type.as_str() {
|
||||
"moss_vl" => "moss-vl",
|
||||
"internvl_chat" => "internvl-2-5",
|
||||
"multi_modality" => "janus-pro",
|
||||
"deepseek_vl_v2" => "deepseek-vl2",
|
||||
"minicpmv" => "minicpmv",
|
||||
"minicpmo" => "minicpmo",
|
||||
"phi4mm" => "phi-4-mm",
|
||||
"deepseek-ocr" => "deepseek-ocr",
|
||||
"unlimited-ocr" => "unlimited-ocr",
|
||||
"paddleocr_vl" => "paddle-ocr",
|
||||
_ => return None,
|
||||
};
|
||||
builtin_template(name)
|
||||
}
|
||||
|
||||
/// Python `get_model_type`: the `model_type` field of the model's `config.json`.
|
||||
fn read_model_type(model_path: &str) -> Option<String> {
|
||||
let config_path = Path::new(model_path).join("config.json");
|
||||
if !config_path.is_file() {
|
||||
return None;
|
||||
}
|
||||
let config: Value = parse_json(
|
||||
&read_to_string(&config_path, "model config").ok()?,
|
||||
&config_path,
|
||||
"model config",
|
||||
)
|
||||
.ok()?;
|
||||
config.get("model_type")?.as_str().map(str::to_owned)
|
||||
}
|
||||
|
||||
fn read_to_string(path: &Path, kind: &'static str) -> Result<String, TemplateError> {
|
||||
std::fs::read_to_string(path).map_err(|source| TemplateError::Read {
|
||||
kind,
|
||||
path: path.to_path_buf(),
|
||||
source,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_json(text: &str, path: &Path, kind: &'static str) -> Result<Value, TemplateError> {
|
||||
serde_json::from_str(text).map_err(|source| TemplateError::Parse {
|
||||
kind,
|
||||
path: path.to_path_buf(),
|
||||
source,
|
||||
})
|
||||
}
|
||||
|
||||
fn set_chat_template(config: &mut Value, chat_template: Value) -> Result<(), TemplateError> {
|
||||
let Some(config) = config.as_object_mut() else {
|
||||
return Err(TemplateError::ConfigNotObject);
|
||||
};
|
||||
config.insert("chat_template".to_string(), chat_template);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn formatter_from_config(config: &Value) -> Result<ChatFormatter, TemplateError> {
|
||||
let template: ChatTemplate = serde_json::from_value(config.clone())
|
||||
.map_err(|source| TemplateError::Config { source })?;
|
||||
if template.chat_template.is_none() {
|
||||
return Err(TemplateError::Missing);
|
||||
}
|
||||
let formatter = PromptFormatter::from_parts(
|
||||
template,
|
||||
ContextMixins::new(&[PromptContextMixin::OaiChat]),
|
||||
true,
|
||||
)
|
||||
.map_err(|error| TemplateError::Renderer {
|
||||
message: error.to_string(),
|
||||
})?;
|
||||
Ok(ChatFormatter::HuggingFace(formatter))
|
||||
}
|
||||
|
||||
/// Port of Python `_load_json_chat_template`: fields mirror `Conversation`
|
||||
/// exactly (missing `sep2`/`image_token`/`audio_token` stay at Python defaults).
|
||||
fn parse_legacy_template(value: &Value, path: &Path) -> Result<LegacySpec, TemplateError> {
|
||||
let object = value
|
||||
.as_object()
|
||||
.ok_or_else(|| TemplateError::LegacyNotObject {
|
||||
path: path.to_path_buf(),
|
||||
})?;
|
||||
|
||||
let required_string = |name: &str| -> Result<String, TemplateError> {
|
||||
object
|
||||
.get(name)
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.ok_or_else(|| TemplateError::LegacyMissingField {
|
||||
path: path.to_path_buf(),
|
||||
field: name.to_string(),
|
||||
})
|
||||
};
|
||||
|
||||
let style = required_string("sep_style")?;
|
||||
if !SUPPORTED_STYLES.contains(&style.as_str()) {
|
||||
return Err(TemplateError::UnknownStyle {
|
||||
path: path.to_path_buf(),
|
||||
style,
|
||||
});
|
||||
}
|
||||
|
||||
// Python `Conversation.stop_str: str | list[str] | None` — the key is
|
||||
// required (`template["stop_str"]` raises KeyError when missing), but an
|
||||
// explicit `null` value means `None`.
|
||||
let stop_str = match object.get("stop_str") {
|
||||
Some(Value::String(value)) => Some(OneOrMany::One(value.clone())),
|
||||
Some(Value::Array(values)) => {
|
||||
let strings = values
|
||||
.iter()
|
||||
.map(Value::as_str)
|
||||
.collect::<Option<Vec<_>>>()
|
||||
.ok_or_else(|| TemplateError::LegacyMissingField {
|
||||
path: path.to_path_buf(),
|
||||
field: "stop_str".to_string(),
|
||||
})?;
|
||||
Some(OneOrMany::Many(
|
||||
strings.into_iter().map(str::to_owned).collect(),
|
||||
))
|
||||
}
|
||||
Some(Value::Null) => None,
|
||||
Some(_) => {
|
||||
return Err(TemplateError::LegacyMissingField {
|
||||
path: path.to_path_buf(),
|
||||
field: "stop_str".to_string(),
|
||||
});
|
||||
}
|
||||
None => {
|
||||
return Err(TemplateError::LegacyMissingField {
|
||||
path: path.to_path_buf(),
|
||||
field: "stop_str".to_string(),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
Ok(LegacySpec {
|
||||
name: required_string("name")?,
|
||||
// Python: `system_template=template["system"] + "\n{system_message}"`.
|
||||
system_template: format!("{}\n{{system_message}}", required_string("system")?),
|
||||
system_message: object
|
||||
.get("system_message")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
roles: (required_string("user")?, required_string("assistant")?),
|
||||
style,
|
||||
sep: object
|
||||
.get("sep")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("\n")
|
||||
.to_string(),
|
||||
sep2: None,
|
||||
stop_str,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
@@ -34,31 +34,34 @@ fn value_error(context: &str, err: impl std::fmt::Display) -> PyErr {
|
||||
pyo3::exceptions::PyValueError::new_err(format!("{context}: {err}"))
|
||||
}
|
||||
|
||||
/// One drained MM result (see [`Server::take_mm`]), consumed by
|
||||
/// `RustServer.build_native_mm` to build the scheduler's
|
||||
/// One drained MM result (see [`Server::take_mm_result`]), consumed by
|
||||
/// `RustMmProcessor.build_output` to build the scheduler's
|
||||
/// `MultimodalProcessorOutput`.
|
||||
#[pyclass(frozen, get_all)]
|
||||
struct MmEncodeResult {
|
||||
/// *Generic.* All items' `pixel_values` concatenated, flat `f32` of logical
|
||||
/// shape `[sum(t*h*w), feature_dim]`; `Some` on the inline (single-rank) path.
|
||||
// General fields.
|
||||
/// All items' `pixel_values` concatenated as flat `f32` with logical shape
|
||||
/// `[sum(t*h*w), feature_dim]`; present on the inline (single-rank) path.
|
||||
features: Option<Py<numpy::PyArray1<f32>>>,
|
||||
/// *Generic.* Per-item POSIX shm segment name holding that item's features
|
||||
/// (`[t*h*w, feature_dim]` f32); `Some` on the TP-broadcast path.
|
||||
/// Per-item POSIX shared-memory segment holding `[t*h*w, feature_dim]` f32
|
||||
/// features; present on the TP-broadcast path.
|
||||
shm_names: Option<Vec<String>>,
|
||||
/// *Generic.* Per-item content hash of the raw source bytes (or the caller's
|
||||
/// `mm_hashes` override), precomputed so the drain never re-hashes.
|
||||
/// Per-item content hash of the raw source bytes, or the caller-provided
|
||||
/// `mm_hashes` override, precomputed so draining never re-hashes.
|
||||
hashes: Vec<u64>,
|
||||
/// *Generic.* Per-item inclusive `(start, end)` placeholder-token span in the
|
||||
/// expanded `input_ids`.
|
||||
/// Per-item inclusive `(start, end)` placeholder-token span in the expanded
|
||||
/// `input_ids`.
|
||||
offsets: Vec<(u32, u32)>,
|
||||
/// *Qwen-VL specific.* Per-item `image_grid_thw` `(t, h, w)` in patch units;
|
||||
/// `t*h*w` is also the item's row count in `features`.
|
||||
|
||||
// Qwen-VL-specific fields.
|
||||
/// Per-item `image_grid_thw` `(t, h, w)` in patch units; `t*h*w` is also the
|
||||
/// item's row count in `features`.
|
||||
grids: Vec<(u32, u32, u32)>,
|
||||
/// *Qwen-VL specific.* M-RoPE position ids, flat `i64` of row-major shape
|
||||
/// `[3, seq_len]` (temporal, height, width rows).
|
||||
/// M-RoPE position ids as flat `i64` with row-major shape `[3, seq_len]`
|
||||
/// (temporal, height, and width rows).
|
||||
mrope: Py<numpy::PyArray1<i64>>,
|
||||
/// *Qwen-VL specific.* M-RoPE delta, `max(mrope) + 1 - seq_len`, that decode
|
||||
/// adds to the plain sequence position.
|
||||
/// M-RoPE delta, `max(mrope) + 1 - seq_len`, added to the plain sequence
|
||||
/// position during decoding.
|
||||
mrope_delta: i64,
|
||||
}
|
||||
|
||||
@@ -93,7 +96,7 @@ impl Server {
|
||||
http_addr = None,
|
||||
to_scheduler_cap = 8192,
|
||||
from_scheduler_cap = 8192,
|
||||
channel_cap = 8192,
|
||||
stage_channel_cap = 8192,
|
||||
cores = None,
|
||||
))]
|
||||
// pyo3 `#[new]` constructor: the wide arg list is the Python-facing boot
|
||||
@@ -104,7 +107,7 @@ impl Server {
|
||||
http_addr: Option<String>,
|
||||
to_scheduler_cap: usize,
|
||||
from_scheduler_cap: usize,
|
||||
channel_cap: usize,
|
||||
stage_channel_cap: usize,
|
||||
cores: Option<Vec<usize>>,
|
||||
) -> PyResult<Self> {
|
||||
// `server_args` already arrived typed (pyo3 rejected any missing/extra/
|
||||
@@ -127,7 +130,7 @@ impl Server {
|
||||
http_api_worker_num: server_args.http_api_worker_num(),
|
||||
to_scheduler_cap,
|
||||
from_scheduler_cap,
|
||||
channel_cap,
|
||||
stage_channel_cap,
|
||||
cores,
|
||||
},
|
||||
server_args: std::sync::Arc::new(server_args),
|
||||
@@ -201,10 +204,10 @@ impl Server {
|
||||
}
|
||||
|
||||
/// Spawn the MM worker pool for the pipeline in `spec` (built from the
|
||||
/// resolved processor config; see `NativeMmHost.resolve_native_spec` and
|
||||
/// resolved processor config; see `RustMmProcessor.resolve_spec` and
|
||||
/// `RustServer._build_mm_spec`). Image-only requests are processed entirely
|
||||
/// in Rust and parked for [`Server::take_mm`]; anything the pipeline cannot
|
||||
/// serve is rejected back to the client — there is no Python fallback.
|
||||
/// in Rust and parked for [`Server::take_mm_result`]; anything the pipeline
|
||||
/// cannot serve is rejected back to the client — there is no Python fallback.
|
||||
fn start_mm_workers(&self, spec: MmSpec, workers: usize) -> PyResult<()> {
|
||||
let ctx = multi_modality::worker::Context::new(
|
||||
spec,
|
||||
@@ -224,7 +227,7 @@ impl Server {
|
||||
/// Runs on the scheduler loop between decode steps, so any per-byte work
|
||||
/// here — memcpy or hashing, tens of MB per image-heavy request — would
|
||||
/// stall every running request's ITL. Hence the worker-precomputed `hashes`.
|
||||
fn take_mm(&self, py: Python<'_>, rid: &str) -> Option<MmEncodeResult> {
|
||||
fn take_mm_result(&self, py: Python<'_>, rid: &str) -> Option<MmEncodeResult> {
|
||||
use numpy::IntoPyArray;
|
||||
|
||||
let res = self.rt.mm_sidecar.take(rid)?;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! Runtime configuration: the rust-server boot knobs
|
||||
//! ([`RustServerServerArgs`]), the scheduler's typed `server_args` handoff
|
||||
//! ([`ServerArgs`] / [`ModelConfig`]), the [`RuntimeConfig`] pairing them for
|
||||
//! `runtime::start`, and the native MM pipeline handoff ([`MmSpec`]).
|
||||
//! `runtime::start`, and the Rust MM pipeline handoff ([`MmSpec`]).
|
||||
//!
|
||||
//! [`ServerArgs`] / [`ModelConfig`] / [`DefaultSamplingParams`] /
|
||||
//! [`DisaggregationMode`] / [`MmSpec`] / [`MmFamily`] / [`MmResample`] are
|
||||
@@ -29,7 +29,7 @@ pub struct RustServerServerArgs {
|
||||
pub http_api_worker_num: usize,
|
||||
pub to_scheduler_cap: usize,
|
||||
pub from_scheduler_cap: usize,
|
||||
pub channel_cap: usize,
|
||||
pub stage_channel_cap: usize,
|
||||
/// CPU core ids the pools pin to (e.g. this rank's NUMA-local cores minus
|
||||
/// the scheduler's reserved launch cores). `None` → run unpinned.
|
||||
pub cores: Option<Vec<usize>>,
|
||||
@@ -42,7 +42,7 @@ impl Default for RustServerServerArgs {
|
||||
http_api_worker_num: 2,
|
||||
to_scheduler_cap: 8192,
|
||||
from_scheduler_cap: 8192,
|
||||
channel_cap: 8192,
|
||||
stage_channel_cap: 8192,
|
||||
cores: None,
|
||||
}
|
||||
}
|
||||
@@ -404,15 +404,15 @@ impl DefaultSamplingParams {
|
||||
}
|
||||
}
|
||||
|
||||
/// The native MM pipeline handoff, built by `RustServer._build_mm_spec` from
|
||||
/// the resolved `NativeMmSpec` and passed to `Server.start_mm_workers`. Same
|
||||
/// The Rust MM pipeline handoff, built by `RustServer._build_mm_spec` from
|
||||
/// the resolved `RustMmSpec` and passed to `Server.start_mm_workers`. Same
|
||||
/// contract as [`ServerArgs`]: every field is a required, typed constructor
|
||||
/// keyword, so a drifted Python caller fails at boot.
|
||||
#[pyo3::pyclass(frozen, from_py_object, module = "sglang.srt.rust_extensions._server")]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct MmSpec {
|
||||
/// Park feature buffers in POSIX shm rather than inline. Set by the Python
|
||||
/// launcher (`NativeMmHost._use_feature_shm`) exactly when the scheduler
|
||||
/// launcher (`RustMmProcessor._use_feature_shm`) exactly when the scheduler
|
||||
/// broadcasts across TP ranks and will unwrap `ShmPointerMMData`.
|
||||
pub feature_shm: bool,
|
||||
/// The family pipeline and its resolved processor parameters.
|
||||
@@ -475,7 +475,7 @@ impl MmSpec {
|
||||
|
||||
/// Which `sglang_mm` family pipeline serves the model — one variant per
|
||||
/// [`sglang_mm::registry::PipelineSpec`] arm. Exposed to Python as an enum
|
||||
/// (`MmFamily.QwenVl`); `NativeMmFamily.name` maps onto it at handoff.
|
||||
/// (`MmFamily.QwenVl`); `RustMmFamily.name` maps onto it at handoff.
|
||||
#[pyo3::pyclass(
|
||||
eq,
|
||||
frozen,
|
||||
@@ -487,9 +487,9 @@ pub enum MmFamily {
|
||||
QwenVl,
|
||||
}
|
||||
|
||||
/// The HF image processor the native resize must reproduce bit-exactly (see
|
||||
/// The HF image processor the Rust resize must reproduce bit-exactly (see
|
||||
/// [`sglang_mm::qwen_vl::Resampler`]). Exposed to Python as an enum
|
||||
/// (`MmResample.AtenU8` / `.Pil`); `NativeMmFamily.image_processors` maps each
|
||||
/// (`MmResample.AtenU8` / `.Pil`); `RustMmFamily.image_processors` maps each
|
||||
/// processor class onto it.
|
||||
#[pyo3::pyclass(
|
||||
eq,
|
||||
|
||||
@@ -72,7 +72,7 @@ fn max_new_tokens_default() -> Option<i64> {
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct SamplingParams {
|
||||
// --- API parameters (set by callers) ---
|
||||
// Output length and stopping.
|
||||
#[serde(default = "max_new_tokens_default")]
|
||||
pub max_new_tokens: Option<i64>,
|
||||
/// API input alias, copied to `stop_strs` then cleared by `normalize`.
|
||||
@@ -85,6 +85,8 @@ pub struct SamplingParams {
|
||||
/// API input alias, copied to `stop_regex_strs` then cleared by `normalize`.
|
||||
#[serde(default)]
|
||||
pub stop_regex: Option<OneOrMany<String>>,
|
||||
|
||||
// Sampling distribution and penalties.
|
||||
#[serde(
|
||||
default = "f64_one::default",
|
||||
deserialize_with = "f64_one::deserialize"
|
||||
@@ -125,6 +127,8 @@ pub struct SamplingParams {
|
||||
deserialize_with = "i64_zero::deserialize"
|
||||
)]
|
||||
pub min_new_tokens: i64,
|
||||
|
||||
// Sequence count and beam search.
|
||||
#[serde(
|
||||
default = "i64_one::default",
|
||||
deserialize_with = "i64_one::deserialize"
|
||||
@@ -134,6 +138,8 @@ pub struct SamplingParams {
|
||||
/// positional wire layout even though the rust path rejects it below.
|
||||
#[serde(default)]
|
||||
pub beam_width: Option<i64>,
|
||||
|
||||
// Structured-output constraints.
|
||||
#[serde(default)]
|
||||
pub json_schema: Option<String>,
|
||||
#[serde(default)]
|
||||
@@ -142,6 +148,8 @@ pub struct SamplingParams {
|
||||
pub ebnf: Option<String>,
|
||||
#[serde(default)]
|
||||
pub structural_tag: Option<String>,
|
||||
|
||||
// Output handling.
|
||||
#[serde(
|
||||
default = "bool_false::default",
|
||||
deserialize_with = "bool_false::deserialize"
|
||||
@@ -164,6 +172,8 @@ pub struct SamplingParams {
|
||||
pub no_stop_trim: bool,
|
||||
#[serde(default)]
|
||||
pub stream_interval: Option<i64>,
|
||||
|
||||
// Logit processing and reproducibility.
|
||||
/// Token id (as a string key, matching Python) → bias. Keys are vocab-bounded
|
||||
/// by [`verify`](Self::verify).
|
||||
#[serde(default)]
|
||||
@@ -175,7 +185,7 @@ pub struct SamplingParams {
|
||||
#[serde(default)]
|
||||
pub custom_params: Option<serde_json::Value>,
|
||||
|
||||
// --- Internal fields (populated by the pipeline below, not API-facing) ---
|
||||
// Normalized internal fields.
|
||||
//
|
||||
// All `skip_deserializing`: they are outputs of `normalize`, and a client that
|
||||
// could set them would be setting the pipeline's own state. `is_normalized` is
|
||||
|
||||
@@ -52,7 +52,7 @@ pub struct Context {
|
||||
pub tokenizer: Option<Arc<dyn TextTokenizer>>,
|
||||
pub sidecar: Sidecar,
|
||||
/// Park feature buffers in POSIX shm. Set by the Python launcher
|
||||
/// (`NativeMmHost._use_feature_shm`) exactly when the scheduler broadcasts
|
||||
/// (`RustMmProcessor._use_feature_shm`) exactly when the scheduler broadcasts
|
||||
/// across TP ranks and will unwrap `ShmPointerMMData`.
|
||||
pub feature_shm: bool,
|
||||
}
|
||||
|
||||
@@ -4,5 +4,7 @@ pub mod channel;
|
||||
pub mod detokenizer;
|
||||
pub mod from_scheduler;
|
||||
pub mod to_scheduler;
|
||||
mod to_scheduler_types;
|
||||
mod to_scheduler_validation;
|
||||
pub mod tokenizer;
|
||||
pub mod wiring;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,892 @@
|
||||
//! Tests for scheduler intake.
|
||||
|
||||
use super::*;
|
||||
use crate::message::request::GenerateRequest;
|
||||
use crate::message::response::ResponseSink;
|
||||
use crate::message::sampling::SamplingParams;
|
||||
use crate::tokenizer_manager::channel::{ToSchedulerRx, to_scheduler};
|
||||
use crate::utils::fsm::RequestState;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// An `Intake` plus its detok-shard receiver, to_scheduler channel consumer (keep alive —
|
||||
/// dropping it closes the channel → false QueueFull), tm inbox sender, and the
|
||||
/// mm-pool receiver (keep alive — dropping it makes mm submits fail).
|
||||
fn make_intake() -> (
|
||||
Intake,
|
||||
flume::Receiver<DetokMsg>,
|
||||
ToSchedulerRx,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
make_intake_with(test_limits())
|
||||
}
|
||||
|
||||
fn make_intake_with_abort(
|
||||
abort_rx: flume::Receiver<AbortSource>,
|
||||
) -> (
|
||||
Intake,
|
||||
flume::Receiver<DetokMsg>,
|
||||
ToSchedulerRx,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
make_intake_inner(test_limits(), abort_rx)
|
||||
}
|
||||
|
||||
fn make_intake_with(
|
||||
limits: Limits,
|
||||
) -> (
|
||||
Intake,
|
||||
flume::Receiver<DetokMsg>,
|
||||
ToSchedulerRx,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx
|
||||
make_intake_inner(limits, abort_rx)
|
||||
}
|
||||
|
||||
fn make_intake_inner(
|
||||
limits: Limits,
|
||||
abort_rx: flume::Receiver<AbortSource>,
|
||||
) -> (
|
||||
Intake,
|
||||
flume::Receiver<DetokMsg>,
|
||||
ToSchedulerRx,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
let (tok_tx, _tok_rx) = flume::unbounded();
|
||||
let (detok_tx, detok_rx) = flume::unbounded();
|
||||
let senders = Senders {
|
||||
tok_manager_tx: flume::unbounded().0,
|
||||
abort_tx: flume::unbounded().0,
|
||||
tokenizer_tx: tok_tx,
|
||||
detokenizer_tx: vec![detok_tx],
|
||||
};
|
||||
let (to_scheduler_tx, consumer) = to_scheduler(16);
|
||||
let (tm_tx, tm_rx) = flume::unbounded();
|
||||
let (mm_tx, mm_rx) = flume::unbounded();
|
||||
// Keep the shutdown sender alive (leak) so its branch never fires — tests
|
||||
// end `run` by dropping `tm_tx`, not by shutdown.
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let intake = Intake::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
to_scheduler_tx,
|
||||
limits,
|
||||
test_mm(mm_tx, true),
|
||||
sd_rx,
|
||||
);
|
||||
(intake, detok_rx, consumer, tm_tx, mm_rx)
|
||||
}
|
||||
|
||||
/// An [`Mm`] over `tx` with a fresh sidecar.
|
||||
fn test_mm(tx: flume::Sender<MmRequest>, enabled: bool) -> Mm {
|
||||
Mm {
|
||||
enabled,
|
||||
tx,
|
||||
sidecar: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Both abort sources do the same two things: drop the detok entry so no
|
||||
/// further chunk can be delivered, and tell the scheduler to stop generating.
|
||||
///
|
||||
/// Neither releases anything, and nothing needs them to. Release ordering used
|
||||
/// to be the delicate part here — `AbortGuard::drop` releasing a rid right
|
||||
/// after enqueuing the abort ordered the SEND, not the EFFECT, so a retry of
|
||||
/// the same rid could `Register` ahead of the stale abort and be torn down by
|
||||
/// it. `Rid::from_client` removes the premise: a retry carries a different
|
||||
/// `Rid`, so no abort in flight can name it.
|
||||
#[test]
|
||||
fn every_abort_source_deregisters_and_stops_the_scheduler() {
|
||||
for source in [
|
||||
AbortSource::Guard("x".into()),
|
||||
AbortSource::Detok("x".into()),
|
||||
] {
|
||||
let (detok_tx, detok_rx) = flume::unbounded::<DetokMsg>();
|
||||
let (to_scheduler_tx, consumer) = to_scheduler(16);
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let mut intake = Intake::new(
|
||||
flume::unbounded().1,
|
||||
flume::unbounded().1,
|
||||
Senders {
|
||||
tok_manager_tx: flume::unbounded().0,
|
||||
abort_tx: flume::unbounded().0,
|
||||
tokenizer_tx: flume::unbounded().0,
|
||||
detokenizer_tx: vec![detok_tx],
|
||||
},
|
||||
to_scheduler_tx,
|
||||
test_limits(),
|
||||
test_mm(flume::unbounded().0, true),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
intake.on_abort(source.clone());
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "x"),
|
||||
"{source:?} must drop the detok entry",
|
||||
);
|
||||
assert_eq!(
|
||||
consumer.drain(8).headers.len(),
|
||||
1,
|
||||
"{source:?} must push an AbortReq so the scheduler stops",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// A context ceiling high enough that only a test which sets one on purpose
|
||||
/// can reach it. `context_len` is mandatory now, so "no ceiling" has to be a
|
||||
/// large number rather than `None`; kept well below `u64::MAX` so the
|
||||
/// `as i64` in the auto-truncate clamp cannot go negative if a future test
|
||||
/// does reach this path.
|
||||
const NO_CONTEXT_CEILING: u64 = 1 << 40;
|
||||
|
||||
/// The default test limits: a real tokenizer, vocab 1000, no context ceiling.
|
||||
/// Spelled out rather than `..Default::default()` — `Limits` deliberately has
|
||||
/// no `Default`, because a zero `vocab_size`/`context_len` would reject every
|
||||
/// request instead of behaving like "unset".
|
||||
fn test_limits() -> Limits {
|
||||
Limits {
|
||||
skip_tokenizer_init: false,
|
||||
vocab_size: 1000,
|
||||
context_len: NO_CONTEXT_CEILING,
|
||||
num_reserved_tokens: 0,
|
||||
allow_auto_truncate: false,
|
||||
enable_return_hidden_states: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn generate_req(id: u64, sampling_params: SamplingParams) -> Request {
|
||||
let (tx, _rx) = mpsc::channel(8);
|
||||
Request {
|
||||
rid: id.to_string().into(),
|
||||
state: RequestState::Received,
|
||||
sink: ResponseSink::Local(tx),
|
||||
kind: RequestKind::Generate(Box::new(GenerateRequest {
|
||||
rid: id.to_string().into(),
|
||||
input_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params,
|
||||
..Default::default()
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// `input + max_new_tokens` past the context window is an actionable 400, not a
|
||||
/// silently truncated 200 (Python `TokenizerManager._validate_one_request`).
|
||||
/// The message names both halves so the client can fix the right one.
|
||||
#[test]
|
||||
fn total_tokens_over_context_is_rejected() {
|
||||
let limits = Limits {
|
||||
context_len: 10,
|
||||
..test_limits()
|
||||
};
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens: Some(100),
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let err = check_total_tokens(&mut g, &limits).unwrap_err();
|
||||
let msg = err.to_string();
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(msg.contains("total of 103 tokens"), "{msg}");
|
||||
assert!(msg.contains("3 tokens from the input"), "{msg}");
|
||||
assert!(msg.contains("100 tokens for the completion"), "{msg}");
|
||||
// Exactly filling the window is allowed (Python compares with `>`).
|
||||
g.sampling_params.max_new_tokens = Some(7);
|
||||
assert!(check_total_tokens(&mut g, &limits).is_ok());
|
||||
assert_eq!(g.sampling_params.max_new_tokens, Some(7), "left alone");
|
||||
}
|
||||
|
||||
/// The reserved slots (eagle draft tokens) count as input, so a request can be
|
||||
/// rejected for them even when the prompt alone would fit.
|
||||
#[test]
|
||||
fn reserved_tokens_count_toward_the_limit() {
|
||||
let limits = Limits {
|
||||
context_len: 10,
|
||||
num_reserved_tokens: 5,
|
||||
..test_limits()
|
||||
};
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens: Some(3), // 3 + 3 fits, but 3 + 5 + 3 does not
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let msg = check_total_tokens(&mut g, &limits).unwrap_err().to_string();
|
||||
assert!(msg.contains("8 tokens from the input"), "{msg}");
|
||||
}
|
||||
|
||||
/// `--allow-auto-truncate` opts into clamping instead of rejecting; with no
|
||||
/// context length, or no `max_new_tokens` cap, there is nothing to check.
|
||||
#[test]
|
||||
fn auto_truncate_clamps_and_unknowns_skip() {
|
||||
let sp = |max_new_tokens| SamplingParams {
|
||||
max_new_tokens,
|
||||
..Default::default()
|
||||
};
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: sp(Some(100)),
|
||||
..Default::default()
|
||||
};
|
||||
let truncating = Limits {
|
||||
context_len: 10,
|
||||
allow_auto_truncate: true,
|
||||
..test_limits()
|
||||
};
|
||||
assert!(check_total_tokens(&mut g, &truncating).is_ok());
|
||||
assert_eq!(g.sampling_params.max_new_tokens, Some(7), "clamped to fit");
|
||||
|
||||
// Unknown context length → no ceiling to enforce.
|
||||
g.sampling_params = sp(Some(100));
|
||||
assert!(check_total_tokens(&mut g, &test_limits()).is_ok());
|
||||
assert_eq!(g.sampling_params.max_new_tokens, Some(100), "untouched");
|
||||
|
||||
// No cap requested → nothing to add to the input length, but the input
|
||||
// itself is still checked (see `input_length_is_checked_unconditionally`).
|
||||
g.sampling_params = sp(None);
|
||||
let roomy = Limits {
|
||||
context_len: 100,
|
||||
..test_limits()
|
||||
};
|
||||
assert!(check_total_tokens(&mut g, &roomy).is_ok());
|
||||
}
|
||||
|
||||
/// `max_new_tokens: null` means "no cap", NOT "skip the checks" — the input
|
||||
/// alone must still fit. Gating the whole function on `max_new_tokens` let an
|
||||
/// over-long prompt through to the scheduler with no error at all.
|
||||
/// Python compares with `>=`: a prompt that exactly fills the window leaves no
|
||||
/// room to generate.
|
||||
#[test]
|
||||
fn input_length_is_checked_unconditionally() {
|
||||
let limits = Limits {
|
||||
context_len: 3,
|
||||
..test_limits()
|
||||
};
|
||||
let req = |max_new_tokens| GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]), // exactly fills a 3-token window
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens,
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
for max_new_tokens in [None, Some(1)] {
|
||||
let err = check_total_tokens(&mut req(max_new_tokens), &limits)
|
||||
.expect_err("input == context_len must be rejected (Python uses >=)");
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("longer than the model's context"));
|
||||
}
|
||||
// One token shorter fits, with or without a cap.
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2]),
|
||||
..Default::default()
|
||||
};
|
||||
g.sampling_params.max_new_tokens = None;
|
||||
assert!(check_total_tokens(&mut g, &limits).is_ok());
|
||||
|
||||
// Under auto-truncate the input is cut to fit instead of rejected.
|
||||
let truncating = Limits {
|
||||
allow_auto_truncate: true,
|
||||
..limits.clone()
|
||||
};
|
||||
let mut g = req(None);
|
||||
assert!(check_total_tokens(&mut g, &truncating).is_ok());
|
||||
assert_eq!(
|
||||
g.input_ids.as_deref(),
|
||||
Some(&[1, 2, 3][..]),
|
||||
"fits at the cap"
|
||||
);
|
||||
}
|
||||
|
||||
/// The clamp runs AFTER `verify` (which happens in `Normalizing`), so lowering
|
||||
/// `max_new_tokens` can leave `min_new_tokens > max_new_tokens`. Nothing
|
||||
/// downstream re-checks — `is_normalized: true` makes the scheduler's own
|
||||
/// verify early-return — so the clamp has to re-assert it here.
|
||||
#[test]
|
||||
fn auto_truncate_cannot_invert_min_and_max_new_tokens() {
|
||||
let limits = Limits {
|
||||
context_len: 10,
|
||||
allow_auto_truncate: true,
|
||||
..test_limits()
|
||||
};
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]), // clamps max_new_tokens to 7
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens: Some(100),
|
||||
min_new_tokens: 50, // …which is below min_new_tokens
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let err = check_total_tokens(&mut g, &limits)
|
||||
.expect_err("a clamp that inverts min/max must 400, not ride the wire");
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("min_new_tokens"), "{err}");
|
||||
|
||||
// A clamp that keeps the invariant still clamps.
|
||||
g.sampling_params.min_new_tokens = 2;
|
||||
g.sampling_params.max_new_tokens = Some(100);
|
||||
assert!(check_total_tokens(&mut g, &limits).is_ok());
|
||||
assert_eq!(g.sampling_params.max_new_tokens, Some(7));
|
||||
}
|
||||
|
||||
/// `return_hidden_states` on a server not launched for it is a 400: the
|
||||
/// scheduler never computes them, so the request would otherwise 200 with
|
||||
/// `meta_info.hidden_states` silently missing.
|
||||
#[test]
|
||||
fn hidden_states_gated_on_server_support() {
|
||||
let req = |want| {
|
||||
let mut r = generate_req(31, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut r.kind {
|
||||
g.return_hidden_states = want;
|
||||
}
|
||||
r
|
||||
};
|
||||
let disabled = test_limits();
|
||||
let err = validate(&mut req(true), &disabled).unwrap_err();
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(
|
||||
err.to_string().contains("--enable-return-hidden-states"),
|
||||
"message must name the flag: {err}"
|
||||
);
|
||||
// Not asking for them (the client sent `false`, or sent nothing and
|
||||
// `into_requests` resolved the default), or asking on a server that
|
||||
// supports them, is fine.
|
||||
assert!(validate(&mut req(false), &disabled).is_ok());
|
||||
let enabled = Limits {
|
||||
enable_return_hidden_states: true,
|
||||
..test_limits()
|
||||
};
|
||||
assert!(validate(&mut req(true), &enabled).is_ok());
|
||||
}
|
||||
|
||||
/// End-to-end through `drive`: an over-context request is rejected on the way
|
||||
/// to the ring, after registration — so it must be deregistered, not leaked.
|
||||
#[test]
|
||||
fn over_context_request_deregisters_and_never_reaches_the_ring() {
|
||||
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake_with(Limits {
|
||||
context_len: 4,
|
||||
..test_limits()
|
||||
});
|
||||
intake.drive(generate_req(
|
||||
33,
|
||||
SamplingParams {
|
||||
max_new_tokens: Some(64),
|
||||
..Default::default()
|
||||
},
|
||||
));
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "33"),
|
||||
"registered before the check",
|
||||
);
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "33"),
|
||||
"must deregister on reject",
|
||||
);
|
||||
assert!(
|
||||
consumer.drain(16).headers.is_empty(),
|
||||
"must not reach the scheduler"
|
||||
);
|
||||
}
|
||||
|
||||
/// A `Detokenize` request terminates at the detok stage, and the shard must
|
||||
/// see its `Register` BEFORE its `Decode` — the shard delivers the result
|
||||
/// through the sink registered under that rid, so a `Decode` that arrives
|
||||
/// unregistered is silently dropped and the caller waits forever. Both
|
||||
/// messages ride one channel from this one thread, which is the FIFO this
|
||||
/// pins. Nothing may reach the scheduler ring.
|
||||
#[test]
|
||||
fn detokenize_flows_register_then_decode_and_skips_the_ring() {
|
||||
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let (tx, mut rx) = mpsc::channel(8);
|
||||
intake.drive(Request {
|
||||
rid: "41".into(),
|
||||
state: RequestState::Received,
|
||||
sink: ResponseSink::Local(tx),
|
||||
kind: RequestKind::Detokenize {
|
||||
token_ids: vec![7, 8, 9],
|
||||
},
|
||||
});
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "41"),
|
||||
"the sink must be registered before the decode job",
|
||||
);
|
||||
assert!(
|
||||
matches!(
|
||||
detok_rx.try_recv(),
|
||||
Ok(DetokMsg::Decode { rid, token_ids })
|
||||
if rid.as_str() == "41" && token_ids == [7, 8, 9]
|
||||
),
|
||||
"the decode job follows, ids intact",
|
||||
);
|
||||
assert!(
|
||||
consumer.drain(16).headers.is_empty(),
|
||||
"must never reach the scheduler"
|
||||
);
|
||||
assert!(
|
||||
rx.try_recv().is_err(),
|
||||
"no response until the shard answers"
|
||||
);
|
||||
}
|
||||
|
||||
/// Negative ids cannot decode (the shard's domain is `&[u32]`): rejected by
|
||||
/// `validate` at `Received` — an `Error` to the sink, and the shard sees
|
||||
/// NOTHING (validation runs before registration, so there is no entry to
|
||||
/// leak and no decode job to drop).
|
||||
#[test]
|
||||
fn detokenize_negative_ids_reject_before_registration() {
|
||||
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let (tx, mut rx) = mpsc::channel(8);
|
||||
intake.drive(Request {
|
||||
rid: "43".into(),
|
||||
state: RequestState::Received,
|
||||
sink: ResponseSink::Local(tx),
|
||||
kind: RequestKind::Detokenize {
|
||||
token_ids: vec![1, -1],
|
||||
},
|
||||
});
|
||||
let Ok(ResponseItem::Error(err)) = rx.try_recv() else {
|
||||
panic!("sink must receive the validation error");
|
||||
};
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("out of range"), "{err}");
|
||||
assert!(detok_rx.try_recv().is_err(), "shard never hears of it");
|
||||
assert!(consumer.drain(16).headers.is_empty());
|
||||
}
|
||||
|
||||
/// A dropped ring push is survivable, and this pins WHY. The ring is bounded,
|
||||
/// so under load the scheduler never learns to stop and keeps generating; its
|
||||
/// chunks then arrive for a rid the detok table no longer holds and are
|
||||
/// dropped. That wastes GPU work but cannot MISDELIVER, because
|
||||
/// `Rid::from_client` guarantees no later request ever answers to that rid.
|
||||
/// The detok entry is dropped either way — that is the half that must not
|
||||
/// depend on the ring.
|
||||
///
|
||||
/// Ring capacity 1: the first abort pushes, the second finds it full.
|
||||
#[test]
|
||||
fn abort_deregisters_even_when_the_ring_push_is_dropped() {
|
||||
let (tok_tx, _tok_rx) = flume::unbounded();
|
||||
let (detok_tx, detok_rx) = flume::unbounded();
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
let senders = Senders {
|
||||
tok_manager_tx: flume::unbounded().0,
|
||||
abort_tx,
|
||||
tokenizer_tx: tok_tx,
|
||||
detokenizer_tx: vec![detok_tx],
|
||||
};
|
||||
let (producer, _consumer) = to_scheduler(1);
|
||||
let (_tm_tx, tm_rx) = flume::unbounded();
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let mut intake = Intake::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
producer,
|
||||
test_limits(),
|
||||
test_mm(flume::unbounded().0, true),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
intake.on_abort(AbortSource::Guard("pushed".into()));
|
||||
intake.on_abort(AbortSource::Guard("dropped".into()));
|
||||
|
||||
// Both deregisters land regardless of whether the ring accepted the push.
|
||||
for expected in ["pushed", "dropped"] {
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == expected),
|
||||
"{expected}: the detok entry must be dropped even when the ring is full",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// The rid keys the detok table and rides on every chunk of every decode step,
|
||||
/// so an unbounded client-supplied one is a recurring cost, not a one-off.
|
||||
#[test]
|
||||
fn oversized_rid_is_rejected() {
|
||||
let mut req = generate_req(51, SamplingParams::default());
|
||||
req.rid = "x".repeat(MAX_RID_LEN + 1).into();
|
||||
let err = validate(&mut req, &test_limits()).expect_err("must be rejected");
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("over the"), "{err}");
|
||||
|
||||
// A uuid-sized rid — what Python mints — is nowhere near the cap.
|
||||
let mut req = generate_req(52, SamplingParams::default());
|
||||
req.rid = "0123456789abcdef0123456789abcdef".into();
|
||||
assert!(validate(&mut req, &test_limits()).is_ok());
|
||||
}
|
||||
|
||||
/// A request rejected BEFORE `register_detok` must not send `Deregister`: the
|
||||
/// handler is a bare `table.remove(&rid)`, so it would evict whatever entry
|
||||
/// holds that key — a concurrent request's sink — leaving that client hung with
|
||||
/// no terminal frame. Python validates before it inserts, so it cannot hit this.
|
||||
#[test]
|
||||
fn pre_registration_failure_does_not_deregister() {
|
||||
// Rejected inside `validate` (out-of-vocab id), which runs before registration.
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(41, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![2_000_000_000]);
|
||||
}
|
||||
intake.drive(req);
|
||||
assert!(
|
||||
detok_rx.try_recv().is_err(),
|
||||
"a pre-registration reject must send NOTHING to the shard — a Deregister \
|
||||
here removes a live request's sink"
|
||||
);
|
||||
|
||||
// A post-registration reject still deregisters (the leak fix stays fixed).
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
intake.drive(generate_req(
|
||||
42,
|
||||
SamplingParams {
|
||||
top_p: 2.0, // rejected by `normalize`, after registration
|
||||
..Default::default()
|
||||
},
|
||||
));
|
||||
assert!(matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. })));
|
||||
assert!(matches!(
|
||||
detok_rx.try_recv(),
|
||||
Ok(DetokMsg::Deregister { .. })
|
||||
));
|
||||
}
|
||||
|
||||
/// A request rejected at normalization (post-register) must not leak: the shard
|
||||
/// sees `Register` then `Deregister`. Regression for RSS growth on bad input.
|
||||
#[test]
|
||||
fn rejected_request_deregisters_from_shard() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
// top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it.
|
||||
let bad = SamplingParams {
|
||||
top_p: 2.0,
|
||||
..Default::default()
|
||||
};
|
||||
intake.drive(generate_req(7, bad));
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "7"),
|
||||
"expected Register for rid 7",
|
||||
);
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "7"),
|
||||
"expected Deregister for rid 7 (leak fix)",
|
||||
);
|
||||
assert!(
|
||||
detok_rx.try_recv().is_err(),
|
||||
"no further shard messages — registration fully cleaned up",
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: an out-of-vocabulary client token id must be rejected at
|
||||
/// with a 400 — passed through, it reaches the embedding lookup
|
||||
/// and kills the scheduler process (`make_intake` bounds vocab at 1000).
|
||||
#[test]
|
||||
fn out_of_vocab_input_ids_rejected() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(21, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![1, 2_000_000_000]);
|
||||
}
|
||||
intake.drive(req);
|
||||
// Rejected before registration: the only shard message is nothing at
|
||||
// all, or a Deregister if registration happened first — never a push.
|
||||
match detok_rx.try_recv() {
|
||||
Err(_) => {}
|
||||
Ok(DetokMsg::Deregister { .. }) => {}
|
||||
Ok(_) => panic!("out-of-vocab request must not be admitted"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Same guard for negative ids and for `token_ids_logprob` entries.
|
||||
#[test]
|
||||
fn negative_and_logprob_token_ids_rejected() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(22, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![-1]);
|
||||
}
|
||||
intake.drive(req);
|
||||
match detok_rx.try_recv() {
|
||||
Err(_) | Ok(DetokMsg::Deregister { .. }) => {}
|
||||
Ok(_) => panic!("negative token id must not be admitted"),
|
||||
}
|
||||
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(23, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.token_ids_logprob = Some(vec![999_999]);
|
||||
}
|
||||
intake.drive(req);
|
||||
match detok_rx.try_recv() {
|
||||
Err(_) | Ok(DetokMsg::Deregister { .. }) => {}
|
||||
Ok(_) => panic!("out-of-vocab token_ids_logprob must not be admitted"),
|
||||
}
|
||||
}
|
||||
|
||||
/// A valid request is registered and handed onward — never deregistered.
|
||||
#[test]
|
||||
fn admitted_request_keeps_registration() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
// Empty map → all sampling defaults, passes normalization.
|
||||
intake.drive(generate_req(9, SamplingParams::default()));
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "9"),
|
||||
"expected Register for rid 9",
|
||||
);
|
||||
assert!(
|
||||
detok_rx.try_recv().is_err(),
|
||||
"admitted request must not be deregistered",
|
||||
);
|
||||
}
|
||||
|
||||
/// A pool return in `Failed` state (failed encode) is rejected via the same
|
||||
/// path and deregistered, not leaked.
|
||||
#[test]
|
||||
fn tokenize_failure_deregisters_via_intake() {
|
||||
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake();
|
||||
// The pool marks a failed encode as `Failed(err)` before returning it.
|
||||
let mut req = generate_req(11, SamplingParams::default());
|
||||
let _ = req
|
||||
.state
|
||||
.apply(Event::Error(Error::Tokenize("boom".into())));
|
||||
tm_tx.send(TmEvent::Tokenized(req)).unwrap();
|
||||
// Close the inbox so the run loop returns after draining the one event.
|
||||
drop(tm_tx);
|
||||
intake.run();
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "11"),
|
||||
"tokenize failure must deregister rid 11",
|
||||
);
|
||||
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
||||
}
|
||||
|
||||
/// An abort deregisters (by the id hashed from the rid string), so a request
|
||||
/// aborted before any terminal chunk can't leak.
|
||||
#[test]
|
||||
fn abort_deregisters_from_shard() {
|
||||
// Aborts arrive on their own unbounded lane now, not the request inbox.
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake_with_abort(abort_rx);
|
||||
abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap();
|
||||
drop(abort_tx);
|
||||
drop(tm_tx);
|
||||
intake.run();
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "rid-13"),
|
||||
"abort must deregister by rid",
|
||||
);
|
||||
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
||||
}
|
||||
|
||||
/// A successful pool return (Queued, ids filled) is pushed to the ring, not
|
||||
/// rejected; its registration is untouched.
|
||||
#[test]
|
||||
fn tokenized_return_pushes_without_deregister() {
|
||||
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(15, SamplingParams::default());
|
||||
// Simulate a successful pool return: ids filled, PreSendValidating.
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![1, 2, 3]);
|
||||
}
|
||||
req.state = RequestState::PreSendValidating;
|
||||
tm_tx.send(TmEvent::Tokenized(req)).unwrap();
|
||||
drop(tm_tx);
|
||||
intake.run();
|
||||
|
||||
// Pushed to the ring; the shard sees nothing.
|
||||
assert!(
|
||||
detok_rx.try_recv().is_err(),
|
||||
"a queued pool-return must be pushed, not touch the shard",
|
||||
);
|
||||
}
|
||||
|
||||
/// If the pool is gone, a request needing tokenization is rejected +
|
||||
/// deregistered, not silently dropped.
|
||||
#[test]
|
||||
fn tokenize_pool_gone_deregisters() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(21, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = None;
|
||||
}
|
||||
intake.drive(req);
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "21"),
|
||||
"expected Register for rid 21",
|
||||
);
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "21"),
|
||||
"pool-gone hand-off must deregister rid 21",
|
||||
);
|
||||
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
||||
}
|
||||
|
||||
/// Build a generate request carrying an image. The parked entry and the
|
||||
/// `MmEncoded` resume path agree on identity via the rid string.
|
||||
fn mm_generate_req(rid: &str) -> Request {
|
||||
let (tx, _rx) = mpsc::channel(8);
|
||||
Request {
|
||||
rid: rid.to_string().into(),
|
||||
state: RequestState::Received,
|
||||
sink: ResponseSink::Local(tx),
|
||||
kind: RequestKind::Generate(Box::new(GenerateRequest {
|
||||
rid: rid.to_string().into(),
|
||||
text: Some("<image> hi".into()),
|
||||
mm: Some(Box::new(crate::message::request::MmData {
|
||||
image_data: Some(rmpv::Value::from("data:image/jpeg;base64,xxxx")),
|
||||
..Default::default()
|
||||
})),
|
||||
..Default::default()
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// An abort while the request is parked for MM cancels it: the pending
|
||||
/// entry is removed, the worker's late result is dropped, and its parked
|
||||
/// sidecar entry is purged — no scheduler work runs for a dead client.
|
||||
#[test]
|
||||
fn abort_cancels_parked_mm_request() {
|
||||
let (mut intake, _detok_rx, consumer, _tm_tx, mm_rx) = make_intake();
|
||||
intake.drive(mm_generate_req("mm-gone"));
|
||||
mm_rx.try_recv().expect("parked to mm pool");
|
||||
|
||||
// The worker parks its result, as it always does before MmEncoded.
|
||||
intake.mm.sidecar.park(
|
||||
"mm-gone".into(),
|
||||
crate::multi_modality::sidecar::MmSidecarEntry {
|
||||
features: crate::multi_modality::sidecar::FeatureStore::Inline(vec![]),
|
||||
grids: vec![],
|
||||
hashes: vec![],
|
||||
offsets: vec![],
|
||||
mrope: vec![],
|
||||
mrope_delta: 0,
|
||||
},
|
||||
);
|
||||
intake.on_abort(AbortSource::Guard("mm-gone".to_string().into()));
|
||||
assert_eq!(consumer.drain(16).headers.len(), 1, "only the AbortReq");
|
||||
|
||||
// The late result must be dropped, not queued, and the sidecar purged.
|
||||
intake.on_mm_encoded("mm-gone".to_string().into(), vec![5, 6]);
|
||||
assert!(
|
||||
consumer.drain(16).headers.is_empty(),
|
||||
"cancelled, not queued"
|
||||
);
|
||||
assert!(intake.mm.sidecar.take("mm-gone").is_none(), "entry purged");
|
||||
}
|
||||
|
||||
/// A multimodal request parks in `Encoding` (submitted to the mm worker
|
||||
/// pool, not the tokenizer pool, not the ring) until `MmEncoded` resumes
|
||||
/// it → ring.
|
||||
#[test]
|
||||
fn mm_request_parks_then_mm_encoded_pushes_to_ring() {
|
||||
let (mut intake, _detok_rx, consumer, _tm_tx, mm_rx) = make_intake();
|
||||
intake.drive(mm_generate_req("mm-1"));
|
||||
|
||||
// Submitted to the mm pool with the typed work item; nothing on the ring yet.
|
||||
let sub = mm_rx.try_recv().expect("mm pool must receive the request");
|
||||
assert_eq!(sub.rid.as_str(), "mm-1");
|
||||
assert_eq!(sub.work.text.as_deref(), Some("<image> hi"));
|
||||
assert!(sub.work.input_ids.is_none(), "no client input_ids");
|
||||
assert_eq!(
|
||||
sub.work.image_data.as_ref().and_then(|v| v.as_str()),
|
||||
Some("data:image/jpeg;base64,xxxx")
|
||||
);
|
||||
assert!(consumer.drain(16).headers.is_empty(), "parked, not queued");
|
||||
|
||||
// The worker returns the final expanded ids → pushed to the ring.
|
||||
intake.on_mm_encoded("mm-1".to_string().into(), vec![5, 6, 7, 8]);
|
||||
let batch = consumer.drain(16);
|
||||
assert_eq!(batch.headers.len(), 1);
|
||||
assert_eq!(
|
||||
batch.lengths,
|
||||
vec![4],
|
||||
"expanded ids ride the columnar cell"
|
||||
);
|
||||
}
|
||||
|
||||
/// A worker failure rejects the parked request (deregister, no ring push).
|
||||
#[test]
|
||||
fn mm_failure_rejects_parked_request() {
|
||||
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
intake.drive(mm_generate_req("mm-2"));
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. })),
|
||||
"registered before parking",
|
||||
);
|
||||
|
||||
intake.on_mm_failed("mm-2".to_string().into(), "bad image".into());
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid })
|
||||
if rid.as_str() == "mm-2"),
|
||||
"mm failure must deregister",
|
||||
);
|
||||
assert!(consumer.drain(16).headers.is_empty(), "nothing queued");
|
||||
}
|
||||
|
||||
/// On a non-multimodal model (`Mm::enabled == false`), image_data is silently
|
||||
/// ignored and the request tokenizes as plain text — the Python
|
||||
/// TokenizerManager behavior when `mm_processor is None`.
|
||||
#[test]
|
||||
fn mm_fields_ignored_when_disabled() {
|
||||
let (tok_tx, tok_rx) = flume::unbounded();
|
||||
let (detok_tx, _detok_rx) = flume::unbounded();
|
||||
let senders = Senders {
|
||||
tok_manager_tx: flume::unbounded().0,
|
||||
abort_tx: flume::unbounded().0,
|
||||
tokenizer_tx: tok_tx,
|
||||
detokenizer_tx: vec![detok_tx],
|
||||
};
|
||||
let (to_scheduler_tx, _consumer) = to_scheduler(16);
|
||||
let (_tm_tx, tm_rx) = flume::unbounded();
|
||||
let (mm_tx, mm_rx) = flume::unbounded();
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
std::mem::forget(abort_tx);
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let mut intake = Intake::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
to_scheduler_tx,
|
||||
test_limits(),
|
||||
test_mm(mm_tx, false),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
intake.drive(mm_generate_req("mm-3"));
|
||||
assert!(
|
||||
mm_rx.try_recv().is_err(),
|
||||
"mm disabled: nothing submitted to the mm channel",
|
||||
);
|
||||
assert!(
|
||||
tok_rx.try_recv().is_ok(),
|
||||
"request must fall through to plain tokenization",
|
||||
);
|
||||
}
|
||||
|
||||
/// A late mm result for a rid that is no longer parked is dropped without
|
||||
/// panicking (e.g. hash-collision overwrite) — regression guard.
|
||||
#[test]
|
||||
fn late_mm_result_is_dropped() {
|
||||
let (mut intake, _detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
intake.on_mm_encoded("ghost".to_string().into(), vec![1]);
|
||||
intake.on_mm_failed("ghost".to_string().into(), "boom".into());
|
||||
assert!(consumer.drain(16).headers.is_empty());
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Configuration and multimodal handles for scheduler intake.
|
||||
|
||||
use crate::message::config::ServerArgs;
|
||||
use crate::message::request::MmRequest;
|
||||
|
||||
/// The intake side of the MM path.
|
||||
#[derive(Clone)]
|
||||
pub struct Mm {
|
||||
/// Whether the model is multimodal. When false, mm fields are silently
|
||||
/// ignored, as the Python `TokenizerManager` does with `mm_processor is
|
||||
/// None`.
|
||||
pub enabled: bool,
|
||||
/// → MM worker pool (spawned via `Server.start_mm_workers`).
|
||||
pub tx: flume::Sender<MmRequest>,
|
||||
/// Results sidecar. Purged here when a late result arrives for a request
|
||||
/// that is no longer parked; otherwise it would leak, since only the
|
||||
/// scheduler drain pops entries.
|
||||
pub sidecar: crate::multi_modality::sidecar::Sidecar,
|
||||
}
|
||||
|
||||
/// Resolved once at boot from the scheduler's `server_args`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Limits {
|
||||
/// Token-ids-in mode: a generate request must arrive already tokenized.
|
||||
pub skip_tokenizer_init: bool,
|
||||
/// `model_config.vocab_size`; bounds client-supplied token ids. A required
|
||||
/// field of the `ServerArgs` schema, so intake can check unconditionally.
|
||||
pub vocab_size: u64,
|
||||
/// `model_config.context_len`, the ceiling for input + `max_new_tokens`.
|
||||
pub context_len: u64,
|
||||
/// Output slots reserved on top of the input (eagle draft tokens).
|
||||
pub num_reserved_tokens: u64,
|
||||
/// Clamp `max_new_tokens` to what fits instead of rejecting the request.
|
||||
pub allow_auto_truncate: bool,
|
||||
/// Whether the server can produce hidden states at all.
|
||||
pub enable_return_hidden_states: bool,
|
||||
}
|
||||
|
||||
impl From<&ServerArgs> for Limits {
|
||||
fn from(sa: &ServerArgs) -> Self {
|
||||
Self {
|
||||
skip_tokenizer_init: sa.skip_tokenizer_init,
|
||||
vocab_size: sa.model_config.vocab_size,
|
||||
context_len: sa.model_config.context_len,
|
||||
num_reserved_tokens: sa.num_reserved_tokens,
|
||||
allow_auto_truncate: sa.allow_auto_truncate,
|
||||
enable_return_hidden_states: sa.enable_return_hidden_states,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
//! Request validation for scheduler intake.
|
||||
|
||||
use crate::message::request::{GenerateRequest, Request, RequestKind};
|
||||
use crate::utils::{
|
||||
error::Error,
|
||||
fsm::{Event, ValidationOutcome},
|
||||
};
|
||||
|
||||
use super::to_scheduler::MAX_RID_LEN;
|
||||
use super::to_scheduler_types::Limits;
|
||||
|
||||
/// `Received → Validating` + admissibility check. Under `skip_tokenizer_init` a
|
||||
/// generate request must already carry token ids (no tokenizer to byte-encode
|
||||
/// text); control requests carry none and are exempt.
|
||||
pub(super) fn validate(req: &mut Request, limits: &Limits) -> Result<(), Error> {
|
||||
let (skip_tokenizer_init, vocab_size) = (limits.skip_tokenizer_init, limits.vocab_size);
|
||||
let _ = req
|
||||
.state
|
||||
.apply(Event::Validated(ValidationOutcome::NeedsTokenize));
|
||||
|
||||
// The rid is the request's identity everywhere downstream: it keys the detok
|
||||
// table, and it rides on EVERY chunk of EVERY decode step. An unbounded
|
||||
// client-supplied rid is therefore a per-step cost, not a one-off. Python's is
|
||||
// a 32-byte uuid hex, so this is generous.
|
||||
// Measured on the CLIENT-facing form: the uniquifier `Rid::from_client` appends
|
||||
// is this server's own overhead, and charging the client for bytes it did not
|
||||
// send would reject a rid exactly at the documented limit.
|
||||
let client_rid_len = req.rid.client_facing().len();
|
||||
if client_rid_len > MAX_RID_LEN {
|
||||
return Err(Error::Validation(format!(
|
||||
"rid is {client_rid_len} bytes, over the {MAX_RID_LEN}-byte limit"
|
||||
)));
|
||||
}
|
||||
if skip_tokenizer_init
|
||||
&& matches!(&req.kind, RequestKind::Generate(g) if !g.already_tokenized())
|
||||
{
|
||||
// `Validation` (400), not `Tokenize` (500): the client sent a request this
|
||||
// server cannot serve, which is their error to fix — Python 400s it too.
|
||||
return Err(Error::Validation(
|
||||
"skip_tokenizer_init is set: request must provide input_ids".into(),
|
||||
));
|
||||
}
|
||||
|
||||
// Client-supplied token ids must be in-vocabulary: an out-of-range id
|
||||
// reaches the embedding lookup and kills the scheduler process, so 400
|
||||
// here instead — mirroring the Python `TokenizerManager` validation.
|
||||
if let RequestKind::Generate(g) = &req.kind {
|
||||
if let Some(ids) = &g.input_ids {
|
||||
for &id in ids {
|
||||
if id < 0 || id as u64 >= vocab_size {
|
||||
return Err(Error::Validation(format!(
|
||||
"input_ids contains out-of-vocabulary token id {id}; \
|
||||
valid range is [0, {vocab_size})"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(ids) = &g.token_ids_logprob {
|
||||
for &id in ids {
|
||||
if id < 0 || id as u64 >= vocab_size {
|
||||
return Err(Error::Validation(format!(
|
||||
"token_ids_logprob contains out-of-vocabulary token id \
|
||||
{id}; valid range is [0, {vocab_size})"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Detokenize ids must fit the shard's `&[u32]` decode domain. No vocab
|
||||
// bound — parity with the retired direct decode service: an unknown id is
|
||||
// the tokenizer's error to report, and nothing here reaches the scheduler's
|
||||
// embedding lookup.
|
||||
if let RequestKind::Detokenize { token_ids } = &req.kind {
|
||||
for &id in token_ids {
|
||||
if u32::try_from(id).is_err() {
|
||||
return Err(Error::Validation(format!("Token ID {id} is out of range")));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The scheduler only computes hidden states when launched for it, so without
|
||||
// this the request would 200 with `meta_info.hidden_states` silently absent
|
||||
// (Python `TokenizerManager._validate_one_request`).
|
||||
if !limits.enable_return_hidden_states
|
||||
&& matches!(&req.kind, RequestKind::Generate(g) if g.return_hidden_states)
|
||||
{
|
||||
return Err(Error::Validation(
|
||||
"The server is not configured to return the hidden states. \
|
||||
Please set `--enable-return-hidden-states` to enable this feature."
|
||||
.into(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The context-window checks that need the tokenized length, mirroring Python
|
||||
/// `TokenizerManager._validate_one_request`: the input alone must fit, and then
|
||||
/// input + `max_new_tokens` must fit. Without them the scheduler silently clamps
|
||||
/// and the client gets a 200 with a truncated completion instead of an actionable
|
||||
/// 400.
|
||||
///
|
||||
/// Under `allow_auto_truncate` both clamp instead of rejecting — the launch flag
|
||||
/// opted into that.
|
||||
pub(super) fn check_total_tokens(g: &mut GenerateRequest, limits: &Limits) -> Result<(), Error> {
|
||||
let max_req_len = limits.context_len;
|
||||
// Python counts the reserved slots as part of the input, so a request can be
|
||||
// rejected for them even when the prompt alone fits.
|
||||
let input_len =
|
||||
g.input_ids.as_ref().map_or(0, |ids| ids.len()) as u64 + limits.num_reserved_tokens;
|
||||
|
||||
// Input length first, and unconditionally: `max_new_tokens: null` means "no
|
||||
// cap", which must not disable this. Python's comparison is `>=` — a prompt
|
||||
// that exactly fills the window leaves no room to generate.
|
||||
if input_len >= max_req_len {
|
||||
if !limits.allow_auto_truncate {
|
||||
return Err(Error::Validation(format!(
|
||||
"The input ({input_len} tokens) is longer than the model's context \
|
||||
length ({max_req_len} tokens)."
|
||||
)));
|
||||
}
|
||||
if let Some(ids) = &mut g.input_ids {
|
||||
ids.truncate(max_req_len as usize);
|
||||
}
|
||||
}
|
||||
let input_len =
|
||||
g.input_ids.as_ref().map_or(0, |ids| ids.len()) as u64 + limits.num_reserved_tokens;
|
||||
|
||||
let Some(max_new_tokens) = g.sampling_params.max_new_tokens else {
|
||||
return Ok(()); // no cap requested → nothing to add to the input length
|
||||
};
|
||||
let total = input_len.saturating_add(max_new_tokens.max(0) as u64);
|
||||
if total <= max_req_len {
|
||||
return Ok(());
|
||||
}
|
||||
if !limits.allow_auto_truncate {
|
||||
return Err(Error::Validation(format!(
|
||||
"Requested token count exceeds the model's maximum context length of \
|
||||
{max_req_len} tokens. You requested a total of {total} tokens: {input_len} \
|
||||
tokens from the input messages and {max_new_tokens} tokens for the \
|
||||
completion. Please reduce the number of tokens in the input messages or \
|
||||
the completion to fit within the limit."
|
||||
)));
|
||||
}
|
||||
let clamped = max_req_len.saturating_sub(input_len) as i64;
|
||||
// Re-check what the clamp can break. `verify` already ran (in Normalizing), so
|
||||
// lowering `max_new_tokens` here can leave `min_new_tokens > max_new_tokens` —
|
||||
// and `is_normalized: true` stops the scheduler from re-verifying, so nothing
|
||||
// downstream would catch it. Python validates before it verifies; we can't
|
||||
// reorder the FSM, so we re-assert the one invariant the clamp can violate.
|
||||
if g.sampling_params.min_new_tokens > clamped {
|
||||
return Err(Error::Validation(format!(
|
||||
"min_new_tokens must be in [0, max_new_tokens({clamped})], got {}",
|
||||
g.sampling_params.min_new_tokens
|
||||
)));
|
||||
}
|
||||
g.sampling_params.max_new_tokens = Some(clamped);
|
||||
Ok(())
|
||||
}
|
||||
@@ -25,7 +25,7 @@ pub enum TmEvent {
|
||||
Tokenized(Request),
|
||||
/// An MM worker finished a request parked in `Encoding`: `input_ids` are the
|
||||
/// final placeholder-expanded prompt ids. The buffers ride the rid-keyed
|
||||
/// sidecar (`Server.take_mm`), not this event.
|
||||
/// sidecar (`Server.take_mm_result`), not this event.
|
||||
MmEncoded { rid: Rid, input_ids: Vec<i32> },
|
||||
/// An MM worker rejected a request parked in `Encoding` (bad media URL,
|
||||
/// unsupported modality, preprocess error, …).
|
||||
|
||||
@@ -51,7 +51,7 @@ pub struct Runtime {
|
||||
/// `skip_tokenizer_init`).
|
||||
pub tokenizer: Option<Arc<dyn tokenizer::TextTokenizer>>,
|
||||
/// MM results parked between a worker's `MmEncoded` and the scheduler drain
|
||||
/// (`Server.take_mm`).
|
||||
/// (`Server.take_mm_result`).
|
||||
pub mm_sidecar: crate::multi_modality::sidecar::Sidecar,
|
||||
/// Worker join handles, joined by `request_shutdown` / `Drop`.
|
||||
threads: Mutex<Vec<JoinHandle<()>>>,
|
||||
@@ -117,18 +117,19 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
|
||||
// --- inter-stage channels ---
|
||||
let (tok_manager_tx, tok_manager_rx) =
|
||||
flume::bounded::<TmEvent>(cfg.rust_server_args.channel_cap);
|
||||
flume::bounded::<TmEvent>(cfg.rust_server_args.stage_channel_cap);
|
||||
let (tokenizer_tx, tokenizer_rx) =
|
||||
flume::bounded::<crate::message::request::Request>(cfg.rust_server_args.channel_cap);
|
||||
flume::bounded::<crate::message::request::Request>(cfg.rust_server_args.stage_channel_cap);
|
||||
// Encoding → MM worker pool. Bounded like the other stage edges so a slow
|
||||
// pool back-pressures instead of buffering unboundedly.
|
||||
let (mm_worker_tx, mm_worker_rx) =
|
||||
flume::bounded::<crate::message::request::MmRequest>(cfg.rust_server_args.channel_cap);
|
||||
let (mm_worker_tx, mm_worker_rx) = flume::bounded::<crate::message::request::MmRequest>(
|
||||
cfg.rust_server_args.stage_channel_cap,
|
||||
);
|
||||
let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num;
|
||||
let mut detokenizer_tx = Vec::with_capacity(detokenizer_worker_num);
|
||||
let mut detokenizer_rx = Vec::with_capacity(detokenizer_worker_num);
|
||||
for _ in 0..detokenizer_worker_num {
|
||||
let (tx, rx) = flume::bounded::<DetokMsg>(cfg.rust_server_args.channel_cap);
|
||||
let (tx, rx) = flume::bounded::<DetokMsg>(cfg.rust_server_args.stage_channel_cap);
|
||||
detokenizer_tx.push(tx);
|
||||
detokenizer_rx.push(rx);
|
||||
}
|
||||
@@ -303,7 +304,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
rt.block_on(api_server::app::serve(
|
||||
listener,
|
||||
senders,
|
||||
cfg.rust_server_args.channel_cap,
|
||||
cfg.rust_server_args.stage_channel_cap,
|
||||
cfg.server_args.clone(),
|
||||
// Response heartbeat watched by `/health_generate`.
|
||||
response_activity,
|
||||
|
||||
Reference in New Issue
Block a user