diff --git a/python/sglang/srt/rust_server/multimodal.py b/python/sglang/srt/rust_server/multimodal.py index e48b601d5..779cd9a7d 100644 --- a/python/sglang/srt/rust_server/multimodal.py +++ b/python/sglang/srt/rust_server/multimodal.py @@ -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 + ), ) diff --git a/python/sglang/srt/rust_server/server.py b/python/sglang/srt/rust_server/server.py index 1e4a9ddc0..64317d04d 100644 --- a/python/sglang/srt/rust_server/server.py +++ b/python/sglang/srt/rust_server/server.py @@ -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 diff --git a/rust/sglang-mm/src/qwen_vl/mod.rs b/rust/sglang-mm/src/qwen_vl/mod.rs index 415848d9c..0985c7769 100644 --- a/rust/sglang-mm/src/qwen_vl/mod.rs +++ b/rust/sglang-mm/src/qwen_vl/mod.rs @@ -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, - /// 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, + /// Per item `[t, h, w]` patch grid. pub grids: Vec<[u32; 3]>, pub hashes: Vec, + /// 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, pub mrope_delta: i64, } -pub fn pack_drain(output: crate::driver::Output) -> Result { +pub fn pack_output(output: crate::driver::Output) -> Result { 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 { ("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, )) } diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index 000ebdd4f..0e090f9bd 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -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 { + fn take_mm_result(&self, py: Python<'_>, rid: &str) -> Option { 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::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/rust/sglang-server/src/multi_modality.rs b/rust/sglang-server/src/multi_modality.rs index 3729b574c..dd189aaef 100644 --- a/rust/sglang-server/src/multi_modality.rs +++ b/rust/sglang-server/src/multi_modality.rs @@ -1,6 +1,6 @@ //! Multimodal worker pool. pub mod payload; +pub mod result_store; mod shm; -pub mod sidecar; pub mod worker; diff --git a/rust/sglang-server/src/multi_modality/sidecar.rs b/rust/sglang-server/src/multi_modality/result_store.rs similarity index 87% rename from rust/sglang-server/src/multi_modality/sidecar.rs rename to rust/sglang-server/src/multi_modality/result_store.rs index fb1bd55a0..75f261e99 100644 --- a/rust/sglang-server/src/multi_modality/sidecar.rs +++ b/rust/sglang-server/src/multi_modality/result_store.rs @@ -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, + /// 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, 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>>); +pub struct MmResultStore(Arc>>); -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 { + pub fn take(&self, rid: &str) -> Option { self.0.lock().unwrap().remove(rid) } pub fn purge(&self, rid: &str) { diff --git a/rust/sglang-server/src/multi_modality/shm.rs b/rust/sglang-server/src/multi_modality/shm.rs index e16c0de47..3b765c6e5 100644 --- a/rust/sglang-server/src/multi_modality/shm.rs +++ b/rust/sglang-server/src/multi_modality/shm.rs @@ -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) diff --git a/rust/sglang-server/src/multi_modality/worker.rs b/rust/sglang-server/src/multi_modality/worker.rs index 280d24c39..3a37ac8f0 100644 --- a/rust/sglang-server/src/multi_modality/worker.rs +++ b/rust/sglang-server/src/multi_modality/worker.rs @@ -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 { } /// Shared state of the mm path, built once at `start_mm_workers`. -pub struct Context { +pub struct MmContext { pub family: Box, /// `None` under `skip_tokenizer_init` (requests must carry `input_ids`). pub tokenizer: Option>, - 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>, - sidecar: Sidecar, + results: MmResultStore, ) -> Result { 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, 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, + /// Back-channel for the workers' `MmEncoded` / `MmFailed` into the + /// to-scheduler loop. + pub tm_tx: flume::Sender, + /// The loaded tokenizer, shared with the tokenizer pool (`None` under + /// `skip_tokenizer_init`). + pub tokenizer: Option>, +} + +/// 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, - tm: flume::Sender, - ctx: Arc, + mm_rx: flume::Receiver, + tm_tx: flume::Sender, + ctx: Arc, } impl MmWorker { pub fn new( - rx: flume::Receiver, - tm: flume::Sender, - ctx: Arc, + mm_rx: flume::Receiver, + tm_tx: flume::Sender, + ctx: Arc, ) -> 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 } } diff --git a/rust/sglang-server/src/tokenizer_manager/to_scheduler.rs b/rust/sglang-server/src/tokenizer_manager/to_scheduler.rs index 71220cd1c..c45ee4248 100644 --- a/rust/sglang-server/src/tokenizer_manager/to_scheduler.rs +++ b/rust/sglang-server/src/tokenizer_manager/to_scheduler.rs @@ -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() { diff --git a/rust/sglang-server/src/tokenizer_manager/to_scheduler_tests.rs b/rust/sglang-server/src/tokenizer_manager/to_scheduler_tests.rs index 85657421e..c1b8139bc 100644 --- a/rust/sglang-server/src/tokenizer_manager/to_scheduler_tests.rs +++ b/rust/sglang-server/src/tokenizer_manager/to_scheduler_tests.rs @@ -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, enabled: bool) -> Mm { - Mm { +/// An [`MmDispatch`] over `tx` with a fresh result store. +fn test_mm(tx: flume::Sender, 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] diff --git a/rust/sglang-server/src/tokenizer_manager/to_scheduler_types.rs b/rust/sglang-server/src/tokenizer_manager/to_scheduler_types.rs index e39d1ec14..a86024bb6 100644 --- a/rust/sglang-server/src/tokenizer_manager/to_scheduler_types.rs +++ b/rust/sglang-server/src/tokenizer_manager/to_scheduler_types.rs @@ -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, - /// 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`. diff --git a/rust/sglang-server/src/tokenizer_manager/wiring.rs b/rust/sglang-server/src/tokenizer_manager/wiring.rs index 4c7235229..9331116d1 100644 --- a/rust/sglang-server/src/tokenizer_manager/wiring.rs +++ b/rust/sglang-server/src/tokenizer_manager/wiring.rs @@ -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 }, /// An MM worker rejected a request parked in `Encoding` (bad media URL, /// unsupported modality, preprocess error, …). diff --git a/rust/sglang-server/src/utils/runtime.rs b/rust/sglang-server/src/utils/runtime.rs index 46f0dd542..1fc359cb6 100644 --- a/rust/sglang-server/src/utils/runtime.rs +++ b/rust/sglang-server/src/utils/runtime.rs @@ -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, - /// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into to_scheduler. - pub from_mm_worker_tx: flume::Sender, - /// The loaded tokenizer, shared with the MM worker path (`None` under - /// `skip_tokenizer_init`). - pub tokenizer: Option>, /// 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>>, /// 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) { + 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 { .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 { .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 { 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)), }) diff --git a/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py b/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py index 40ec519f8..4e37bcb0a 100644 --- a/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py +++ b/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py @@ -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.""" diff --git a/test/registered/unit/multimodal/rust/qwen/test_rust_mm_processor.py b/test/registered/unit/multimodal/rust/qwen/test_rust_mm_processor.py index 0e9070470..6f408a92d 100644 --- a/test/registered/unit/multimodal/rust/qwen/test_rust_mm_processor.py +++ b/test/registered/unit/multimodal/rust/qwen/test_rust_mm_processor.py @@ -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"]) diff --git a/test/registered/unit/multimodal/rust/shared/test_build_rust_mm_output.py b/test/registered/unit/multimodal/rust/shared/test_wrap_encoded.py similarity index 95% rename from test/registered/unit/multimodal/rust/shared/test_build_rust_mm_output.py rename to test/registered/unit/multimodal/rust/shared/test_wrap_encoded.py index fde7e42dd..35dd929ca 100644 --- a/test/registered/unit/multimodal/rust/shared/test_build_rust_mm_output.py +++ b/test/registered/unit/multimodal/rust/shared/test_wrap_encoded.py @@ -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."""