[mm] refactor mm code for rust tokenizer manager (#34660)

Co-authored-by: Rain Jiang <rain-jiang@outlook.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Kan Wu
2026-09-04 03:03:25 -07:00
committed by GitHub
co-authored by Rain Jiang Cursor Claude Fable 5
parent b42569a0f1
commit 12735c2d76
16 changed files with 200 additions and 161 deletions
+28 -23
View File
@@ -23,7 +23,7 @@ class RustMmSpec(msgspec.Struct, frozen=True, kw_only=True):
consumed by the Rust worker pool (as the typed extension ``MmSpec``, see
:meth:`RustServer._build_mm_spec`), the ``_multimodal`` parity API
(:meth:`rust_json`) and the drain adapter
(:meth:`RustMmProcessor.build_output`)."""
(:meth:`RustMmProcessor.wrap_encoded`)."""
family: str
feature_shm: bool
@@ -119,7 +119,7 @@ class RustMmProcessor:
TokenizerManager would build — not to process requests (the Rust worker pool
does that, GIL-free) but as the source of truth
:meth:`resolve_spec` resolves the pipeline parameters from. At drain
time :meth:`build_output` wraps the Rust-produced buffers into the
time :meth:`wrap_encoded` wraps the Rust-produced buffers into the
scheduler's ``MultimodalProcessorOutput``.
There is no Python fallback: a model without a Rust MM spec fails at launch,
@@ -174,17 +174,17 @@ class RustMmProcessor:
)
if family is None:
return None
ip = getattr(self._processor, "image_processor", None)
resample = family.image_processors.get(type(ip).__name__)
image_processor = getattr(self._processor, "image_processor", None)
resample = family.image_processors.get(type(image_processor).__name__)
if resample is None:
return None
# The Rust pipeline always resizes, rescales by 1/255 and normalizes;
# Rust's fused normalize constants assume that factor. Anything else
# would silently produce different features.
stages = ("do_resize", "do_rescale", "do_normalize")
if not all(getattr(ip, stage, True) for stage in stages):
if not all(getattr(image_processor, stage, True) for stage in stages):
return None
if getattr(ip, "rescale_factor", None) != 1 / 255:
if getattr(image_processor, "rescale_factor", None) != 1 / 255:
return None
# `--mm-process-config {"image": {...}}`: only pixel-limit overrides are
@@ -193,25 +193,27 @@ class RustMmProcessor:
if not set(image_overrides) <= {"min_pixels", "max_pixels"}:
return None
size = getattr(ip, "size", None) or {}
size = getattr(image_processor, "size", None) or {}
min_pixels = image_overrides.get(
"min_pixels", getattr(ip, "min_pixels", None) or size.get("shortest_edge")
"min_pixels",
getattr(image_processor, "min_pixels", None) or size.get("shortest_edge"),
)
max_pixels = image_overrides.get(
"max_pixels", getattr(ip, "max_pixels", None) or size.get("longest_edge")
"max_pixels",
getattr(image_processor, "max_pixels", None) or size.get("longest_edge"),
)
try:
spec = RustMmSpec(
family=family.name,
feature_shm=self._use_feature_shm(),
image_token_id=hf_config.image_token_id,
patch_size=ip.patch_size,
merge_size=ip.merge_size,
temporal_patch_size=ip.temporal_patch_size,
patch_size=image_processor.patch_size,
merge_size=image_processor.merge_size,
temporal_patch_size=image_processor.temporal_patch_size,
min_pixels=int(min_pixels),
max_pixels=int(max_pixels),
image_mean=tuple(float(x) for x in ip.image_mean),
image_std=tuple(float(x) for x in ip.image_std),
image_mean=tuple(float(x) for x in image_processor.image_mean),
image_std=tuple(float(x) for x in image_processor.image_std),
resample=resample,
vision_start_token_id=getattr(hf_config, "vision_start_token_id", None),
vision_end_token_id=getattr(hf_config, "vision_end_token_id", None),
@@ -247,10 +249,11 @@ class RustMmProcessor:
)
@staticmethod
def build_output(spec: RustMmSpec, entry):
"""Drain-time adapter: wrap the Rust-produced buffers of one ``MmEncodeResult``
into the scheduler's ``MultimodalProcessorOutput``. Wrapping only — load,
resize, patchify, token expansion and M-RoPE all ran in Rust.
def wrap_encoded(spec: RustMmSpec, encoded):
"""Drain-time adapter: wrap the Rust-produced buffers of one
``MmEncodedResult`` into the scheduler's ``MultimodalProcessorOutput``.
Wrapping only — load, resize, patchify, token expansion and M-RoPE all
ran in Rust.
Runs on the scheduler loop, so it must stay copy-free *and* hash-free:
``take_mm_result``'s numpy arrays own the Rust buffers, ``torch.from_numpy`` just
@@ -267,13 +270,13 @@ class RustMmProcessor:
MultimodalProcessorOutput,
)
shm_names = entry.shm_names
shm_names = encoded.shm_names
if shm_names is None:
features = torch.from_numpy(entry.features.reshape(-1, spec.feature_dim))
features = torch.from_numpy(encoded.features.reshape(-1, spec.feature_dim))
items = []
row = 0
for index, ((t, h, w), item_hash, offset) in enumerate(
zip(entry.grids, entry.hashes, entry.offsets)
zip(encoded.grids, encoded.hashes, encoded.offsets)
):
n = t * h * w
if shm_names is None:
@@ -314,6 +317,8 @@ class RustMmProcessor:
im_start_id=spec.vision_start_token_id,
im_end_id=spec.vision_end_token_id,
video_token_id=spec.video_token_id,
mrope_positions=torch.from_numpy(entry.mrope.reshape(3, -1)),
mrope_position_delta=torch.tensor([[entry.mrope_delta]], dtype=torch.long),
mrope_positions=torch.from_numpy(encoded.mrope.reshape(3, -1)),
mrope_position_delta=torch.tensor(
[[encoded.mrope_delta]], dtype=torch.long
),
)
+7 -8
View File
@@ -205,14 +205,13 @@ class RustServer:
obj.input_ids = ids
pos += nbytes
if self.mm_spec is not None and isinstance(obj, TokenizedGenerateReqInput):
# The buffers rode the Rust sidecar, parked before the ring push;
# wrapping them into tensors is the only Python step of the Rust
# path. `None` for a text-only request on a multimodal model.
mm_result = self.server.take_mm_result(obj.rid)
if mm_result is not None:
obj.mm_inputs = RustMmProcessor.build_output(
self.mm_spec, mm_result
)
# The buffers were parked in the Rust result store before the
# ring push; wrapping them into tensors is the only Python step
# of the Rust path. `None` for a text-only request on a
# multimodal model.
encoded = self.server.take_mm_result(obj.rid)
if encoded is not None:
obj.mm_inputs = RustMmProcessor.wrap_encoded(self.mm_spec, encoded)
out.append(obj)
return out
+27 -18
View File
@@ -344,31 +344,36 @@ pub fn mrope_image_only(
/// The qwen scheduler-drain shape, extracted from the generic driver
/// [`Output`](crate::driver::Output). Shared by `sglang-server`'s MM worker
/// and the parity binding so the mapping can't drift; replaced by a generic
/// named-tensor handoff once a second family needs a different shape.
pub struct QwenDrain {
/// and the parity binding so the mapping can't drift. TODO(mm-families):
/// replace with a generic named-tensor handoff once a second family needs a
/// different shape.
pub struct QwenPackedOutput {
pub input_ids: Vec<i32>,
/// All items' `pixel_values`, concatenated in prompt order.
/// All items' `pixel_values`, concatenated in prompt order; flattened
/// `[Σ t·h·w, 3·temporal_patch_size·patch_size²]`.
pub features: Vec<f32>,
/// Per item `[t, h, w]` patch grid.
pub grids: Vec<[u32; 3]>,
pub hashes: Vec<u64>,
/// Per item inclusive token range in `input_ids`.
pub offsets: Vec<(u32, u32)>,
/// Flattened row-major `[3, input_len]` M-RoPE positions.
pub mrope: Vec<i64>,
pub mrope_delta: i64,
}
pub fn pack_drain(output: crate::driver::Output) -> Result<QwenDrain, String> {
pub fn pack_output(output: crate::driver::Output) -> Result<QwenPackedOutput, String> {
use crate::pipeline::PositionOutput;
let PositionOutput::MRope { positions, delta } = output.positions else {
return Err("qwen_vl drain: expected M-RoPE positions".into());
return Err("qwen_vl pack: expected M-RoPE positions".into());
};
let mut features = Vec::new();
let mut grids = Vec::with_capacity(output.items.len());
let mut hashes = Vec::with_capacity(output.items.len());
for item in output.items {
let TensorData::F32(pixel_values) = item.feature.data else {
return Err("qwen_vl drain: expected f32 feature".into());
return Err("qwen_vl pack: expected f32 feature".into());
};
features.extend(pixel_values);
let grid = item
@@ -378,11 +383,11 @@ pub fn pack_drain(output: crate::driver::Output) -> Result<QwenDrain, String> {
("image_grid_thw", TensorData::I64(v)) => Some(v),
_ => None,
})
.ok_or("qwen_vl drain: missing image_grid_thw")?;
.ok_or("qwen_vl pack: missing image_grid_thw")?;
grids.push([grid[0] as u32, grid[1] as u32, grid[2] as u32]);
hashes.push(item.hash);
}
Ok(QwenDrain {
Ok(QwenPackedOutput {
input_ids: output.input_ids,
features,
grids,
@@ -504,23 +509,27 @@ mod python {
input_ids,
images,
};
let drain = py
let packed = py
.detach(move || {
let family = crate::registry::pipeline_from_spec(&spec_json)?;
let output = crate::driver::process(family.as_ref(), input, |_| {
Err("native parity API requires input_ids".into())
})?;
pack_drain(output)
pack_output(output)
})
.map_err(PyValueError::new_err)?;
Ok((
drain.input_ids,
drain.features.into_pyarray(py),
drain.grids.into_iter().map(|[t, h, w]| (t, h, w)).collect(),
drain.hashes,
drain.offsets,
drain.mrope.into_pyarray(py),
drain.mrope_delta,
packed.input_ids,
packed.features.into_pyarray(py),
packed
.grids
.into_iter()
.map(|[t, h, w]| (t, h, w))
.collect(),
packed.hashes,
packed.offsets,
packed.mrope.into_pyarray(py),
packed.mrope_delta,
))
}
+12 -17
View File
@@ -6,7 +6,7 @@
//! ([`ServerArgs`] and its parts, constructed by keyword from Python; their
//! `#[pyclass]`es and constructors live in `message::config`), [`Server`]
//! (boot, `recv_requests`/`wait_request`, `push_*`, MM handoff, shutdown),
//! [`RequestBatch`] and [`MmEncodeResult`]. Everything behind that boundary —
//! [`RequestBatch`] and [`MmEncodedResult`]. Everything behind that boundary —
//! receiving requests, encoding multimodal inputs, tokenizing, detokenizing,
//! SSE streaming, and so on — is implemented purely in Rust and never touches
//! a `PyObject`.
@@ -29,10 +29,10 @@ use crate::utils::startup::{listen_addr, value_error};
use crate::utils::{logging, runtime};
/// One drained MM result (see [`Server::take_mm_result`]), consumed by
/// `RustMmProcessor.build_output` to build the scheduler's
/// `RustMmProcessor.wrap_encoded` to build the scheduler's
/// `MultimodalProcessorOutput`.
#[pyclass(frozen, get_all)]
struct MmEncodeResult {
struct MmEncodedResult {
// 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.
@@ -199,14 +199,9 @@ impl Server {
/// 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,
self.rt.tokenizer.clone(),
self.rt.mm_sidecar.clone(),
)
.map_err(|e| value_error("mm spec", e))?;
self.rt.spawn_mm_pool(workers, std::sync::Arc::new(ctx));
Ok(())
self.rt
.start_mm_workers(spec, workers)
.map_err(|e| value_error("mm spec", e))
}
/// Pop the MM result for `rid` — parked strictly before the request reached
@@ -217,22 +212,22 @@ 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_result(&self, py: Python<'_>, rid: &str) -> Option<MmEncodeResult> {
fn take_mm_result(&self, py: Python<'_>, rid: &str) -> Option<MmEncodedResult> {
use numpy::IntoPyArray;
let res = self.rt.mm_sidecar.take(rid)?;
let res = self.rt.mm_results.take(rid)?;
let (features, shm_names) = match res.features {
multi_modality::sidecar::FeatureStore::Inline(v) => {
multi_modality::result_store::FeatureStore::Inline(v) => {
(Some(v.into_pyarray(py).unbind()), None)
}
// The segments — and the duty to unlink — move to Python here;
// `materialize()` unlinks after the post-broadcast clone on each rank.
multi_modality::sidecar::FeatureStore::Shm(segments) => (
multi_modality::result_store::FeatureStore::Shm(segments) => (
None,
Some(segments.into_iter().map(|s| s.into_name()).collect()),
),
};
Some(MmEncodeResult {
Some(MmEncodedResult {
features,
shm_names,
grids: res.grids.iter().map(|g| (g[0], g[1], g[2])).collect(),
@@ -276,6 +271,6 @@ fn _server(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<MmSpec>()?;
m.add_class::<Server>()?;
m.add_class::<RequestBatch>()?;
m.add_class::<MmEncodeResult>()?;
m.add_class::<MmEncodedResult>()?;
Ok(())
}
+1 -1
View File
@@ -1,6 +1,6 @@
//! Multimodal worker pool.
pub mod payload;
pub mod result_store;
mod shm;
pub mod sidecar;
pub mod worker;
@@ -7,17 +7,22 @@ use std::sync::{Arc, Mutex};
use super::shm::{ShmSegment, shm_name};
/// One parked result: the buffers the drain-time Python adapter needs (the
/// expanded `input_ids` travel separately, via `TmEvent::MmEncoded`). The qwen
/// drain shape (`sglang_mm::qwen_vl::pack_drain`); generalizes to a
/// named-tensor handoff once a family needs a different one.
/// expanded `input_ids` travel separately, via `TmEvent::MmEncoded`).
///
/// TODO(mm-families): these fields are the shape the only current family
/// (qwen_vl) produces; generalize to a named-tensor handoff when a family
/// needs a different one.
///
/// Constructed from outside the module only by tests; the worker parks every
/// real entry itself.
pub struct MmSidecarEntry {
pub struct MmEncodedEntry {
pub features: FeatureStore,
/// Per item `[t, h, w]` patch grid.
pub grids: Vec<[u32; 3]>,
pub hashes: Vec<u64>,
/// Per item inclusive token range in the expanded prompt.
pub offsets: Vec<(u32, u32)>,
/// Flattened row-major `[3, input_len]` M-RoPE positions.
pub mrope: Vec<i64>,
pub mrope_delta: i64,
}
@@ -37,13 +42,13 @@ pub enum FeatureStore {
/// strictly before `MmEncoded`, [`take`](Self::take) at the drain,
/// [`purge`](Self::purge) for requests that die while parked.
#[derive(Clone, Default)]
pub struct Sidecar(Arc<Mutex<HashMap<String, MmSidecarEntry>>>);
pub struct MmResultStore(Arc<Mutex<HashMap<String, MmEncodedEntry>>>);
impl Sidecar {
pub fn park(&self, rid: String, entry: MmSidecarEntry) {
impl MmResultStore {
pub fn park(&self, rid: String, entry: MmEncodedEntry) {
self.0.lock().unwrap().insert(rid, entry);
}
pub fn take(&self, rid: &str) -> Option<MmSidecarEntry> {
pub fn take(&self, rid: &str) -> Option<MmEncodedEntry> {
self.0.lock().unwrap().remove(rid)
}
pub fn purge(&self, rid: &str) {
+1 -1
View File
@@ -86,7 +86,7 @@ pub(super) fn shm_name(item: usize) -> String {
format!("sglmm-{}-{n}-{item}", std::process::id())
}
/// Test helper shared with the sidecar's parking tests.
/// Test helper shared with the result store's parking tests.
#[cfg(test)]
pub(super) fn shm_path(name: &str) -> std::path::PathBuf {
std::path::Path::new("/dev/shm").join(name)
+47 -29
View File
@@ -3,7 +3,7 @@
use std::sync::Arc;
use super::sidecar::{FeatureStore, MmSidecarEntry, Sidecar, park_features_in_shm};
use super::result_store::{FeatureStore, MmEncodedEntry, MmResultStore, park_features_in_shm};
use crate::message::config::MmSpec;
use crate::message::ids::Rid;
use crate::message::request::MmRequest;
@@ -46,27 +46,27 @@ fn parse_caller_hash(entry: &str) -> Option<u64> {
}
/// Shared state of the mm path, built once at `start_mm_workers`.
pub struct Context {
pub struct MmContext {
pub family: Box<dyn sglang_mm::pipeline::MmFamilyProcessor>,
/// `None` under `skip_tokenizer_init` (requests must carry `input_ids`).
pub tokenizer: Option<Arc<dyn TextTokenizer>>,
pub sidecar: Sidecar,
pub results: MmResultStore,
/// Park feature buffers in POSIX shm. Set by the Python launcher
/// (`RustMmProcessor._use_feature_shm`) exactly when the scheduler broadcasts
/// across TP ranks and will unwrap `ShmPointerMMData`.
pub feature_shm: bool,
}
impl Context {
impl MmContext {
pub fn new(
spec: MmSpec,
tokenizer: Option<Arc<dyn TextTokenizer>>,
sidecar: Sidecar,
results: MmResultStore,
) -> Result<Self, String> {
Ok(Self {
family: sglang_mm::registry::build_pipeline(spec.pipeline)?,
tokenizer,
sidecar,
results,
feature_shm: spec.feature_shm,
})
}
@@ -75,7 +75,7 @@ impl Context {
/// Run the pipeline for one request. `Ok` returns the final expanded ids, the
/// buffers already parked; `Err` rejects the request back to the client.
fn process(
ctx: &Context,
ctx: &MmContext,
rid: &Rid,
mut work: crate::message::request::MmWorkItem,
) -> Result<Vec<i32>, String> {
@@ -87,42 +87,60 @@ fn process(
})?;
tokenizer.encode(text).map_err(|error| error.to_string())
})?;
let mut drain = sglang_mm::qwen_vl::pack_drain(output)?;
apply_caller_hashes(&mut drain.hashes, &caller_hashes);
// TODO(mm-families): the one family-specific call in this worker — dispatch
// on the spec's `family` (as `registry::build_pipeline` does) once a
// second family lands.
let mut packed = sglang_mm::qwen_vl::pack_output(output)?;
apply_caller_hashes(&mut packed.hashes, &caller_hashes);
let features = if ctx.feature_shm {
park_features_in_shm(&drain.features, &drain.grids)
park_features_in_shm(&packed.features, &packed.grids)
} else {
FeatureStore::Inline(drain.features)
FeatureStore::Inline(packed.features)
};
ctx.sidecar.park(
ctx.results.park(
rid.as_str().to_owned(),
MmSidecarEntry {
MmEncodedEntry {
features,
grids: drain.grids,
hashes: drain.hashes,
offsets: drain.offsets,
mrope: drain.mrope,
mrope_delta: drain.mrope_delta,
grids: packed.grids,
hashes: packed.hashes,
offsets: packed.offsets,
mrope: packed.mrope,
mrope_delta: packed.mrope_delta,
},
);
Ok(drain.input_ids)
Ok(packed.input_ids)
}
/// One MM worker, spawned via `Runtime::spawn_mm_pool` (which owns the
/// Boot-time wiring of the MM path, held privately by the `Runtime` for the
/// late pool spawn (`Runtime::start_mm_workers`, once Python has resolved
/// the spec).
pub struct MmWiring {
/// Requests parked in `Encoding`, drained by the worker pool. Stays empty
/// for non-multimodal models — nothing routes to it.
pub mm_rx: flume::Receiver<MmRequest>,
/// Back-channel for the workers' `MmEncoded` / `MmFailed` into the
/// to-scheduler loop.
pub tm_tx: flume::Sender<TmEvent>,
/// The loaded tokenizer, shared with the tokenizer pool (`None` under
/// `skip_tokenizer_init`).
pub tokenizer: Option<Arc<dyn TextTokenizer>>,
}
/// One MM worker, spawned via `Runtime::start_mm_workers` (which owns the
/// pinning policy for this pool — see its docs).
pub struct MmWorker {
rx: flume::Receiver<MmRequest>,
tm: flume::Sender<TmEvent>,
ctx: Arc<Context>,
mm_rx: flume::Receiver<MmRequest>,
tm_tx: flume::Sender<TmEvent>,
ctx: Arc<MmContext>,
}
impl MmWorker {
pub fn new(
rx: flume::Receiver<MmRequest>,
tm: flume::Sender<TmEvent>,
ctx: Arc<Context>,
mm_rx: flume::Receiver<MmRequest>,
tm_tx: flume::Sender<TmEvent>,
ctx: Arc<MmContext>,
) -> Self {
Self { rx, tm, ctx }
Self { mm_rx, tm_tx, ctx }
}
}
@@ -131,7 +149,7 @@ impl Runnable for MmWorker {
/// shutdown). One request at a time, so the pool size bounds MM
/// concurrency; an error rejects the request back to the client.
fn run(self) {
while let Ok(req) = self.rx.recv() {
while let Ok(req) = self.mm_rx.recv() {
let rid = req.rid;
let event = match process(&self.ctx, &rid, req.work) {
Ok(input_ids) => {
@@ -143,7 +161,7 @@ impl Runnable for MmWorker {
TmEvent::MmFailed { rid, message }
}
};
if self.tm.send(event).is_err() {
if self.tm_tx.send(event).is_err() {
return; // to-scheduler gone: shutdown
}
}
@@ -11,7 +11,7 @@ use crate::message::request::{MmRequest, Request, RequestKind, SchedulerRequest}
use crate::message::response::ResponseItem;
use crate::runtime::Runnable;
use crate::tokenizer_manager::channel::ToSchedulerTx;
pub use crate::tokenizer_manager::to_scheduler_types::{Limits, Mm};
pub use crate::tokenizer_manager::to_scheduler_types::{Limits, MmDispatch};
use crate::tokenizer_manager::to_scheduler_validation::{
check_total_tokens, validate, validate_input_ids,
};
@@ -36,7 +36,7 @@ pub struct Intake {
senders: Senders,
to_scheduler_tx: ToSchedulerTx,
limits: Limits,
mm: Mm,
mm: MmDispatch,
/// Requests parked in `Encoding` while an MM worker processes their media;
/// resumed by `MmEncoded` / `MmFailed`. Only this thread touches it, so no
/// lock.
@@ -51,7 +51,7 @@ impl Intake {
senders: Senders,
to_scheduler_tx: ToSchedulerTx,
limits: Limits,
mm: Mm,
mm: MmDispatch,
shutdown: flume::Receiver<()>,
) -> Self {
Self {
@@ -126,7 +126,7 @@ impl Intake {
}
// A rejected request never reaches the scheduler drain, so purge any
// parked MM result (no-op for the common non-mm request).
self.mm.sidecar.purge(req.rid.as_str());
self.mm.results.purge(req.rid.as_str());
let _ = req.state.apply(Event::Error(err.clone()));
let _ = req.sink.try_send(ResponseItem::Error(err)); // client may be gone
if registered {
@@ -402,7 +402,7 @@ impl Intake {
let Some(mut req) = self.pending_mm.remove(&rid) else {
tracing::debug!(rid = %rid, "mm result for unknown/finished request; dropped");
// It will never reach the scheduler drain, so purge or leak.
self.mm.sidecar.purge(rid.as_str());
self.mm.results.purge(rid.as_str());
return;
};
if let RequestKind::Generate(g) = &mut req.kind {
@@ -432,8 +432,8 @@ impl Intake {
/// ([`Rid::from_client`]), so no later request can ever answer to it.
///
/// A request parked in `pending_mm` is cancelled here, so the worker's late
/// result lands in `on_mm_encoded`'s no-entry branch and purges the sidecar —
/// no generation runs for output nobody will read.
/// result lands in `on_mm_encoded`'s no-entry branch and purges the parked
/// result — no generation runs for output nobody will read.
fn on_abort(&mut self, source: AbortSource) {
let rid = source.rid().clone();
if self.pending_mm.remove(&rid).is_some() {
@@ -84,12 +84,12 @@ fn make_intake_inner(
(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 {
/// An [`MmDispatch`] over `tx` with a fresh result store.
fn test_mm(tx: flume::Sender<MmRequest>, enabled: bool) -> MmDispatch {
MmDispatch {
enabled,
tx,
sidecar: Default::default(),
results: Default::default(),
}
}
@@ -781,7 +781,7 @@ fn mm_generate_req(rid: &str) -> Request {
/// 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.
/// result-store 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();
@@ -789,10 +789,10 @@ fn abort_cancels_parked_mm_request() {
mm_rx.try_recv().expect("parked to mm pool");
// The worker parks its result, as it always does before MmEncoded.
intake.mm.sidecar.park(
intake.mm.results.park(
"mm-gone".into(),
crate::multi_modality::sidecar::MmSidecarEntry {
features: crate::multi_modality::sidecar::FeatureStore::Inline(vec![]),
crate::multi_modality::result_store::MmEncodedEntry {
features: crate::multi_modality::result_store::FeatureStore::Inline(vec![]),
grids: vec![],
hashes: vec![],
offsets: vec![],
@@ -803,13 +803,13 @@ fn abort_cancels_parked_mm_request() {
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.
// The late result must be dropped, not queued, and the parked result 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");
assert!(intake.mm.results.take("mm-gone").is_none(), "entry purged");
}
/// A multimodal request parks in `Encoding` (submitted to the mm worker
@@ -861,7 +861,7 @@ fn mm_failure_rejects_parked_request() {
assert!(consumer.drain(16).headers.is_empty(), "nothing queued");
}
/// On a non-multimodal model (`Mm::enabled == false`), image_data is silently
/// On a non-multimodal model (`MmDispatch::enabled == false`), image_data is silently
/// ignored and the request tokenizes as plain text — the Python
/// TokenizerManager behavior when `mm_processor is None`.
#[test]
@@ -3,19 +3,19 @@
use crate::message::config::ServerArgs;
use crate::message::request::MmRequest;
/// The intake side of the MM path.
/// Dispatches multimodal requests onto the MM worker channel.
#[derive(Clone)]
pub struct Mm {
pub struct MmDispatch {
/// 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
/// Parked results. 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,
pub results: crate::multi_modality::result_store::MmResultStore,
}
/// Resolved once at boot from the scheduler's `server_args`.
@@ -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_result`), not this event.
/// result store (`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, …).
+29 -23
View File
@@ -9,7 +9,7 @@
//! * To_scheduler — 1 thread driving the FSM
//! * From_scheduler — 1 thread draining the scheduler → detok shards
//! * MM workers — K unpinned OS threads, spawned late via
//! [`Runtime::spawn_mm_pool`] (multimodal models only)
//! [`Runtime::start_mm_workers`] (multimodal models only)
//!
//! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling
//! axum's worker threads.
@@ -41,18 +41,11 @@ pub trait Runnable: Send + 'static {
pub struct Runtime {
pub to_scheduler_rx: ToSchedulerRx,
pub from_scheduler_tx: FromSchedulerTx,
/// Requests parked in `Encoding`, drained by the MM worker pool
/// (`Server.start_mm_workers`). Stays empty for non-multimodal models —
/// request never routes to it.
pub to_mm_worker_rx: flume::Receiver<crate::message::request::MmRequest>,
/// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into to_scheduler.
pub from_mm_worker_tx: flume::Sender<TmEvent>,
/// The loaded tokenizer, shared with the MM worker path (`None` under
/// `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_result`).
pub mm_sidecar: crate::multi_modality::sidecar::Sidecar,
pub mm_results: crate::multi_modality::result_store::MmResultStore,
/// Wiring for the late-spawned MM pool ([`Runtime::start_mm_workers`]).
mm_wiring: crate::multi_modality::worker::MmWiring,
/// Worker join handles, joined by `request_shutdown` / `Drop`.
threads: Mutex<Vec<JoinHandle<()>>>,
/// The single shutdown sender.
@@ -64,23 +57,34 @@ pub struct Runtime {
const SHUTDOWN_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
impl Runtime {
/// Spawn `workers` `mm-worker-{i}` threads into the shutdown join set —
/// late, once Python has built the mm spec (`Server::start_mm_workers`).
/// Build the family context from `spec` and spawn `workers`
/// `mm-worker-{i}` threads into the shutdown join set — late, once Python
/// has built the mm spec (`Server.start_mm_workers`).
///
/// Deliberately unpinned: the threads inherit the launch thread's affinity,
/// already narrowed by `RustServer.launch` to the server cores, so bursty
/// MM preprocessing floats over that whole set (rather than owning cores
/// that idle between bursts) and never preempts the scheduler's reserved
/// cores.
pub fn spawn_mm_pool(&self, workers: usize, ctx: Arc<crate::multi_modality::worker::Context>) {
pub fn start_mm_workers(
&self,
spec: crate::message::config::MmSpec,
workers: usize,
) -> Result<(), String> {
let ctx = Arc::new(crate::multi_modality::worker::MmContext::new(
spec,
self.mm_wiring.tokenizer.clone(),
self.mm_results.clone(),
)?);
let mut threads = self.threads.lock().unwrap();
spawn_pool("mm-worker", None, workers.max(1), &mut threads, |_| {
crate::multi_modality::worker::MmWorker::new(
self.to_mm_worker_rx.clone(),
self.from_mm_worker_tx.clone(),
self.mm_wiring.mm_rx.clone(),
self.mm_wiring.tm_tx.clone(),
ctx.clone(),
)
});
Ok(())
}
/// Stop the runtime and join every worker thread (with a bounded wait).
@@ -165,7 +169,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
.map(|t| Arc::new(tokenizer::DynamoTokenizer::new(t.clone())) as _);
// Shared: MM workers park, the Python drain pops.
let mm_sidecar: crate::multi_modality::sidecar::Sidecar = Default::default();
let mm_results: crate::multi_modality::result_store::MmResultStore = Default::default();
// --- Detokenizer shards (pinned, CPU bound) ---
{
@@ -250,10 +254,10 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
.and_then(|p| p.tm.get(1).or_else(|| p.tm.first()).copied())
.map(|c| vec![c]);
let limits = tokenizer_manager::to_scheduler::Limits::from(&*cfg.server_args);
let mm = tokenizer_manager::to_scheduler::Mm {
let mm = tokenizer_manager::to_scheduler::MmDispatch {
enabled: cfg.server_args.model_is_multimodal(),
tx: mm_worker_tx,
sidecar: mm_sidecar.clone(),
results: mm_results.clone(),
};
let mut parts = Some((tok_manager_rx, to_scheduler_tx)); // moved into the single worker
let shutdown_rx = shutdown_rx.clone();
@@ -318,10 +322,12 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
Ok(Runtime {
to_scheduler_rx,
from_scheduler_tx,
to_mm_worker_rx: mm_worker_rx,
from_mm_worker_tx: tok_manager_tx,
tokenizer: text_tokenizer,
mm_sidecar,
mm_results,
mm_wiring: crate::multi_modality::worker::MmWiring {
mm_rx: mm_worker_rx,
tm_tx: tok_manager_tx,
tokenizer: text_tokenizer,
},
threads: Mutex::new(threads),
shutdown_tx: Mutex::new(Some(shutdown_tx)),
})
@@ -1,7 +1,7 @@
"""End-to-end parity at the scheduler-input boundary.
`test_preprocess.py` pins the `preprocess` binding; this drives the whole native
path — the `process_mm` driver, then `RustMmProcessor.build_output` — and
path — the `process_mm` driver, then `RustMmProcessor.wrap_encoded` — and
compares every field the scheduler reads against the Python `mm_processor`.
Bitwise, for both HF backends: the Rust resize clones PIL's fixed-point bicubic
and ATen's uint8 antialias kernel, so whichever one a server is configured with
@@ -82,9 +82,9 @@ class TestQwenE2eParity(CustomTestCase):
ids, features, grids, hashes, offsets, mrope, delta = DRIVER(
PROMPT_PER_IMAGE * len(sources), sources, spec.rust_json()
)
# The shape of Rust's MmEncodeResult, inline transport (test_build_output
# pins the shm shape).
handoff = SimpleNamespace(
# The shape of Rust's MmEncodedResult, inline transport
# (test_wrap_encoded pins the shm shape).
encoded = SimpleNamespace(
features=features,
shm_names=None,
grids=grids,
@@ -93,7 +93,7 @@ class TestQwenE2eParity(CustomTestCase):
mrope=mrope,
mrope_delta=delta,
)
return snapshot(ids, RustMmProcessor.build_output(spec, handoff))
return snapshot(ids, RustMmProcessor.wrap_encoded(spec, encoded))
def run_python(self, sources):
"""The reference path: the Python `mm_processor` the scheduler would use."""
@@ -49,6 +49,8 @@ class TestQwenRustMmHashes(CustomTestCase):
def setUp(self):
from sglang.srt.managers.multimodal_processor import import_processors
# The hash helper builds RustMmProcessor via __new__ (no __init__),
# so processors must be registered here for resolve_spec's lookup.
import_processors("sglang.srt.multimodal.processors")
self.processor = make_processor(self, PROCESSOR_CONFIGS["qwen2_5_vl"])
@@ -1,4 +1,4 @@
"""``RustMmProcessor.build_output``: the drain-time
"""``RustMmProcessor.wrap_encoded``: the drain-time
wrapping contracts tensors are zero-copy views over the Rust-owned buffers, and
pad values come from worker-precomputed hashes, since the scheduler loop must
never hash features. Synthetic buffers, so this needs no Rust extension."""
@@ -23,7 +23,7 @@ from sglang.srt.rust_server.multimodal import ( # noqa: E402
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
class TestBuildRustMmOutput(CustomTestCase):
class TestWrapEncoded(CustomTestCase):
def setUp(self):
# feature_dim == 3 * temporal_patch_size * patch_size**2 == 6.
self.spec = RustMmSpec(
@@ -53,9 +53,9 @@ class TestBuildRustMmOutput(CustomTestCase):
def build(self):
features = np.arange(30, dtype=np.float32)
output = RustMmProcessor.build_output(
output = RustMmProcessor.wrap_encoded(
self.spec,
SimpleNamespace( # the shape of Rust's MmEncodeResult
SimpleNamespace( # the shape of Rust's MmEncodedResult
grids=self.GRIDS,
hashes=self.HASHES,
offsets=self.OFFSETS,
@@ -102,7 +102,7 @@ class TestBuildRustMmOutput(CustomTestCase):
)
class TestBuildRustMmOutputShm(TestBuildRustMmOutput):
class TestWrapEncodedShm(TestWrapEncoded):
"""The shm entry shape (TP>1): features arrive as named POSIX segments, and
each item becomes a ``ShmPointerMMData`` stub whose ``materialize()`` yields
that item's slice — and unlinks, taking the cleanup duty exactly once."""