diff --git a/.github/workflows/pr-test-rust-exts.yml b/.github/workflows/pr-test-rust-exts.yml new file mode 100644 index 000000000..ce78da39d --- /dev/null +++ b/.github/workflows/pr-test-rust-exts.yml @@ -0,0 +1,70 @@ +name: PR Test (Rust extensions) + +# In-crate checks for the Rust extensions bundled with the `sglang` wheel +# (rust/sglang-mm). The pyo3 bindings themselves are exercised by the Python +# unit suite (base-a-test-cpu), which builds the extensions from source; this +# job covers the pure-Rust core, exactly as sglang-server links it, plus lints. +# +# Scoped to sglang-mm on purpose: sglang-server and sglang-grpc have no +# in-crate test suite yet, so a `rust/**` trigger would spawn a job that does +# not actually cover the changed crate. Widen the paths together with the job. + +on: + push: + branches: [ main ] + paths: + - "rust/sglang-mm/**" + - "rust/Cargo.toml" + - ".github/workflows/pr-test-rust-exts.yml" + pull_request: + branches: [ main ] + types: [opened, synchronize, reopened, labeled] + paths: + - "rust/sglang-mm/**" + - "rust/Cargo.toml" + - ".github/workflows/pr-test-rust-exts.yml" + workflow_dispatch: + +concurrency: + group: rust-exts-${{ github.ref }} + cancel-in-progress: true + +jobs: + sglang-mm-unit: + if: | + github.event_name != 'pull_request' || + (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) || + (github.event.action == 'labeled' && github.event.label.name == 'run-ci') + runs-on: ubuntu-latest + timeout-minutes: 15 + defaults: + run: + working-directory: rust/sglang-mm + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + # The workspace root owns Cargo.lock and target/. + workspaces: rust + + # rustfmt and non-test clippy for rust/ are already gated by the Lint + # workflow's pre-commit hooks; `--all-targets` is the part it misses, so + # this lints the test code too rather than duplicating that gate. + - name: cargo clippy (test targets) + run: | + rustup component add clippy + cargo clippy --all-targets -- -D warnings + + # Default features are exactly what sglang-server links: no pyo3, no + # rayon. `rlib_is_single_threaded` only runs in this shape. + - name: cargo test (rlib shape — no pyo3, no rayon) + run: cargo test + + # The wheel's shape (minus pyo3, whose extension-module crates cannot + # link test binaries). Both shapes must pass: `common::par` has two + # implementations and results are required to be identical. + - name: cargo test (parallel shape) + run: cargo test --features parallel diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d25216620..dbc61a358 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -130,10 +130,12 @@ repos: pass_filenames: false - id: clippy-rust-workspace name: clippy rust/ workspace (auto-fix) - # Two steps: apply machine-applicable fixes, then fail on anything left - # (combining --fix with -D warnings can abort the fix phase). protoc is - # not required: sglang-grpc's build.rs falls back to a vendored binary. - entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings' + # Three steps: apply machine-applicable fixes, then fail on anything + # left (combining --fix with -D warnings can abort the fix phase), then + # cover sglang-mm's PyO3 bindings + rayon fan-out — both sit behind + # non-default features, so `--workspace` alone never compiles them. protoc + # is not required: sglang-grpc's build.rs falls back to a vendored binary. + entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings && cargo clippy -p sglang-mm --features python,parallel --lib -- -D warnings' language: system files: ^rust/.*\.rs$ pass_filenames: false diff --git a/python/setup.py b/python/setup.py index 735c846b6..0f81e61c8 100644 --- a/python/setup.py +++ b/python/setup.py @@ -6,6 +6,7 @@ crate whose Cargo.toml declares [package.metadata.sglang] python-module = "sglang.srt.._core" # import path inside the wheel debug = false # optional RustExtension knob + features = ["python"] # optional cargo features to enable is built as a PyO3 extension module at that import path. Adding a new extension crate therefore needs no pyproject changes — declare the metadata in the crate. @@ -115,6 +116,9 @@ def _discovered_rust_extensions(): path=package["manifest_path"], binding=Binding.PyO3, debug=sglang_meta.get("debug"), + # Crates that gate their PyO3 bindings behind a non-default + # feature (so the pure-Rust core stays pyo3-free) declare it here. + features=sglang_meta.get("features"), ) ) if not extensions: diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 3e3d564bb..164c87904 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -3679,6 +3679,9 @@ dependencies = [ "numpy", "pyo3", "rayon", + "serde", + "serde_json", + "ureq", ] [[package]] diff --git a/rust/sglang-mm/Cargo.toml b/rust/sglang-mm/Cargo.toml index b8b8a5fff..16b3ae3e4 100644 --- a/rust/sglang-mm/Cargo.toml +++ b/rust/sglang-mm/Cargo.toml @@ -8,9 +8,12 @@ license.workspace = true # Consumed by python/setup.py: registers this crate as a PyO3 extension module # of the main sglang wheel at the given import path. debug = false keeps # editable installs on release builds (image preprocessing is perf-sensitive). +# features: the wheel build needs the PyO3 bindings, which are NOT default (see +# [features] below). [package.metadata.sglang] python-module = "sglang.srt.multimodal._core" debug = false +features = ["python", "parallel"] [lib] # Unique per-crate artifact name (the shared workspace target/ dir cannot hold @@ -18,14 +21,45 @@ debug = false # name comes from the `#[pymodule]` entry point and the setuptools-rust # `target` in python/pyproject.toml, which renames the built artifact. name = "sglang_mm_core" -crate-type = ["cdylib"] +# cdylib: the PyO3 module (`sglang.srt.multimodal._core`). +# rlib: pure-Rust core linked by sglang-server's native MM path. +crate-type = ["cdylib", "rlib"] + +[features] +# Both features are deliberately NOT default: cargo unifies a package's +# features across everything selected in one invocation, so anything default-on +# here would leak back into sglang-server's copy of this crate under +# `cargo build --workspace` even though the server depends on it with +# `default-features = false`. Opting in explicitly keeps the pure-Rust core +# pure no matter how the workspace is built; the wheel asks for both via +# `[package.metadata.sglang] features` above. +default = [] +# Every PyO3 binding. +python = ["dep:pyo3", "dep:numpy"] +# Fan work out onto crate-owned rayon pools (see `common::par`). Off for the +# rlib: sglang-server provides concurrency across requests and owns its core +# budget, so preprocessing runs inline on the calling thread and rayon is not +# linked at all. Orthogonal to `python` so both shapes stay testable. +parallel = ["dep:rayon"] [dependencies] -pyo3 = { workspace = true } +pyo3 = { workspace = true, optional = true } +numpy = { version = "0.29", optional = true } base64 = "0.22" blake3 = "1" half = "2.4" -image = { version = "0.25", default-features = false, features = ["jpeg", "png"] } -numpy = "0.29" -rayon = "1.10" +# The formats the Python (PIL) path commonly accepts; all pure-Rust decoders. +image = { version = "0.25", default-features = false, features = [ + "jpeg", + "png", + "webp", + "gif", + "bmp", +] } +rayon = { version = "1.10", optional = true } +serde = { workspace = true } +serde_json = { workspace = true } +# Native media fetch (stage 1 of the server MM pipeline); rustls so the +# extension never links the system OpenSSL. +ureq = { version = "2", default-features = false, features = ["tls", "gzip"] } diff --git a/rust/sglang-mm/README.md b/rust/sglang-mm/README.md index 11a4e6531..b4d528694 100644 --- a/rust/sglang-mm/README.md +++ b/rust/sglang-mm/README.md @@ -1,24 +1,117 @@ # sglang-mm Rust-accelerated multimodal preprocessing for SGLang. Fused image decode, -resize, patchify, normalize, and content hash — all parallel and GIL-released. +fetch, resize, patchify, normalize, and content hash — all parallel and +GIL-released. -Compiled as `sglang.srt.multimodal._core` via setuptools-rust when installing sglang. +Built two ways: + +- **PyO3 extension** `sglang.srt.multimodal._core` (features `python,parallel`, + requested by the wheel build) via setuptools-rust when installing sglang — + used by Python processors and parity tests. +- **Pure-Rust `rlib`** (default features, i.e. neither) linked by + `sglang-server`'s MM worker path — that copy needs no pyo3, no libpython, and + no rayon: it spawns no threads and runs inline on the calling thread, because + the server supplies concurrency across requests and pins its own cores. + `tests/rlib_is_single_threaded.rs` guards that from the outside. ## Architecture ``` src/ -├── lib.rs # PyO3 module root (_core) -├── registry.rs # ImageProcessorSpec trait + ProcessorRegistry +├── lib.rs # module root; PyO3 module (_core) feature-gated +├── pipeline.rs # the server-pipeline contract: MmFamilyProcessor +│ # trait + the carriers (Tensor, TokenLayout, ...) +├── driver.rs # model-independent request driver (fetch → +│ # decode → process_item → layout → positions) +├── registry.rs # ImageProcessorSpec registry (Python-facing) +│ # + pipeline_from_spec (family factory) ├── common/ -│ ├── mod.rs # thread pool, image decode, SHA256 hash, base64 -│ ├── resize.rs # PIL-exact Lanczos resize +│ ├── mod.rs # thread pool, image decode, content hash, base64 +│ ├── fetch.rs # media source → bytes (data:/base64/file/http) +│ ├── par.rs # the only fan-out seam (rayon, or inline) +│ ├── resize.rs # PIL-exact Lanczos + Bicubic resize +│ ├── token_layout.rs # TokenLayout mechanics (apply_layout + helpers) │ └── transforms.rs # reusable primitives: normalize, pad, extract_patches └── / - └── mod.rs # model-specific processor + └── mod.rs # model-specific processor (inkling, qwen_vl, ...) ``` +## Server pipeline architecture + +`sglang-server`'s MM workers process an image request entirely in Rust. +`driver::process` runs the same fixed steps for every model family: + +``` +MmInput { text?, input_ids?, images } + 1. per image: fetch_bytes (inline, sequential — see Design notes), then + fanned out via common::par: + content hash → decode_rgb → family.process_item() + → ProcessedItem { feature, aux, geometry } + 2. family.layout(input_ids, geometries) → TokenLayout + apply_layout: expanded input_ids + per-item (start, end) offsets + 3. family.positions(len, offsets, geoms) → Rope1D | MRope + 4. Output { input_ids, items: [{feature, aux, hash}], offsets, positions } +``` + +The driver owns these steps and their failure semantics — any `Err` at any +step rejects the request as a 400 (there is no Python fallback path). A +model family fills in only the `family.*` calls, by implementing +`MmFamilyProcessor` (`pipeline.rs`): it describes its data, it never runs +the request. With qwen as the example: + +- **`process_item`** — one decoded image → `ProcessedItem`: + - `feature`: the model's feature tensor. Qwen: `pixel_values`, from + smart_resize → bicubic → normalize → patchify. The item identity is the + driver's hash of the raw encoded source bytes, taken before decode — the + same role as Python's `hash_feature`, but a different algorithm over + different input, so never comparable across paths. + - `aux`: named tensors for the model runner. Qwen: `image_grid_thw`; + other families: `image_sizes`, `tgt_sizes`, ... (Python: + `model_specific_data`). + - `geometry`: whatever this family's `layout`/`positions` need later. + Qwen: the `[t, h, w]` patch grid. +- **`layout`** — how the prompt expands, described as a value. Example: the + prompt `[A, , B]` with one 4-token image becomes + + ``` + [Text(0..1), Media { item: 0, Repeat( × 4) }, Text(2..3)] + ``` + + which the driver expands to `[A, , , , , B]` with + offsets `[(1, 4)]`. Qwen builds this with the `layout_by_placeholder` + helper; families that interleave tile markers or row separators + (internvl/minicpm-style) use `Explicit` id sequences instead. Expansion, + offsets, and position inputs all derive from this one value, so a family + cannot get them out of sync. +- **`positions`** — `Rope1D` (default: the scheduler needs nothing extra) + or `MRope` (qwen's image-only fast path). +- **`capabilities`** — which modalities the family accepts; the server + rejects everything else per family. + +Why not give each family the whole request, like Python's per-family +`process_mm_data_async` override? In the server core, every request must +resolve to exactly one accept/reject with its buffers parked in order — +that invariant only holds structurally if the driver owns the flow. + +Two things stay in Python permanently: HF config parsing (a family is +configured by a spec JSON of already-resolved params, selected via +`registry::pipeline_from_spec`) and the thin drain adapter mapping +feature/aux tensors to model kwargs. The carriers grow by need, not +speculation: `DecodedMedia` gains a variant per modality (video/audio), +`Geometry` per family style (tile sets), `TensorData` per dtype. + +Supported families: `qwen_vl` (Qwen2-VL / 2.5-VL / 3-VL / 3.5; images only). +Adding one = a `MmFamilyProcessor` impl in `src//mod.rs` plus a +`family` arm in `pipeline_from_spec`. + +`common::fetch` matches the Python `get_image_bytes` semantics +(`REQUEST_TIMEOUT` env, `HTTP(S)_PROXY` / `ALL_PROXY` / `NO_PROXY` including +IPv4-CIDR and `host:port` entries) with two deliberate differences: every +source form is capped at 64 MiB — plus 64 items / 256 MiB per request in the +driver — and `file://` URLs actually work (the Python helper passes the +un-stripped URL to `open()`). + ## Python API ```python @@ -28,7 +121,8 @@ from sglang.srt.multimodal._core import common, inkling common.resize_rgb(arr, out_w, out_h) common.scaled_dims(w, h, rescale_frac, rescale_cap) common.image_decode_rgb(bytes) # -> (h, w, ndarray) -common.data_hash(bytes) # -> u64 SHA256 +common.content_hash(bytes) # -> u64 (blake3, truncated) +common.fetch_bytes(source) # -> bytes (data:/base64/file/http) common.base64_decode(str) # -> bytes # Model-specific @@ -44,8 +138,8 @@ inkling.patchify_rgb(arr, patch_size) ```rust use crate::common; +use crate::common::par; use crate::registry::ImageProcessorSpec; -use rayon::prelude::*; pub struct MyModelProcessor; @@ -61,14 +155,14 @@ impl ImageProcessorSpec for MyModelProcessor { rescale_frac: Option, rescale_cap: Option, ) -> Result, u64)>, String> { - common::pool().install(|| { - datas.par_iter().map(|data| { - let hash = common::sha256_u64(data); - let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?; - // Use common::transforms::* or model-specific logic - let patches = my_patchify(&rgb, h, w, patch_size); - Ok((h, w, patches, hash)) - }).collect() + // Always fan out through `par`, never rayon directly: that is what + // keeps the rlib build rayon-free (see Design notes). + par::try_map(datas, |data| { + let hash = common::content_hash_u64(data); + let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?; + // Use common::transforms::* or model-specific logic + let patches = my_patchify(&rgb, h, w, patch_size); + Ok((h, w, patches, hash)) }) } } @@ -93,9 +187,26 @@ impl ImageProcessorSpec for MyModelProcessor { ## Design notes -- Thread pool capped at `min(8, cores)`. Override: `SGL_MM_RS_THREADS`. -- PNG decode is bit-exact vs PIL; JPEG may differ by ±1 LSB. -- Lanczos resize is a bit-exact clone of PIL's fixed-point implementation. +- All fan-out goes through `common::par`, so whether this crate owns threads is + decided by the `parallel` feature alone. With it on: CPU pool capped at + `min(8, cores)` (override `SGL_MM_RS_THREADS`). With it off: no rayon, no + threads, everything inline. Output is bit-identical either way — the fan-outs + are order-preserving maps and writes into disjoint slices, never reductions. + Note that sizing a pool to 1 is *not* the same as off: `install` blocks the + caller and would serialize every concurrent request in the process. +- Media fetch is blocking I/O and deliberately never enters the CPU pool; it + runs inline and sequentially in `driver::process`. Contract: callers on a + fixed worker pool (sglang-server) must resolve I/O-backed string sources — + URLs *and* file paths (a network mount can hang far longer than any HTTP + timeout) — on their own I/O layer and pass bytes, so workers never block on + I/O. `data:`/base64 sources are pure CPU and stay on the worker. +- PNG decode is bit-exact vs PIL; JPEG may differ by ±1 LSB. WebP/GIF/BMP also + decode (GIF: first frame); their parity is not bit-audited. Samples deeper + than 8 bits are rejected rather than rescaled (PIL clips instead). +- Lanczos and Bicubic resize are bit-exact clones of PIL's fixed-point + implementations. +- `common::content_hash_u64` is blake3, *not* Python's SHA-256 + `mm_utils.data_hash`. Hashes are consistent within one path only. ## Build @@ -104,17 +215,23 @@ Automatically built when installing sglang: pip install -e "python" ``` -Or standalone for development: +Or standalone for development (the PyO3 bindings are behind a non-default +feature — see `[features]` in `Cargo.toml` for why): ```bash cd rust/sglang-mm pip install maturin -maturin develop --release +maturin develop --release --features python ``` ## Test ```bash -python bench/generate_golden.py # regenerate fixtures -pytest bench/test_golden.py # regression tests +cd rust/sglang-mm +cargo test --no-default-features # pure-Rust unit tests (CI: pr-test-rust-exts) +python tests/generate_golden.py # regenerate fixtures +pytest tests/test_golden.py # regression tests python bench/bench_parity.py # parity + benchmark ``` + +Scheduler-boundary parity tests against the real HF processors live in +`test/registered/unit/multimodal/rust/`. diff --git a/rust/sglang-mm/bench/bench_parity.py b/rust/sglang-mm/bench/bench_parity.py index 5da717c7d..78fab9371 100644 --- a/rust/sglang-mm/bench/bench_parity.py +++ b/rust/sglang-mm/bench/bench_parity.py @@ -5,7 +5,7 @@ import numpy as np import torch from PIL import Image -import sglang.srt.multimodal._core.inkling +from sglang.srt.multimodal._core import inkling as _rs_inkling from sglang.srt.multimodal.inkling.image_processing import ( IMAGE_MEAN, IMAGE_STD, @@ -30,12 +30,12 @@ def rs_patchify(arr: np.ndarray) -> torch.Tensor: h, w, _ = arr.shape nph = (h + PS - 1) // PS npw = w // PS + 1 - bits = sglang.srt.multimodal._core.inkling.patchify_rgb(arr, PS) + bits = _rs_inkling.patchify_rgb(arr, PS) return torch.from_numpy(bits).view(torch.bfloat16).reshape(nph * npw, PS, PS, 3) def rs_decode_patchify(data: bytes) -> torch.Tensor: - h, w, bits = sglang.srt.multimodal._core.inkling.decode_patchify(data, PS) + h, w, bits = _rs_inkling.decode_patchify(data, PS) nph = (h + PS - 1) // PS npw = w // PS + 1 return torch.from_numpy(bits).view(torch.bfloat16).reshape(nph * npw, PS, PS, 3) @@ -114,7 +114,7 @@ def bench(): rescale_image_max_upscaled_long_edge=None, ) rs_decode_patchify(jpeg) - sglang.srt.multimodal._core.inkling.decode_patchify_batch([jpeg] * 5, PS) + _rs_inkling.decode_patchify_batch([jpeg] * 5, PS) def run(label, fn, iters=n, images_per_call=1): t0, c0 = time.perf_counter(), time.process_time() @@ -137,9 +137,7 @@ def bench(): w_rs, c_rs = run("rust decode_patchify", lambda: rs_decode_patchify(jpeg)) w_rb, c_rb = run( "rust decode_patchify_batch (5 imgs/call)", - lambda: sglang.srt.multimodal._core.inkling.decode_patchify_batch( - [jpeg] * 5, PS - ), + lambda: _rs_inkling.decode_patchify_batch([jpeg] * 5, PS), iters=max(n // 5, 5), images_per_call=5, ) diff --git a/rust/sglang-mm/src/common/fetch.rs b/rust/sglang-mm/src/common/fetch.rs new file mode 100644 index 000000000..b391db2b2 --- /dev/null +++ b/rust/sglang-mm/src/common/fetch.rs @@ -0,0 +1,331 @@ +//! Stage 1 of the server MM pipeline: resolve one media source to raw bytes. +//! +//! Mirrors the Python `get_image_bytes` source handling (and its precedence): +//! raw bytes, `http(s)://` (bounded download, `REQUEST_TIMEOUT` and the proxy +//! env vars like Python), `file://` / absolute path, `data:` URL, else bare +//! base64. + +use std::io::Read; +use std::sync::OnceLock; + +use base64::Engine; + +/// Cap on any single resolved payload — HTTP, file, or base64 — so no source +/// form can exhaust memory (the Python path has no such cap; oversized +/// payloads reject the request here). +pub const MAX_FETCH_BYTES: u64 = 64 << 20; + +/// Resolve one string-typed image source into raw encoded-image bytes. +/// An `Err` rejects the request, matching the Python per-request +/// exception → 400. +pub fn fetch_bytes(src: &str) -> Result, String> { + if src.starts_with("http://") || src.starts_with("https://") { + return http_get(src); + } + if let Some(path) = src.strip_prefix("file://") { + return read_file(path); + } + if src.starts_with('/') { + return read_file(src); + } + if let Some(rest) = src.strip_prefix("data:") { + let encoded = rest + .split_once(',') + .ok_or_else(|| "media fetch: malformed data: URL".to_string())? + .1; + return b64(encoded); + } + // Python treats any other string as bare base64. + b64(src) +} + +/// Bounded read: never trusts metadata, so huge and non-regular files +/// (`/dev/zero`) hit the cap instead of exhausting memory. +fn read_file(path: &str) -> Result, String> { + let file = std::fs::File::open(path).map_err(|e| format!("media fetch: {path}: {e}"))?; + read_capped(file, path) +} + +fn read_capped(reader: impl Read, what: &str) -> Result, String> { + let mut buf = Vec::new(); + reader + .take(MAX_FETCH_BYTES + 1) + .read_to_end(&mut buf) + .map_err(|e| format!("media fetch: read {what}: {e}"))?; + if buf.len() as u64 > MAX_FETCH_BYTES { + return Err(format!( + "media fetch: {what}: exceeds {MAX_FETCH_BYTES} bytes" + )); + } + Ok(buf) +} + +fn b64(encoded: &str) -> Result, String> { + // Slightly laxer than Python's `pybase64.b64decode(validate=True)`: + // surrounding whitespace (e.g. a trailing newline) is trimmed here. + let encoded = encoded.trim(); + // Reject by encoded length before allocating the decode buffer. + if encoded.len() as u64 / 4 * 3 > MAX_FETCH_BYTES { + return Err(format!( + "media fetch: base64 payload exceeds {MAX_FETCH_BYTES} bytes" + )); + } + base64::engine::general_purpose::STANDARD + .decode(encoded.as_bytes()) + .map_err(|e| format!("media fetch: base64 decode: {e}")) +} + +/// Shared pooled agent honoring `HTTP_PROXY`/`HTTPS_PROXY`/`ALL_PROXY`, as the +/// Python `requests` session does. +fn http_agent() -> &'static ureq::Agent { + static AGENT: OnceLock = OnceLock::new(); + AGENT.get_or_init(|| ureq::AgentBuilder::new().try_proxy_from_env(true).build()) +} + +/// Companion agent that ignores the proxy env vars, for hosts matched by +/// `NO_PROXY`. ureq has no `NO_PROXY` support of its own, and silently sending +/// an internal image host through a corporate proxy breaks deployments that +/// work on the Python path, so the match is applied here. +fn direct_agent() -> &'static ureq::Agent { + static AGENT: OnceLock = OnceLock::new(); + AGENT.get_or_init(|| ureq::AgentBuilder::new().build()) +} + +/// The host component of an `http(s)://` URL (lowercased, without userinfo) +/// plus its explicit port, if any — what `NO_PROXY` entries are matched +/// against. +fn host_port_of(url: &str) -> Option<(String, Option)> { + let rest = url + .strip_prefix("http://") + .or_else(|| url.strip_prefix("https://"))?; + let authority = rest.split(['/', '?', '#']).next()?; + let host = authority.rsplit_once('@').map_or(authority, |(_, h)| h); + // Bracketed IPv6 literal, else split off a trailing `:port`. + let (host, port) = match host.strip_prefix('[') { + Some(v6) => match v6.split_once(']') { + Some((h, p)) => (h, p.strip_prefix(':')), + None => (v6, None), + }, + None => match host.split_once(':') { + Some((h, p)) => (h, Some(p)), + None => (host, None), + }, + }; + let port = port.and_then(|p| p.parse().ok()); + (!host.is_empty()).then(|| (host.to_ascii_lowercase(), port)) +} + +fn bypasses_proxy(host: &str, port: Option) -> bool { + ["no_proxy", "NO_PROXY"] + .iter() + .find_map(|key| std::env::var(key).ok()) + .is_some_and(|list| no_proxy_matches(&list, host, port)) +} + +/// `NO_PROXY` semantics as `requests` implements them: comma-separated +/// entries; `*` matches everything; an IPv4 CIDR entry matches an IPv4 host in +/// that network (requests supports IPv4 networks only); a `host:port` entry +/// matches only that explicit port; otherwise an entry matches a host that +/// equals it or is a subdomain of it (leading dots ignored). Kept pure so it +/// is testable without mutating process-global env. +fn no_proxy_matches(no_proxy: &str, host: &str, port: Option) -> bool { + no_proxy.split(',').any(|entry| { + let entry = entry.trim().trim_start_matches('.').to_ascii_lowercase(); + if entry.is_empty() { + return false; + } + if entry == "*" { + return true; + } + if let Some((net, bits)) = parse_ipv4_cidr(&entry) { + return host + .parse::() + .is_ok_and(|ip| in_ipv4_network(ip, net, bits)); + } + let (entry_host, entry_port) = match entry.rsplit_once(':') { + Some((h, p)) if p.bytes().all(|b| b.is_ascii_digit()) => (h, p.parse::().ok()), + _ => (entry.as_str(), None), + }; + if entry_port.is_some() && entry_port != port { + return false; + } + host == entry_host || host.ends_with(&format!(".{entry_host}")) + }) +} + +fn parse_ipv4_cidr(entry: &str) -> Option<(std::net::Ipv4Addr, u32)> { + let (net, bits) = entry.split_once('/')?; + Some((net.parse().ok()?, bits.parse().ok().filter(|b| *b <= 32)?)) +} + +fn in_ipv4_network(ip: std::net::Ipv4Addr, net: std::net::Ipv4Addr, bits: u32) -> bool { + let mask = u32::MAX.checked_shl(32 - bits).unwrap_or(0); + u32::from(ip) & mask == u32::from(net) & mask +} + +fn http_get(url: &str) -> Result, String> { + // Python: `int(os.getenv("REQUEST_TIMEOUT", "3"))` seconds per image GET. + let timeout = std::env::var("REQUEST_TIMEOUT") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(3); + let agent = match host_port_of(url) { + Some((host, port)) if bypasses_proxy(&host, port) => direct_agent(), + _ => http_agent(), + }; + let resp = agent + .get(url) + .timeout(std::time::Duration::from_secs(timeout)) + .call() + .map_err(|e| format!("media fetch: GET {url}: {e}"))?; + read_capped(resp.into_reader(), url) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn data_url_and_bare_base64_decode() { + let b64 = base64::engine::general_purpose::STANDARD.encode(b"hello"); + assert_eq!( + fetch_bytes(&format!("data:image/png;base64,{b64}")).unwrap(), + b"hello" + ); + assert_eq!(fetch_bytes(&b64).unwrap(), b"hello"); + } + + #[test] + fn bad_base64_fails() { + assert!(fetch_bytes("!!not-base64!!").is_err()); + } + + #[test] + fn missing_file_fails() { + assert!(fetch_bytes("file:///definitely/not/here.jpg").is_err()); + assert!(fetch_bytes("/definitely/not/here.jpg").is_err()); + } + + /// A non-regular file must hit the byte cap, not exhaust memory. + #[test] + fn unbounded_file_capped() { + let err = fetch_bytes("/dev/zero").err().unwrap(); + assert!(err.contains("exceeds"), "{err}"); + } + + /// Oversized base64 is rejected from its encoded length, before decoding. + #[test] + fn oversized_base64_rejected() { + let encoded = "A".repeat((MAX_FETCH_BYTES / 3 * 4 + 8) as usize); + let err = fetch_bytes(&encoded).err().unwrap(); + assert!(err.contains("exceeds"), "{err}"); + } + + #[test] + fn host_parsing_strips_userinfo_and_path() { + assert_eq!( + host_port_of("http://Example.COM/a/b.png").unwrap(), + ("example.com".into(), None) + ); + assert_eq!( + host_port_of("https://u:p@images.internal:8443/x").unwrap(), + ("images.internal".into(), Some(8443)) + ); + assert_eq!( + host_port_of("http://[::1]:8080/x.png").unwrap(), + ("::1".into(), Some(8080)) + ); + assert_eq!( + host_port_of("http://host?q=1").unwrap(), + ("host".into(), None) + ); + assert!(host_port_of("data:image/png;base64,AAA").is_none()); + } + + /// `NO_PROXY` must bypass the proxy for exact hosts and subdomains but not + /// for lookalike suffixes — sending an internal host to a corporate proxy + /// is a silent failure that works fine on the Python path. + #[test] + fn no_proxy_matches_host_and_subdomains_only() { + let list = " .internal ,localhost"; + assert!(no_proxy_matches(list, "images.internal", None)); + assert!(no_proxy_matches(list, "internal", None)); + assert!(no_proxy_matches(list, "localhost", None)); + assert!(!no_proxy_matches(list, "notinternal", None)); + assert!(!no_proxy_matches(list, "example.com", None)); + assert!(no_proxy_matches("*", "anything.example.com", None)); + // An empty or all-empty list must not bypass everything. + assert!(!no_proxy_matches("", "example.com", None)); + assert!(!no_proxy_matches(" , ", "example.com", None)); + } + + /// IPv4 CIDR entries match IP-literal hosts, as `requests` does. + #[test] + fn no_proxy_matches_ipv4_cidr() { + assert!(no_proxy_matches("10.0.0.0/8", "10.1.2.3", None)); + assert!(!no_proxy_matches("10.0.0.0/8", "11.1.2.3", None)); + assert!(no_proxy_matches("192.168.1.0/24", "192.168.1.77", Some(80))); + assert!(!no_proxy_matches("192.168.1.0/24", "192.168.2.1", None)); + assert!(no_proxy_matches("0.0.0.0/0", "8.8.8.8", None)); + // CIDR entries never match hostnames. + assert!(!no_proxy_matches("10.0.0.0/8", "example.com", None)); + } + + /// `host:port` entries match only that explicit port. + #[test] + fn no_proxy_matches_host_with_port() { + assert!(no_proxy_matches("internal:8443", "internal", Some(8443))); + assert!(!no_proxy_matches("internal:8443", "internal", Some(80))); + assert!(!no_proxy_matches("internal:8443", "internal", None)); + assert!(no_proxy_matches( + "internal:8443", + "img.internal", + Some(8443) + )); + // A port-free entry matches any port. + assert!(no_proxy_matches("internal", "internal", Some(8443))); + } + + /// End-to-end HTTP download against a local one-shot server, and the + /// capped rejection of an oversized response. + #[test] + fn http_download_and_cap() { + let serve = |body: Vec, content_length: u64| { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let handle = std::thread::spawn(move || { + use std::io::{BufRead, Write}; + let (stream, _) = listener.accept().unwrap(); + let mut reader = std::io::BufReader::new(stream); + let mut line = String::new(); + while reader.read_line(&mut line).unwrap() > 2 { + line.clear(); // headers until the blank line + } + let mut stream = reader.into_inner(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {content_length}\r\n\r\n" + ) + .unwrap(); + stream.write_all(&body).unwrap(); + }); + (addr, handle) + }; + + let (addr, handle) = serve(b"tiny image".to_vec(), 10); + assert_eq!( + fetch_bytes(&format!("http://{addr}/img.png")).unwrap(), + b"tiny image" + ); + handle.join().unwrap(); + + // A response over the cap is rejected without buffering it all. + let over = MAX_FETCH_BYTES + 2; + let (addr, handle) = serve(vec![0u8; over as usize], over); + let err = fetch_bytes(&format!("http://{addr}/big.png")) + .err() + .unwrap(); + assert!(err.contains("exceeds"), "{err}"); + drop(handle); // server thread may die on the closed socket; don't join + } +} diff --git a/rust/sglang-mm/src/common/mod.rs b/rust/sglang-mm/src/common/mod.rs index 792971985..21593c61c 100644 --- a/rust/sglang-mm/src/common/mod.rs +++ b/rust/sglang-mm/src/common/mod.rs @@ -1,12 +1,19 @@ +pub mod fetch; +pub mod par; pub mod resize; +pub mod token_layout; pub mod transforms; +#[cfg(feature = "parallel")] use std::sync::OnceLock; -use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods}; -use pyo3::exceptions::PyValueError; -use pyo3::prelude::*; - +/// CPU pool: decode, resize, patchify, hash. Sized to cores (capped at 8) +/// because the work is compute-bound — never run blocking I/O on it, or one +/// request's remote fetches stall every other request's preprocessing. +/// +/// Only exists under the `parallel` feature; reach it through [`par`], never +/// directly, so the rayon-less build stays compiling. +#[cfg(feature = "parallel")] pub fn pool() -> &'static rayon::ThreadPool { static POOL: OnceLock = OnceLock::new(); POOL.get_or_init(|| { @@ -23,13 +30,27 @@ pub fn pool() -> &'static rayon::ThreadPool { }) } -pub fn sha256_u64(data: &[u8]) -> u64 { +/// Content hash for cache/dedup identity: blake3 truncated to its first 8 +/// bytes, big-endian. Deliberately *not* Python's `mm_utils.data_hash` (which +/// is SHA-256 truncated the same way) — hashes are consistent within a path, +/// never comparable across the Rust and Python paths. +pub fn content_hash_u64(data: &[u8]) -> u64 { let digest = blake3::hash(data); u64::from_be_bytes(digest.as_bytes()[..8].try_into().unwrap()) } pub fn decode_rgb(data: &[u8]) -> Result<(Vec, usize, usize), String> { + use image::ColorType; + let img = image::load_from_memory(data).map_err(|e| format!("image decode: {e}"))?; + // >8-bit samples: PIL clips to 255 where `to_rgb8` would rescale. Refuse + // rather than silently diverge from the Python (PIL) pipeline. + if !matches!( + img.color(), + ColorType::L8 | ColorType::La8 | ColorType::Rgb8 | ColorType::Rgba8 + ) { + return Err(format!("image decode: unsupported color {:?}", img.color())); + } let rgb = img.to_rgb8(); let (w, h) = rgb.dimensions(); Ok((rgb.into_raw(), h as usize, w as usize)) @@ -48,87 +69,144 @@ pub fn decode_rescale( Ok((resize::resize_lanczos_rgb(&rgb, h, w, th, tw), th, tw)) } -// --- Python-exposed functions --- +// --- Python bindings (feature-gated: absent from the pure-Rust rlib) --- -#[pyfunction] -pub fn resize_rgb<'py>( - py: Python<'py>, - arr: PyReadonlyArray3<'py, u8>, - out_w: usize, - out_h: usize, -) -> PyResult>> { - if out_w == 0 || out_h == 0 { - return Err(PyValueError::new_err("output size must be positive")); +#[cfg(feature = "python")] +mod python { + use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods}; + use pyo3::exceptions::PyValueError; + use pyo3::prelude::*; + use pyo3::types::PyBytes; + + use super::{decode_rgb, resize}; + + #[pyfunction] + pub fn resize_rgb<'py>( + py: Python<'py>, + arr: PyReadonlyArray3<'py, u8>, + out_w: usize, + out_h: usize, + ) -> PyResult>> { + if out_w == 0 || out_h == 0 { + return Err(PyValueError::new_err("output size must be positive")); + } + let shape = arr.shape(); + let (h, w, c) = (shape[0], shape[1], shape[2]); + if c != 3 { + return Err(PyValueError::new_err(format!( + "expected HWC RGB array with 3 channels, got {c}" + ))); + } + let data = arr + .as_slice() + .map_err(|_| PyValueError::new_err("array must be C-contiguous"))? + .to_vec(); + let out = py.detach(move || resize::resize_lanczos_rgb(&data, h, w, out_h, out_w)); + Ok(out.into_pyarray(py)) } - let shape = arr.shape(); - let (h, w, c) = (shape[0], shape[1], shape[2]); - if c != 3 { - return Err(PyValueError::new_err(format!( - "expected HWC RGB array with 3 channels, got {c}" - ))); + + #[pyfunction] + #[pyo3(signature = (w, h, rescale_frac=None, rescale_cap=None))] + pub fn scaled_dims( + w: usize, + h: usize, + rescale_frac: Option, + rescale_cap: Option, + ) -> (usize, usize) { + resize::scaled_dims(w, h, rescale_frac, rescale_cap) + } + + #[pyfunction] + pub fn image_decode_rgb<'py>( + py: Python<'py>, + data: Vec, + ) -> PyResult<(usize, usize, Bound<'py, PyArray1>)> { + let (rgb, h, w) = py + .detach(move || decode_rgb(&data)) + .map_err(PyValueError::new_err)?; + Ok((h, w, rgb.into_pyarray(py))) + } + + /// Named `content_hash`, not `data_hash`, so it is not mistaken for + /// `sglang.srt.managers.mm_utils.data_hash` (SHA-256); this is blake3. + #[pyfunction] + pub fn content_hash(py: Python<'_>, data: Vec) -> u64 { + py.detach(move || super::content_hash_u64(&data)) + } + + #[pyfunction] + pub fn fetch_bytes<'py>(py: Python<'py>, source: String) -> PyResult> { + let data = py + .detach(move || super::fetch::fetch_bytes(&source)) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + Ok(PyBytes::new(py, &data)) + } + + #[pyfunction] + pub fn base64_decode<'py>( + py: Python<'py>, + encoded: &str, + ) -> PyResult> { + use base64::Engine; + let decoded = py + .detach(|| { + base64::engine::general_purpose::STANDARD + .decode(encoded) + .map_err(|e| format!("base64 decode error: {e}")) + }) + .map_err(PyValueError::new_err)?; + Ok(pyo3::types::PyBytes::new(py, &decoded)) + } + + pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> { + let m = PyModule::new(parent.py(), "common")?; + m.add_function(wrap_pyfunction!(resize_rgb, &m)?)?; + m.add_function(wrap_pyfunction!(scaled_dims, &m)?)?; + m.add_function(wrap_pyfunction!(image_decode_rgb, &m)?)?; + m.add_function(wrap_pyfunction!(content_hash, &m)?)?; + m.add_function(wrap_pyfunction!(fetch_bytes, &m)?)?; + m.add_function(wrap_pyfunction!(base64_decode, &m)?)?; + parent.add_submodule(&m)?; + Ok(()) } - let data = arr - .as_slice() - .map_err(|_| PyValueError::new_err("array must be C-contiguous"))? - .to_vec(); - let out = - py.detach(move || pool().install(|| resize::resize_lanczos_rgb(&data, h, w, out_h, out_w))); - Ok(out.into_pyarray(py)) } -#[pyfunction] -#[pyo3(signature = (w, h, rescale_frac=None, rescale_cap=None))] -pub fn scaled_dims( - w: usize, - h: usize, - rescale_frac: Option, - rescale_cap: Option, -) -> (usize, usize) { - resize::scaled_dims(w, h, rescale_frac, rescale_cap) -} +#[cfg(feature = "python")] +pub use python::register; -#[pyfunction] -pub fn image_decode_rgb<'py>( - py: Python<'py>, - data: Vec, -) -> PyResult<(usize, usize, Bound<'py, PyArray1>)> { - let (rgb, h, w) = py - .detach(move || decode_rgb(&data)) - .map_err(PyValueError::new_err)?; - Ok((h, w, rgb.into_pyarray(py))) -} +#[cfg(test)] +mod tests { + use super::decode_rgb; + use image::ImageFormat; -#[pyfunction] -pub fn data_hash(py: Python<'_>, data: Vec) -> u64 { - py.detach(move || { - let digest = blake3::hash(&data); - u64::from_be_bytes(digest.as_bytes()[..8].try_into().unwrap()) - }) -} + fn encode(img: &image::DynamicImage, fmt: ImageFormat) -> Vec { + let mut buf = std::io::Cursor::new(Vec::new()); + img.write_to(&mut buf, fmt).unwrap(); + buf.into_inner() + } -#[pyfunction] -pub fn base64_decode<'py>( - py: Python<'py>, - encoded: &str, -) -> PyResult> { - use base64::Engine; - let decoded = py - .detach(|| { - base64::engine::general_purpose::STANDARD - .decode(encoded) - .map_err(|e| format!("base64 decode error: {e}")) - }) - .map_err(PyValueError::new_err)?; - Ok(pyo3::types::PyBytes::new(py, &decoded)) -} + /// Formats the Python (PIL) path accepts must decode, not reject. + #[test] + fn decodes_webp_gif_bmp() { + let img = image::DynamicImage::ImageRgb8(image::RgbImage::from_fn(6, 4, |x, y| { + image::Rgb([x as u8 * 40, y as u8 * 60, 7]) + })); + for fmt in [ImageFormat::WebP, ImageFormat::Gif, ImageFormat::Bmp] { + let (rgb, h, w) = decode_rgb(&encode(&img, fmt)).unwrap(); + assert_eq!((h, w), (4, 6), "{fmt:?}"); + assert_eq!(rgb.len(), 4 * 6 * 3, "{fmt:?}"); + } + } -pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> { - let m = PyModule::new(parent.py(), "common")?; - m.add_function(wrap_pyfunction!(resize_rgb, &m)?)?; - m.add_function(wrap_pyfunction!(scaled_dims, &m)?)?; - m.add_function(wrap_pyfunction!(image_decode_rgb, &m)?)?; - m.add_function(wrap_pyfunction!(data_hash, &m)?)?; - m.add_function(wrap_pyfunction!(base64_decode, &m)?)?; - parent.add_submodule(&m)?; - Ok(()) + /// Samples deeper than 8 bits stay rejected (PIL clips; we refuse). + #[test] + fn deep_png_rejected() { + let img = image::DynamicImage::ImageRgb16(image::ImageBuffer::from_pixel( + 2, + 2, + image::Rgb([65535u16, 0, 0]), + )); + let err = decode_rgb(&encode(&img, ImageFormat::Png)).err().unwrap(); + assert!(err.contains("unsupported color"), "{err}"); + } } diff --git a/rust/sglang-mm/src/common/par.rs b/rust/sglang-mm/src/common/par.rs new file mode 100644 index 000000000..dc226b53e --- /dev/null +++ b/rust/sglang-mm/src/common/par.rs @@ -0,0 +1,93 @@ +//! The crate's only parallelism seam. +//! +//! Every fan-out in the crate goes through the functions below, so whether +//! this crate owns worker threads at all is decided in exactly one place: the +//! `parallel` cargo feature. +//! +//! * **feature on** (the PyO3 extension): work is fanned out on the crate's +//! rayon pool. A Python processor calls in from one or two worker threads +//! with the GIL released, so intra-call parallelism is the whole point. +//! * **feature off** (the pure-Rust `rlib` that `sglang-server` links): rayon +//! is not even a dependency, and everything runs inline on the calling +//! thread. A server supplies concurrency across requests and owns its own +//! core budget (it pins threads via `core_affinity`), so a library that +//! silently spawns its own pools would fight it. +//! +//! Note that "sequential" here means *inline on the caller*, not a one-thread +//! pool: `ThreadPool::install` injects work into the pool and blocks the +//! caller, so sizing a pool to 1 would serialize every concurrent request in +//! the process instead of just declining to fan out. +//! +//! Results are identical either way — the fan-outs are order-preserving maps +//! and writes into disjoint slices, never reductions. + +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +/// Map `items`, short-circuiting on the first error. Output order matches input +/// order. CPU-bound work: decode, resize, patchify, hash. +#[cfg(feature = "parallel")] +pub fn try_map<'a, T, R, E>( + items: &'a [T], + f: impl Fn(&'a T) -> Result + Send + Sync, +) -> Result, E> +where + T: Send + Sync, + R: Send, + E: Send, +{ + super::pool().install(|| items.par_iter().map(f).collect()) +} + +#[cfg(not(feature = "parallel"))] +pub fn try_map<'a, T, R, E>( + items: &'a [T], + f: impl Fn(&'a T) -> Result + Send + Sync, +) -> Result, E> +where + T: Send + Sync, + R: Send, + E: Send, +{ + items.iter().map(f).collect() +} + +/// Apply `f(chunk_index, chunk)` over disjoint `chunk_size`-element windows of +/// `buf`. The final chunk is short when `chunk_size` does not divide the length. +#[cfg(feature = "parallel")] +pub fn for_chunks_mut( + buf: &mut [T], + chunk_size: usize, + f: impl Fn(usize, &mut [T]) + Send + Sync, +) { + super::pool().install(|| { + buf.par_chunks_mut(chunk_size) + .enumerate() + .for_each(|(index, chunk)| f(index, chunk)); + }); +} + +#[cfg(not(feature = "parallel"))] +pub fn for_chunks_mut( + buf: &mut [T], + chunk_size: usize, + f: impl Fn(usize, &mut [T]) + Send + Sync, +) { + for (index, chunk) in buf.chunks_mut(chunk_size).enumerate() { + f(index, chunk); + } +} + +/// Run `f` with the CPU pool already entered, so nested [`for_chunks_mut`] +/// calls inside it reuse this entry instead of injecting a job each. Use it to +/// wrap a multi-stage leaf (e.g. the two passes of a separable resize) that +/// would otherwise pay per-stage pool entry. +#[cfg(feature = "parallel")] +pub fn in_pool(f: impl FnOnce() -> R + Send) -> R { + super::pool().install(f) +} + +#[cfg(not(feature = "parallel"))] +pub fn in_pool(f: impl FnOnce() -> R + Send) -> R { + f() +} diff --git a/rust/sglang-mm/src/common/resize.rs b/rust/sglang-mm/src/common/resize.rs index 5a9ebb88f..cb8277933 100644 --- a/rust/sglang-mm/src/common/resize.rs +++ b/rust/sglang-mm/src/common/resize.rs @@ -1,7 +1,33 @@ -use rayon::prelude::*; +use super::par; const PRECISION_BITS: i32 = 32 - 8 - 2; +/// Resampling filters, bit-exact clones of PIL's kernels. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Filter { + /// support 3.0 — PIL `LANCZOS`. + Lanczos, + /// support 2.0, a = -0.5 — PIL `BICUBIC` (≈ torchvision antialiased + /// bicubic, which the HF "fast" image processors use). + Bicubic, +} + +impl Filter { + fn support(self) -> f64 { + match self { + Filter::Lanczos => 3.0, + Filter::Bicubic => 2.0, + } + } + + fn eval(self, x: f64) -> f64 { + match self { + Filter::Lanczos => lanczos(x), + Filter::Bicubic => bicubic(x), + } + } +} + fn sinc(x: f64) -> f64 { if x == 0.0 { return 1.0; @@ -18,16 +44,28 @@ fn lanczos(x: f64) -> f64 { } } +fn bicubic(x: f64) -> f64 { + const A: f64 = -0.5; + let x = x.abs(); + if x < 1.0 { + ((A + 2.0) * x - (A + 3.0)) * x * x + 1.0 + } else if x < 2.0 { + (((x - 5.0) * x + 8.0) * x - 4.0) * A + } else { + 0.0 + } +} + struct Coeffs { bounds: Vec<(usize, usize)>, kk: Vec, ksize: usize, } -fn precompute_coeffs(in_size: usize, out_size: usize) -> Coeffs { +fn precompute_coeffs(in_size: usize, out_size: usize, filter: Filter) -> Coeffs { let scale = in_size as f64 / out_size as f64; let filterscale = if scale < 1.0 { 1.0 } else { scale }; - let support = 3.0 * filterscale; + let support = filter.support() * filterscale; let ksize = support.ceil() as usize * 2 + 1; let ss = 1.0 / filterscale; @@ -47,7 +85,7 @@ fn precompute_coeffs(in_size: usize, out_size: usize) -> Coeffs { let k = &mut kkf[xx * ksize..(xx + 1) * ksize]; let mut ww = 0.0f64; for (x, kv) in k[..count].iter_mut().enumerate() { - let w = lanczos((x as f64 + xmin as f64 - center + 0.5) * ss); + let w = filter.eval((x as f64 + xmin as f64 - center + 0.5) * ss); *kv = w; ww += w; } @@ -86,32 +124,30 @@ fn clip8(v: i32) -> u8 { fn resample_horizontal(src: &[u8], h: usize, w: usize, out_w: usize, c: &Coeffs) -> Vec { let mut out = vec![0u8; h * out_w * 3]; - out.par_chunks_mut(out_w * 3) - .enumerate() - .for_each(|(y, row)| { - let src_row = &src[y * w * 3..(y + 1) * w * 3]; - for xx in 0..out_w { - let (xmin, count) = c.bounds[xx]; - let k = &c.kk[xx * c.ksize..xx * c.ksize + count]; - let mut s = [1i32 << (PRECISION_BITS - 1); 3]; - for (x, &coef) in k.iter().enumerate() { - let p = (xmin + x) * 3; - s[0] += src_row[p] as i32 * coef; - s[1] += src_row[p + 1] as i32 * coef; - s[2] += src_row[p + 2] as i32 * coef; - } - let o = xx * 3; - row[o] = clip8(s[0]); - row[o + 1] = clip8(s[1]); - row[o + 2] = clip8(s[2]); + par::for_chunks_mut(&mut out, out_w * 3, |y, row| { + let src_row = &src[y * w * 3..(y + 1) * w * 3]; + for xx in 0..out_w { + let (xmin, count) = c.bounds[xx]; + let k = &c.kk[xx * c.ksize..xx * c.ksize + count]; + let mut s = [1i32 << (PRECISION_BITS - 1); 3]; + for (x, &coef) in k.iter().enumerate() { + let p = (xmin + x) * 3; + s[0] += src_row[p] as i32 * coef; + s[1] += src_row[p + 1] as i32 * coef; + s[2] += src_row[p + 2] as i32 * coef; } - }); + let o = xx * 3; + row[o] = clip8(s[0]); + row[o + 1] = clip8(s[1]); + row[o + 2] = clip8(s[2]); + } + }); out } fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec { let mut out = vec![0u8; out_h * w * 3]; - out.par_chunks_mut(w * 3).enumerate().for_each(|(yy, row)| { + par::for_chunks_mut(&mut out, w * 3, |yy, row| { let (ymin, count) = c.bounds[yy]; let k = &c.kk[yy * c.ksize..yy * c.ksize + count]; for x in 0..w { @@ -131,25 +167,51 @@ fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec out } -pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec { +/// PIL-exact separable resize of a flat HWC RGB buffer with the given filter. +/// +/// Enters the fan-out pool once for both passes; the per-row `for_chunks_mut` +/// calls inside then reuse that entry rather than injecting a job per pass. +pub fn resize_rgb_filter( + src: &[u8], + h: usize, + w: usize, + out_h: usize, + out_w: usize, + filter: Filter, +) -> Vec { + par::in_pool(move || resize_passes(src, h, w, out_h, out_w, filter)) +} + +fn resize_passes( + src: &[u8], + h: usize, + w: usize, + out_h: usize, + out_w: usize, + filter: Filter, +) -> Vec { let need_h = out_w != w; let need_v = out_h != h; if need_h && need_v { - let ch = precompute_coeffs(w, out_w); + let ch = precompute_coeffs(w, out_w, filter); let tmp = resample_horizontal(src, h, w, out_w, &ch); - let cv = precompute_coeffs(h, out_h); + let cv = precompute_coeffs(h, out_h, filter); resample_vertical(&tmp, out_w, out_h, &cv) } else if need_h { - let ch = precompute_coeffs(w, out_w); + let ch = precompute_coeffs(w, out_w, filter); resample_horizontal(src, h, w, out_w, &ch) } else if need_v { - let cv = precompute_coeffs(h, out_h); + let cv = precompute_coeffs(h, out_h, filter); resample_vertical(src, w, out_h, &cv) } else { src.to_vec() } } +pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec { + resize_rgb_filter(src, h, w, out_h, out_w, Filter::Lanczos) +} + pub fn scaled_dims(w: usize, h: usize, frac: Option, cap: Option) -> (usize, usize) { let Some(frac) = frac else { return (w, h); diff --git a/rust/sglang-mm/src/common/token_layout.rs b/rust/sglang-mm/src/common/token_layout.rs new file mode 100644 index 000000000..0c824a61a --- /dev/null +++ b/rust/sglang-mm/src/common/token_layout.rs @@ -0,0 +1,239 @@ +//! Token-layout mechanics for the server MM pipeline. +//! +//! Families describe their prompt geometry as a [`TokenLayout`] value +//! (`pipeline.rs`); [`apply_layout`] applies it mechanically. Expanding the +//! already-tokenized prompt means non-media tokens can never drift from a +//! retokenize (the `SGLANG_MM_AVOID_RETOKENIZE` idea, unconditional here). + +use crate::pipeline::{Segment, TokenLayout, TokenPattern}; + +/// The expanded prompt plus, per media item (indexed as in the layout), the +/// inclusive `(start, end)` token range it occupies — the Python +/// `get_mm_items_offset` convention. +pub struct ExpandedPrompt { + pub input_ids: Vec, + pub offsets: Vec<(u32, u32)>, +} + +/// Apply a family's [`TokenLayout`] to the original prompt. +/// +/// The point of the layout being data is that a family cannot get expansion, +/// offsets, and positions out of sync, so this validates the whole contract +/// rather than just indexing safely: +/// * text ranges are in bounds, ascending, and non-overlapping; +/// * together with the media placeholders they cover every source token +/// exactly once — a family that forgets a tail segment must not silently +/// truncate the prompt; +/// * every one of the `n_items` media items is placed exactly once; +/// * no item expands to zero tokens (which would have no representable offset). +pub fn apply_layout( + src: &[i32], + layout: &TokenLayout, + n_items: usize, +) -> Result { + let mut out = Vec::new(); + let mut offsets: Vec> = vec![None; n_items]; + // Source tokens consumed so far: `Text` copies them, a `Media` segment + // replaces exactly the one placeholder token sitting at this position. + let mut consumed = 0usize; + for segment in &layout.segments { + match segment { + Segment::Text(range) => { + let text = src + .get(range.clone()) + .ok_or_else(|| format!("layout: text range {range:?} out of bounds"))?; + if range.start != consumed { + return Err(format!( + "layout: text range {range:?} does not resume at source index {consumed}" + )); + } + consumed = range.end; + out.extend_from_slice(text); + } + Segment::Media { item, pattern } => { + if consumed >= src.len() { + return Err(format!( + "layout: media item {item} has no source placeholder at index {consumed}" + )); + } + consumed += 1; + let start = out.len() as u32; + let n = match pattern { + TokenPattern::Repeat { id, n } => { + out.resize(out.len() + n, *id); + *n + } + TokenPattern::Explicit(ids) => { + out.extend_from_slice(ids); + ids.len() + } + }; + if n == 0 { + return Err(format!("layout: media item {item} expands to zero tokens")); + } + let slot = offsets + .get_mut(*item) + .ok_or_else(|| format!("layout: media item {item} out of range"))?; + if slot.replace((start, start + n as u32 - 1)).is_some() { + return Err(format!("layout: media item {item} placed twice")); + } + } + } + } + if consumed != src.len() { + return Err(format!( + "layout: covers {consumed} of {} source token(s)", + src.len() + )); + } + let offsets = offsets + .into_iter() + .enumerate() + .map(|(i, slot)| slot.ok_or_else(|| format!("layout: media item {i} not placed"))) + .collect::, _>>()?; + Ok(ExpandedPrompt { + input_ids: out, + offsets, + }) +} + +/// Build the simplest layout: each occurrence of `placeholder_id` in `ids` +/// becomes `counts[i]` copies (i-th occurrence ↔ i-th media item). Errs when +/// the occurrence count and `counts` disagree. +pub fn layout_by_placeholder( + ids: &[i32], + placeholder_id: i32, + counts: &[usize], +) -> Result { + let found = ids.iter().filter(|&&id| id == placeholder_id).count(); + if found != counts.len() { + return Err(format!( + "prompt has {found} media placeholder(s) but {} media item(s)", + counts.len() + )); + } + let mut segments = Vec::new(); + let mut text_start = 0; + let mut item = 0; + for (pos, &id) in ids.iter().enumerate() { + if id == placeholder_id { + if text_start < pos { + segments.push(Segment::Text(text_start..pos)); + } + segments.push(Segment::Media { + item, + pattern: TokenPattern::Repeat { + id: placeholder_id, + n: counts[item], + }, + }); + item += 1; + text_start = pos + 1; + } + } + if text_start < ids.len() { + segments.push(Segment::Text(text_start..ids.len())); + } + Ok(TokenLayout { segments }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn expand(ids: &[i32], placeholder: i32, counts: &[usize]) -> Result { + apply_layout( + ids, + &layout_by_placeholder(ids, placeholder, counts)?, + counts.len(), + ) + } + + #[test] + fn expands_in_order_with_inclusive_offsets() { + // [7, PAD, 8, PAD, 9] with counts [2, 3] + let e = expand(&[7, 1, 8, 1, 9], 1, &[2, 3]).unwrap(); + assert_eq!(e.input_ids, vec![7, 1, 1, 8, 1, 1, 1, 9]); + assert_eq!(e.offsets, vec![(1, 2), (4, 6)]); + } + + #[test] + fn count_mismatch_errs() { + assert!(expand(&[7, 1, 9], 1, &[2, 3]).is_err()); + assert!(expand(&[7, 1, 1, 9], 1, &[2]).is_err()); + } + + #[test] + fn zero_count_errs() { + assert!(expand(&[7, 1, 9], 1, &[0]).is_err()); + } + + #[test] + fn no_placeholders_no_items_ok() { + let e = expand(&[7, 8], 1, &[]).unwrap(); + assert_eq!(e.input_ids, vec![7, 8]); + assert!(e.offsets.is_empty()); + } + + #[test] + fn explicit_patterns_and_placement_validation() { + // Structured expansion: marker tokens around the item span. + let layout = TokenLayout { + segments: vec![ + Segment::Text(0..1), + Segment::Media { + item: 0, + pattern: TokenPattern::Explicit(vec![90, 5, 5, 91]), + }, + Segment::Text(2..3), + ], + }; + let e = apply_layout(&[7, 1, 9], &layout, 1).unwrap(); + assert_eq!(e.input_ids, vec![7, 90, 5, 5, 91, 9]); + assert_eq!(e.offsets, vec![(1, 4)]); + + // Every item must be placed exactly once; ranges must be in bounds. + let missing = TokenLayout { + segments: vec![Segment::Text(0..3)], + }; + assert!(apply_layout(&[7, 1, 9], &missing, 1).is_err()); + let out_of_bounds = TokenLayout { + segments: vec![Segment::Text(0..4)], + }; + assert!(apply_layout(&[7, 1, 9], &out_of_bounds, 0).is_err()); + } + + /// A family that skips, repeats, or reorders source tokens would silently + /// serve a truncated or scrambled prompt; the layout must reject it instead. + #[test] + fn incomplete_or_disordered_coverage_errs() { + let media = || Segment::Media { + item: 0, + pattern: TokenPattern::Repeat { id: 5, n: 2 }, + }; + let cases = [ + // Dropped tail: [7, PAD, 9] expanded without the trailing 9. + vec![Segment::Text(0..1), media()], + // Dropped head. + vec![media(), Segment::Text(2..3)], + // Gap in the middle (source index 1 never consumed). + vec![Segment::Text(0..1), Segment::Text(2..3), media()], + // Duplicated source span. + vec![ + Segment::Text(0..1), + Segment::Text(0..1), + media(), + Segment::Text(2..3), + ], + // Out of order. + vec![Segment::Text(2..3), media(), Segment::Text(0..1)], + ]; + for (i, segments) in cases.into_iter().enumerate() { + let layout = TokenLayout { segments }; + assert!( + apply_layout(&[7, 1, 9], &layout, 1).is_err(), + "case {i} should be rejected" + ); + } + } +} diff --git a/rust/sglang-mm/src/driver.rs b/rust/sglang-mm/src/driver.rs new file mode 100644 index 000000000..9b70a52ea --- /dev/null +++ b/rust/sglang-mm/src/driver.rs @@ -0,0 +1,280 @@ +//! Shared multimodal request driver for the server (pure-Rust) pipeline. +//! +//! Owns the request control flow — parallel fan-out, layout application, +//! failure semantics — while every model decision lives behind +//! [`MmFamilyProcessor`] (see `pipeline.rs`). Families produce data; they +//! cannot alter orchestration. + +use crate::common::{self, fetch, par, token_layout}; +use crate::pipeline::{DecodedMedia, MmFamilyProcessor, PositionOutput, ProcessedItem}; + +/// Per-request bounds: together with [`fetch::MAX_FETCH_BYTES`] they cap what +/// one request can make the pipeline buffer. +pub const MAX_ITEMS_PER_REQUEST: usize = 64; +pub const MAX_REQUEST_BYTES: u64 = 256 << 20; + +/// One raw image source from the request. +#[derive(Debug)] +pub enum ImageSource { + /// `data:`/base64/file/http — resolved by [`fetch::fetch_bytes`]. + String(String), + /// Already-raw encoded image bytes. + Bytes(Vec), +} + +/// Typed multimodal request input. The server's message layer owns the wire +/// format and parses its payload into this before calling [`process`]. +pub struct MmInput { + pub text: Option, + pub input_ids: Option>, + pub images: Vec, +} + +/// One processed media item at the request boundary. +pub struct OutputItem { + pub feature: crate::pipeline::Tensor, + pub aux: crate::pipeline::NamedTensors, + /// [`common::content_hash_u64`] of the raw encoded source bytes — the same + /// identity role as the Python path's `hash_feature`, but a different + /// algorithm, so hashes are consistent within the server pipeline and + /// never comparable across the two paths. + pub hash: u64, +} + +/// The per-request result parked for the scheduler drain. +pub struct Output { + pub input_ids: Vec, + /// In prompt order; `offsets[i]` is `items[i]`'s inclusive token range. + pub items: Vec, + pub offsets: Vec<(u32, u32)>, + pub positions: PositionOutput, +} + +/// Resolve one image source to raw encoded bytes, borrowing when the request +/// already carries them. +fn resolve(source: &ImageSource) -> Result, String> { + match source { + ImageSource::String(source) => Ok(fetch::fetch_bytes(source)?.into()), + ImageSource::Bytes(bytes) => Ok(bytes.as_slice().into()), + } +} + +/// Run one request through the pipeline. Any `Err` rejects the request back +/// to the client — including inputs merely outside the pipeline's scope +/// (video/audio, precomputed features, undecodable images), since there is +/// no Python fallback path. +pub fn process( + family: &dyn MmFamilyProcessor, + input: MmInput, + tokenize: impl FnOnce(&str) -> Result, String>, +) -> Result { + if input.images.is_empty() { + return Err("multimodal request without image sources".into()); + } + if input.images.len() > MAX_ITEMS_PER_REQUEST { + return Err(format!( + "multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items" + )); + } + // Stage 1 (fetch) is blocking I/O and runs inline, sequentially — never on + // the CPU pool, where a slow URL or file read would starve decode/resize + // for other requests. Contract: callers on a fixed worker pool (the + // server) must resolve I/O-backed string sources — URLs and file paths — + // on their own I/O layer and pass `Bytes`. + let mut fetched: Vec> = Vec::with_capacity(input.images.len()); + let mut total: u64 = 0; + for source in &input.images { + let bytes = resolve(source)?; + total += bytes.len() as u64; + if total > MAX_REQUEST_BYTES { + return Err(format!( + "multimodal request exceeds {MAX_REQUEST_BYTES} total media bytes" + )); + } + fetched.push(bytes); + } + let processed: Vec<(ProcessedItem, u64)> = + par::try_map(&fetched, |bytes| -> Result<(ProcessedItem, u64), String> { + let hash = common::content_hash_u64(bytes); + // Inputs PIL accepts but decode_rgb refuses (e.g. 16-bit PNG) + // error here and reject the request. + let (rgb, height, width) = common::decode_rgb(bytes)?; + let item = family.process_item(&DecodedMedia::Image { rgb, height, width })?; + Ok((item, hash)) + })?; + + let input_ids = match input.input_ids { + Some(input_ids) if !input_ids.is_empty() => input_ids, + _ => { + let text = input + .text + .as_deref() + .ok_or("multimodal request without text or input_ids")?; + tokenize(text)? + } + }; + let geometries = processed + .iter() + .map(|(item, _)| item.geometry.clone()) + .collect::>(); + let layout = family.layout(&input_ids, &geometries)?; + let expanded = token_layout::apply_layout(&input_ids, &layout, processed.len())?; + let positions = family.positions(expanded.input_ids.len(), &expanded.offsets, &geometries)?; + + Ok(Output { + input_ids: expanded.input_ids, + items: processed + .into_iter() + .map(|(item, hash)| OutputItem { + feature: item.feature, + aux: item.aux, + hash, + }) + .collect(), + offsets: expanded.offsets, + positions, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::registry::pipeline_from_spec; + + const SPEC: &str = r#"{"family":"qwen_vl","image_token_id":1,"patch_size":2, + "merge_size":2,"temporal_patch_size":2,"min_pixels":4, + "max_pixels":1073741824,"image_mean":[0.0,0.0,0.0],"image_std":[1.0,1.0,1.0]}"#; + + fn png(w: u32, h: u32) -> Vec { + let img = image::RgbImage::from_fn(w, h, |x, y| image::Rgb([x as u8, y as u8, 7])); + let mut buf = std::io::Cursor::new(Vec::new()); + img.write_to(&mut buf, image::ImageFormat::Png).unwrap(); + buf.into_inner() + } + + #[test] + fn processes_typed_image_request() { + let family = pipeline_from_spec(SPEC).unwrap(); + let input = MmInput { + text: None, + input_ids: Some(vec![7, 1, 8]), + images: vec![ImageSource::Bytes(png(8, 8))], + }; + let out = process(family.as_ref(), input, |_| Err("no tokenizer".into())).unwrap(); + // 8x8, factor 4 → grid [1, 4, 4] → 16 patches / merge² = 4 tokens. + assert_eq!(out.input_ids, vec![7, 1, 1, 1, 1, 8]); + assert_eq!(out.offsets, vec![(1, 4)]); + let item = &out.items[0]; + assert_eq!(item.feature.shape, [16, 3 * 2 * 2 * 2]); + assert_eq!(item.aux[0].0, "image_grid_thw"); + let crate::pipeline::PositionOutput::MRope { positions, .. } = &out.positions else { + panic!("qwen emits mrope") + }; + assert_eq!(positions.len(), 3 * out.input_ids.len()); + } + + /// String sources — HTTP URL, `file://`, and bare path — all resolve to + /// the same bytes and flow through the full pipeline. + #[test] + fn fetches_url_and_file_sources() { + let png = png(8, 8); + let dir = std::env::temp_dir().join(format!("sglang-mm-driver-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("img.png"); + std::fs::write(&path, &png).unwrap(); + + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let body = png.clone(); + let server = std::thread::spawn(move || { + use std::io::{BufRead, Write}; + let (stream, _) = listener.accept().unwrap(); + let mut reader = std::io::BufReader::new(stream); + let mut line = String::new(); + while reader.read_line(&mut line).unwrap() > 2 { + line.clear(); // headers until the blank line + } + let mut stream = reader.into_inner(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", + body.len() + ) + .unwrap(); + stream.write_all(&body).unwrap(); + }); + + let family = pipeline_from_spec(SPEC).unwrap(); + let input = MmInput { + text: None, + input_ids: Some(vec![7, 1, 1, 1, 8]), + images: vec![ + ImageSource::String(format!("http://{addr}/img.png")), + ImageSource::String(format!("file://{}", path.display())), + ImageSource::String(path.display().to_string()), + ], + }; + let out = process(family.as_ref(), input, |_| unreachable!()).unwrap(); + server.join().unwrap(); + std::fs::remove_dir_all(&dir).ok(); + + assert_eq!(out.items.len(), 3); + // Identical source bytes → identical content hashes. + assert_eq!(out.items[0].hash, out.items[1].hash); + assert_eq!(out.items[1].hash, out.items[2].hash); + assert_eq!(out.input_ids.len(), 5 + 3 * 3); // each placeholder → 4 tokens + } + + #[test] + fn per_request_caps_enforced() { + let family = pipeline_from_spec(SPEC).unwrap(); + let too_many = MmInput { + text: None, + input_ids: Some(vec![1]), + images: (0..=MAX_ITEMS_PER_REQUEST) + .map(|_| ImageSource::Bytes(vec![])) + .collect(), + }; + let err = process(family.as_ref(), too_many, |_| unreachable!()) + .err() + .unwrap(); + assert!(err.contains("media items"), "{err}"); + + let chunk = (MAX_REQUEST_BYTES / 2 + 1) as usize; + let too_big = MmInput { + text: None, + input_ids: Some(vec![1, 1]), + images: vec![ + ImageSource::Bytes(vec![0; chunk]), + ImageSource::Bytes(vec![0; chunk]), + ], + }; + let err = process(family.as_ref(), too_big, |_| unreachable!()) + .err() + .unwrap(); + assert!(err.contains("total media bytes"), "{err}"); + } + + #[test] + fn image_free_and_mismatched_requests_rejected() { + let family = pipeline_from_spec(SPEC).unwrap(); + let no_images = MmInput { + text: None, + input_ids: Some(vec![7, 1]), + images: vec![], + }; + let err = process(family.as_ref(), no_images, |_| unreachable!()) + .err() + .unwrap(); + assert!(err.contains("image sources")); + let no_placeholder = MmInput { + text: None, + input_ids: Some(vec![7, 8]), + images: vec![ImageSource::Bytes(png(8, 8))], + }; + let err = process(family.as_ref(), no_placeholder, |_| unreachable!()) + .err() + .unwrap(); + assert!(err.contains("placeholder")); + } +} diff --git a/rust/sglang-mm/src/inkling/mod.rs b/rust/sglang-mm/src/inkling/mod.rs index d953d0e8b..2003c6374 100644 --- a/rust/sglang-mm/src/inkling/mod.rs +++ b/rust/sglang-mm/src/inkling/mod.rs @@ -1,21 +1,8 @@ use std::sync::OnceLock; -use half::bf16; -use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods}; -use pyo3::exceptions::PyValueError; -use pyo3::prelude::*; -use rayon::prelude::*; - use crate::common; - -/// `(height, width, patches_as_u16_bits)` for one decoded image. -type Patches = (usize, usize, Vec); -/// [`Patches`] plus the image content hash. -type HashedPatches = (usize, usize, Vec, u64); -/// [`Patches`] with the patch data as a numpy array bound to `'py`. -type PyPatches<'py> = (usize, usize, Bound<'py, PyArray1>); -/// [`PyPatches`] plus the image content hash. -type PyHashedPatches<'py> = (usize, usize, Bound<'py, PyArray1>, u64); +use crate::common::par; +use half::bf16; const MEAN: [f32; 3] = [ 0.48145466f64 as f32, @@ -90,9 +77,7 @@ fn patchify_into(arr: &[u8], h: usize, w: usize, ps: usize, out: &mut [u16]) { } }; - common::pool().install(|| { - out.par_chunks_mut(row_elems).enumerate().for_each(body); - }); + par::for_chunks_mut(out, row_elems, |index, chunk| body((index, chunk))); } fn patchify_alloc(arr: &[u8], h: usize, w: usize, ps: usize) -> Vec { @@ -102,119 +87,6 @@ fn patchify_alloc(arr: &[u8], h: usize, w: usize, ps: usize) -> Vec { out } -fn check_ps(ps: usize) -> PyResult<()> { - if ps == 0 { - return Err(PyValueError::new_err( - "patch_size must be greater than zero", - )); - } - Ok(()) -} - -#[pyfunction] -fn patchify_rgb<'py>( - py: Python<'py>, - arr: PyReadonlyArray3<'py, u8>, - patch_size: usize, -) -> PyResult>> { - check_ps(patch_size)?; - let shape = arr.shape(); - let (h, w, c) = (shape[0], shape[1], shape[2]); - if c != 3 { - return Err(PyValueError::new_err(format!( - "expected HWC RGB array with 3 channels, got {c}" - ))); - } - let data = arr - .as_slice() - .map_err(|_| PyValueError::new_err("array must be C-contiguous"))? - .to_vec(); - let out = py.detach(move || patchify_alloc(&data, h, w, patch_size)); - Ok(out.into_pyarray(py)) -} - -#[pyfunction] -#[pyo3(signature = (data, patch_size, rescale_frac=None, rescale_cap=None))] -fn decode_patchify<'py>( - py: Python<'py>, - data: Vec, - patch_size: usize, - rescale_frac: Option, - rescale_cap: Option, -) -> PyResult<(usize, usize, Bound<'py, PyArray1>)> { - check_ps(patch_size)?; - let (h, w, out) = py - .detach(move || { - common::pool().install(|| { - let (rgb, h, w) = common::decode_rescale(&data, rescale_frac, rescale_cap)?; - Ok::<_, String>((h, w, patchify_alloc(&rgb, h, w, patch_size))) - }) - }) - .map_err(PyValueError::new_err)?; - Ok((h, w, out.into_pyarray(py))) -} - -#[pyfunction] -#[pyo3(signature = (datas, patch_size, rescale_frac=None, rescale_cap=None))] -fn decode_patchify_batch<'py>( - py: Python<'py>, - datas: Vec>, - patch_size: usize, - rescale_frac: Option, - rescale_cap: Option, -) -> PyResult>> { - check_ps(patch_size)?; - let results: Vec> = py.detach(move || { - common::pool().install(|| { - datas - .par_iter() - .map(|data| { - let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?; - Ok((h, w, patchify_alloc(&rgb, h, w, patch_size))) - }) - .collect() - }) - }); - results - .into_iter() - .map(|r| { - let (h, w, v) = r.map_err(PyValueError::new_err)?; - Ok((h, w, v.into_pyarray(py))) - }) - .collect() -} - -#[pyfunction] -#[pyo3(signature = (datas, patch_size, rescale_frac=None, rescale_cap=None))] -fn preprocess_images<'py>( - py: Python<'py>, - datas: Vec>, - patch_size: usize, - rescale_frac: Option, - rescale_cap: Option, -) -> PyResult>> { - check_ps(patch_size)?; - let results: Vec> = py.detach(move || { - common::pool().install(|| { - datas - .par_iter() - .map(|data| { - let hash = common::sha256_u64(data); - let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?; - Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash)) - }) - .collect() - }) - }); - results - .into_iter() - .map(|r| { - let (h, w, v, hash) = r.map_err(PyValueError::new_err)?; - Ok((h, w, v.into_pyarray(py), hash)) - }) - .collect() -} - /// Struct implementing ImageProcessorSpec for Inkling. pub struct InklingProcessor; @@ -233,44 +105,158 @@ impl crate::registry::ImageProcessorSpec for InklingProcessor { if patch_size == 0 { return Err("patch_size must be greater than zero".into()); } - common::pool().install(|| { - datas - .par_iter() - .map(|data| { - let hash = common::sha256_u64(data); - let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?; - Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash)) - }) - .collect() + par::try_map(datas, |data| { + let hash = common::content_hash_u64(data); + let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?; + Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash)) }) } } -#[pyfunction] -#[pyo3(signature = (arr, raw_bytes, patch_size, rescale_frac=None, rescale_cap=None))] -fn rescale_patchify_hash<'py>( - py: Python<'py>, - arr: PyReadonlyArray3<'py, u8>, - raw_bytes: &[u8], - patch_size: usize, - rescale_frac: Option, - rescale_cap: Option, -) -> PyResult<(usize, usize, Bound<'py, PyArray1>, u64)> { - check_ps(patch_size)?; - let shape = arr.shape(); - let (h, w, c) = (shape[0], shape[1], shape[2]); - if c != 3 { - return Err(PyValueError::new_err(format!( - "expected HWC RGB array with 3 channels, got {c}" - ))); +// --- Python bindings (feature-gated: absent from the pure-Rust rlib) --- + +#[cfg(feature = "python")] +mod python { + use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods}; + use pyo3::exceptions::PyValueError; + use pyo3::prelude::*; + + use super::patchify_alloc; + use crate::common; + use crate::common::par; + + /// `(height, width, patches_as_u16_bits)` for one decoded image. + type Patches = (usize, usize, Vec); + /// [`Patches`] plus the image content hash. + type HashedPatches = (usize, usize, Vec, u64); + /// [`Patches`] with the patch data as a numpy array bound to `'py`. + type PyPatches<'py> = (usize, usize, Bound<'py, PyArray1>); + /// [`PyPatches`] plus the image content hash. + type PyHashedPatches<'py> = (usize, usize, Bound<'py, PyArray1>, u64); + + fn check_ps(ps: usize) -> PyResult<()> { + if ps == 0 { + return Err(PyValueError::new_err( + "patch_size must be greater than zero", + )); + } + Ok(()) } - let hash = common::sha256_u64(raw_bytes); - let rgb = arr - .as_slice() - .map_err(|_| PyValueError::new_err("array must be C-contiguous"))? - .to_vec(); - let (oh, ow, out) = py.detach(move || { - common::pool().install(|| { + + #[pyfunction] + fn patchify_rgb<'py>( + py: Python<'py>, + arr: PyReadonlyArray3<'py, u8>, + patch_size: usize, + ) -> PyResult>> { + check_ps(patch_size)?; + let shape = arr.shape(); + let (h, w, c) = (shape[0], shape[1], shape[2]); + if c != 3 { + return Err(PyValueError::new_err(format!( + "expected HWC RGB array with 3 channels, got {c}" + ))); + } + let data = arr + .as_slice() + .map_err(|_| PyValueError::new_err("array must be C-contiguous"))? + .to_vec(); + let out = py.detach(move || patchify_alloc(&data, h, w, patch_size)); + Ok(out.into_pyarray(py)) + } + + #[pyfunction] + #[pyo3(signature = (data, patch_size, rescale_frac=None, rescale_cap=None))] + fn decode_patchify<'py>( + py: Python<'py>, + data: Vec, + patch_size: usize, + rescale_frac: Option, + rescale_cap: Option, + ) -> PyResult<(usize, usize, Bound<'py, PyArray1>)> { + check_ps(patch_size)?; + let (h, w, out) = py + .detach(move || { + let (rgb, h, w) = common::decode_rescale(&data, rescale_frac, rescale_cap)?; + Ok::<_, String>((h, w, patchify_alloc(&rgb, h, w, patch_size))) + }) + .map_err(PyValueError::new_err)?; + Ok((h, w, out.into_pyarray(py))) + } + + #[pyfunction] + #[pyo3(signature = (datas, patch_size, rescale_frac=None, rescale_cap=None))] + fn decode_patchify_batch<'py>( + py: Python<'py>, + datas: Vec>, + patch_size: usize, + rescale_frac: Option, + rescale_cap: Option, + ) -> PyResult>> { + check_ps(patch_size)?; + let results: Vec = py + .detach(move || { + par::try_map(&datas, |data| -> Result { + let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?; + Ok((h, w, patchify_alloc(&rgb, h, w, patch_size))) + }) + }) + .map_err(PyValueError::new_err)?; + results + .into_iter() + .map(|(h, w, v)| Ok((h, w, v.into_pyarray(py)))) + .collect() + } + + #[pyfunction] + #[pyo3(signature = (datas, patch_size, rescale_frac=None, rescale_cap=None))] + fn preprocess_images<'py>( + py: Python<'py>, + datas: Vec>, + patch_size: usize, + rescale_frac: Option, + rescale_cap: Option, + ) -> PyResult>> { + check_ps(patch_size)?; + let results: Vec = py + .detach(move || { + par::try_map(&datas, |data| -> Result { + let hash = common::content_hash_u64(data); + let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?; + Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash)) + }) + }) + .map_err(PyValueError::new_err)?; + results + .into_iter() + .map(|(h, w, v, hash)| Ok((h, w, v.into_pyarray(py), hash))) + .collect() + } + + #[pyfunction] + #[pyo3(signature = (arr, raw_bytes, patch_size, rescale_frac=None, rescale_cap=None))] + fn rescale_patchify_hash<'py>( + py: Python<'py>, + arr: PyReadonlyArray3<'py, u8>, + raw_bytes: &[u8], + patch_size: usize, + rescale_frac: Option, + rescale_cap: Option, + ) -> PyResult<(usize, usize, Bound<'py, PyArray1>, u64)> { + check_ps(patch_size)?; + let shape = arr.shape(); + let (h, w, c) = (shape[0], shape[1], shape[2]); + if c != 3 { + return Err(PyValueError::new_err(format!( + "expected HWC RGB array with 3 channels, got {c}" + ))); + } + let hash = common::content_hash_u64(raw_bytes); + let rgb = arr + .as_slice() + .map_err(|_| PyValueError::new_err("array must be C-contiguous"))? + .to_vec(); + let (oh, ow, out) = py.detach(move || { let (tw, th) = common::resize::scaled_dims(w, h, rescale_frac, rescale_cap); let (rgb, h, w) = if (tw, th) != (w, h) { ( @@ -282,18 +268,21 @@ fn rescale_patchify_hash<'py>( (rgb, h, w) }; (h, w, patchify_alloc(&rgb, h, w, patch_size)) - }) - }); - Ok((oh, ow, out.into_pyarray(py), hash)) + }); + Ok((oh, ow, out.into_pyarray(py), hash)) + } + + pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> { + let m = PyModule::new(parent.py(), "inkling")?; + m.add_function(wrap_pyfunction!(patchify_rgb, &m)?)?; + m.add_function(wrap_pyfunction!(decode_patchify, &m)?)?; + m.add_function(wrap_pyfunction!(decode_patchify_batch, &m)?)?; + m.add_function(wrap_pyfunction!(preprocess_images, &m)?)?; + m.add_function(wrap_pyfunction!(rescale_patchify_hash, &m)?)?; + parent.add_submodule(&m)?; + Ok(()) + } } -pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> { - let m = PyModule::new(parent.py(), "inkling")?; - m.add_function(wrap_pyfunction!(patchify_rgb, &m)?)?; - m.add_function(wrap_pyfunction!(decode_patchify, &m)?)?; - m.add_function(wrap_pyfunction!(decode_patchify_batch, &m)?)?; - m.add_function(wrap_pyfunction!(preprocess_images, &m)?)?; - m.add_function(wrap_pyfunction!(rescale_patchify_hash, &m)?)?; - parent.add_submodule(&m)?; - Ok(()) -} +#[cfg(feature = "python")] +pub use python::register; diff --git a/rust/sglang-mm/src/lib.rs b/rust/sglang-mm/src/lib.rs index 92d9dd04a..1865d879e 100644 --- a/rust/sglang-mm/src/lib.rs +++ b/rust/sglang-mm/src/lib.rs @@ -1,12 +1,26 @@ -mod common; -mod inkling; +//! sglang-mm: Rust-accelerated multimodal preprocessing for SGLang. +//! +//! Built two ways: +//! * PyO3 extension `sglang.srt.multimodal._core` (feature `python`, default), +//! used by Python processors (e.g. Inkling) and by parity tests. +//! * Pure-Rust `rlib` (`default-features = false`), linked by `sglang-server`'s +//! MM worker path — no pyo3 in that dependency graph. + +pub mod common; +pub mod driver; +pub mod inkling; +pub mod pipeline; +pub mod qwen_vl; pub mod registry; +#[cfg(feature = "python")] use pyo3::prelude::*; +#[cfg(feature = "python")] #[pymodule] fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { common::register(m)?; inkling::register(m)?; + qwen_vl::register(m)?; Ok(()) } diff --git a/rust/sglang-mm/src/pipeline.rs b/rust/sglang-mm/src/pipeline.rs new file mode 100644 index 000000000..a6afe04b1 --- /dev/null +++ b/rust/sglang-mm/src/pipeline.rs @@ -0,0 +1,132 @@ +//! The model-family seam of the server MM pipeline. +//! +//! Design rule: **families produce data, the driver owns control flow.** A +//! family never sees the request loop, the thread pool, or the failure +//! protocol — it implements [`MmFamilyProcessor`], turning decoded media into +//! named tensors and describing its prompt geometry as a [`TokenLayout`] +//! value. `driver::process` applies the layout mechanically, so expansion, +//! per-item offsets, and position inputs all derive from one declarative +//! structure and every family gets identical failure semantics for free. + +/// Typed tensor payload. Grows a variant per dtype actually produced by a +/// family — not speculatively. +pub enum TensorData { + F32(Vec), + I64(Vec), +} + +pub struct Tensor { + pub shape: Vec, + pub data: TensorData, +} + +/// Named auxiliary tensors that reach the model runner as kwargs — the Rust +/// analogue of Python's `MultimodalDataItem.model_specific_data`, e.g. qwen's +/// `image_grid_thw`. +pub type NamedTensors = Vec<(String, Tensor)>; + +/// One decoded media item handed to [`MmFamilyProcessor::process_item`]. +/// Grows a variant per modality as families that need it are ported. +pub enum DecodedMedia { + /// HWC u8 RGB. + Image { + rgb: Vec, + height: usize, + width: usize, + }, +} + +/// Family-internal geometry of one processed item, consumed by +/// [`MmFamilyProcessor::layout`] / [`MmFamilyProcessor::positions`]. Grows a +/// variant per family style; the driver never interprets it. +#[derive(Clone, Debug)] +pub enum Geometry { + /// `[t, h, w]` patch grid (`t` = 1 for still images). + Grid([u32; 3]), +} + +/// One processed media item, mirroring Python's `MultimodalDataItem`: the +/// primary feature tensor, named auxiliary tensors, and the geometry the +/// family's own `layout`/`positions` hooks need. +pub struct ProcessedItem { + /// The model's feature tensor for this item (qwen: `pixel_values`). + pub feature: Tensor, + pub aux: NamedTensors, + pub geometry: Geometry, +} + +/// The tokens one media item occupies in the expanded prompt. +pub enum TokenPattern { + /// N copies of one placeholder id (qwen-style). + Repeat { id: i32, n: usize }, + /// An explicit id sequence — tile markers, row separators, wrapper + /// tokens (minicpm/internvl-style structured expansions). + Explicit(Vec), +} + +/// One span of the expanded prompt. +pub enum Segment { + /// Copy `src` (a range into the original ids) verbatim. + Text(std::ops::Range), + /// Media item `item`'s token span. + Media { item: usize, pattern: TokenPattern }, +} + +/// Prompt geometry as data: the family *describes* the expansion, the driver +/// *applies* it (`common::token_layout::apply_layout`) — deriving final input ids +/// and per-item offsets, and validating that every item is placed exactly +/// once. +pub struct TokenLayout { + pub segments: Vec, +} + +/// Modalities a family accepts; the server's message layer rejects anything +/// a family does not declare. +#[derive(Clone, Copy, Debug, Default)] +pub struct Capabilities { + pub video: bool, + pub audio: bool, +} + +/// Position scheme of the expanded prompt. +pub enum PositionOutput { + /// Plain sequential positions — the scheduler needs nothing extra. + Rope1D, + /// M-RoPE: flattened row-major `[3, input_len]` positions + the position + /// delta (`max + 1 - input_len`). + MRope { positions: Vec, delta: i64 }, +} + +/// The per-model-family hooks of the server pipeline. Adding a family = +/// implementing this in `src//mod.rs` and adding its `family` arm to +/// [`crate::registry::pipeline_from_spec`]. All parameters come from the +/// runtime spec JSON (resolved from the HF config on the Python side); +/// nothing is hardcoded per model. +pub trait MmFamilyProcessor: Send + Sync { + /// Modalities beyond images this family accepts. Default: images only. + fn capabilities(&self) -> Capabilities { + Capabilities::default() + } + + /// Preprocess one decoded media item: the model's HF processor + /// equivalent (resize/tile/normalize/patchify → named tensors) plus the + /// geometry `layout`/`positions` will need. + fn process_item(&self, media: &DecodedMedia) -> Result; + + /// Describe how the prompt expands around the processed items (in + /// prompt order). Sees the full original prompt and all items, so + /// structured schemes (tile markers, separators) are expressible. + fn layout(&self, input_ids: &[i32], items: &[Geometry]) -> Result; + + /// Positions for the expanded prompt. Families without a custom scheme + /// keep the default. + fn positions( + &self, + input_len: usize, + offsets: &[(u32, u32)], + items: &[Geometry], + ) -> Result { + let _ = (input_len, offsets, items); + Ok(PositionOutput::Rope1D) + } +} diff --git a/rust/sglang-mm/src/qwen_vl/mod.rs b/rust/sglang-mm/src/qwen_vl/mod.rs new file mode 100644 index 000000000..d32ac4756 --- /dev/null +++ b/rust/sglang-mm/src/qwen_vl/mod.rs @@ -0,0 +1,645 @@ +//! Qwen VL family (Qwen2-VL / 2.5-VL / 3-VL / 3.5) server-pipeline image processor. +//! +//! Pure-Rust equivalent of the HF `Qwen2VLImageProcessor` pipeline the Python +//! `QwenVLImageProcessor` drives: `smart_resize` → bicubic resize → rescale + +//! normalize → patchify into `[grid_h*grid_w, C*tps*ps*ps]` (HF flatten order: +//! patches by `(gh/m, gw/m, m, m)`, features by `(C, tps, ps, ps)`, temporal +//! copies duplicated for stills) — plus the image-only M-RoPE fast path. +//! All parameters come from the runtime spec; nothing is hardcoded per model. + +use crate::common::{par, resize, token_layout}; +use crate::pipeline::{ + DecodedMedia, Geometry, MmFamilyProcessor, PositionOutput, ProcessedItem, Tensor, TensorData, + TokenLayout, +}; + +const MAX_RATIO: f64 = 200.0; + +/// One media item's placement for M-RoPE: inclusive token range + patch grid. +pub struct MropeItem { + pub start: u32, + pub end: u32, + pub grid: [u32; 3], +} + +/// Resolved processor params, deserialized from the Python-side spec JSON +/// (unknown fields like `family` are ignored here). +#[derive(Clone, Debug, serde::Deserialize)] +pub struct QwenVlSpec { + pub image_token_id: i32, + pub patch_size: usize, + pub merge_size: usize, + pub temporal_patch_size: usize, + pub min_pixels: usize, + pub max_pixels: usize, + pub image_mean: [f32; 3], + pub image_std: [f32; 3], +} + +pub struct QwenVlProcessor { + spec: QwenVlSpec, + /// Per-channel u8 → normalized-f32 lookup: `(v/255 - mean) / std`. + lut: [[f32; 256]; 3], +} + +impl QwenVlProcessor { + pub fn new(spec: QwenVlSpec) -> Result { + if spec.patch_size == 0 || spec.merge_size == 0 || spec.temporal_patch_size == 0 { + return Err("qwen_vl spec: sizes must be positive".into()); + } + let lut = core::array::from_fn(|c| { + core::array::from_fn(|v| (v as f32 / 255.0 - spec.image_mean[c]) / spec.image_std[c]) + }); + Ok(Self { spec, lut }) + } + + pub fn from_spec_json(json: &str) -> Result { + let spec: QwenVlSpec = + serde_json::from_str(json).map_err(|e| format!("qwen_vl spec: {e}"))?; + Self::new(spec) + } + + fn factor(&self) -> usize { + self.spec.patch_size * self.spec.merge_size + } + + /// HF flatten: patches ordered `(gh/m, gw/m, m, m)`, features `(C, tps, + /// ps, ps)`; parallel over merged-block rows. + fn patchify(&self, rgb: &[u8], h: usize, w: usize) -> Vec { + let (ps, m, tps) = ( + self.spec.patch_size, + self.spec.merge_size, + self.spec.temporal_patch_size, + ); + let (gh, gw) = (h / ps, w / ps); + let dim = 3 * tps * ps * ps; + let block_row = gw * m * dim; // one merged-block row of patches + let mut out = vec![0.0f32; gh * gw * dim]; + + par::for_chunks_mut(&mut out, block_row, |i, chunk| { + let mut p = 0; + for j in 0..gw / m { + for mh in 0..m { + for mw in 0..m { + let y0 = (i * m + mh) * ps; + let x0 = (j * m + mw) * ps; + let patch = &mut chunk[p * dim..(p + 1) * dim]; + for c in 0..3 { + let ch = &mut patch[c * tps * ps * ps..]; + for py in 0..ps { + let src = ((y0 + py) * w + x0) * 3 + c; + for px in 0..ps { + ch[py * ps + px] = self.lut[c][rgb[src + px * 3] as usize]; + } + } + // Temporal copies of a still are duplicates. + let (t0, rest) = ch.split_at_mut(ps * ps); + for t in 0..tps - 1 { + rest[t * ps * ps..(t + 1) * ps * ps].copy_from_slice(t0); + } + } + p += 1; + } + } + } + }); + out + } +} + +impl QwenVlProcessor { + fn tokens_per_image(&self, grid: &[u32; 3]) -> usize { + (grid[0] as usize * grid[1] as usize * grid[2] as usize) + / (self.spec.merge_size * self.spec.merge_size) + } +} + +impl MmFamilyProcessor for QwenVlProcessor { + fn process_item(&self, media: &DecodedMedia) -> Result { + let DecodedMedia::Image { rgb, height, width } = media; + let (h, w) = (*height, *width); + let (th, tw) = smart_resize( + h, + w, + self.factor(), + self.spec.min_pixels, + self.spec.max_pixels, + )?; + let resized; + let data = if (th, tw) != (h, w) { + resized = resize::resize_rgb_filter(rgb, h, w, th, tw, resize::Filter::Bicubic); + &resized + } else { + rgb.as_slice() + }; + let (gh, gw) = (th / self.spec.patch_size, tw / self.spec.patch_size); + // `smart_resize` guarantees both: dims are positive and divisible by + // `patch_size * merge_size`. `patchify` indexes on that (and the `dim` + // division below needs a non-empty grid), so fail loudly rather than + // panic if a future spec change breaks the guarantee. + if gh == 0 || gw == 0 || gh % self.spec.merge_size != 0 || gw % self.spec.merge_size != 0 { + return Err(format!( + "qwen_vl: patch grid {gh}x{gw} is empty or not a multiple of \ + merge_size {}", + self.spec.merge_size + )); + } + let pixel_values = self.patchify(data, th, tw); + let dim = pixel_values.len() / (gh * gw); + Ok(ProcessedItem { + feature: Tensor { + shape: vec![gh * gw, dim], + data: TensorData::F32(pixel_values), + }, + aux: vec![( + "image_grid_thw".to_string(), + Tensor { + shape: vec![3], + data: TensorData::I64(vec![1, gh as i64, gw as i64]), + }, + )], + geometry: Geometry::Grid([1, gh as u32, gw as u32]), + }) + } + + fn layout(&self, input_ids: &[i32], items: &[Geometry]) -> Result { + let counts = items + .iter() + .map(|Geometry::Grid(grid)| self.tokens_per_image(grid)) + .collect::>(); + token_layout::layout_by_placeholder(input_ids, self.spec.image_token_id, &counts) + } + + fn positions( + &self, + input_len: usize, + offsets: &[(u32, u32)], + items: &[Geometry], + ) -> Result { + let mrope_items = offsets + .iter() + .zip(items) + .map(|(&(start, end), Geometry::Grid(grid))| MropeItem { + start, + end, + grid: *grid, + }) + .collect::>(); + let (positions, delta) = mrope_image_only(input_len, &mrope_items, self.spec.merge_size)?; + Ok(PositionOutput::MRope { positions, delta }) + } +} + +/// Python-`round()` (round-half-to-even), which `round_by_factor` relies on. +fn round_half_even(x: f64) -> f64 { + if (x - x.trunc()).abs() == 0.5 { + (x / 2.0).round() * 2.0 + } else { + x.round() + } +} + +/// The Qwen `smart_resize`: dims divisible by `factor`, total pixels within +/// `[min_pixels, max_pixels]`, aspect ratio preserved as closely as possible. +pub fn smart_resize( + height: usize, + width: usize, + factor: usize, + min_pixels: usize, + max_pixels: usize, +) -> Result<(usize, usize), String> { + let (h, w) = (height as f64, width as f64); + if height == 0 || width == 0 { + return Err("empty image".into()); + } + let ratio = h.max(w) / h.min(w); + if ratio > MAX_RATIO { + return Err(format!( + "absolute aspect ratio must be smaller than {MAX_RATIO}, got {ratio}" + )); + } + let f = factor as f64; + let mut h_bar = ((round_half_even(h / f) * f) as usize).max(factor); + let mut w_bar = ((round_half_even(w / f) * f) as usize).max(factor); + if h_bar * w_bar > max_pixels { + let beta = (h * w / max_pixels as f64).sqrt(); + h_bar = ((h / beta / f).floor() * f) as usize; + w_bar = ((w / beta / f).floor() * f) as usize; + } else if h_bar * w_bar < min_pixels { + let beta = (min_pixels as f64 / (h * w)).sqrt(); + h_bar = ((h * beta / f).ceil() * f) as usize; + w_bar = ((w * beta / f).ceil() * f) as usize; + } + // The downscale branch floors without a lower clamp (as Python does), so a + // very thin image against a small `max_pixels` can floor a side to 0. + // Python then fails inside PIL's resize; here it would reach the resize + // coefficient math (overflow panic in debug, garbage in release) and the + // `dim = len / (gh * gw)` division, so reject it as a request error. + if h_bar == 0 || w_bar == 0 { + return Err(format!( + "smart_resize: {height}x{width} degenerates to {h_bar}x{w_bar} at \ + max_pixels={max_pixels}; image is too thin for this pixel budget" + )); + } + Ok((h_bar, w_bar)) +} + +/// Image-only M-RoPE fast path (the image branch of +/// `MRotaryEmbedding.get_rope_index`, identical across Qwen generations): +/// text runs sequentially on all three rows; each image spans `(t, h/m, w/m)` +/// index grids; positions advance by `max(t, h/m, w/m)` past an image. +/// Returns flattened row-major `[3, input_len]` positions and the delta +/// (`max + 1 - input_len`). `items` must be in prompt order. +pub fn mrope_image_only( + input_len: usize, + items: &[MropeItem], + merge_size: usize, +) -> Result<(Vec, i64), String> { + let len = input_len; + let mut pos = vec![0i64; 3 * len]; + let fill_text = |st: usize, n: usize, base: i64, pos: &mut [i64]| { + for k in 0..n { + let v = base + k as i64; + pos[st + k] = v; + pos[len + st + k] = v; + pos[2 * len + st + k] = v; + } + }; + let mut st = 0usize; + let mut next_pos = 0i64; + for item in items { + let (start, end) = (item.start as usize, item.end as usize); + if start < st || end >= len { + return Err(format!( + "mrope: item range ({start},{end}) out of order/bounds" + )); + } + fill_text(st, start - st, next_pos, &mut pos); + next_pos += (start - st) as i64; + + let t = item.grid[0] as usize; + let gh = item.grid[1] as usize / merge_size; + let gw = item.grid[2] as usize / merge_size; + if t * gh * gw != end - start + 1 { + return Err("mrope: token span does not match grid".into()); + } + for ti in 0..t { + for hi in 0..gh { + for wi in 0..gw { + let idx = start + (ti * gh + hi) * gw + wi; + pos[idx] = next_pos + ti as i64; + pos[len + idx] = next_pos + hi as i64; + pos[2 * len + idx] = next_pos + wi as i64; + } + } + } + next_pos += (t.max(gh).max(gw)) as i64; + st = end + 1; + } + if st < len { + fill_text(st, len - st, next_pos, &mut pos); + } + let max = pos.iter().copied().max().unwrap_or(-1); + Ok((pos, max + 1 - len as i64)) +} + +/// 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 { + pub input_ids: Vec, + /// All items' `pixel_values`, concatenated in prompt order. + pub features: Vec, + pub grids: Vec<[u32; 3]>, + pub hashes: Vec, + pub offsets: Vec<(u32, u32)>, + pub mrope: Vec, + pub mrope_delta: i64, +} + +pub fn pack_drain(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()); + }; + 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()); + }; + features.extend(pixel_values); + let grid = item + .aux + .into_iter() + .find_map(|(name, tensor)| match (name.as_str(), tensor.data) { + ("image_grid_thw", TensorData::I64(v)) => Some(v), + _ => None, + }) + .ok_or("qwen_vl drain: missing image_grid_thw")?; + grids.push([grid[0] as u32, grid[1] as u32, grid[2] as u32]); + hashes.push(item.hash); + } + Ok(QwenDrain { + input_ids: output.input_ids, + features, + grids, + hashes, + offsets: output.offsets, + mrope: positions, + mrope_delta: delta, + }) +} + +// --- Python bindings (parity tests drive the exact server pipeline) --- + +#[cfg(feature = "python")] +mod python { + use numpy::{IntoPyArray, PyArray1}; + use pyo3::exceptions::PyValueError; + use pyo3::prelude::*; + + use super::*; + use crate::pipeline::TensorData; + + /// `(pixel_values flat f32, (t, h, w))` for one preprocessed image. + type PyProcessedImage<'py> = (Bound<'py, PyArray1>, (u32, u32, u32)); + /// Full native pipeline output at the scheduler boundary: + /// `(input_ids, features, grids, hashes, offsets, mrope, mrope_delta)`. + type PyNativeOutput<'py> = ( + Vec, + Bound<'py, PyArray1>, + Vec<(u32, u32, u32)>, + Vec, + Vec<(u32, u32)>, + Bound<'py, PyArray1>, + i64, + ); + + /// Run the full native image path on encoded image bytes: + /// decode → smart_resize → bicubic → normalize → patchify. + /// Returns `(pixel_values flat f32, (t, h, w))`. + #[pyfunction] + fn preprocess<'py>( + py: Python<'py>, + data: Vec, + spec_json: &str, + ) -> PyResult> { + let proc = QwenVlProcessor::from_spec_json(spec_json).map_err(PyValueError::new_err)?; + let out = py + .detach(move || { + let (rgb, height, width) = crate::common::decode_rgb(&data)?; + proc.process_item(&DecodedMedia::Image { rgb, height, width }) + }) + .map_err(PyValueError::new_err)?; + let Geometry::Grid([t, h, w]) = out.geometry; + let TensorData::F32(pixel_values) = out.feature.data else { + return Err(PyValueError::new_err("qwen_vl: expected f32 feature")); + }; + Ok((pixel_values.into_pyarray(py), (t, h, w))) + } + + #[pyfunction] + fn smart_resize_py( + height: usize, + width: usize, + factor: usize, + min_pixels: usize, + max_pixels: usize, + ) -> PyResult<(usize, usize)> { + smart_resize(height, width, factor, min_pixels, max_pixels).map_err(PyValueError::new_err) + } + + /// `(positions flat [3*input_len], delta)` for image-only requests; + /// `items` = [(start, end_inclusive, t, h, w), ...] in prompt order. + #[pyfunction] + fn mrope_image_only_py<'py>( + py: Python<'py>, + input_len: usize, + items: Vec<(u32, u32, u32, u32, u32)>, + merge_size: usize, + ) -> PyResult<(Bound<'py, PyArray1>, i64)> { + let items: Vec = items + .into_iter() + .map(|(start, end, t, h, w)| MropeItem { + start, + end, + grid: [t, h, w], + }) + .collect(); + let (pos, delta) = + mrope_image_only(input_len, &items, merge_size).map_err(PyValueError::new_err)?; + Ok((pos.into_pyarray(py), delta)) + } + + /// One image source: a `str` (data:/base64/file/http, resolved by + /// `common::fetch`) or raw encoded `bytes`. + #[derive(FromPyObject)] + enum PyImageSource { + Str(String), + Bytes(Vec), + } + + /// Drive the same typed native Qwen request pipeline used by + /// `sglang-server` (whose message layer owns the wire-payload parsing). + #[pyfunction] + #[pyo3(signature = (input_ids, images, spec_json))] + fn process_native_mm<'py>( + py: Python<'py>, + input_ids: Option>, + images: Vec, + spec_json: String, + ) -> PyResult> { + let images = images + .into_iter() + .map(|source| match source { + PyImageSource::Str(s) => crate::driver::ImageSource::String(s), + PyImageSource::Bytes(b) => crate::driver::ImageSource::Bytes(b), + }) + .collect(); + let input = crate::driver::MmInput { + text: None, + input_ids, + images, + }; + let drain = 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) + }) + .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, + )) + } + + pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> { + let m = PyModule::new(parent.py(), "qwen_vl")?; + m.add_function(wrap_pyfunction!(preprocess, &m)?)?; + m.add_function(wrap_pyfunction!(smart_resize_py, &m)?)?; + m.add_function(wrap_pyfunction!(mrope_image_only_py, &m)?)?; + m.add_function(wrap_pyfunction!(process_native_mm, &m)?)?; + parent.add_submodule(&m)?; + Ok(()) + } +} + +#[cfg(feature = "python")] +pub use python::register; + +#[cfg(test)] +mod tests { + use super::*; + + fn spec() -> QwenVlSpec { + QwenVlSpec { + image_token_id: 1, + patch_size: 2, + merge_size: 2, + temporal_patch_size: 2, + min_pixels: 4, + max_pixels: 1 << 30, + image_mean: [0.0; 3], + image_std: [1.0; 3], + } + } + + #[test] + fn smart_resize_matches_python_reference() { + // Values from the Python `smart_resize` (qwen_vl.py) run offline. + assert_eq!( + smart_resize(1365, 2048, 28, 3136, 12845056).unwrap(), + (1372, 2044) + ); + assert_eq!( + smart_resize(100, 100, 28, 3136, 12845056).unwrap(), + (112, 112) + ); + // Downscale branch: 4000x3000 exceeds 1280*28*28 → floor_by_factor. + assert_eq!( + smart_resize(3000, 4000, 28, 3136, 1003520).unwrap(), + (840, 1148) + ); + // Upscale branch: tiny image below min_pixels → ceil_by_factor. + assert_eq!(smart_resize(20, 20, 28, 3136, 12845056).unwrap(), (56, 56)); + // Qwen3.5 factors (patch 16 * merge 2, min 65536, max 16777216). + assert_eq!( + smart_resize(1365, 2048, 32, 65536, 16777216).unwrap(), + (1376, 2048) + ); + // Banker's rounding tie: 48/32 = 1.5 rounds to 2 (even), not 1. + assert_eq!(smart_resize(4000, 48, 32, 4, 1 << 30).unwrap(), (4000, 64)); + // Extreme aspect ratio rejected. + assert!(smart_resize(10000, 10, 28, 3136, 12845056).is_err()); + } + + /// A thin image against a small `max_pixels` floors one side to 0. That + /// used to reach the resize coefficient math and panic on a worker thread + /// (`attempt to multiply with overflow`) instead of rejecting the request. + #[test] + fn degenerate_target_is_rejected_not_panicked() { + // Aspect ratio 200 is exactly at MAX_RATIO, so it passes that guard; + // 10 / beta then floors to 0 with factor 28. + assert!(smart_resize(10, 2000, 28, 3136, 3136).is_err()); + + let mut spec = spec(); + spec.patch_size = 14; + spec.min_pixels = 3136; + spec.max_pixels = 3136; + let proc = QwenVlProcessor::new(spec).unwrap(); + let err = proc + .process_item(&DecodedMedia::Image { + rgb: vec![0u8; 10 * 2000 * 3], + height: 10, + width: 2000, + }) + .err() + .expect("degenerate geometry must be an Err, never a panic"); + assert!(err.contains("smart_resize"), "unexpected error: {err}"); + } + + /// The server's message layer gates modalities on what a family declares, + /// so a family gaining video/audio support must not silently inherit the + /// images-only default. + #[test] + fn qwen_declares_images_only() { + let caps = QwenVlProcessor::new(spec()).unwrap().capabilities(); + assert!(!caps.video && !caps.audio); + } + + #[test] + fn patchify_layout_matches_hf_order() { + // 4x8 image, ps=2, m=2, tps=2 → gh=2, gw=4, dim=3*2*2*2=24. + // Pixel value encodes its (y, x): v = y*16 + x*2 (fits u8). + let (h, w) = (4usize, 8usize); + let mut rgb = vec![0u8; h * w * 3]; + for y in 0..h { + for x in 0..w { + for c in 0..3 { + rgb[(y * w + x) * 3 + c] = (y * 16 + x * 2 + c) as u8; + } + } + } + let proc = QwenVlProcessor::new(spec()).unwrap(); + let pv = proc.patchify(&rgb, h, w); + let dim = 24; // 3 * tps * ps * ps + assert_eq!(pv.len(), 2 * 4 * dim); + + // Patch order (gh/m=1, gw/m=2, m, m): patch 0 = block(0,0) offset (0,0), + // patch 1 = (0,0)+(0,1) → x0=2, patch 2 = (0,0)+(1,0) → y0=2, + // patch 4 = block(0,1) → x0=4. + let lut = |y: usize, x: usize, c: usize| ((y * 16 + x * 2 + c) as f32) / 255.0; + // patch 1, channel 0, t=0, (py=0, px=0) → pixel (0, 2). + assert_eq!(pv[dim], lut(0, 2, 0)); + // patch 2, channel 0, t=0, (0,0) → pixel (2, 0). + assert_eq!(pv[2 * dim], lut(2, 0, 0)); + // patch 4, channel 0 → pixel (0, 4). + assert_eq!(pv[4 * dim], lut(0, 4, 0)); + // Temporal duplicate: t=1 block equals t=0 block. + let ps2 = 4; // ps*ps + assert_eq!(pv[dim + ps2], pv[dim]); + // Channel 1 block of patch 0 → same pixel, c=1. + assert_eq!(pv[2 * ps2], lut(0, 0, 1)); // c stride = tps*ps*ps = 8 + } + + #[test] + fn mrope_image_only_matches_reference() { + // 3 text tokens, image of grid [1, 4, 6] (m=2 → 2x3 = 6 tokens), 2 text. + // input: [T T T I I I I I I T T], len 11. + let items = [MropeItem { + start: 3, + end: 8, + grid: [1, 4, 6], + }]; + let (pos, delta) = mrope_image_only(11, &items, 2).unwrap(); + let len = 11; + // Text prefix 0..3: all rows 0,1,2. + for k in 0..3 { + assert_eq!( + (pos[k], pos[len + k], pos[2 * len + k]), + (k as i64, k as i64, k as i64) + ); + } + // Image tokens: t=0, h in 0..2, w in 0..3, +3 offset. + assert_eq!((pos[3], pos[len + 3], pos[2 * len + 3]), (3, 3, 3)); + assert_eq!((pos[4], pos[len + 4], pos[2 * len + 4]), (3, 3, 4)); + assert_eq!((pos[6], pos[len + 6], pos[2 * len + 6]), (3, 4, 3)); + // Text tail resumes at 3 + max(1,2,3) = 6. + assert_eq!((pos[9], pos[len + 9], pos[2 * len + 9]), (6, 6, 6)); + assert_eq!((pos[10], pos[len + 10], pos[2 * len + 10]), (7, 7, 7)); + // delta = max + 1 - len = 7 + 1 - 11. + assert_eq!(delta, -3); + } +} diff --git a/rust/sglang-mm/src/registry.rs b/rust/sglang-mm/src/registry.rs index 411ffa5bd..061dbc9f3 100644 --- a/rust/sglang-mm/src/registry.rs +++ b/rust/sglang-mm/src/registry.rs @@ -1,7 +1,13 @@ -//! Model processor registry. +//! Model processor registries. //! -//! Each model implements `ImageProcessorSpec` and registers itself. The Python -//! layer looks up a processor by model name at init time. +//! Two registries live here: +//! * [`ImageProcessorSpec`] / [`ProcessorRegistry`] — the Python-facing batch +//! preprocess interface (e.g. Inkling), looked up by name at init time. +//! * [`pipeline_from_spec`] — the pure-Rust request pipeline `sglang-server`'s +//! MM workers drive. Each model family implements +//! [`crate::pipeline::MmFamilyProcessor`] in `src//mod.rs`; the Python +//! side selects one by serializing a spec +//! (`{"family": ..., resolved processor params}`). /// `(height, width, patches_as_u16_bits, content_hash)` for one image. pub type PreprocessedImage = (usize, usize, Vec, u64); @@ -59,3 +65,23 @@ pub fn default_registry() -> ProcessorRegistry { reg.register(Box::new(crate::inkling::InklingProcessor)); reg } +// --- Server (pure-Rust) request pipeline --- + +/// Build a family processor from the Python-side spec JSON. `Err` on an +/// unknown family or malformed spec — the caller treats that as "no Rust +/// pipeline". +pub fn pipeline_from_spec( + json: &str, +) -> Result, String> { + #[derive(serde::Deserialize)] + struct Header { + family: String, + } + let header: Header = serde_json::from_str(json).map_err(|e| format!("mm spec: {e}"))?; + match header.family.as_str() { + "qwen_vl" => Ok(Box::new(crate::qwen_vl::QwenVlProcessor::from_spec_json( + json, + )?)), + other => Err(format!("unknown mm family: {other}")), + } +} diff --git a/rust/sglang-mm/tests/generate_golden.py b/rust/sglang-mm/tests/generate_golden.py index c1cf5ffde..b26eccca1 100644 --- a/rust/sglang-mm/tests/generate_golden.py +++ b/rust/sglang-mm/tests/generate_golden.py @@ -6,7 +6,7 @@ import numpy as np import torch from PIL import Image -sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "bench")) from bench_parity import PS, make_photo_like, ref_patchify OUT = ( diff --git a/rust/sglang-mm/tests/rlib_is_single_threaded.rs b/rust/sglang-mm/tests/rlib_is_single_threaded.rs new file mode 100644 index 000000000..22da13520 --- /dev/null +++ b/rust/sglang-mm/tests/rlib_is_single_threaded.rs @@ -0,0 +1,68 @@ +//! The pure-Rust `rlib` that `sglang-server` links must not own worker +//! threads: the server supplies concurrency across requests and pins its own +//! cores, so a library spawning pools behind its back would fight it. +//! +//! Guarding this from the outside (thread count of the process) rather than by +//! inspecting the code, so it stays true no matter how the fan-out seam in +//! `common::par` is refactored. Runs only in the default (rayon-less) build; +//! under `--features parallel` the pools are expected. + +#![cfg(not(feature = "parallel"))] + +use sglang_mm_core::driver::{ImageSource, MmInput, process}; +use sglang_mm_core::registry::pipeline_from_spec; + +const SPEC: &str = r#"{"family":"qwen_vl","image_token_id":1,"patch_size":14, + "merge_size":2,"temporal_patch_size":2,"min_pixels":3136, + "max_pixels":12845056,"image_mean":[0.0,0.0,0.0],"image_std":[1.0,1.0,1.0]}"#; + +fn thread_names() -> Vec { + std::fs::read_dir("/proc/self/task") + .expect("procfs") + .filter_map(|entry| { + let comm = entry.ok()?.path().join("comm"); + Some(std::fs::read_to_string(comm).ok()?.trim().to_string()) + }) + .collect() +} + +fn png(w: u32, h: u32) -> Vec { + let img = image::RgbImage::from_fn(w, h, |x, y| image::Rgb([x as u8, y as u8, 7])); + let mut buf = std::io::Cursor::new(Vec::new()); + img.write_to(&mut buf, image::ImageFormat::Png).unwrap(); + buf.into_inner() +} + +#[test] +fn processing_a_request_spawns_no_worker_threads() { + let before = thread_names().len(); + let family = pipeline_from_spec(SPEC).unwrap(); + + // Two images, so the per-item fan-out seam is exercised, not bypassed. + let out = process( + family.as_ref(), + MmInput { + text: None, + input_ids: Some(vec![7, 1, 8, 1, 9]), + images: vec![ + ImageSource::Bytes(png(112, 112)), + ImageSource::Bytes(png(84, 140)), + ], + }, + |_| Err("no tokenizer".into()), + ) + .expect("request should succeed"); + assert_eq!(out.items.len(), 2); + + let after = thread_names(); + let spawned: Vec<&String> = after.iter().filter(|t| t.starts_with("sgl-mm")).collect(); + assert!( + spawned.is_empty(), + "rlib build spawned crate-owned worker threads: {spawned:?}" + ); + assert_eq!( + after.len(), + before, + "rlib build changed the process thread count: {after:?}" + ); +} diff --git a/rust/sglang-mm/tests/test_golden.py b/rust/sglang-mm/tests/test_golden.py index 24baa18d5..1ec5e6f72 100644 --- a/rust/sglang-mm/tests/test_golden.py +++ b/rust/sglang-mm/tests/test_golden.py @@ -4,7 +4,7 @@ import os import numpy as np import pytest -import sglang.srt.multimodal._core.inkling +from sglang.srt.multimodal._core import inkling as _rs_inkling GOLDEN_DIR = os.environ.get( "INKLING_MM_GOLDEN_DIR", @@ -20,9 +20,7 @@ def bf16_bits_to_f32(bits: np.ndarray) -> np.ndarray: @pytest.mark.parametrize("path", GOLDENS, ids=[os.path.basename(p) for p in GOLDENS]) def test_patchify_rgb_bit_exact(path): g = np.load(path) - got = sglang.srt.multimodal._core.inkling.patchify_rgb( - g["arr"], int(g["patch_size"]) - ) + got = _rs_inkling.patchify_rgb(g["arr"], int(g["patch_size"])) np.testing.assert_array_equal(got, g["bits"].reshape(-1)) @@ -30,9 +28,7 @@ def test_patchify_rgb_bit_exact(path): def test_decode_patchify_png_bit_exact(path): g = np.load(path) h_ref, w_ref = g["arr"].shape[:2] - h, w, got = sglang.srt.multimodal._core.inkling.decode_patchify( - g["png"].tobytes(), int(g["patch_size"]) - ) + h, w, got = _rs_inkling.decode_patchify(g["png"].tobytes(), int(g["patch_size"])) assert (h, w) == (h_ref, w_ref) np.testing.assert_array_equal(got, g["bits"].reshape(-1)) @@ -41,9 +37,7 @@ def test_batch_matches_single(): gs = [np.load(p) for p in GOLDENS] data = [g["png"].tobytes() for g in gs] ps = int(gs[0]["patch_size"]) - for (h, w, bits), g in zip( - sglang.srt.multimodal._core.inkling.decode_patchify_batch(data, ps), gs - ): + for (h, w, bits), g in zip(_rs_inkling.decode_patchify_batch(data, ps), gs): np.testing.assert_array_equal(bits, g["bits"].reshape(-1)) diff --git a/rust/sglang-mm/tests/test_hash_fetch.py b/rust/sglang-mm/tests/test_hash_fetch.py index 3ac8bc241..76d32d55a 100644 --- a/rust/sglang-mm/tests/test_hash_fetch.py +++ b/rust/sglang-mm/tests/test_hash_fetch.py @@ -11,11 +11,15 @@ import soundfile as sf import torch from PIL import Image -sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "bench")) from bench_parity import make_photo_like from sglang.srt.managers.mm_utils import data_hash, hash_feature +from sglang.srt.multimodal._core import common as _rs_common from sglang.srt.multimodal.inkling import InklingProcessor +from sglang.srt.multimodal.inkling.image_processing_rust import ( + InklingRustImageProcessor, +) from sglang.srt.multimodal.processors import inkling as prc @@ -39,7 +43,9 @@ def make_proc(): proc.IMAGE_TOKEN_ID = 100 proc.AUDIO_TOKEN_ID = 101 proc.AUDIO_END_TOKEN_ID = 102 - proc.inkling_processor = InklingProcessor() + proc.inkling_processor = InklingProcessor( + image_processor=InklingRustImageProcessor() + ) return proc @@ -50,14 +56,19 @@ aud = wav_bytes() out = proc.assemble([1, 100, 2, 101, 3], [img], [aud]) img_item = next(i for i in out.mm_items if i.modality.name == "IMAGE") aud_item = next(i for i in out.mm_items if i.modality.name == "AUDIO") -assert img_item.hash == data_hash(img), "image hash != data_hash(raw bytes)" -assert aud_item.hash == data_hash(aud), "audio hash != data_hash(raw bytes)" -print(f" assemble: image hash={img_item.hash:#x} audio hash={aud_item.hash:#x} OK") +# The rust image path stores its raw-bytes content hash (blake3) eagerly; +# non-rust items get hashed lazily from the feature at set_pad_value time. +assert img_item.hash == _rs_common.content_hash(img), "image hash != rust content_hash" +assert aud_item.hash is None, "audio hash expected to be lazy (feature-based)" +print(f" assemble: image hash={img_item.hash:#x} OK") h0 = img_item.hash img_item.set_pad_value() assert img_item.hash == h0 and img_item.pad_value is not None print(f" set_pad_value: hash preserved, pad_value={img_item.pad_value} OK") +aud_item.set_pad_value() +assert aud_item.hash is not None and aud_item.pad_value is not None +print(" set_pad_value: audio feature-hash filled OK") out2 = proc.assemble([1, 100, 2], [img], []) assert out2.mm_items[0].hash == h0 @@ -73,14 +84,8 @@ out3 = asyncio.run( assert all(i.hash == h0 for i in out3.mm_items), "data: URL roundtrip hash mismatch" print(" process_mm_data_async: concurrent resolve + hash OK") -orig = prc._resolve_media_item -prc._resolve_media_item = lambda it: (time.sleep(0.3), orig(it))[1] -t0 = time.perf_counter() -asyncio.run(prc._resolve_media_items([data_url] * 8)) -elapsed = time.perf_counter() - t0 -prc._resolve_media_item = orig -assert elapsed < 1.2, f"8x 0.3s resolves took {elapsed:.2f}s; expected ~0.3s" -print(f" concurrency: 8 x 0.3s resolves in {elapsed:.2f}s OK") +# (The old `_resolve_media_items` concurrency helper is gone; resolution now +# happens inline in `process_mm_data_async`, covered by the roundtrip above.) imgs_5 = [png_bytes(make_photo_like(1080, 1920, seed=s)) for s in range(5)] feats = [ diff --git a/rust/sglang-mm/tests/test_integration.py b/rust/sglang-mm/tests/test_integration.py index cfb620c09..85486f908 100644 --- a/rust/sglang-mm/tests/test_integration.py +++ b/rust/sglang-mm/tests/test_integration.py @@ -6,10 +6,13 @@ import numpy as np import torch from PIL import Image -sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "bench")) from bench_parity import make_photo_like -import sglang.srt.multimodal.inkling.image_processing as ip +from sglang.srt.multimodal.inkling.image_processing import InklingImageProcessor +from sglang.srt.multimodal.inkling.image_processing_rust import ( + InklingRustImageProcessor, +) def encode(arr, fmt): @@ -21,17 +24,14 @@ def encode(arr, fmt): def run(images, use_rs: bool, rescale: bool): - ip._rs_module = None - os.environ["SGLANG_RS_MM_PREPROCESS"] = "1" if use_rs else "0" kwargs = ( {} if rescale else {"rescale_image_frac": None, "rescale_image_max_upscaled_long_edge": None} ) - proc = ip.InklingImageProcessor(patch_size=40, **kwargs) - out = proc.preprocess(images) - assert (ip._rs_module is not False) == use_rs, "rust module gating mismatch" - return out + cls = InklingRustImageProcessor if use_rs else InklingImageProcessor + proc = cls(patch_size=40, **kwargs) + return proc.preprocess(images) def compare(tag, images, expect_exact, rescale=False): @@ -59,7 +59,7 @@ arr1 = make_photo_like(1080, 1920, seed=1) arr2 = make_photo_like(720, 1280, seed=2) arr3 = make_photo_like(480, 640, seed=3) -print("=== integration: InklingImageProcessor env-gated rust path ===") +print("=== integration: Inkling rust vs python processor parity ===") compare("single PNG", [encode(arr1, "PNG")], expect_exact=True) compare("single JPEG", [encode(arr1, "JPEG")], expect_exact=False) compare( diff --git a/rust/sglang-mm/tests/test_resize_parity.py b/rust/sglang-mm/tests/test_resize_parity.py index 23fc6c0ef..e09f93a2d 100644 --- a/rust/sglang-mm/tests/test_resize_parity.py +++ b/rust/sglang-mm/tests/test_resize_parity.py @@ -6,7 +6,8 @@ import numpy as np import pytest from PIL import Image -import sglang.srt.multimodal._core.inkling +from sglang.srt.multimodal._core import common as _rs_common +from sglang.srt.multimodal._core import inkling as _rs_inkling def py_scaled_dims( @@ -41,9 +42,7 @@ def pil_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray: def rs_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray: - return sglang.srt.multimodal._core.inkling.resize_rgb(arr, tw, th).reshape( - th, tw, 3 - ) + return _rs_common.resize_rgb(arr, tw, th).reshape(th, tw, 3) CASES = [ @@ -74,9 +73,9 @@ def test_scaled_dims_sweep(): sizes += [(2048, 1024), (2049, 100), (1024, 2048), (1, 1), (4096, 4096)] for frac, cap in [(2.0, 2048), (1.5, 2048), (3.0, None), (None, None), (2.0, 1)]: for w, h in sizes: - assert sglang.srt.multimodal._core.inkling.scaled_dims( + assert _rs_common.scaled_dims(w, h, frac, cap) == py_scaled_dims( w, h, frac, cap - ) == py_scaled_dims(w, h, frac, cap), ( + ), ( w, h, frac, @@ -93,12 +92,10 @@ def test_decode_patchify_rescaled_matches_pil_pipeline(): arr = rng.integers(0, 256, (1080, 1920, 3), dtype=np.uint8) buf = io.BytesIO() Image.fromarray(arr).save(buf, format="PNG") - h, w, bits = sglang.srt.multimodal._core.inkling.decode_patchify( - buf.getvalue(), 40, 2.0, 2048 - ) + h, w, bits = _rs_inkling.decode_patchify(buf.getvalue(), 40, 2.0, 2048) assert (w, h) == py_scaled_dims(1920, 1080, 2.0, 2048) ref_arr = pil_resize(arr, w, h) - ref_bits = sglang.srt.multimodal._core.inkling.patchify_rgb(ref_arr, 40) + ref_bits = _rs_inkling.patchify_rgb(ref_arr, 40) np.testing.assert_array_equal(bits, ref_bits) assert torch.from_numpy(bits).view(torch.bfloat16).shape[0] > 0 diff --git a/test/registered/unit/multimodal/rust/_mm_rust_utils.py b/test/registered/unit/multimodal/rust/_mm_rust_utils.py new file mode 100644 index 000000000..b8cc6c759 --- /dev/null +++ b/test/registered/unit/multimodal/rust/_mm_rust_utils.py @@ -0,0 +1,92 @@ +"""Shared fixtures for the native Rust multimodal suites. + +Imported via ``sys.path`` from the sibling suites (unittest runs these files by +path, so a package-relative import would break ``python ``); the module +name is deliberately specific so it cannot shadow another suite's helpers on the +process-global ``sys.path``. +""" + +import io +import json + +import numpy as np +from PIL import Image + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import is_in_ci + +register_cpu_ci( + est_time=0, suite="base-a-test-cpu", disabled="Rust multimodal test helpers" +) + + +def load_core(): + """The Rust ``_core`` extension, or ``None`` (→ skip) when not built + locally. In CI a missing extension is a hard failure, never a silent + skip — the CPU suite builds it from source.""" + try: + from sglang.srt.multimodal import _core + + return _core + except ImportError: + if is_in_ci(): + raise + return None + + +IMAGE_TOKEN_ID = 900 +VISION_START_ID = 901 +VISION_END_ID = 902 +VIDEO_TOKEN_ID = 903 + +PROCESSOR_CONFIGS = { + "qwen2_vl": dict( + patch_size=14, + merge_size=2, + temporal_patch_size=2, + min_pixels=56 * 56, + max_pixels=28 * 28 * 1280, + image_mean=[0.48145466, 0.4578275, 0.40821073], + image_std=[0.26862954, 0.26130258, 0.27577711], + ), + "qwen2_5_vl": dict( + patch_size=14, + merge_size=2, + temporal_patch_size=2, + min_pixels=56 * 56, + max_pixels=28 * 28 * 1280, + image_mean=[0.5] * 3, + image_std=[0.5] * 3, + ), + "qwen3_5": dict( + patch_size=16, + merge_size=2, + temporal_patch_size=2, + min_pixels=65536, + max_pixels=16777216, + image_mean=[0.5] * 3, + image_std=[0.5] * 3, + ), +} + + +def make_image(width, height, seed=0): + rng = np.random.default_rng(seed) + y, x = np.mgrid[0:height, 0:width] + base = np.stack( + (x * 255 / max(width - 1, 1), y * 255 / max(height - 1, 1), (x + y) % 256), + axis=-1, + ) + return Image.fromarray( + np.clip(base + rng.integers(0, 24, base.shape), 0, 255).astype(np.uint8) + ) + + +def image_bytes(width, height, seed=0): + buffer = io.BytesIO() + make_image(width, height, seed).save(buffer, format="PNG") + return buffer.getvalue() + + +def spec_json(config, image_token_id=IMAGE_TOKEN_ID): + return json.dumps({"family": "qwen_vl", "image_token_id": image_token_id, **config}) diff --git a/test/registered/unit/multimodal/rust/inkling/test_bindings.py b/test/registered/unit/multimodal/rust/inkling/test_bindings.py new file mode 100644 index 000000000..23ab5a939 --- /dev/null +++ b/test/registered/unit/multimodal/rust/inkling/test_bindings.py @@ -0,0 +1,74 @@ +"""Smoke coverage for the ``_core.inkling`` PyO3 bindings. + +Covers the ``#[pyfunction]``s in ``rust/sglang-mm/src/inkling/mod.rs`` +(``patchify_rgb`` / ``decode_patchify`` / ``decode_patchify_batch`` / +``preprocess_images`` / ``rescale_patchify_hash``) plus +``_core.common.content_hash``. + +The inkling bindings are otherwise exercised only by the GPU e2e model test; +this pins the binding surface (signatures, dtypes, cross-binding consistency) +in the CPU suite so a rework of the extension can't silently break them. +""" + +import sys +import unittest +from pathlib import Path + +import numpy as np + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _mm_rust_utils import image_bytes, load_core # noqa: E402 + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +CORE = load_core() +PATCH_SIZE = 16 + + +@unittest.skipUnless(CORE, "sglang-mm extension not built") +class TestInklingBindings(CustomTestCase): + def test_bindings_are_consistent(self): + data = image_bytes(50, 34) + h, w, patches = CORE.inkling.decode_patchify(data, PATCH_SIZE) + self.assertEqual((h, w), (34, 50)) + patches = np.asarray(patches) + self.assertEqual(patches.dtype, np.uint16) + # One padded patch column, ceil rows (the inkling grid convention). + expected_len = -(-h // PATCH_SIZE) * (w // PATCH_SIZE + 1) * PATCH_SIZE**2 * 3 + self.assertEqual(patches.size, expected_len) + + # patchify_rgb on the decoded array must match the fused decode path. + dh, dw, rgb = CORE.common.image_decode_rgb(data) + arr = np.asarray(rgb).reshape(dh, dw, 3) + np.testing.assert_array_equal( + np.asarray(CORE.inkling.patchify_rgb(arr, PATCH_SIZE)), patches + ) + + # Batch and hashed variants agree with the single-image call. + [(bh, bw, batch)] = CORE.inkling.decode_patchify_batch([data], PATCH_SIZE) + self.assertEqual((bh, bw), (h, w)) + np.testing.assert_array_equal(np.asarray(batch), patches) + [(ph, pw, pre, phash)] = CORE.inkling.preprocess_images([data], PATCH_SIZE) + self.assertEqual((ph, pw), (h, w)) + np.testing.assert_array_equal(np.asarray(pre), patches) + self.assertEqual(phash, CORE.common.content_hash(data)) + + rh, rw, rpatches, rhash = CORE.inkling.rescale_patchify_hash( + arr, data, PATCH_SIZE + ) + self.assertEqual((rh, rw), (h, w)) + np.testing.assert_array_equal(np.asarray(rpatches), patches) + self.assertEqual(rhash, phash) + + def test_invalid_inputs_rejected(self): + with self.assertRaises(ValueError): + CORE.inkling.decode_patchify(b"junk", PATCH_SIZE) + with self.assertRaises(ValueError): + CORE.inkling.decode_patchify(image_bytes(16, 16), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/rust/qwen/test_driver.py b/test/registered/unit/multimodal/rust/qwen/test_driver.py new file mode 100644 index 000000000..59ad187e5 --- /dev/null +++ b/test/registered/unit/multimodal/rust/qwen/test_driver.py @@ -0,0 +1,109 @@ +"""Native driver error paths: out-of-scope and malformed inputs are rejected. + +Covers ``process`` in ``rust/sglang-mm/src/driver.rs`` (via the +``_core.qwen_vl.process_native_mm`` binding). The wire-payload parsing that +feeds this driver (modality/shape rejection) lives in ``sglang-server``'s +message layer and is tested with the integration PR. + +There is no Python fallback path, so the server rejects every driver error +back to the client as a 400; the message must say why (placeholder mismatch, +undecodable image, missing prompt). This pins that contract for each +rejection class. +""" + +import io +import sys +import unittest +from pathlib import Path + +import numpy as np +from PIL import Image + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _mm_rust_utils import ( # noqa: E402 + IMAGE_TOKEN_ID, + PROCESSOR_CONFIGS, + VISION_END_ID, + VISION_START_ID, + image_bytes, + load_core, + spec_json, +) + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +QWEN_CORE = getattr(load_core(), "qwen_vl", None) +SPEC = spec_json(PROCESSOR_CONFIGS["qwen2_5_vl"]) +IMAGE_IDS = [7, VISION_START_ID, IMAGE_TOKEN_ID, VISION_END_ID, 8] + + +def gif_bytes(): + buffer = io.BytesIO() + Image.fromarray(np.zeros((16, 16, 3), dtype=np.uint8)).save(buffer, format="GIF") + return buffer.getvalue() + + +@unittest.skipUnless( + QWEN_CORE and hasattr(QWEN_CORE, "process_native_mm"), + "sglang-mm native Qwen driver not built", +) +class TestNativeDriverErrorPaths(CustomTestCase): + def assert_rejected(self, input_ids, images, pattern, spec=SPEC): + with self.assertRaisesRegex(ValueError, pattern): + QWEN_CORE.process_native_mm(input_ids, images, spec) + + def test_degenerate_geometry_rejected_not_panicked(self): + """A thin image against a tight ``max_pixels`` floors a side of the + smart_resize target to 0. That used to panic on a worker thread inside + the resize coefficient math; a Rust panic surfaces as ``PanicException`` + (a ``BaseException``), not ``ValueError``, so this asserts the request + is rejected cleanly rather than crashing the pipeline.""" + config = dict(PROCESSOR_CONFIGS["qwen2_5_vl"], min_pixels=3136, max_pixels=3136) + self.assert_rejected( + IMAGE_IDS, + [image_bytes(2000, 10)], + "smart_resize", + spec=spec_json(config), + ) + + def test_placeholder_count_mismatches_rejected(self): + cases = { + "no placeholder": ([7, 8], [image_bytes(80, 80)]), + "more images": (IMAGE_IDS, [image_bytes(80, 80), image_bytes(88, 80, 1)]), + "more placeholders": (IMAGE_IDS + IMAGE_IDS, [image_bytes(80, 80)]), + } + for name, (ids, images) in cases.items(): + with self.subTest(case=name): + self.assert_rejected(ids, images, "placeholder") + + def test_undecodable_images_rejected(self): + # Corrupt bytes are outside the native decoder's scope; the server + # rejects them as a 400. + self.assert_rejected(IMAGE_IDS, [b"junk"], "decode") + + def test_gif_serves_through_the_pipeline(self): + """GIF moved from rejected to served when the pure-Rust webp/gif/bmp + decoders were enabled; this pins the accept side of that contract flip + (the reject side used to be asserted here and broke in CI).""" + _, _, grids, _, offsets, _, _ = QWEN_CORE.process_native_mm( + IMAGE_IDS, [gif_bytes()], SPEC + ) + self.assertEqual(len(grids), 1) + self.assertEqual(len(offsets), 1) + + def test_missing_text_and_input_ids_rejected(self): + for input_ids in (None, []): + with self.subTest(input_ids=input_ids): + self.assert_rejected( + input_ids, [image_bytes(80, 80)], "without text or input_ids" + ) + + def test_image_free_request_rejected(self): + self.assert_rejected(IMAGE_IDS, [], "image sources") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/rust/qwen/test_preprocess.py b/test/registered/unit/multimodal/rust/qwen/test_preprocess.py new file mode 100644 index 000000000..d9843a045 --- /dev/null +++ b/test/registered/unit/multimodal/rust/qwen/test_preprocess.py @@ -0,0 +1,105 @@ +"""Qwen native image preprocessing parity against Transformers. + +Covers ``QwenVlProcessor::process_item`` and ``smart_resize`` in +``rust/sglang-mm/src/qwen_vl/mod.rs`` (via the ``_core.qwen_vl.preprocess`` +and ``smart_resize_py`` bindings), against the HF Qwen2-VL image processors +and the Python ``smart_resize``. + +Both HF processors are pinned, because they resample differently and only one +of them is what a server actually runs. On transformers 5.x +``Qwen2VLImageProcessor`` is the torchvision path (the ``Fast`` suffix was +dropped) and is what ``AutoImageProcessor`` hands SGLang by default; +``Qwen2VLImageProcessorPil`` is the PIL path, reachable via +``--disable-fast-image-processor``. The Rust resize is a bit-exact clone of +PIL's fixed-point kernel, so the PIL processor is asserted exactly and the +torchvision one carries the cross-implementation envelope. +""" + +import sys +import unittest +from pathlib import Path + +import numpy as np + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _mm_rust_utils import ( # noqa: E402 + PROCESSOR_CONFIGS, + image_bytes, + load_core, + make_image, + spec_json, +) + +register_cpu_ci(est_time=20, suite="base-a-test-cpu") + +QWEN_CORE = getattr(load_core(), "qwen_vl", None) + + +SIZES = ((640, 480), (1024, 683), (50, 40), (300, 301)) + + +@unittest.skipUnless(QWEN_CORE, "sglang-mm Qwen binding not built") +class TestQwenImagePreprocess(CustomTestCase): + def _assert_matches(self, processor, max_diff, mean_diff): + for family, config in PROCESSOR_CONFIGS.items(): + hf = processor(**config) + for index, size in enumerate(SIZES): + with self.subTest(family=family, size=size): + image = make_image(*size, seed=index) + actual, grid = QWEN_CORE.preprocess( + image_bytes(*size, seed=index), spec_json(config) + ) + expected = hf(images=[image], return_tensors="pt") + self.assertEqual(grid, tuple(expected.image_grid_thw[0].tolist())) + diff = np.abs( + np.asarray(actual).reshape(expected.pixel_values.shape) + - expected.pixel_values.numpy() + ) + # LessEqual, not Less: max_diff=0.0 is a real bound here. + self.assertLessEqual(diff.max(), max_diff) + self.assertLessEqual(diff.mean(), mean_diff) + + def test_features_match_pil_processor_exactly(self): + """Against the PIL processor the native path is bit-exact, so this is + asserted with zero tolerance: every stage (smart_resize geometry, + the fixed-point bicubic kernel, rescale/normalize, HF patch order) is + pinned, and any drift in any of them shows up here rather than being + absorbed by a tolerance.""" + from transformers.models.qwen2_vl.image_processing_pil_qwen2_vl import ( + Qwen2VLImageProcessorPil, + ) + + self._assert_matches(Qwen2VLImageProcessorPil, max_diff=0.0, mean_diff=0.0) + + def test_features_match_torchvision_processor_within_envelope(self): + """The torchvision processor is what a default server runs, so its + divergence is bounded separately — it is a different antialiased-bicubic + implementation, which the bit-exact PIL assertion above says nothing + about. Measured worst case is max 0.030 / mean 6.7e-5 (≈2 u8 levels + after normalize with the qwen2_vl std).""" + from transformers.models.qwen2_vl.image_processing_qwen2_vl import ( + Qwen2VLImageProcessor, + ) + + self._assert_matches(Qwen2VLImageProcessor, max_diff=0.035, mean_diff=1e-3) + + def test_smart_resize_matches_python(self): + from sglang.srt.multimodal.processors.qwen_vl import smart_resize + + cases = ( + (1365, 2048, 28, 3136, 12845056), + (3000, 4000, 28, 3136, 1003520), + (20, 20, 28, 3136, 12845056), + (1365, 2048, 32, 65536, 16777216), + (4000, 48, 32, 4, 1 << 30), + ) + for case in cases: + with self.subTest(case=case): + self.assertEqual(QWEN_CORE.smart_resize_py(*case), smart_resize(*case)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/rust/qwen/test_token_layout_mrope.py b/test/registered/unit/multimodal/rust/qwen/test_token_layout_mrope.py new file mode 100644 index 000000000..409b26282 --- /dev/null +++ b/test/registered/unit/multimodal/rust/qwen/test_token_layout_mrope.py @@ -0,0 +1,104 @@ +"""Qwen placeholder expansion, offsets, and M-RoPE parity. + +Covers ``layout_by_placeholder`` / ``apply_layout`` in +``rust/sglang-mm/src/common/token_layout.rs`` and ``mrope_image_only`` in +``rust/sglang-mm/src/qwen_vl/mod.rs`` (via the +``_core.qwen_vl.process_native_mm`` and ``mrope_image_only_py`` +bindings), against ``BaseMultimodalProcessor`` expansion/offsets and +``MRotaryEmbedding.get_rope_index``. +""" + +import sys +import unittest +from pathlib import Path + +import numpy as np +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _mm_rust_utils import ( # noqa: E402 + IMAGE_TOKEN_ID, + PROCESSOR_CONFIGS, + VIDEO_TOKEN_ID, + VISION_END_ID, + VISION_START_ID, + image_bytes, + load_core, + spec_json, +) + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +QWEN_CORE = getattr(load_core(), "qwen_vl", None) + + +@unittest.skipUnless( + QWEN_CORE and hasattr(QWEN_CORE, "process_native_mm"), + "sglang-mm native Qwen driver not built", +) +class TestQwenPromptGeometry(CustomTestCase): + def test_placeholder_expansion_and_offsets(self): + from sglang.srt.multimodal.processors.base_processor import ( + BaseMultimodalProcessor, + ) + + config = PROCESSOR_CONFIGS["qwen2_5_vl"] + for image_count in (1, 2): + ids = [7] + for _ in range(image_count): + ids.extend((VISION_START_ID, IMAGE_TOKEN_ID, VISION_END_ID, 8)) + images = [image_bytes(96 + 8 * i, 80, i) for i in range(image_count)] + with self.subTest(image_count=image_count): + actual_ids, _, grids, _, offsets, _, _ = QWEN_CORE.process_native_mm( + ids, images, spec_json(config) + ) + counts = [t * h * w // config["merge_size"] ** 2 for t, h, w in grids] + expected_ids = BaseMultimodalProcessor._expand_input_ids( + ids, counts, IMAGE_TOKEN_ID + ) + self.assertEqual(actual_ids, expected_ids) + self.assertEqual( + offsets, + BaseMultimodalProcessor.get_mm_items_offset( + torch.tensor(expected_ids), IMAGE_TOKEN_ID + ), + ) + + def test_mrope_matches_model_reference(self): + """The single Rust image-only M-RoPE must match ``get_rope_index`` + for every native Qwen family (image-only makes them coincide).""" + from sglang.srt.layers.rotary_embedding import MRotaryEmbedding + + grids = [(1, 4, 6), (1, 6, 4)] + ids, items = [10], [] + for grid in grids: + ids.extend((VISION_START_ID, IMAGE_TOKEN_ID)) + start = len(ids) - 1 + ids.extend([IMAGE_TOKEN_ID] * (np.prod(grid) // 4 - 1)) + items.append((start, len(ids) - 1, *grid)) + ids.extend((VISION_END_ID, 11)) + + actual, delta = QWEN_CORE.mrope_image_only_py(len(ids), items, 2) + actual = np.asarray(actual).reshape(3, -1) + for model_type in ("qwen2_vl", "qwen2_5_vl", "qwen3_vl", "qwen3_5"): + with self.subTest(model_type=model_type): + expected, expected_delta = MRotaryEmbedding.get_rope_index( + spatial_merge_size=2, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_ID, + model_type=model_type, + tokens_per_second=2 if model_type == "qwen2_5_vl" else None, + input_ids=torch.tensor(ids).unsqueeze(0), + image_grid_thw=torch.tensor(grids), + video_grid_thw=None, + ) + np.testing.assert_array_equal(actual, expected.squeeze(1).numpy()) + self.assertEqual(delta, int(expected_delta.item())) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/rust/shared/test_fetch.py b/test/registered/unit/multimodal/rust/shared/test_fetch.py new file mode 100644 index 000000000..14a8b8031 --- /dev/null +++ b/test/registered/unit/multimodal/rust/shared/test_fetch.py @@ -0,0 +1,83 @@ +"""Model-independent parity tests for Rust multimodal source loading. + +Covers ``fetch_bytes`` in ``rust/sglang-mm/src/common/fetch.rs`` (via the +``_core.common.fetch_bytes`` binding), against the Python reference +``sglang.srt.utils.common.get_image_bytes``. +""" + +import base64 +import http.server +import sys +import tempfile +import threading +import unittest +from pathlib import Path + +from sglang.srt.utils.common import get_image_bytes +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _mm_rust_utils import load_core # noqa: E402 + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +CORE = load_core() +FETCH = CORE and CORE.common.fetch_bytes + + +@unittest.skipUnless(FETCH, "sglang-mm fetch binding not built") +class TestRustMediaSourceLoading(CustomTestCase): + DATA = b"native-mm-source" + + def test_inline_sources(self): + encoded = base64.b64encode(self.DATA).decode() + for source in (encoded, f"data:application/octet-stream;base64,{encoded}"): + with self.subTest(source=source[:8]): + self.assertEqual(bytes(FETCH(source)), get_image_bytes(source)) + + def test_file_sources(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "image.png" + path.write_bytes(self.DATA) + self.assertEqual(bytes(FETCH(str(path))), get_image_bytes(str(path))) + # `file://` is asserted against the payload, not the Python helper, + # on purpose: `get_image_bytes` passes the un-stripped URL straight + # to `open()`, so the reference raises here. The native path strips + # the scheme and succeeds — a deliberate divergence, not an omission. + with self.assertRaises(OSError): + get_image_bytes(path.as_uri()) + self.assertEqual(bytes(FETCH(path.as_uri())), self.DATA) + + def test_http_source(self): + data = self.DATA + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.end_headers() + self.wfile.write(data) + + def log_message(self, *_): + pass + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + url = f"http://127.0.0.1:{server.server_port}/image" + self.assertEqual(bytes(FETCH(url)), get_image_bytes(url)) + finally: + server.shutdown() + server.server_close() + thread.join() + + def test_invalid_sources_fail(self): + for source in ("not base64!", "/definitely/missing/image.png"): + with self.subTest(source=source): + with self.assertRaises(ValueError): + FETCH(source) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/rust/shared/test_image_decode.py b/test/registered/unit/multimodal/rust/shared/test_image_decode.py new file mode 100644 index 000000000..2b5d80c26 --- /dev/null +++ b/test/registered/unit/multimodal/rust/shared/test_image_decode.py @@ -0,0 +1,102 @@ +"""Model-independent image decode parity for native Rust MM. + +Covers ``decode_rgb`` in ``rust/sglang-mm/src/common/mod.rs`` (via the +``_core.common.image_decode_rgb`` binding), against PIL's +``Image.open(...).convert("RGB")``. +""" + +import io +import sys +import unittest +from pathlib import Path + +import numpy as np +from PIL import Image + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _mm_rust_utils import load_core # noqa: E402 + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +CORE = load_core() +DECODE = CORE and CORE.common.image_decode_rgb + + +def encode(image, fmt, **kwargs): + buffer = io.BytesIO() + image.save(buffer, format=fmt, **kwargs) + return buffer.getvalue() + + +@unittest.skipUnless(DECODE, "sglang-mm decode binding not built") +class TestRustImageDecode(CustomTestCase): + def assert_matches_pil(self, data, tolerance=0): + height, width, pixels = DECODE(data) + expected = np.asarray(Image.open(io.BytesIO(data)).convert("RGB")) + self.assertEqual((height, width), expected.shape[:2]) + actual = np.asarray(pixels).reshape(expected.shape) + diff = np.abs(actual.astype(int) - expected.astype(int)) + self.assertLessEqual(diff.max(), tolerance) + + def test_png_modes_match_pil(self): + rgb = np.random.default_rng(1).integers(0, 256, (19, 23, 3), dtype=np.uint8) + cases = [ + ("RGB", Image.fromarray(rgb)), + ("L", Image.fromarray(rgb[..., 0])), + ("RGBA", Image.fromarray(np.dstack((rgb, rgb[..., 0])))), + ("P", Image.fromarray(rgb).quantize(colors=16)), + ] + for mode, image in cases: + with self.subTest(mode=mode): + self.assert_matches_pil(encode(image, "PNG")) + + def test_jpeg_modes_match_with_decoder_tolerance(self): + rgb = np.random.default_rng(2).integers(0, 256, (31, 29, 3), dtype=np.uint8) + image = Image.fromarray(rgb) + exif = image.getexif() + exif[274] = 6 # EXIF orientation: neither decoder applies it + cases = [ + ("RGB", encode(image, "JPEG")), + ("L", encode(Image.fromarray(rgb[..., 0]), "JPEG")), + ("CMYK", encode(image.convert("CMYK"), "JPEG")), + ("EXIF-rotated", encode(image, "JPEG", exif=exif)), + ] + for mode, data in cases: + with self.subTest(mode=mode): + self.assert_matches_pil(data, tolerance=3) + + def test_lossless_formats_match_pil_exactly(self): + """GIF/BMP/lossless-WebP joined the native decoder with the pure-Rust + webp/gif/bmp enablement. Their reconstruction is exact by format spec, + so parity with PIL is pinned at zero tolerance like PNG.""" + rgb = np.random.default_rng(4).integers(0, 256, (17, 21, 3), dtype=np.uint8) + image = Image.fromarray(rgb) + cases = { + "gif": encode(image.quantize(colors=64), "GIF"), + "bmp": encode(image, "BMP"), + "webp-lossless": encode(image, "WEBP", lossless=True), + } + for name, data in cases.items(): + with self.subTest(fmt=name): + self.assert_matches_pil(data) + + def test_unsupported_inputs_fail(self): + gray16 = Image.fromarray(np.zeros((9, 9), dtype=np.uint16)) + cases = { + "corrupt": b"not an image", + # >8-bit depths must error (the request is then rejected — there + # is no Python fallback) — never silently diverge from PIL's + # clipping. + "png16": encode(gray16, "PNG"), + } + for name, data in cases.items(): + with self.subTest(input=name): + with self.assertRaises(ValueError): + DECODE(data) + + +if __name__ == "__main__": + unittest.main()