Skip to content

[mm] rust-server: native multimodal processing for Qwen VL (integrate sglang-mm, e2e) - #32365

Merged
hnyls2002 merged 33 commits into
mainfrom
kan/rust-server-native-mm
Aug 6, 2026
Merged

hnyls2002 merged 33 commits into
mainfrom
kan/rust-server-native-mm

Conversation

@sherlockwu

@sherlockwu sherlockwu commented Jul 25, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

Second half of the native-MM stack: integrate the sglang-mm pipeline (#32364) into the Rust tokenizer manager, so SGLANG_RUST_SERVER=1 serves Qwen VL end to end. Image preprocessing runs entirely in Rust — no Python mm_processor in the request path, and no Python fallback (unsupported models fail at launch; unsupported inputs are rejected per request).

Design

HTTP /generate ──► tm-ingress FSM ──[bounded mm channel: rid + msgpack payload]──► MM worker pool
                        ▲                                                          (N Rust threads)
                        │                                                               │
                        └──── TmEvent::MmEncoded { rid, final input_ids } ◄─────────────┤ (ids only)
                        │                                                               │
                   ingress ring                                    rid-keyed sidecar ◄──┘ (features)
                        │                                                    │
                        ▼                                                    ▼
                 Python scheduler drain ───────── take_mm(rid): zero-copy ───┘

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_ids back to ingress. The scheduler drain pops the sidecar by rid, copy-free and hash-free: numpy takes ownership of the Rust vectors, torch.from_numpy wraps them into MultimodalProcessorOutput.

Semantics:

  • Sidecar entry is stored strictly before MmEncoded, so a drained request always finds its buffers; late results for rejected/aborted requests are purged.
  • Worker errors (bad URL, undecodable image, video/audio) → MmFailed → 400/500; a full mm channel back-pressures like a full ingress ring (503).
  • Pool size = --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, driving sglang-mm's driver::process with the server tokenizer); message/request.rs + message/mm_payload.rs (image/video/audio_data wire fields, batch fan-out mirroring Python _normalize_*_data); ingress.rs + fsm.rs (Encoding stage: park in pending_mm, resume on MmEncoded/MmFailed); lib.rs (pyo3 start_mm_workers / take_mm).
  • Python (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 in NATIVE_MM_FAMILIES, one entry per Rust pipeline_from_spec arm. NativeMmHost.build_native_mm() — the zero-copy drain adapter.

Testing

  • E2E (GPU, new): test_rust_native_mm_e2e.py — live SGLANG_RUST_SERVER=1 server: 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.
  • CPU units: qwen/test_native_mm_host.py (scheduler input identical to the Python mm_processor output), 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).
  • Rust units: cargo test -p sglang-server (238 tests); workspace cargo check clean.
image

CI States

Latest PR Test (Base): 🚫 Run #31049346794
Latest PR Test (Extra): ❌ Run #31049346320

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@github-actions github-actions Bot added documentation Improvements or additions to documentation dependencies Pull requests that update a dependency file npu labels Jul 25, 2026
sherlockwu added a commit that referenced this pull request Jul 25, 2026
…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>
@sherlockwu
sherlockwu force-pushed the kan/rust-server-native-mm branch 4 times, most recently from 41997a2 to 9123bc5 Compare July 26, 2026 09:44
@sherlockwu
sherlockwu changed the base branch from rainj-me/rust-server to kan/sglang-mm-on-rust-server July 26, 2026 09:44
@sherlockwu
sherlockwu force-pushed the kan/rust-server-native-mm branch 2 times, most recently from fa72bd6 to e4178da Compare July 28, 2026 06:07
@sherlockwu
sherlockwu marked this pull request as ready for review July 28, 2026 09:57
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Comment thread python/sglang/srt/managers/rust_server.py Outdated
Comment thread python/sglang/srt/managers/rust_server.py Outdated
@AgainstEntropy

Copy link
Copy Markdown
Collaborator

btw I encountered a bug when running sglang.bench_serving locally. The cause is that RustServer._build_server_args overwrites server_args["tokenizer_path"] in-place (from model_id to local snapshot path), which is reflected in the returns of /model_info.
Then get_tokenizer routes the HF tokenizer path ending in .json wrongly to TiktokenTokenizer, resulting in KeyError: 'regular_tokens'.

This is out of scope if this PR, so I will figure it out from the get_tokenizer side. Nothing needed here, just want to surface and mark the potential robustness issue.

@sherlockwu

Copy link
Copy Markdown
Collaborator Author

btw I encountered a bug when running sglang.bench_serving locally. The cause is that RustServer._build_server_args overwrites server_args["tokenizer_path"] in-place (from model_id to local snapshot path), which is reflected in the returns of /model_info. Then get_tokenizer routes the HF tokenizer path ending in .json wrongly to TiktokenTokenizer, resulting in KeyError: 'regular_tokens'.

This is out of scope if this PR, so I will figure it out from the get_tokenizer side. Nothing needed here, just want to surface and mark the potential robustness issue.

yes, this is relevant to: #33219

sherlockwu added a commit that referenced this pull request Aug 2, 2026
…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>
@sherlockwu
sherlockwu force-pushed the kan/rust-server-native-mm branch from 60c2595 to be45af4 Compare August 2, 2026 12:33
sherlockwu and others added 27 commits August 5, 2026 21:25
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

dependencies Pull requests that update a dependency file documentation Improvements or additions to documentation high priority Multi-modal multi-modal language model npu run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants