Repository navigation
[mm] rust-server: native multimodal processing for Qwen VL (integrate sglang-mm, e2e) - #32365
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
…envelope - Add CPU-CI smoke coverage for the _core.inkling bindings (previously only exercised by the GPU e2e model test). - Tighten the qwen feature-parity max-diff envelope to 0.035: the Rust bicubic resize is bit-exact vs PIL, so the tolerance only needs to absorb HF-internal noise (measured worst case 0.030). - Use VISION_END_ID instead of the literal 902. - Drop _fixtures.py: its only consumer (test_scheduler_input_parity) lives in the stacked integration PR #32365, which already carries this file. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
41997a2 to
9123bc5
Compare
fa72bd6 to
e4178da
Compare
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
btw I encountered a bug when running This is out of scope if this PR, so I will figure it out from the |
yes, this is relevant to: #33219 |
…envelope - Add CPU-CI smoke coverage for the _core.inkling bindings (previously only exercised by the GPU e2e model test). - Tighten the qwen feature-parity max-diff envelope to 0.035: the Rust bicubic resize is bit-exact vs PIL, so the tolerance only needs to absorb HF-internal noise (measured worst case 0.030). - Use VISION_END_ID instead of the literal 902. - Drop _fixtures.py: its only consumer (test_scheduler_input_parity) lives in the stacked integration PR #32365, which already carries this file. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
60c2595 to
be45af4
Compare
The name now says it is the live-server end-to-end suite (vlm-dir convention, cf. test_token_id_retokenize_e2e), distinct from the CPU units and the MMMU accuracy gate. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…adcast_pyobj
Fixes the TP>1 image-TTFT regression vs the Python TokenizerManager
(Qwen3.5-397B TP4: +44% mean TTFT at c=64, +80% at c=256; reproduced on
Qwen3.5-35B-A3B TP4).
Root cause, established by per-rank instrumentation: the sidecar delivers
feature tensors zero-copy into rank 0 only, so the ~20 MB/image buffer rode
`broadcast_pyobj` (pickle -> gloo -> unpickle) to ranks 1..n on the request's
critical path. The TP-sharded ViT has a collective in every block, so rank 0
started ~31 ms before the last rank and stalled at the first allreduce for
exactly that skew (python arm: all ranks start within 2 ms — its
TokenizerManager ships features as `ShmPointerMMData` stubs and every rank
materializes in parallel). Preprocessing was never the problem: native inline
is 12.9 ms vs HF's 14.6 ms per 720p image.
The fix reuses the Python path's transport: when the scheduler will broadcast
across TP ranks (and the receiver's `unwrap_shm_features` gates pass), the MM
worker parks each item's feature buffer in a named POSIX segment — written
off the scheduler loop — and `take_mm` hands Python the names. The drain
wraps each name in a `ShmPointerMMData` stub (post-`__setstate__` form: rank
0 never pickle-roundtrips its own copy), so `broadcast_pyobj` carries ~100
bytes and the receiver's existing barrier + materialize does the parallel
fan-out. Single-rank serving keeps the zero-copy inline path.
Lifecycle: segment ownership transfers exactly once — worker creates,
`take_mm` disarms the Rust `Drop` (which otherwise unlinks when a parked
result is purged: abort/late-result paths), Python's `materialize()` unlinks
after the post-broadcast clone. shm failure (e.g. /dev/shm full) degrades to
the inline path, mirroring `_wrap_shm_or_inline`.
Measured on Qwen3.5-35B-A3B-FP8 TP4, c=1, one 720p image, p50 of 24:
text TTFT image TTFT increment per-rank ViT
python TM 14.3 87.7 73-84 34/34/34/34
rust TM before 14.5 111.9 97-108 62/33/55/37
rust TM + shm fan-out 14.3 83.2 68.9 35/37/37/36
Rank starts realign and the rust arm now beats python on image requests too.
Tests: cargo -p sglang-server 142 (4 new: segment roundtrip/drop-unlink,
into_name ownership transfer, per-grid slicing, degenerate-shape fallback);
CPU suites 24 passed / 86 subtests (test_build_native_mm grows a shm-shape
mirror class); GPU e2e 4/4 (TP1 inline path).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`_partition_cores` gave the rust pools every allowed core minus the two launch cores — on a multi-GPU host that is the whole NUMA node, which the sibling TP ranks' processes share. MM preprocessing bursts (decode + resize + patchify, ~15 ms CPU per image) could then land on — and preempt — a sibling's CUDA-launch thread; with the TP-sharded ViT synchronizing at every block, one preempted rank inflates every rank's forward. Measured on Qwen3.5-35B-A3B TP4 (1×720p image/request): the unbounded mask cost the worst sibling ~20 ms of ViT wall time per image; bounding the pools removed it. The budget is `max(8, mm_workers + 4)` cores: one per MM worker (the only CPU-hot pool threads) plus headroom for the I/O-shaped tokenizer / ingress / egress / api threads, with the floor covering text-only serving. Rest of the node stays with the scheduler ranks. Validated on the 35B TP4 image sweep (rust arm, vs the unbounded mask): c=16 -2.1%, c=64 -5.3%, c=256 -7.1% mean TTFT; c=1 p50 unchanged (80.2 vs 83.2 ms probe) — the mask only ever mattered under concurrent load. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The sglang-mm driver hashes the raw encoded source (32364 review), not the feature tensor; pin the drain hashes to that convention. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
NativeMmHost built the full Python mm_processor stack only to type-check it in resolve_native_spec — every spec parameter comes from the HF image processor and hf_config. Expose the selection half of get_mm_processor as get_mm_processor_cls and gate on the resolved CLASS instead: same semantics (an SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE override or the transformers-backend fallback still disables the native path) with no processor constructed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
resolve_native_spec() returned the Rust JSON while side-effect-writing self._native, an untyped dict superset the drain adapter read later. Make resolution pure: it returns a frozen msgspec.Struct; the Rust wire is serialized from its subset (rust_json), feature_dim becomes a derived property, and build_native_mm is a static function of (spec, entry). RustServer holds the spec — NativeMmHost is launch-scoped only. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…rkers The MM worker pool is fixed, core-pinned CPU capacity; URL sources used to resolve inside driver::process, so a few slow image hosts could occupy the whole pool and a multi-image request paid n * REQUEST_TIMEOUT serially. The /generate handler now downloads http(s) image sources before dispatch: spawn_blocking around sglang-mm's fetch_bytes (one owner for proxy/timeout/ cap semantics) under a global 32-permit semaphore, all images of a batch concurrently. The bytes ride out-of-band (MmData/MmRequest::prefetched) — the client's opaque msgpack payload is never re-encoded — and mm_payload::parse swaps them back in by walk order. A network source that reaches an MM worker without prefetched bytes is now a hard error, so the workers-never-do-I/O contract is enforced rather than assumed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A file path blocks an MM worker just like a URL — worse on a hanging network mount, which has no timeout. is_network_source widens to is_io_source (http(s)/file:///bare path); file reads now resolve on the API runtime through the same capped fetch_bytes, and data:/base64 stay on the worker (pure CPU). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ress Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ocess take_mm_work() moves the worker-owned fields out of the parked request (text is cloned: the scheduler header still needs it; input_ids is replaced by the expanded ids). Kills the positional wire contract and a serialize+parse copy per request. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
HeapBytes for rmpv::Value (recursive: payload bytes + node overhead) feeds the same check_broadcast_budget as sampling_params, so a 16 MiB base64 scalar broadcast to 4096 prompts is a 400, not a 64 GiB abort. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
on_abort removes the pending_mm entry, so the worker's late result hits on_mm_encoded's no-entry branch and the parked buffers are purged — no generation runs for a dead client. MmResult renamed MmSidecarEntry (the expanded input_ids travel separately, and it isn't a Result); the raw Arc<Mutex<HashMap>> becomes a Sidecar type owning park/take/purge. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mirrors the Python TokenizerManager exactly: hex hashes override the computed source-byte hashes so external routers' keys align with the prefix cache; honored on single requests only (the Python batch __getitem__ drops the field), and mismatch/malformed entries warn and fall back per item rather than rejecting. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Named fields across the Rust/Python boundary so the handoff schema is explicit; the drain adapter and its tests read attributes instead of unpacking positions. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two shifts invalidated the GIF probe in test_unsupported_format_is_rejected: sglang-mm's webp/gif/bmp enablement (d49ef4c), and — in the server binary — feature unification through dynamo-parsers → openai-harmony, which depends on image with default features and so re-enables EVERY default format (tiff/avif/ico/exr/...) despite sglang-mm's minimal list. PCX is a format PIL writes but the image crate does not recognize at all, so the rejection path stays covered in every build shape. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Python's `int(hex_hash, 16)` takes hex of any width, but the Rust path parsed into u64 and treated anything wider as malformed — so a router sending a full SHA-256 silently fell back to the computed hash and lost its prefix-cache alignment. Parse any width and keep the low 64 bits: only the low 30 bits are observable, via `_compute_pad_value`, so a wide digest yields the same pad value Python's arbitrary-precision int would. Empty or non-hex entries still warn and fall back. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`prefetch_all` resolved every source with only the per-source 64 MiB cap, and the driver's per-request guards (64 items, 256 MiB) ran once all of it was resident — 64 valid 64 MiB URLs buffered ~4 GiB before being rejected. Both budgets now gate the prefetch: the item count rejects before a single byte is fetched, and the sources of one request share a `ByteBudget`. Streamed sources charge it chunk by chunk, so overflow rejects mid-download and `try_join_all` drops the rest; releasing what a short read did not use is what keeps this from over-rejecting a legal 64 x 1 MiB request. Base64 payloads are already resident in the request body and cannot amplify, so they charge once decoded. The driver keeps its checks as the backstop for callers without a prefetch layer, and because base64 sources are counted only there. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`GenerateReqInput` declares batched `mm_hashes` shapes, but its `__getitem__` never forwards the field to sub-requests, so Python answers 200 with hashes the caller did not ask for. Honoring the batch form only here would give the two servers different prefix-cache keys for the same body, so this rejects it with a 400 instead: the field exists to align a caller's keys, and silently substituting others is the failure that is hard to notice. A deliberate divergence, and the only one — single requests behave exactly as before, and an absent or empty field still passes. Implementing the batch contract is a `GenerateReqInput` change, for both paths at once. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Publicizing `_resolve_local_or_cached_file` as `resolve_local_or_cached_file` (so the rust-server launch path can share it) missed the one caller outside `hf_transformers`: `models/gemma3_causal.py` still imported the private name, so importing it raised ImportError. The model registry then fell back to `TransformersMultiModalForCausalLM`, which cannot find gemma-3's vision tower parameters — `test_vlm_models.py` failed to launch `google/gemma-3-4b-it` with "No module or parameter named 'model.vision_tower.vision_model'". Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Comment pass over the whole native-MM stack: keep the "why" that isn't
inferable from the code, cut the restatements of it. The same three facts
(no Python fallback, sidecar parked before MmEncoded, drain must stay
copy-free) were spelled out in four or five places each; each now lives
once, where it is enforced.
Code cleanups, no behavior change:
- `Ingress::new` takes one `Mm { enabled, tx, sidecar }` instead of three
positional args, which drops the `too_many_arguments` allow and shortens
every test call site.
- `NativeMmSpec.rust_json` derives its field set from `__struct_fields__`
minus `DRAIN_ONLY` rather than restating all ten by hand.
- `build_native_mm` reads `entry.*` directly instead of unpacking into
seven locals first.
- `test_build_native_mm` shares one handoff builder between the inline and
shm cases via a `transport()` hook, and loses a test override that only
called `super()`.
Also fixes stale references the review commits left behind: `take_mm`'s doc
still described the 7-tuple it no longer returns, `prefetch`'s module doc
pointed at `mm_payload::parse` (now `to_mm_input`), and the ingress tests
still called the worker pool a "bridge".
The native path was bit-exact only against `Qwen2VLImageProcessorPil`. On transformers 5.x the default `Qwen2VLImageProcessor` is the torchvision path, and against it the features drifted by 1 u8 level (2/255 after normalize) on ~0.4% of pixels — so a default server's native features were not what its own Python path would produce. Two causes, both now reproduced instead of tolerated: - Resize. ATen's uint8 antialias kernel shares PIL's geometry, bicubic kernel and per-pass u8 rounding, and differs only in quantizing the weights to i16 at the widest non-saturating precision instead of PIL's fixed 22-bit i32. `Resample::AtenU8` implements that; verified dispatch-independent (identical under ATEN_CPU_CAPABILITY=default/avx2/avx512), so it is a stable target. - Normalize. The fast processor folds the rescale into mean/std before normalizing (`_fuse_mean_std_and_rescale_factor`), which lands on a different f32 for 128 of the 256 u8 inputs. `normalize_lut` now picks the matching form. `NativeMmSpec.resample` carries which processor to mirror, and the gate maps each admitted class to it. The gate also now rejects a processor that disables resize/rescale/normalize or uses a rescale factor other than 1/255 — those would have silently produced different features, and the fused constants assume it. Tests: `test_e2e_parity.py` asserts the whole scheduler-input boundary bitwise for both backends; `test_preprocess.py` and `test_native_mm_host.py` drop their tolerances entirely; `test_resize_parity.py` covers all three resamplers over the fixed cases plus a random sweep, since the `aten_u8` precision varies with the scale factor. `resize_rgb` takes a `resample` argument to make that testable. Also lets the gate accept `Qwen2VLImageProcessorPil`, which it previously refused — `--disable-fast-image-processor` + `SGLANG_RUST_SERVER=1` hard-failed launch (smoke-tested on a live server).
Replace the qwen-specific gate constants (NATIVE_QWEN_MODEL_TYPES, NATIVE_IMAGE_PROCESSORS, the hard-coded family string and processor identity check) with a NATIVE_MM_FAMILIES registry mirroring the Rust pipeline_from_spec arms: supporting a new model family is one declaration plus its Rust arm. native_mm_family_for is the shared selection seam. test_native_mm_gate pins the negative half the suites never covered: a non-Qwen multimodal model (Inkling) keeps its Python processor and matches no native family, and matching is by class identity so an SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE override stays off the native path. Also drops a redundant rebinding in into_requests. Co-authored-by: Cursor <cursoragent@cursor.com>
Readability only, same coverage: PROMPT_PER_IMAGE names the fixture-vocab prompt, both_paths splits into run_native/run_python (each naming the production path it mirrors), source forms carry names so a failing subTest says which transport broke, and the PIL subclass adds the transport-only test instead of the torchvision subclass nulling an inherited one. Co-authored-by: Cursor <cursoragent@cursor.com>
A different approach lands in the next PR; the Rust hf-hub fallback (with the env-var patch) covers the default paths. Co-authored-by: Cursor <cursoragent@cursor.com>
The MM workers were the one stage with a bespoke thread-spawn loop. MmWorker now implements Runnable and Runtime::spawn_mm_pool (replacing adopt_threads) spawns it through threads::spawn_pool, deliberately unpinned: the pool keeps inheriting the Python-narrowed server-core mask so bursty MM preprocessing floats instead of owning idle cores. Also correct the stale "blocking threads are unpinned" comments: the api runtime's on_thread_start fires for blocking-pool threads too and round-robins them onto the api core set, which is the desired policy. Co-authored-by: Cursor <cursoragent@cursor.com>
Failed(Error) is read by ingress now, so the lint suppress is unused. Co-authored-by: Cursor <cursoragent@cursor.com>
Motivation
Second half of the native-MM stack: integrate the
sglang-mmpipeline (#32364) into the Rust tokenizer manager, soSGLANG_RUST_SERVER=1serves Qwen VL end to end. Image preprocessing runs entirely in Rust — no Pythonmm_processorin the request path, and no Python fallback (unsupported models fail at launch; unsupported inputs are rejected per request).Design
Key decision: MM output features never pass through ingress. Feature tensors are tens of MB per image-heavy request; routing them through the single-threaded ingress loop would stall admission, and per-byte work at drain time would inflate every running request's ITL. So a worker parks the big buffers (features, grids, hashes, M-RoPE) in a rid-keyed sidecar and sends only the placeholder-expanded
input_idsback to ingress. The scheduler drain pops the sidecar by rid, copy-free and hash-free: numpy takes ownership of the Rust vectors,torch.from_numpywraps them intoMultimodalProcessorOutput.Semantics:
MmEncoded, so a drained request always finds its buffers; late results for rejected/aborted requests are purged.MmFailed→ 400/500; a full mm channel back-pressures like a full ingress ring (503).--mm-processor-worker-num(0 → 8); workers inherit the server-core affinity mask, never the scheduler's reserved cores.Modifications
rust/sglang-server:mm.rs(worker pool + sidecar, drivingsglang-mm'sdriver::processwith the server tokenizer);message/request.rs+message/mm_payload.rs(image/video/audio_datawire fields, batch fan-out mirroring Python_normalize_*_data);ingress.rs+fsm.rs(Encodingstage: park inpending_mm, resume onMmEncoded/MmFailed);lib.rs(pyo3start_mm_workers/take_mm).rust_server.py,io_struct.py):NativeMmHost.resolve_native_spec()— launch gate building the Rust pipeline spec from the resolved HF processor; anything it can't mirror exactly fails launch fast. Native families are declared inNATIVE_MM_FAMILIES, one entry per Rustpipeline_from_specarm.NativeMmHost.build_native_mm()— the zero-copy drain adapter.Testing
test_rust_native_mm_e2e.py— liveSGLANG_RUST_SERVER=1server: single/multi-image generate (URL + base64), rejection of unsupported/corrupt images without crashing;test_rust_native_mm_mmmu.py— MMMU accuracy gate (≥ 0.30 on a fixed 100-sample subset, matching the Python path) and asserts the native pipeline engaged.qwen/test_native_mm_host.py(scheduler input identical to the Pythonmm_processoroutput),shared/test_build_native_mm.py(zero-copy + no-hashing drain contracts),shared/test_native_mm_gate.py(family selection by class identity; non-native models keep their Python processor).cargo test -p sglang-server(238 tests); workspacecargo checkclean.CI States
Latest PR Test (Base): 🚫 Run #31049346794
Latest PR Test (Extra): ❌ Run #31049346320