[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:
co-authored by
Rain Jiang
Cursor
Claude Fable 5
parent
b42569a0f1
commit
12735c2d76
@@ -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
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
))
|
||||
}
|
||||
|
||||
|
||||
@@ -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,6 +1,6 @@
|
||||
//! Multimodal worker pool.
|
||||
|
||||
pub mod payload;
|
||||
pub mod result_store;
|
||||
mod shm;
|
||||
pub mod sidecar;
|
||||
pub mod worker;
|
||||
|
||||
+13
-8
@@ -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) {
|
||||
@@ -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)
|
||||
|
||||
@@ -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, …).
|
||||
|
||||
@@ -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"])
|
||||
|
||||
|
||||
+5
-5
@@ -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."""
|
||||
Reference in New Issue
Block a user