From 32e5d788bd4c8a4ec61dce728adc1b5670530c90 Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Wed, 5 Aug 2026 23:35:50 -0700 Subject: [PATCH] [mm] rust-server: native multimodal processing for Qwen VL (integrate sglang-mm, e2e) (#32365) Co-authored-by: Claude Fable 5 Co-authored-by: Cursor --- python/sglang/srt/managers/io_struct.py | 5 +- python/sglang/srt/managers/mm_utils.py | 4 +- .../srt/managers/multimodal_processor.py | 41 +- python/sglang/srt/managers/rust_server.py | 382 +++++++++++++++- .../sglang/srt/managers/tokenizer_manager.py | 8 +- rust/Cargo.lock | 4 + rust/sglang-mm/src/common/fetch.rs | 132 +++++- rust/sglang-mm/src/common/mod.rs | 18 +- rust/sglang-mm/src/common/resize.rs | 114 +++-- rust/sglang-mm/src/qwen_vl/mod.rs | 61 ++- rust/sglang-mm/tests/test_resize_parity.py | 58 ++- rust/sglang-server/Cargo.toml | 11 + rust/sglang-server/src/api_server.rs | 1 + .../src/api_server/native_api.rs | 7 +- rust/sglang-server/src/api_server/prefetch.rs | 199 +++++++++ rust/sglang-server/src/fsm.rs | 13 +- rust/sglang-server/src/lib.rs | 65 +++ rust/sglang-server/src/message.rs | 6 +- rust/sglang-server/src/message/mm_payload.rs | 222 ++++++++++ rust/sglang-server/src/message/request.rs | 335 +++++++++++++- rust/sglang-server/src/mm.rs | 413 ++++++++++++++++++ rust/sglang-server/src/runtime.rs | 56 ++- rust/sglang-server/src/runtime/config.rs | 11 + rust/sglang-server/src/tokenizer_manager.rs | 9 +- .../src/tokenizer_manager/ingress.rs | 353 +++++++++++++-- .../unit/multimodal/rust/_mm_rust_utils.py | 11 +- .../unit/multimodal/rust/qwen/_fixtures.py | 102 +++++ .../multimodal/rust/qwen/test_e2e_parity.py | 162 +++++++ .../rust/qwen/test_native_mm_host.py | 105 +++++ .../multimodal/rust/qwen/test_preprocess.py | 48 +- .../rust/shared/test_build_native_mm.py | 166 +++++++ .../rust/shared/test_native_mm_gate.py | 61 +++ .../rust/shared/test_partition_cores.py | 54 +++ .../registered/vlm/test_rust_native_mm_e2e.py | 129 ++++++ .../vlm/test_rust_native_mm_mmmu.py | 168 +++++++ 35 files changed, 3356 insertions(+), 178 deletions(-) create mode 100644 rust/sglang-server/src/api_server/prefetch.rs create mode 100644 rust/sglang-server/src/message/mm_payload.rs create mode 100644 rust/sglang-server/src/mm.rs create mode 100644 test/registered/unit/multimodal/rust/qwen/_fixtures.py create mode 100644 test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py create mode 100644 test/registered/unit/multimodal/rust/qwen/test_native_mm_host.py create mode 100644 test/registered/unit/multimodal/rust/shared/test_build_native_mm.py create mode 100644 test/registered/unit/multimodal/rust/shared/test_native_mm_gate.py create mode 100644 test/registered/unit/multimodal/rust/shared/test_partition_cores.py create mode 100644 test/registered/vlm/test_rust_native_mm_e2e.py create mode 100644 test/registered/vlm/test_rust_native_mm_mmmu.py diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 928348eb5..b2b63b1fd 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -2260,7 +2260,10 @@ def unwrap_from_pickle(obj: Optional[object]) -> Optional[object]: return None if _USE_PICKLE_IPC: return obj - assert isinstance(obj, PickleWrapper) + if not isinstance(obj, PickleWrapper): + # Already materialized: the embedded Rust server attaches in-process + # objects (native-MM `mm_inputs`) without a pickle hop. + return obj return pickle.loads(obj.data) diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index aebb16253..5e8e2f57d 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1336,11 +1336,11 @@ def _get_is_default_transport(): global _is_default_tensor_transport if _is_default_tensor_transport is None: from sglang.srt.managers.tokenizer_manager import ( - _determine_tensor_transport_mode, + determine_tensor_transport_mode, ) _is_default_tensor_transport = ( - _determine_tensor_transport_mode(get_server_args()) == "default" + determine_tensor_transport_mode(get_server_args()) == "default" ) return _is_default_tensor_transport diff --git a/python/sglang/srt/managers/multimodal_processor.py b/python/sglang/srt/managers/multimodal_processor.py index 0554ed34e..4af1a0a6d 100644 --- a/python/sglang/srt/managers/multimodal_processor.py +++ b/python/sglang/srt/managers/multimodal_processor.py @@ -41,14 +41,9 @@ def import_processors(package_name: str, overwrite: bool = False): PROCESSOR_MAPPING[arch] = cls -def get_mm_processor( - hf_config, - server_args: ServerArgs, - processor, - transport_mode, - model_config=None, - **kwargs, -) -> BaseMultimodalProcessor: +def get_mm_processor_cls(hf_config, server_args: ServerArgs, model_config=None): + """The class :func:`get_mm_processor` would instantiate, or ``None`` when the + architecture has no registered processor.""" model_impl = str(getattr(server_args, "model_impl", "auto")).lower() uses_transformers_backend = model_impl == "transformers" if model_impl == "auto" and model_config is not None: @@ -64,20 +59,30 @@ def get_mm_processor( if not uses_transformers_backend or getattr( processor_cls, "supports_transformers_backend", False ): - return processor_cls( - hf_config, server_args, processor, transport_mode, **kwargs - ) + return processor_cls if uses_transformers_backend: from sglang.srt.multimodal.processors.transformers_auto import ( TransformersAutoMultimodalProcessor, ) - return TransformersAutoMultimodalProcessor( - hf_config, server_args, processor, transport_mode, **kwargs - ) + return TransformersAutoMultimodalProcessor - raise ValueError( - f"No processor registered for architecture: {hf_config.architectures}.\n" - f"Registered architectures: {[model_cls.__name__ for model_cls in PROCESSOR_MAPPING.keys()]}" - ) + return None + + +def get_mm_processor( + hf_config, + server_args: ServerArgs, + processor, + transport_mode, + model_config=None, + **kwargs, +) -> BaseMultimodalProcessor: + processor_cls = get_mm_processor_cls(hf_config, server_args, model_config) + if processor_cls is None: + raise ValueError( + f"No processor registered for architecture: {hf_config.architectures}.\n" + f"Registered architectures: {[model_cls.__name__ for model_cls in PROCESSOR_MAPPING.keys()]}" + ) + return processor_cls(hf_config, server_args, processor, transport_mode, **kwargs) diff --git a/python/sglang/srt/managers/rust_server.py b/python/sglang/srt/managers/rust_server.py index 4c947ae9f..3275baeec 100644 --- a/python/sglang/srt/managers/rust_server.py +++ b/python/sglang/srt/managers/rust_server.py @@ -10,14 +10,17 @@ scheduler holds an `Optional[RustServer]` and delegates to it. from __future__ import annotations +import importlib import logging import os from array import array from itertools import chain -from typing import TYPE_CHECKING, Any, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Optional, Tuple import msgspec +from sglang.srt.environ import envs +from sglang.srt.managers.io_struct import TokenizedGenerateReqInput from sglang.srt.managers.utils import ( MsgpackDecodeError, compute_num_reserved_tokens, @@ -31,13 +34,315 @@ from sglang.srt.utils.flatten import ( from sglang.version import __version__ if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig from sglang.srt.managers.io_struct import BatchTokenIDOutput from sglang.srt.managers.scheduler import Scheduler from sglang.srt.server._core import Server + from sglang.srt.server_args import ServerArgs logger = logging.getLogger(__name__) +class NativeMmSpec(msgspec.Struct, frozen=True, kw_only=True): + """Resolved parameters of the native Rust MM pipeline for one model, + consumed by the Rust worker pool (:meth:`rust_json`) and the drain + adapter (:meth:`NativeMmHost.build_native_mm`).""" + + family: str + feature_shm: bool + image_token_id: int + patch_size: int + merge_size: int + temporal_patch_size: int + min_pixels: int + max_pixels: int + image_mean: Tuple[float, ...] + image_std: Tuple[float, ...] + # Which HF processor the Rust resize must reproduce bit-exactly, from + # `NativeMmHost.NATIVE_IMAGE_PROCESSORS`. + resample: str + vision_start_token_id: Optional[int] + vision_end_token_id: Optional[int] + video_token_id: Optional[int] + + # Used by the drain adapter only; every other field goes to Rust. + DRAIN_ONLY = ("vision_start_token_id", "vision_end_token_id", "video_token_id") + + @property + def feature_dim(self) -> int: + return 3 * self.temporal_patch_size * self.patch_size * self.patch_size + + def rust_json(self) -> str: + """The subset `sglang_mm::registry::pipeline_from_spec` parses.""" + fields = (f for f in self.__struct_fields__ if f not in self.DRAIN_ONLY) + return msgspec.json.encode({f: getattr(self, f) for f in fields}).decode() + + +class NativeMmFamily(msgspec.Struct, frozen=True, kw_only=True): + """The Python half of one Rust MM family (an arm of + `sglang_mm::registry::pipeline_from_spec`): which models it serves. + Supporting a new model family = one entry in :data:`NATIVE_MM_FAMILIES` + plus its Rust arm — the launch gate is data-driven.""" + + name: str + # The registered Python mm-processor the native pipeline replaces, as + # "module:Class". Compared by identity, so an + # SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE override still disables the native path. + mm_processor: str + # Model types whose image-only M-RoPE matches the family's fast path. + model_types: FrozenSet[str] + # HF image processors the native resize reproduces bit-exactly, each mapped + # to the `resample` the Rust pipeline must use (see `NativeMmSpec.resample`). + image_processors: Dict[str, str] + + def serves(self, mm_processor_cls: Any, model_type: Optional[str]) -> bool: + module_name, _, class_name = self.mm_processor.partition(":") + cls = getattr(importlib.import_module(module_name), class_name) + return mm_processor_cls is cls and model_type in self.model_types + + +NATIVE_MM_FAMILIES: Tuple[NativeMmFamily, ...] = ( + NativeMmFamily( + name="qwen_vl", + mm_processor="sglang.srt.multimodal.processors.qwen_vl:QwenVLImageProcessor", + model_types=frozenset( + ( + "qwen2_vl", + "qwen2_5_vl", + "qwen3_vl", + "qwen3_vl_moe", + "qwen3_5", + "qwen3_5_moe", + ) + ), + image_processors={ + "Qwen2VLImageProcessor": "aten_u8", + "Qwen2VLImageProcessorFast": "aten_u8", + "Qwen2VLImageProcessorPil": "pil", + }, + ), +) + + +def native_mm_family_for( + mm_processor_cls: Any, model_type: Optional[str] +) -> Optional[NativeMmFamily]: + """The declared family serving this model, or ``None`` — which + :meth:`RustServer.launch` turns into a hard error (no Python fallback).""" + return next( + (f for f in NATIVE_MM_FAMILIES if f.serves(mm_processor_cls, model_type)), None + ) + + +class NativeMmHost: + """Builds and validates the native Rust MM pipeline for one model. + + Construction registers the same ``mm_processor`` mapping the Python + TokenizerManager would build — not to process requests (the Rust worker pool + does that, GIL-free) but as the source of truth + :meth:`resolve_native_spec` resolves the pipeline parameters from. At drain + time :meth:`build_native_mm` wraps the Rust-produced buffers into the + scheduler's ``MultimodalProcessorOutput``. + + There is no Python fallback: a model without a native spec fails at launch, + and inputs outside the pipeline's scope are rejected per request. + """ + + # Rust mm-worker threads when --mm-processor-worker-num is 0. They are + # GIL-free, so unlike the Python processor pool more than one always helps. + AUTO_MM_WORKERS = 8 + + def __init__( + self, + *, + server_args: ServerArgs, + model_config: ModelConfig, + processor: Any = None, + ): + # Lazy: this class exists only for multimodal models under + # SGLANG_RUST_SERVER. + from sglang.srt.managers.multimodal_processor import import_processors + from sglang.srt.managers.tokenizer_manager import get_processor_wrapper + + self.server_args = server_args + self.model_config = model_config + # Worker threads == max concurrently-processed mm requests. + self.mm_workers = server_args.mm_processor_worker_num or self.AUTO_MM_WORKERS + + # The mapping the Python TokenizerManager builds in + # init_tokenizer_and_processor. The caller's already-loaded HF + # AutoProcessor is reused when available (identical construction args). + import_processors("sglang.srt.multimodal.processors") + if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get(): + import_processors(mm_process_pkg, overwrite=True) + self._processor = processor or get_processor_wrapper(self.server_args) + + def resolve_native_spec(self) -> Optional[NativeMmSpec]: + """The :class:`NativeMmSpec` for this model, or ``None`` when it has no + native pipeline (the launch gate turns that into a hard error). + + Carries only resolved settings — patch geometry, pixel limits, + normalization, token ids — never the HF config, and is conservative by + design: an unrecognized knob disables the native path rather than being + approximated.""" + from sglang.srt.managers.multimodal_processor import get_mm_processor_cls + + hf_config = self.model_config.hf_config + mm_processor_cls = get_mm_processor_cls( + hf_config, self.server_args, model_config=self.model_config + ) + family = native_mm_family_for( + mm_processor_cls, getattr(hf_config, "model_type", None) + ) + if family is None: + return None + ip = getattr(self._processor, "image_processor", None) + resample = family.image_processors.get(type(ip).__name__) + if resample is None: + return None + # The native pipeline always resizes, rescales by 1/255 and normalizes; + # Rust's fused normalize constants assume that factor. Anything else + # would silently produce different features. + stages = ("do_resize", "do_rescale", "do_normalize") + if not all(getattr(ip, stage, True) for stage in stages): + return None + if getattr(ip, "rescale_factor", None) != 1 / 255: + return None + + # `--mm-process-config {"image": {...}}`: only pixel-limit overrides are + # mirrored natively, anything else disables the pipeline. + image_overrides = dict( + (self.server_args.mm_process_config or {}).get("image", {}) + ) + if not set(image_overrides) <= {"min_pixels", "max_pixels"}: + return None + + size = getattr(ip, "size", None) or {} + min_pixels = image_overrides.get( + "min_pixels", getattr(ip, "min_pixels", None) or size.get("shortest_edge") + ) + max_pixels = image_overrides.get( + "max_pixels", getattr(ip, "max_pixels", None) or size.get("longest_edge") + ) + try: + spec = NativeMmSpec( + family=family.name, + feature_shm=self._use_feature_shm(), + image_token_id=hf_config.image_token_id, + patch_size=ip.patch_size, + merge_size=ip.merge_size, + temporal_patch_size=ip.temporal_patch_size, + min_pixels=int(min_pixels), + max_pixels=int(max_pixels), + image_mean=tuple(float(x) for x in ip.image_mean), + image_std=tuple(float(x) for x in ip.image_std), + resample=resample, + vision_start_token_id=getattr(hf_config, "vision_start_token_id", None), + vision_end_token_id=getattr(hf_config, "vision_end_token_id", None), + video_token_id=getattr(hf_config, "video_token_id", None), + ) + except (AttributeError, TypeError): # missing/odd processor attrs + return None + logger.info("rust server: native MM pipeline enabled (family=%s)", family.name) + return spec + + def _use_feature_shm(self) -> bool: + """Whether to park feature buffers in POSIX shm rather than inline. + + On exactly when the drained request is broadcast across TP ranks *and* + the receiver's ``unwrap_shm_features`` will materialize the stubs (its + gates: non-default tensor transport, no ``skip_tokenizer_init``). + + Inline, the whole ~20 MB/image buffer rides ``broadcast_pyobj`` serially + on the scheduler loop, so ranks 1..n start the TP-sharded ViT ~30 ms + after rank 0 and every rank then stalls that long at the first + collective. With shm the broadcast carries a ~100-byte stub and all ranks + map in parallel — the transport the Python TokenizerManager already uses. + Single-rank serving stays inline, where shm would only add a copy. + """ + from sglang.srt.managers.tokenizer_manager import ( + determine_tensor_transport_mode, + ) + + return ( + self.server_args.tp_size > 1 + and determine_tensor_transport_mode(self.server_args) != "default" + and not self.server_args.skip_tokenizer_init + ) + + @staticmethod + def build_native_mm(spec: NativeMmSpec, entry): + """Drain-time adapter: wrap the Rust-produced buffers of one ``MmHandoff`` + into the scheduler's ``MultimodalProcessorOutput``. Wrapping only — load, + resize, patchify, token expansion and M-RoPE all ran in Rust. + + Runs on the scheduler loop, so it must stay copy-free *and* hash-free: + ``take_mm``'s numpy arrays own the Rust buffers, ``torch.from_numpy`` just + views them, and each item's ``hash`` is worker-precomputed so + ``set_pad_value`` skips ``hash_feature``. Any per-byte work here — memcpy, + sha256, tens of MB per image-heavy request — measurably inflates every + running request's inter-token latency.""" + import torch + + from sglang.srt.managers.mm_utils import ShmPointerMMData + from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalProcessorOutput, + ) + + shm_names = entry.shm_names + if shm_names is None: + features = torch.from_numpy(entry.features.reshape(-1, spec.feature_dim)) + items = [] + row = 0 + for index, ((t, h, w), item_hash, offset) in enumerate( + zip(entry.grids, entry.hashes, entry.offsets) + ): + n = t * h * w + if shm_names is None: + feature = features[row : row + n] + else: + # The worker parked this item's buffer in a named POSIX + # segment (see `_use_feature_shm`). Build the stub in its + # post-`__setstate__` form: rank 0 never pickle-roundtrips its + # own copy, and `materialize()` needs the mapped view. + # Ownership of the unlink moved here with `take_mm`. + feature = ShmPointerMMData.__new__(ShmPointerMMData) + feature.__setstate__( + { + "shm_name": shm_names[index], + "shape": (n, spec.feature_dim), + "dtype": torch.float32, + "precomputed_hash": item_hash, + } + ) + items.append( + MultimodalDataItem( + modality=Modality.IMAGE, + feature=feature, + hash=item_hash, + offsets=[tuple(offset)], + model_specific_data={ + "image_grid_thw": torch.tensor([[t, h, w]], dtype=torch.long) + }, + ) + ) + row += n + if envs.SGLANG_MM_PRECOMPUTE_HASH.get(): + for item in items: + item.set_pad_value() + return MultimodalProcessorOutput( + mm_items=items, + im_token_id=spec.image_token_id, + im_start_id=spec.vision_start_token_id, + im_end_id=spec.vision_end_token_id, + video_token_id=spec.video_token_id, + mrope_positions=torch.from_numpy(entry.mrope.reshape(3, -1)), + mrope_position_delta=torch.tensor([[entry.mrope_delta]], dtype=torch.long), + ) + + class RustServer: """Owns the embedded multi-threaded Rust server (``sglang_server.Server``). @@ -45,8 +350,14 @@ class RustServer: all implemented as Rust threads in scheduler process. """ - def __init__(self, server: Server, max_per_poll: int = 256): + def __init__( + self, + server: Server, + mm_spec: Optional[NativeMmSpec] = None, + max_per_poll: int = 256, + ): self.server = server + self.mm_spec = mm_spec self._max_per_poll = max_per_poll @classmethod @@ -83,7 +394,13 @@ class RustServer: if dp_rank is not None: http_addr = f"{server_args.host}:{server_args.port + dp_rank}" - launch_cores, server_cores = cls._partition_cores() + launch_cores, server_cores = cls._partition_cores( + mm_workers=( + (server_args.mm_processor_worker_num or NativeMmHost.AUTO_MM_WORKERS) + if scheduler.model_config.is_multimodal + else 0 + ) + ) server = Server( # None -> run unpinned; the list carries the pinning decision. @@ -92,6 +409,40 @@ class RustServer: server_args_json=cls._build_server_args(scheduler), ) + # Multimodal models must have a native Rust pipeline — there is no Python + # fallback. + mm_spec = None + if scheduler.model_config.is_multimodal: + # New threads inherit the spawning thread's affinity, and this launch + # thread still holds the full mask. Narrow it first so every MM thread + # created below (the processor's executors, the Rust MM workers) stays + # off the scheduler's reserved cores, where MM preprocessing would + # preempt the scheduler loop and inflate inter-token latency. + if server_cores is not None: + try: + os.sched_setaffinity(0, set(server_cores)) + except OSError as e: + logger.warning( + "rust server: cannot confine mm threads to server cores: %s", e + ) + mm_host = NativeMmHost( + server_args=server_args, + model_config=scheduler.model_config, + processor=scheduler.processor, + ) + mm_spec = mm_host.resolve_native_spec() + if mm_spec is None: + supported = sorted( + set(chain.from_iterable(f.model_types for f in NATIVE_MM_FAMILIES)) + ) + raise RuntimeError( + "SGLANG_RUST_SERVER=1: no native Rust MM pipeline for " + f"model_type={scheduler.model_config.hf_config.model_type!r} " + f"(supported: {', '.join(supported)}; " + "images only). Unset SGLANG_RUST_SERVER to serve this model." + ) + server.start_mm_workers(mm_spec.rust_json(), mm_host.mm_workers) + # Narrow the scheduler thread only after the server threads are launched. if launch_cores is not None: try: @@ -111,7 +462,7 @@ class RustServer: dp_note, ) - return cls(server) + return cls(server, mm_spec=mm_spec) def wait_ingress(self, timeout_ms: int) -> None: """Block until a request is pushed into the in-process ring or the timeout @@ -162,6 +513,13 @@ class RustServer: ids.frombytes(ids_view[pos : pos + nbytes]) obj.input_ids = ids pos += nbytes + if self.mm_spec is not None and isinstance(obj, TokenizedGenerateReqInput): + # The buffers rode the Rust sidecar, parked before the ring push; + # wrapping them into tensors is the only Python step of the native + # path. `None` for a text-only request on a multimodal model. + native = self.server.take_mm(obj.rid) + if native is not None: + obj.mm_inputs = NativeMmHost.build_native_mm(self.mm_spec, native) out.append(obj) return out @@ -374,7 +732,9 @@ class RustServer: return msgspec.json.encode(server_args, enc_hook=str).decode("utf-8") @staticmethod - def _partition_cores() -> Tuple[Optional[List[int]], Optional[List[int]]]: + def _partition_cores( + mm_workers: int = 0, + ) -> Tuple[Optional[List[int]], Optional[List[int]]]: """Split this rank's allowed cores into ``(launch_cores, server_cores)``. Pure computation — no affinity is changed here. Both sets are a subset @@ -403,7 +763,17 @@ class RustServer: # effectively serial) and never take more than a quarter of the cores. reserve = min(2, len(allowed) // 4) launch_cores = allowed[:reserve] - server_cores = allowed[reserve:] + # Bound the pool instead of taking the whole remainder: this rank's + # allowed cores are usually the entire NUMA node, shared with the sibling + # TP ranks' processes, so an unbounded mask lets MM preprocessing bursts + # preempt a sibling's CUDA-launch thread and inflate every rank's forward + # through the TP collectives. Measured on Qwen3.5-35B TP4 at one 720p + # image per request: ~20 ms of ViT wall time on the worst sibling, gone + # once bounded. The budget covers the CPU-hot threads (MM workers, plus + # the I/O-shaped tokenizer/ingress/egress/api ones that are rarely all hot + # at once) and leaves the rest of the node to the scheduler ranks. + pool_budget = max(8, mm_workers + 4) + server_cores = allowed[reserve : reserve + pool_budget] logger.info( "rust server cores=%s, scheduler launch cores=%s", server_cores, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 80351f330..f3525b2ef 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -469,8 +469,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): import_processors("sglang.srt.multimodal.processors") if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get(): import_processors(mm_process_pkg, overwrite=True) - _processor = _get_processor_wrapper(server_args) - transport_mode = _determine_tensor_transport_mode(self.server_args) + _processor = get_processor_wrapper(server_args) + transport_mode = determine_tensor_transport_mode(self.server_args) # We want to parallelize the image pre-processing so we create an executor for it # We create mm_processor for any skip_tokenizer_init to make sure we still encode @@ -3699,7 +3699,7 @@ async def print_exception_wrapper(func): sys.exit(1) -def _get_processor_wrapper(server_args): +def get_processor_wrapper(server_args): try: processor = get_processor( server_args.tokenizer_path, @@ -3730,7 +3730,7 @@ def _get_processor_wrapper(server_args): return processor -def _determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode: +def determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode: is_cross_node = server_args.dist_init_addr if is_cross_node: diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 164c87904..74d9a404b 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -3690,6 +3690,7 @@ version = "0.1.0" dependencies = [ "async-stream", "axum 0.8.9", + "bytemuck", "bytes", "core_affinity", "dynamo-parsers", @@ -3700,12 +3701,15 @@ dependencies = [ "futures", "hf-hub", "itertools", + "libc", + "numpy", "pyo3", "regex-syntax", "rmp-serde", "rmpv", "serde", "serde_json", + "sglang-mm", "socket2 0.6.5", "thiserror", "tokio", diff --git a/rust/sglang-mm/src/common/fetch.rs b/rust/sglang-mm/src/common/fetch.rs index b391db2b2..a5d8ceb3b 100644 --- a/rust/sglang-mm/src/common/fetch.rs +++ b/rust/sglang-mm/src/common/fetch.rs @@ -15,49 +15,111 @@ use base64::Engine; /// payloads reject the request here). pub const MAX_FETCH_BYTES: u64 = 64 << 20; +/// Charge granularity of a streaming read: the most an in-flight source can +/// over-charge a shared [`ByteBudget`] by. +const CHUNK_BYTES: u64 = 256 << 10; + +/// A byte allowance shared by every source of one request, charged *as they +/// stream*, so concurrent fetches stop at their combined size rather than each +/// stopping at [`MAX_FETCH_BYTES`]. +#[derive(Debug)] +pub struct ByteBudget(std::sync::atomic::AtomicU64); + +impl ByteBudget { + pub fn new(total: u64) -> Self { + Self(std::sync::atomic::AtomicU64::new(total)) + } + + /// Claim `n` bytes, or `Err` once the allowance is spent. + fn claim(&self, n: u64) -> Result<(), ()> { + use std::sync::atomic::Ordering::{AcqRel, Acquire}; + self.0 + .fetch_update(AcqRel, Acquire, |left| left.checked_sub(n)) + .map(|_| ()) + .map_err(|_| ()) + } + + /// Give back bytes claimed for a chunk but not filled by the read. + fn release(&self, n: u64) { + self.0.fetch_add(n, std::sync::atomic::Ordering::AcqRel); + } +} + /// 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> { + fetch_bytes_budgeted(src, &ByteBudget::new(MAX_FETCH_BYTES)) +} + +/// [`fetch_bytes`] against a caller-owned allowance, for resolving several +/// sources under one whole-request bound. [`MAX_FETCH_BYTES`] still caps each. +pub fn fetch_bytes_budgeted(src: &str, budget: &ByteBudget) -> Result, String> { if src.starts_with("http://") || src.starts_with("https://") { - return http_get(src); + return http_get(src, budget); } if let Some(path) = src.strip_prefix("file://") { - return read_file(path); + return read_file(path, budget); } if src.starts_with('/') { - return read_file(src); + return read_file(src, budget); } 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); + return charge_decoded(b64(encoded)?, budget); } // Python treats any other string as bare base64. - b64(src) + charge_decoded(b64(src)?, budget) +} + +/// Base64 payloads are already resident in the request body — they cannot +/// amplify the way a download can, so they charge once decoded, not per chunk. +fn charge_decoded(decoded: Vec, budget: &ByteBudget) -> Result, String> { + budget + .claim(decoded.len() as u64) + .map_err(|()| over_budget("base64 payload"))?; + Ok(decoded) +} + +fn over_budget(what: &str) -> String { + format!("media fetch: {what}: exceeds the request media byte budget") } /// 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> { +fn read_file(path: &str, budget: &ByteBudget) -> Result, String> { let file = std::fs::File::open(path).map_err(|e| format!("media fetch: {path}: {e}"))?; - read_capped(file, path) + read_capped(file, path, budget) } -fn read_capped(reader: impl Read, what: &str) -> Result, String> { +/// Read to EOF, charging `budget` per chunk, so an oversized source stops +/// mid-stream instead of going fully resident first. +fn read_capped(mut reader: impl Read, what: &str, budget: &ByteBudget) -> Result, String> { + let too_big = || format!("media fetch: {what}: exceeds {MAX_FETCH_BYTES} bytes"); 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" - )); + loop { + // `+ 1`: read one byte past the cap, so oversized is detected, not truncated. + let want = CHUNK_BYTES.min(MAX_FETCH_BYTES + 1 - buf.len() as u64); + if want == 0 { + return Err(too_big()); + } + budget.claim(want).map_err(|()| over_budget(what))?; + let read = reader + .by_ref() + .take(want) + .read_to_end(&mut buf) + .map_err(|e| format!("media fetch: read {what}: {e}"))? as u64; + budget.release(want - read); + if buf.len() as u64 > MAX_FETCH_BYTES { + return Err(too_big()); + } + if read < want { + return Ok(buf); // short read == EOF + } } - Ok(buf) } fn b64(encoded: &str) -> Result, String> { @@ -163,7 +225,7 @@ fn in_ipv4_network(ip: std::net::Ipv4Addr, net: std::net::Ipv4Addr, bits: u32) - u32::from(ip) & mask == u32::from(net) & mask } -fn http_get(url: &str) -> Result, String> { +fn http_get(url: &str, budget: &ByteBudget) -> Result, String> { // Python: `int(os.getenv("REQUEST_TIMEOUT", "3"))` seconds per image GET. let timeout = std::env::var("REQUEST_TIMEOUT") .ok() @@ -178,7 +240,7 @@ fn http_get(url: &str) -> Result, String> { .timeout(std::time::Duration::from_secs(timeout)) .call() .map_err(|e| format!("media fetch: GET {url}: {e}"))?; - read_capped(resp.into_reader(), url) + read_capped(resp.into_reader(), url, budget) } #[cfg(test)] @@ -221,6 +283,38 @@ mod tests { assert!(err.contains("exceeds"), "{err}"); } + /// One budget spans sources: each fits alone, the set does not. + #[test] + fn shared_budget_spans_sources() { + let payload = base64::engine::general_purpose::STANDARD.encode([7u8; 4096]); + let budget = ByteBudget::new(6144); + assert_eq!(fetch_bytes_budgeted(&payload, &budget).unwrap().len(), 4096); + let err = fetch_bytes_budgeted(&payload, &budget).err().unwrap(); + assert!(err.contains("request media byte budget"), "{err}"); + } + + /// Unused claims come back, so small sources fit in a budget their + /// worst-case sizes would have exhausted. + #[test] + fn short_reads_release_their_claim() { + let path = std::env::temp_dir().join(format!("sglang-budget-{}", std::process::id())); + std::fs::write(&path, [0u8; 1024]).unwrap(); + let src = path.display().to_string(); + let budget = ByteBudget::new(CHUNK_BYTES + 4096); + for _ in 0..4 { + assert_eq!(fetch_bytes_budgeted(&src, &budget).unwrap().len(), 1024); + } + std::fs::remove_file(&path).ok(); + } + + /// The per-source cap holds even under a larger shared budget. + #[test] + fn per_source_cap_survives_a_large_budget() { + let budget = ByteBudget::new(MAX_FETCH_BYTES * 4); + let err = fetch_bytes_budgeted("/dev/zero", &budget).err().unwrap(); + assert!(err.contains(&format!("exceeds {MAX_FETCH_BYTES}")), "{err}"); + } + #[test] fn host_parsing_strips_userinfo_and_path() { assert_eq!( diff --git a/rust/sglang-mm/src/common/mod.rs b/rust/sglang-mm/src/common/mod.rs index 21593c61c..529aa04d7 100644 --- a/rust/sglang-mm/src/common/mod.rs +++ b/rust/sglang-mm/src/common/mod.rs @@ -80,16 +80,32 @@ mod python { use super::{decode_rgb, resize}; + /// `resample` names the implementation to reproduce: `"pil_lanczos"` (the + /// inkling default), `"pil_bicubic"`, or `"aten_u8"` (torchvision's uint8 + /// antialias bicubic). Exposed so the bit-exactness tests can cover each. #[pyfunction] + #[pyo3(signature = (arr, out_w, out_h, resample="pil_lanczos"))] pub fn resize_rgb<'py>( py: Python<'py>, arr: PyReadonlyArray3<'py, u8>, out_w: usize, out_h: usize, + resample: &str, ) -> PyResult>> { if out_w == 0 || out_h == 0 { return Err(PyValueError::new_err("output size must be positive")); } + let resample = match resample { + "pil_lanczos" => resize::Resample::Pil(resize::Filter::Lanczos), + "pil_bicubic" => resize::Resample::Pil(resize::Filter::Bicubic), + "aten_u8" => resize::Resample::AtenU8, + other => { + return Err(PyValueError::new_err(format!( + "unknown resample {other:?}; expected \"pil_lanczos\", \ + \"pil_bicubic\" or \"aten_u8\"" + ))); + } + }; let shape = arr.shape(); let (h, w, c) = (shape[0], shape[1], shape[2]); if c != 3 { @@ -101,7 +117,7 @@ mod python { .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)); + let out = py.detach(move || resize::resize_rgb(&data, h, w, out_h, out_w, resample)); Ok(out.into_pyarray(py)) } diff --git a/rust/sglang-mm/src/common/resize.rs b/rust/sglang-mm/src/common/resize.rs index cb8277933..bbe923d6f 100644 --- a/rust/sglang-mm/src/common/resize.rs +++ b/rust/sglang-mm/src/common/resize.rs @@ -1,17 +1,52 @@ use super::par; -const PRECISION_BITS: i32 = 32 - 8 - 2; +/// PIL's `PRECISION_BITS` for 8-bit images: weights quantized to i32. +const PIL_PRECISION_BITS: u32 = 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). + /// support 2.0, a = -0.5 — PIL `BICUBIC`. Bicubic, } +/// A resampler reproduced bit-exactly. Both share PIL's geometry, kernels and +/// per-pass u8 rounding, and differ only in how the weights are quantized. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Resample { + /// PIL `Image.resize`, i32 weights. + Pil(Filter), + /// ATen's uint8 antialias bicubic — torchvision `resize(antialias=True)` on + /// a uint8 tensor. i16 weights, so it rounds unlike `Pil(Bicubic)`. + AtenU8, +} + +impl Resample { + fn filter(self) -> Filter { + match self { + Resample::Pil(filter) => filter, + Resample::AtenU8 => Filter::Bicubic, + } + } + + /// Fixed-point precision for one axis's already-normalized weights. ATen + /// (`_compute_weights_precision`) takes the widest that stays inside i16. + fn precision(self, weights: &[f64]) -> u32 { + match self { + Resample::Pil(_) => PIL_PRECISION_BITS, + Resample::AtenU8 => { + let wmax = weights.iter().fold(0.0f64, |m, w| m.max(w.abs())); + (1..PIL_PRECISION_BITS) + .take_while(|&p| (0.5 + wmax * (1u64 << p) as f64) < (1 << 15) as f64) + .last() + .unwrap_or(1) + } + } + } +} + impl Filter { fn support(self) -> f64 { match self { @@ -60,9 +95,11 @@ struct Coeffs { bounds: Vec<(usize, usize)>, kk: Vec, ksize: usize, + prec: u32, } -fn precompute_coeffs(in_size: usize, out_size: usize, filter: Filter) -> Coeffs { +fn precompute_coeffs(in_size: usize, out_size: usize, resample: Resample) -> Coeffs { + let filter = resample.filter(); let scale = in_size as f64 / out_size as f64; let filterscale = if scale < 1.0 { 1.0 } else { scale }; let support = filter.support() * filterscale; @@ -97,7 +134,8 @@ fn precompute_coeffs(in_size: usize, out_size: usize, filter: Filter) -> Coeffs bounds[xx] = (xmin as usize, count); } - let factor = (1i64 << PRECISION_BITS) as f64; + let prec = resample.precision(&kkf); + let factor = (1i64 << prec) as f64; let kk = kkf .iter() .map(|&v| { @@ -108,17 +146,22 @@ fn precompute_coeffs(in_size: usize, out_size: usize, filter: Filter) -> Coeffs } }) .collect(); - Coeffs { bounds, kk, ksize } + Coeffs { + bounds, + kk, + ksize, + prec, + } } #[inline] -fn clip8(v: i32) -> u8 { - if v >= 1 << (PRECISION_BITS + 8) { +fn clip8(v: i32, prec: u32) -> u8 { + if v >= 1 << (prec + 8) { 255 } else if v <= 0 { 0 } else { - (v >> PRECISION_BITS) as u8 + (v >> prec) as u8 } } @@ -129,7 +172,7 @@ fn resample_horizontal(src: &[u8], h: usize, w: usize, out_w: usize, c: &Coeffs) 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]; + let mut s = [1i32 << (c.prec - 1); 3]; for (x, &coef) in k.iter().enumerate() { let p = (xmin + x) * 3; s[0] += src_row[p] as i32 * coef; @@ -137,9 +180,9 @@ fn resample_horizontal(src: &[u8], h: usize, w: usize, out_w: usize, c: &Coeffs) 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]); + row[o] = clip8(s[0], c.prec); + row[o + 1] = clip8(s[1], c.prec); + row[o + 2] = clip8(s[2], c.prec); } }); out @@ -151,7 +194,7 @@ fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec let (ymin, count) = c.bounds[yy]; let k = &c.kk[yy * c.ksize..yy * c.ksize + count]; for x in 0..w { - let mut s = [1i32 << (PRECISION_BITS - 1); 3]; + let mut s = [1i32 << (c.prec - 1); 3]; for (y, &coef) in k.iter().enumerate() { let p = ((ymin + y) * w + x) * 3; s[0] += src[p] as i32 * coef; @@ -159,27 +202,27 @@ fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec s[2] += src[p + 2] as i32 * coef; } let o = x * 3; - row[o] = clip8(s[0]); - row[o + 1] = clip8(s[1]); - row[o + 2] = clip8(s[2]); + row[o] = clip8(s[0], c.prec); + row[o + 1] = clip8(s[1], c.prec); + row[o + 2] = clip8(s[2], c.prec); } }); out } -/// PIL-exact separable resize of a flat HWC RGB buffer with the given filter. +/// Separable resize of a flat HWC RGB buffer, bit-exact against `resample`. /// /// 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( +pub fn resize_rgb( src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize, - filter: Filter, + resample: Resample, ) -> Vec { - par::in_pool(move || resize_passes(src, h, w, out_h, out_w, filter)) + par::in_pool(move || resize_passes(src, h, w, out_h, out_w, resample)) } fn resize_passes( @@ -188,28 +231,23 @@ fn resize_passes( w: usize, out_h: usize, out_w: usize, - filter: Filter, + resample: Resample, ) -> 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, filter); - let tmp = resample_horizontal(src, h, w, out_w, &ch); - 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, filter); - resample_horizontal(src, h, w, out_w, &ch) - } else if need_v { - let cv = precompute_coeffs(h, out_h, filter); - resample_vertical(src, w, out_h, &cv) - } else { - src.to_vec() + // Per-axis coefficients — and, under `AtenU8`, a per-axis precision. + let coeffs = |in_size, out_size| precompute_coeffs(in_size, out_size, resample); + match (out_w != w, out_h != h) { + (true, true) => { + let tmp = resample_horizontal(src, h, w, out_w, &coeffs(w, out_w)); + resample_vertical(&tmp, out_w, out_h, &coeffs(h, out_h)) + } + (true, false) => resample_horizontal(src, h, w, out_w, &coeffs(w, out_w)), + (false, true) => resample_vertical(src, w, out_h, &coeffs(h, out_h)), + (false, false) => 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) + resize_rgb(src, h, w, out_h, out_w, Resample::Pil(Filter::Lanczos)) } pub fn scaled_dims(w: usize, h: usize, frac: Option, cap: Option) -> (usize, usize) { diff --git a/rust/sglang-mm/src/qwen_vl/mod.rs b/rust/sglang-mm/src/qwen_vl/mod.rs index d32ac4756..55933dce3 100644 --- a/rust/sglang-mm/src/qwen_vl/mod.rs +++ b/rust/sglang-mm/src/qwen_vl/mod.rs @@ -34,21 +34,60 @@ pub struct QwenVlSpec { pub max_pixels: usize, pub image_mean: [f32; 3], pub image_std: [f32; 3], + #[serde(default)] + pub resample: Resampler, +} + +/// The HF image processor the pipeline must match bit-exactly. Defaults to the +/// one a default server runs. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Resampler { + /// `Qwen2VLImageProcessor` / `…Fast` — torchvision on a uint8 tensor. + #[default] + AtenU8, + /// `Qwen2VLImageProcessorPil`, behind `--disable-fast-image-processor`. + Pil, +} + +impl From for resize::Resample { + fn from(r: Resampler) -> Self { + match r { + Resampler::AtenU8 => resize::Resample::AtenU8, + Resampler::Pil => resize::Resample::Pil(resize::Filter::Bicubic), + } + } } pub struct QwenVlProcessor { spec: QwenVlSpec, - /// Per-channel u8 → normalized-f32 lookup: `(v/255 - mean) / std`. + /// Per-channel u8 → normalized-f32 lookup; see [`normalize_lut`]. lut: [[f32; 256]; 3], } +/// `1 / rescale_factor`; `resolve_native_spec` rejects any other factor. +const INV_RESCALE: f32 = 255.0; + +/// u8 → normalized f32, rounded as the mirrored processor rounds. The slow one +/// rescales then normalizes; the fast one folds the rescale into mean/std first +/// (`_fuse_mean_std_and_rescale_factor`), which differs on 128 of the 256 inputs. +fn normalize_lut(resample: Resampler, mean: f32, std: f32) -> [f32; 256] { + match resample { + Resampler::Pil => core::array::from_fn(|v| (v as f32 / INV_RESCALE - mean) / std), + Resampler::AtenU8 => { + let (mean, std) = (mean * INV_RESCALE, std * INV_RESCALE); + core::array::from_fn(|v| (v as f32 - mean) / std) + } + } +} + 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]) + normalize_lut(spec.resample, spec.image_mean[c], spec.image_std[c]) }); Ok(Self { spec, lut }) } @@ -127,7 +166,7 @@ impl MmFamilyProcessor for QwenVlProcessor { )?; let resized; let data = if (th, tw) != (h, w) { - resized = resize::resize_rgb_filter(rgb, h, w, th, tw, resize::Filter::Bicubic); + resized = resize::resize_rgb(rgb, h, w, th, tw, self.spec.resample.into()); &resized } else { rgb.as_slice() @@ -513,6 +552,22 @@ mod tests { max_pixels: 1 << 30, image_mean: [0.0; 3], image_std: [1.0; 3], + resample: Resampler::default(), + } + } + + /// The fused and unfused normalize forms are not interchangeable: with + /// mean = std = 0.5 they disagree on 128 of the 256 u8 inputs, so picking + /// the wrong one silently costs bit-exactness with the HF processor. + #[test] + fn normalize_lut_differs_per_resampler() { + let pil = normalize_lut(Resampler::Pil, 0.5, 0.5); + let aten = normalize_lut(Resampler::AtenU8, 0.5, 0.5); + assert_eq!(pil.iter().zip(aten).filter(|(p, a)| *p != a).count(), 128); + // Both still span [-1, 1] — this is rounding, not a scale error. + for lut in [pil, aten] { + assert_eq!(lut[0], -1.0); + assert_eq!(lut[255], 1.0); } } diff --git a/rust/sglang-mm/tests/test_resize_parity.py b/rust/sglang-mm/tests/test_resize_parity.py index e09f93a2d..407dede65 100644 --- a/rust/sglang-mm/tests/test_resize_parity.py +++ b/rust/sglang-mm/tests/test_resize_parity.py @@ -34,15 +34,34 @@ def py_scaled_dims( return scale(width), scale(height) -def pil_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray: - return np.array( - Image.fromarray(arr).resize((tw, th), resample=Image.Resampling.LANCZOS), - dtype=np.uint8, +def pil_resize(arr: np.ndarray, tw: int, th: int, filter=Image.Resampling.LANCZOS): + return np.array(Image.fromarray(arr).resize((tw, th), resample=filter), np.uint8) + + +def tv_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray: + """torchvision's uint8 antialias bicubic — ATen's fixed-point kernel.""" + import torch + from torchvision.transforms.v2 import functional as F + + tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) + out = F.resize( + tensor, [th, tw], interpolation=F.InterpolationMode.BICUBIC, antialias=True ) + return out[0].permute(1, 2, 0).numpy() -def rs_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray: - return _rs_common.resize_rgb(arr, tw, th).reshape(th, tw, 3) +# Every resampler the Rust resize claims, and its reference: `aten_u8` for a +# default server, `pil_bicubic` for --disable-fast-image-processor, `pil_lanczos` +# for inkling. +REFERENCES = { + "pil_lanczos": lambda a, tw, th: pil_resize(a, tw, th, Image.Resampling.LANCZOS), + "pil_bicubic": lambda a, tw, th: pil_resize(a, tw, th, Image.Resampling.BICUBIC), + "aten_u8": tv_resize, +} + + +def rs_resize(arr, tw: int, th: int, resample: str = "pil_lanczos") -> np.ndarray: + return _rs_common.resize_rgb(arr, tw, th, resample).reshape(th, tw, 3) CASES = [ @@ -58,13 +77,36 @@ CASES = [ ] +@pytest.mark.parametrize("resample", sorted(REFERENCES)) @pytest.mark.parametrize( "h,w,th,tw", CASES, ids=[f"{h}x{w}->{th}x{tw}" for h, w, th, tw in CASES] ) -def test_resize_bit_exact(h, w, th, tw): +def test_resize_bit_exact(h, w, th, tw, resample): rng = np.random.default_rng(h * 10000 + w) arr = rng.integers(0, 256, (h, w, 3), dtype=np.uint8) - np.testing.assert_array_equal(rs_resize(arr, tw, th), pil_resize(arr, tw, th)) + np.testing.assert_array_equal( + rs_resize(arr, tw, th, resample), REFERENCES[resample](arr, tw, th) + ) + + +@pytest.mark.parametrize("resample", sorted(REFERENCES)) +def test_resize_bit_exact_random_sweep(resample): + """`aten_u8`'s weight precision varies with the scale factor, so the fixed + cases above are not enough coverage on their own.""" + rng = np.random.default_rng(7) + for h, w, th, tw in rng.integers(1, 200, (40, 4)): + arr = rng.integers(0, 256, (h, w, 3), dtype=np.uint8) + np.testing.assert_array_equal( + rs_resize(arr, tw, th, resample), + REFERENCES[resample](arr, tw, th), + err_msg=f"{h}x{w}->{th}x{tw} under {resample}", + ) + + +def test_unknown_resample_rejected(): + arr = np.zeros((4, 4, 3), dtype=np.uint8) + with pytest.raises(ValueError, match="unknown resample"): + _rs_common.resize_rgb(arr, 2, 2, "nearest") def test_scaled_dims_sweep(): diff --git a/rust/sglang-server/Cargo.toml b/rust/sglang-server/Cargo.toml index d0e2afc20..2d7f514e3 100644 --- a/rust/sglang-server/Cargo.toml +++ b/rust/sglang-server/Cargo.toml @@ -43,14 +43,25 @@ dynamo-parsers = "7.0.1" dynamo-protocols = "5.1.0" dynamo-renderer = "5.0.0" flume = "0.12.0" +# Safe POD slice casts (feature buffers viewed as bytes for the shm copy). +bytemuck = "1" itertools = "0.14" hf-hub = { version = "0.4", default-features = false } +# POSIX shm for the MM feature fan-out (`mm::ShmSegment`). +libc = "0.2" +# Same major as the workspace pyo3: the zero-copy MM drain (`take_mm`) moves +# Rust vectors into numpy arrays. +numpy = "0.29.0" rmp-serde = "1" rmpv = { version = "1", features = ["with-serde"] } # Pinned EXACTLY: this crate's accepted grammar defines the # "anything Rust admits, Python can compile" invariant in `message::sampling`. # A minor bump can widen it and silently reopen a scheduler-killing hole. regex-syntax = "=0.8.11" +# The pure-Rust core of the MM pipeline. `default-features = false` drops the +# pyo3 bindings so it links as a plain rlib; renamed so `use sglang_mm::…` reads +# naturally while the crate keeps its own artifact name in the shared target/. +sglang_mm = { package = "sglang-mm", path = "../sglang-mm", default-features = false } [dev-dependencies] # `Router::oneshot` for handler-level router tests. diff --git a/rust/sglang-server/src/api_server.rs b/rust/sglang-server/src/api_server.rs index ff3b2cf4e..c078549d1 100644 --- a/rust/sglang-server/src/api_server.rs +++ b/rust/sglang-server/src/api_server.rs @@ -10,6 +10,7 @@ mod log; mod native_api; mod openai; mod pd_bootstrap; +mod prefetch; mod submit; use std::sync::Arc; diff --git a/rust/sglang-server/src/api_server/native_api.rs b/rust/sglang-server/src/api_server/native_api.rs index 6d53b9b89..bbe3a92b0 100644 --- a/rust/sglang-server/src/api_server/native_api.rs +++ b/rust/sglang-server/src/api_server/native_api.rs @@ -154,7 +154,7 @@ async fn generate( let stream = body.stream; // Fan `text`/`input_ids`/`sampling_params` (scalar or list) into per-request // payloads. `is_batch` = list form → the response is a JSON array. - let (payloads, is_batch) = match body.into_requests() { + let (mut payloads, is_batch) = match body.into_requests() { Ok(v) => v, // The error carries its own status (a bad batch is `Validation` → 400). Err(e) => { @@ -162,6 +162,11 @@ async fn generate( return pre_submit_error(code, &e.to_string(), stream); } }; + // Media I/O (URL downloads, file reads) happens here, on the API runtime + // — never on the MM worker pool (see `prefetch`). + if let Err(e) = super::prefetch::prefetch_all(&mut payloads).await { + return pre_submit_error(StatusCode::BAD_REQUEST, &e, stream); + } if !is_batch { // `into_requests` guarantees exactly one payload for a non-batch body. let payload = payloads diff --git a/rust/sglang-server/src/api_server/prefetch.rs b/rust/sglang-server/src/api_server/prefetch.rs new file mode 100644 index 000000000..4d5af2436 --- /dev/null +++ b/rust/sglang-server/src/api_server/prefetch.rs @@ -0,0 +1,199 @@ +//! Resolve I/O-backed media sources on the API runtime, before MM dispatch. +//! +//! The MM worker pool is fixed, core-pinned CPU capacity: a slow image host — or +//! a file on a hanging network mount — must never occupy it, and a request's +//! images must download concurrently, not in `n * REQUEST_TIMEOUT`. URLs and +//! file paths resolve here through `sglang-mm`'s `fetch_bytes_budgeted` (one +//! owner for proxy/timeout/cap semantics) and ride out-of-band as +//! [`crate::message::MmData::prefetched`], which +//! [`crate::message::mm_payload::to_mm_input`] swaps back in. + +use std::sync::Arc; + +use bytes::Bytes; +use sglang_mm::common::fetch::{ByteBudget, fetch_bytes_budgeted}; +use sglang_mm::driver::{MAX_ITEMS_PER_REQUEST, MAX_REQUEST_BYTES}; +use tokio::sync::Semaphore; + +use crate::message::mm_payload::{io_sources, item_count}; +use crate::message::{GenerateRequest, MmData}; + +/// Global bound on concurrent media fetches across all in-flight requests; +/// excess acquisitions queue on the semaphore without holding a thread. +static PERMITS: Semaphore = Semaphore::const_new(32); + +/// Fill [`MmData::prefetched`] for every request, all fetches across the batch +/// concurrent. Any failure rejects the call (a 400, as on the Python path). +/// +/// The driver's budgets ([`MAX_ITEMS_PER_REQUEST`], [`MAX_REQUEST_BYTES`]) are +/// enforced *here* rather than in `sglang_mm::driver::process`, where 64 sources +/// of 64 MiB would already be resident. The driver keeps its own checks as the +/// backstop for callers without a prefetch layer. +pub async fn prefetch_all(requests: &mut [GenerateRequest]) -> Result<(), String> { + // The item budget rejects before a single byte is fetched. + let plan = |mm: &Option>| -> Result, String> { + let Some(image_data) = mm.as_deref().and_then(|m| m.image_data.as_ref()) else { + return Ok(Vec::new()); + }; + if item_count(image_data) > MAX_ITEMS_PER_REQUEST { + return Err(format!( + "multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items" + )); + } + Ok(io_sources(image_data)) + }; + let plans = requests + .iter() + .map(|r| plan(&r.mm)) + .collect::, String>>()?; + let fetches = plans + .into_iter() + .map(|sources| fetch_ordered(sources, MAX_REQUEST_BYTES)); + let fetched = futures::future::try_join_all(fetches).await?; + for (req, bytes) in requests.iter_mut().zip(fetched) { + if !bytes.is_empty() { + req.mm.as_mut().expect("sources came from mm").prefetched = bytes; + } + } + Ok(()) +} + +/// Resolve one request's sources concurrently (globally bounded), in order, +/// against one shared `total_bytes` allowance. Overflow rejects mid-download and +/// `try_join_all` drops the rest, so queued sources never start. +async fn fetch_ordered(sources: Vec, total_bytes: u64) -> Result, String> { + let budget = Arc::new(ByteBudget::new(total_bytes)); + futures::future::try_join_all(sources.into_iter().map(|src| { + let budget = Arc::clone(&budget); + async move { + let _permit = PERMITS.acquire().await.expect("semaphore never closed"); + // Blocking I/O: parks a lazily-spawned blocking-pool thread, never + // an API worker. Those threads are pinned round-robin over the api + // core set (see `on_thread_start` in `runtime::start`) — off the + // CPU-bound stages, and mostly I/O-parked, so sharing is fine. + tokio::task::spawn_blocking(move || fetch_bytes_budgeted(&src, &budget)) + .await + .map_err(|e| format!("media prefetch: {e}"))? + .map(Bytes::from) + } + })) + .await +} + +#[cfg(test)] +mod tests { + use rmpv::Value; + + use super::*; + + fn serve(bodies: Vec>) -> std::net::SocketAddr { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + std::thread::spawn(move || { + for body in bodies { + 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(); + } + }); + addr + } + + fn mm_request(image_data: Value) -> GenerateRequest { + GenerateRequest { + mm: Some(Box::new(MmData { + image_data: Some(image_data), + ..Default::default() + })), + ..Default::default() + } + } + + /// URLs and file paths resolve concurrently into `prefetched` in source + /// order; CPU-only sources and mm-free requests are untouched. + #[tokio::test] + async fn resolves_io_sources() { + let addr = serve(vec![b"one".to_vec(), b"two".to_vec()]); + let path = std::env::temp_dir().join(format!("sglang-prefetch-{}", std::process::id())); + std::fs::write(&path, b"zzz").unwrap(); + let mut requests = vec![ + mm_request(Value::Array(vec![ + Value::from(format!("http://{addr}/a.png")), + Value::from("data:image/png;base64,x"), + Value::from(format!("http://{addr}/b.png")), + Value::from(path.display().to_string()), + ])), + GenerateRequest::default(), + ]; + prefetch_all(&mut requests).await.unwrap(); + std::fs::remove_file(&path).ok(); + let fetched = &requests[0].mm.as_ref().unwrap().prefetched; + // The one-shot server answers in accept order, so contents may swap + // between the two URLs; all three bodies must arrive. + let mut got: Vec<&[u8]> = fetched.iter().map(|b| b.as_ref()).collect(); + got.sort(); + assert_eq!(got, vec![b"one".as_ref(), b"two".as_ref(), b"zzz".as_ref()]); + assert!(requests[1].mm.is_none()); + } + + #[tokio::test] + async fn failed_download_rejects() { + let mut requests = vec![mm_request(Value::from("http://127.0.0.1:1/nope.png"))]; + let err = prefetch_all(&mut requests).await.err().unwrap(); + assert!(err.contains("media fetch"), "{err}"); + } + + /// The item budget rejects before any source is touched: all of these would + /// fail to fetch, so a fetch error would prove fetching started. + #[tokio::test] + async fn item_budget_rejects_before_fetching() { + let sources: Vec = (0..=MAX_ITEMS_PER_REQUEST) + .map(|i| Value::from(format!("/definitely/not/here-{i}.png"))) + .collect(); + let mut requests = vec![mm_request(Value::Array(sources))]; + let err = prefetch_all(&mut requests).await.err().unwrap(); + assert_eq!( + err, + format!("multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items") + ); + assert!(requests[0].mm.as_ref().unwrap().prefetched.is_empty()); + } + + /// Sources legal alone but collectively over the limit are rejected while + /// downloading, not once every body is resident. + #[tokio::test] + async fn byte_budget_is_shared_across_sources() { + let addr = serve(vec![vec![b'a'; 4096], vec![b'b'; 4096]]); + let sources = vec![ + format!("http://{addr}/a.png"), + format!("http://{addr}/b.png"), + ]; + // Room for one body, not both. + let err = fetch_ordered(sources, 6144).await.err().unwrap(); + assert!(err.contains("request media byte budget"), "{err}"); + } + + /// ...and a fitting set still fetches: the budget never over-rejects. + #[tokio::test] + async fn byte_budget_admits_a_fitting_request() { + let addr = serve(vec![vec![b'a'; 4096], vec![b'b'; 4096]]); + let sources = vec![ + format!("http://{addr}/a.png"), + format!("http://{addr}/b.png"), + ]; + let fetched = fetch_ordered(sources, MAX_REQUEST_BYTES).await.unwrap(); + assert_eq!(fetched.iter().map(|b| b.len()).sum::(), 8192); + } +} diff --git a/rust/sglang-server/src/fsm.rs b/rust/sglang-server/src/fsm.rs index 2e54f7b6a..5777891dd 100644 --- a/rust/sglang-server/src/fsm.rs +++ b/rust/sglang-server/src/fsm.rs @@ -14,8 +14,6 @@ use crate::error::Error; -// `Failed(Error)` carries the cause for observability even where it isn't read -// back yet; `EncodeDone` belongs to the deferred Encoder edge. #[derive(Debug, Clone)] pub enum RequestState { Received, @@ -41,8 +39,8 @@ pub enum RequestState { /// Outcome of validation, selecting the ingress branch. #[derive(Debug, Clone, Copy)] pub enum ValidationOutcome { - /// Has multimodal inputs → Encoding. Deferred: no encoder yet. - #[allow(dead_code)] + /// Has multimodal inputs → Encoding, where an MM worker runs the native + /// pipeline and returns the final expanded `input_ids`. HasMultimodal, /// Plain text → Tokenizing. NeedsTokenize, @@ -52,7 +50,6 @@ pub enum ValidationOutcome { /// Events that drive transitions. Each variant maps 1:1 to an edge in the /// design's transition table. -#[allow(dead_code)] // EncodeDone is the deferred Encoder edge. #[derive(Debug)] pub enum Event { // --- ingress --- @@ -125,7 +122,11 @@ impl RequestState { (Normalizing, Validated(HasMultimodal)) => Encoding, (Normalizing, Validated(NeedsTokenize)) => Tokenizing, (Normalizing, Validated(AlreadyTokenized)) => PreSendValidating, - (Encoding, EncodeDone) => Tokenizing, + // The MM worker returns the *final* placeholder-expanded input_ids, + // so an encoded request skips the tokenizer pool — but not the + // pre-send checks: expanded image tokens count against the same + // input + max_new_tokens ceiling as tokenized text. + (Encoding, EncodeDone) => PreSendValidating, // Every ingress branch funnels through the pre-send checks, so they // run exactly once per request no matter how it got its ids. (Tokenizing, TokenizeDone) => PreSendValidating, diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index 767a14c5c..b2aab2c54 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -17,6 +17,7 @@ mod error; mod fsm; mod ids; mod message; +mod mm; mod ring; mod runtime; mod tokenizer; @@ -31,6 +32,21 @@ use pyo3::types::PyBytes; use crate::runtime::{Runtime, RuntimeConfig}; +/// One drained MM result (see [`Server::take_mm`]). Exactly one of +/// `features`/`shm_names` is `Some`: inline features for single-rank serving +/// (zero-copy into numpy), or one POSIX segment name per item when the scheduler +/// broadcasts across TP ranks and Python wraps each in a `ShmPointerMMData`. +#[pyclass(frozen, get_all)] +struct MmHandoff { + features: Option>>, + shm_names: Option>, + grids: Vec<(u32, u32, u32)>, + hashes: Vec, + offsets: Vec<(u32, u32)>, + mrope: Py>, + mrope_delta: i64, +} + /// Columnar ingress batch handed to Python by [`Server::recv_requests`]. /// `frozen`: immutable snapshot, so field access never contends on a borrow. #[pyclass(frozen, get_all)] @@ -199,6 +215,54 @@ impl Server { self.push_frame(py, crate::message::frame_egress_error(rid, message)) } + /// Spawn the MM worker pool for the pipeline in `spec_json` (built from the + /// resolved processor config; see `NativeMmHost.resolve_native_spec`). + /// Image-only requests are processed entirely in Rust and parked for + /// [`Server::take_mm`]; anything the pipeline cannot serve is rejected back to + /// the client — there is no Python fallback. + fn start_mm_workers(&self, spec_json: &str, workers: usize) -> PyResult<()> { + let ctx = mm::Context::new( + spec_json, + self.rt.tokenizer.clone(), + self.rt.mm_sidecar.clone(), + ) + .map_err(PyErr::new::)?; + self.rt.spawn_mm_pool(workers, std::sync::Arc::new(ctx)); + Ok(()) + } + + /// Pop the MM result for `rid` — parked strictly before the request reached + /// the ingress ring — or `None` if there is none. The numeric buffers become + /// 1-D numpy arrays that take **ownership** of the Rust vectors, no copy. + /// + /// Runs on the scheduler loop (`RustServer.drain`, under the GIL) between + /// decode steps, so any per-byte work here — memcpy or hashing, tens of MB + /// per image-heavy request — would stall every running request's ITL. Hence + /// the worker-precomputed `hashes`. + fn take_mm(&self, py: Python<'_>, rid: &str) -> Option { + use numpy::IntoPyArray; + + let res = self.rt.mm_sidecar.take(rid)?; + let (features, shm_names) = match res.features { + mm::FeatureStore::Inline(v) => (Some(v.into_pyarray(py).unbind()), None), + // The segments — and the duty to unlink — move to Python here; + // `materialize()` unlinks after the post-broadcast clone on each rank. + mm::FeatureStore::Shm(segments) => ( + None, + Some(segments.into_iter().map(|s| s.into_name()).collect()), + ), + }; + Some(MmHandoff { + features, + shm_names, + grids: res.grids.iter().map(|g| (g[0], g[1], g[2])).collect(), + hashes: res.hashes, + offsets: res.offsets, + mrope: res.mrope.into_pyarray(py).unbind(), + mrope_delta: res.mrope_delta, + }) + } + /// Signal all threads to stop (best effort). fn shutdown(&self) { self.rt.request_shutdown(); @@ -245,5 +309,6 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { .try_init(); m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/rust/sglang-server/src/message.rs b/rust/sglang-server/src/message.rs index 74af4a41d..5747e926c 100644 --- a/rust/sglang-server/src/message.rs +++ b/rust/sglang-server/src/message.rs @@ -11,6 +11,7 @@ mod egress; mod finish_reason; mod io_struct; +pub mod mm_payload; mod request; mod sampling; mod types; @@ -22,7 +23,10 @@ pub use egress::{ }; pub use finish_reason::Matched; pub(crate) use io_struct::{AbortReq, ControlRequest, GetInternalStateReq}; -pub use request::{GenerateBody, GenerateRequest, RequestKind}; +pub use request::{GenerateBody, GenerateRequest, MmRequest, MmWorkItem, RequestKind}; +// Constructed directly only by tests: `api_server::prefetch` fills its +// `prefetched` field, everything else gets it packed inside a `GenerateRequest`. +pub use request::MmData; pub(crate) use sampling::{SamplingParams, SamplingParamsInput}; pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds}; diff --git a/rust/sglang-server/src/message/mm_payload.rs b/rust/sglang-server/src/message/mm_payload.rs new file mode 100644 index 000000000..397f88034 --- /dev/null +++ b/rust/sglang-server/src/message/mm_payload.rs @@ -0,0 +1,222 @@ +//! Convert a parked request's [`MmWorkItem`] into the typed [`MmInput`] the +//! `sglang-mm` driver consumes — an in-process handoff, nothing serialized. +//! +//! Every `Err` rejects the request back to the client; the message says whether +//! the input is malformed or merely outside the pipeline's scope (video/audio, +//! precomputed features, …). + +use bytes::Bytes; +use rmpv::Value; +use sglang_mm::driver::{ImageSource, MmInput}; + +use super::request::MmWorkItem; + +/// True for sources the API layer must resolve before MM dispatch: I/O — network +/// *or* disk, since a network mount can hang past any HTTP timeout — never runs +/// on the fixed MM worker pool (see `api_server::prefetch`). `data:` and bare +/// base64 are pure CPU and stay on the worker. Lives next to `collect_images` so +/// the prefetch walk and the parse walk cannot drift. +pub fn is_io_source(src: &str) -> bool { + src.starts_with("http://") + || src.starts_with("https://") + || src.starts_with("file://") + || src.starts_with('/') +} + +/// The I/O-backed sources of an `image_data` value, in `collect_images` order. +pub fn io_sources(value: &Value) -> Vec { + let mut out = Vec::new(); + let mut walk = |value: &Value| { + if let Some(src) = value.as_str().filter(|s| is_io_source(s)) { + out.push(src.to_owned()); + } + }; + if let Value::Array(values) = value { + values.iter().for_each(&mut walk); + } else { + walk(value); + } + out +} + +/// How many media items an `image_data` value contributes, walked the way +/// [`collect_images`] walks it, so the item budget can reject before fetching. +pub fn item_count(value: &Value) -> usize { + match value { + Value::Nil => 0, + Value::Array(values) => values.iter().map(item_count).sum(), + _ => 1, + } +} + +/// I/O-backed sources are swapped for their `work.prefetched` bytes (in +/// [`io_sources`] order); one left without an entry is an internal error here, +/// never a fetch. +pub fn to_mm_input(work: MmWorkItem) -> Result { + let present = |v: &Option| v.as_ref().is_some_and(value_present); + if present(&work.video_data) || present(&work.audio_data) { + return Err("unsupported modality: video/audio input".into()); + } + let mut images = Vec::new(); + if let Some(image_data) = &work.image_data { + collect_images(image_data, &mut work.prefetched.iter(), &mut images)?; + } + if images.is_empty() { + return Err("no raw image sources in mm input".into()); + } + Ok(MmInput { + text: work.text, + input_ids: work.input_ids, + images, + }) +} + +fn collect_images( + value: &Value, + prefetched: &mut std::slice::Iter, + out: &mut Vec, +) -> Result<(), String> { + match value { + Value::Nil => Ok(()), + Value::String(value) => { + let value = value + .as_str() + .ok_or_else(|| "non-utf8 image source".to_string())?; + if is_io_source(value) { + let bytes = prefetched + .next() + .ok_or_else(|| "I/O-backed image source was not prefetched".to_string())?; + out.push(ImageSource::Bytes(bytes.to_vec())); + } else { + out.push(ImageSource::String(value.to_owned())); + } + Ok(()) + } + Value::Binary(value) => { + out.push(ImageSource::Bytes(value.clone())); + Ok(()) + } + Value::Array(values) => { + for value in values { + match value { + Value::String(_) | Value::Binary(_) | Value::Nil => { + collect_images(value, prefetched, out)? + } + _ => { + return Err("unsupported image_data shape: nested/typed item".into()); + } + } + } + Ok(()) + } + _ => Err("unsupported image_data shape".into()), + } +} + +/// Rust mirror of Python `has_valid_data`: `nil` and (recursively) empty or +/// all-nil lists don't count as multimodal input. Shared with the ingress +/// `has_multimodal` check so routing and parsing cannot drift. +pub fn value_present(value: &Value) -> bool { + match value { + Value::Nil => false, + Value::Array(values) => values.iter().any(value_present), + _ => true, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn image_work(image: Value) -> MmWorkItem { + MmWorkItem { + text: Some("prompt".into()), + image_data: Some(image), + ..Default::default() + } + } + + #[test] + fn converts_string_and_list_images() { + let one = to_mm_input(image_work(Value::from("data:image/png;base64,x"))).unwrap(); + assert_eq!(one.images.len(), 1); + let many = to_mm_input(image_work(Value::Array(vec![ + Value::from("a"), + Value::from("b"), + ]))) + .unwrap(); + assert_eq!(many.images.len(), 2); + } + + #[test] + fn unsupported_modalities_and_shapes_rejected() { + let video = MmWorkItem { + video_data: Some(Value::from("video.mp4")), + ..Default::default() + }; + assert!(to_mm_input(video).err().unwrap().contains("video/audio")); + + let dict = Value::Map(vec![(Value::from("format"), Value::from("x"))]); + assert!( + to_mm_input(image_work(Value::Array(vec![dict]))) + .err() + .unwrap() + .contains("image_data shape") + ); + } + + #[test] + fn empty_video_audio_lists_are_not_modalities() { + // Mirrors Python `has_valid_data`: nil / empty lists don't count. + let work = MmWorkItem { + input_ids: Some(vec![1]), + image_data: Some(Value::from("a")), + video_data: Some(Value::Array(vec![])), + audio_data: Some(Value::Array(vec![Value::Array(vec![])])), + ..Default::default() + }; + assert_eq!(to_mm_input(work).unwrap().images.len(), 1); + } + + /// I/O-backed sources (URLs, file paths) take their prefetched bytes in walk + /// order; one left unfetched errors, so no I/O can reach an MM worker. + #[test] + fn io_sources_use_prefetched_bytes() { + let image = Value::Array(vec![ + Value::from("http://a/x.png"), + Value::from("data:image/png;base64,x"), + Value::from("/mnt/nfs/y.png"), + ]); + assert_eq!(io_sources(&image), vec!["http://a/x.png", "/mnt/nfs/y.png"]); + + let mut work = image_work(image.clone()); + work.prefetched = vec![Bytes::from_static(b"aa"), Bytes::from_static(b"bb")]; + let input = to_mm_input(work).unwrap(); + let as_bytes = |i: usize| match &input.images[i] { + ImageSource::Bytes(b) => b.as_slice(), + other => panic!("expected bytes, got {other:?}"), + }; + assert_eq!(as_bytes(0), b"aa"); + assert_eq!(as_bytes(2), b"bb"); + assert!(matches!(&input.images[1], ImageSource::String(_))); + + let err = to_mm_input(image_work(image)).err().unwrap(); + assert!(err.contains("not prefetched"), "{err}"); + } + + #[test] + fn image_free_work_rejected() { + assert!( + to_mm_input(image_work(Value::Nil)) + .err() + .unwrap() + .contains("no raw image sources") + ); + assert!( + to_mm_input(MmWorkItem::default()) + .err() + .unwrap() + .contains("no raw image sources") + ); + } +} diff --git a/rust/sglang-server/src/message/request.rs b/rust/sglang-server/src/message/request.rs index fa7abee3f..dc35dd94b 100644 --- a/rust/sglang-server/src/message/request.rs +++ b/rust/sglang-server/src/message/request.rs @@ -106,6 +106,23 @@ pub struct GenerateBody { pub routed_dp_rank: Option, #[serde(default)] pub disagg_prefill_dp_rank: Option, + // Multimodal inputs, permissive `Value` so any shape Python's + // `GenerateReqInput` accepts (URL / base64 / list / list-of-lists) parses. + // `into_requests` fans them out per the Python + // `_normalize_{image,video,audio}_data` batch rules. + #[serde(default)] + pub image_data: Option, + /// Caller-supplied per-item content hashes (hex) overriding the computed + /// ones, so an external router's keys align with the prefix cache. Single + /// requests only: Python declares the batched shapes but `__getitem__` never + /// forwards them, so a batch is rejected here rather than answered with + /// hashes it did not ask for. + #[serde(default)] + pub mm_hashes: Option, + #[serde(default)] + pub video_data: Option, + #[serde(default)] + pub audio_data: Option, } impl GenerateBody { @@ -135,6 +152,10 @@ impl GenerateBody { decode_tp_size, routed_dp_rank, disagg_prefill_dp_rank, + image_data, + video_data, + audio_data, + mm_hashes, // Unported `GenerateReqInput` fields land here and are dropped, as they // are on the Python path. .. @@ -320,10 +341,26 @@ impl GenerateBody { let bootstrap_pair_keys = flatten_column(fan_out(bootstrap_pair_key, n, "bootstrap_pair_key")?); let decode_tp_sizes = flatten_column(fan_out(decode_tp_size, n, "decode_tp_size")?); + // `mm_hashes` has no batch form: honoring it only here would give the two + // servers different prefix-cache keys for the same body. Reject instead of + // dropping it silently as Python does — the field exists to align a + // caller's keys, so ignoring it returns subtly wrong ones. + if is_batch && mm_value_present(&mm_hashes) { + return Err(Error::Validation( + "mm_hashes is not supported for batch requests; send one request per prompt".into(), + )); + } + // Multimodal columns; see `split_mm_column` for the Python parity rules. + let images = split_mm_column(image_data, n, is_batch, MmBroadcast::WrapInList) + .map_err(|e| Error::Validation(format!("image_data: {e}")))?; + let videos = split_mm_column(video_data, n, is_batch, MmBroadcast::AsIs) + .map_err(|e| Error::Validation(format!("video_data: {e}")))?; + let audios = split_mm_column(audio_data, n, is_batch, MmBroadcast::AsIs) + .map_err(|e| Error::Validation(format!("audio_data: {e}")))?; // Every column above is exactly `n` long, so zip them by value: each // request takes ownership of its cell, with no indexing or bounds checks. - let requests = izip!( + let mut requests: Vec = izip!( rids, texts, id_lists, @@ -338,6 +375,9 @@ impl GenerateBody { bootstrap_rooms, bootstrap_pair_keys, decode_tp_sizes, + images, + videos, + audios, ) .map( |( @@ -355,6 +395,9 @@ impl GenerateBody { bootstrap_room, bootstrap_pair_key, decode_tp_size, + image_data, + video_data, + audio_data, )| GenerateRequest { rid, text, @@ -381,13 +424,126 @@ impl GenerateBody { decode_tp_size, routed_dp_rank, disagg_prefill_dp_rank, + mm: pack_mm(image_data, video_data, audio_data), }, ) .collect(); + // Single requests only (batches rejected above). Malformed entries are + // dropped here and warned about in `mm::apply_caller_hashes`, never a 400. + if !is_batch + && let (Some(rmpv::Value::Array(vals)), Some(req)) = (mm_hashes, requests.first_mut()) + && let Some(mm) = req.mm.as_deref_mut() + { + mm.mm_hashes = vals + .iter() + .filter_map(|v| v.as_str().map(str::to_owned)) + .collect(); + } Ok((requests, is_batch)) } } +/// Box the per-item mm values, `None` when the item has none — the common +/// text-only case keeps `GenerateRequest` slim. +fn pack_mm( + image_data: Option, + video_data: Option, + audio_data: Option, +) -> Option> { + if image_data.is_none() && video_data.is_none() && audio_data.is_none() { + return None; + } + Some(Box::new(MmData { + image_data, + video_data, + audio_data, + ..Default::default() + })) +} + +/// How a scalar mm value broadcasts across a batch: images become a one-image +/// list per item (`[[img]] * num` in Python `_normalize_image_data`), +/// video/audio broadcast bare (`[v] * num` in `_normalize_video_data`). +#[derive(Clone, Copy)] +enum MmBroadcast { + WrapInList, + AsIs, +} + +/// Fan one mm field into per-item values, mirroring Python's +/// `_normalize_{image,video,audio}_data`: +/// * `None` / empty list → `None` for every item; +/// * single request → the raw value passes through (the processor wraps a +/// non-list into a one-element list); +/// * batch + scalar → broadcast to every item, per `MmBroadcast`; +/// * batch + list → per item, length must equal the batch size. +fn split_mm_column( + v: Option, + n: usize, + is_batch: bool, + broadcast: MmBroadcast, +) -> Result>, String> { + let Some(v) = v else { + return Ok(vec![None; n]); + }; + if v.is_nil() { + return Ok(vec![None; n]); + } + if !is_batch { + return Ok(vec![Some(v)]); + } + match v { + rmpv::Value::Array(items) if items.is_empty() => Ok(vec![None; n]), + rmpv::Value::Array(items) => { + if items.len() != n { + return Err(format!( + "list length {} does not match batch size {n}", + items.len() + )); + } + Ok(items.into_iter().map(Some).collect()) + } + scalar => { + // A broadcast deep-clones once per prompt — same blow-up as + // sampling_params above, so bound the product before any clone. + check_broadcast_budget(scalar.heap_bytes(), n, "value").map_err(|e| e.to_string())?; + Ok(match broadcast { + MmBroadcast::WrapInList => vec![Some(rmpv::Value::Array(vec![scalar])); n], + MmBroadcast::AsIs => vec![Some(scalar); n], + }) + } + } +} + +/// One request handed to the MM worker pool: the rid to correlate the result, +/// plus the owned inputs from [`GenerateRequest::take_mm_work`]. +#[derive(Debug)] +pub struct MmRequest { + pub rid: crate::ids::Rid, + pub work: MmWorkItem, +} + +/// The parked request's fields the MM worker owns; converted to the driver input +/// by [`super::mm_payload::to_mm_input`]. +#[derive(Debug, Default)] +pub struct MmWorkItem { + pub text: Option, + pub input_ids: Option>, + pub image_data: Option, + pub video_data: Option, + pub audio_data: Option, + /// See [`MmData::prefetched`]. + pub prefetched: Vec, + /// See [`GenerateBody::mm_hashes`]. + pub mm_hashes: Vec, +} + +/// Whether an optional mm field counts as multimodal input, via the same +/// `value_present` the MM worker's payload parser uses. +fn mm_value_present(v: &Option) -> bool { + v.as_ref().is_some_and(super::mm_payload::value_present) +} + /// Request variant — selects the ingress branch, scheduler wire message, and /// egress shape. Each owns its body, so generate/control fields stay type-separate. #[derive(Debug)] @@ -475,6 +631,26 @@ pub struct GenerateRequest { /// so these are pure passthrough for the scheduler/LB protocol. pub routed_dp_rank: Option, pub disagg_prefill_dp_rank: Option, + /// Multimodal inputs, carried opaquely. Consumed by the Encoding stage, + /// which ships them to the MM worker pool; never read by the tokenizer or + /// serialized onto the scheduler header. Boxed so the common text-only + /// request doesn't grow every `Request` moved between stages. + pub mm: Option>, +} + +/// The opaque multimodal fields of one request (see [`GenerateRequest::mm`]). +#[derive(Debug, Default)] +pub struct MmData { + pub image_data: Option, + pub video_data: Option, + pub audio_data: Option, + /// Bytes of `image_data`'s I/O-backed sources, resolved by + /// `api_server::prefetch` in `mm_payload::io_sources` order so MM workers + /// never block on I/O. Out-of-band: the values above stay as the client + /// sent them. + pub prefetched: Vec, + /// See [`GenerateBody::mm_hashes`]; applied by the MM worker. + pub mm_hashes: Vec, } impl GenerateRequest { @@ -483,11 +659,33 @@ impl GenerateRequest { self.input_ids.as_ref().is_some_and(|v| !v.is_empty()) } - /// Multimodal detection hook. Deferred (Encoder stubbed): always false until mm - /// fields are wired in. - #[allow(dead_code)] + /// True when the request carries a usable multimodal payload — the mirror of + /// Python `GenerateReqInput.contains_mm_input()`. pub fn has_multimodal(&self) -> bool { - false + self.mm.as_ref().is_some_and(|mm| { + mm_value_present(&mm.image_data) + || mm_value_present(&mm.video_data) + || mm_value_present(&mm.audio_data) + }) + } + + /// Carve out the MM worker's inputs: `text` is cloned (the scheduler header + /// still needs it), `input_ids` is taken (the expanded ids replace it), and + /// the mm values move wholesale. + pub fn take_mm_work(&mut self) -> MmWorkItem { + let mut work = MmWorkItem { + text: self.text.clone(), + input_ids: self.input_ids.take(), + ..Default::default() + }; + if let Some(m) = self.mm.as_deref_mut() { + work.image_data = m.image_data.take(); + work.video_data = m.video_data.take(); + work.audio_data = m.audio_data.take(); + work.prefetched = std::mem::take(&mut m.prefetched); + work.mm_hashes = std::mem::take(&mut m.mm_hashes); + } + work } pub fn encode_header(&self) -> Result { @@ -539,6 +737,23 @@ impl HeapBytes for Option { self.as_ref().map_or(0, HeapBytes::heap_bytes) } } +impl HeapBytes for rmpv::Value { + fn heap_bytes(&self) -> usize { + use rmpv::Value; + const NODE: usize = std::mem::size_of::(); + match self { + Value::String(s) => s.as_bytes().len(), + Value::Binary(b) => b.len(), + Value::Ext(_, b) => b.len(), + Value::Array(items) => items.iter().map(|v| NODE + v.heap_bytes()).sum(), + Value::Map(entries) => entries + .iter() + .map(|(k, v)| 2 * NODE + k.heap_bytes() + v.heap_bytes()) + .sum(), + _ => 0, + } + } +} /// Collapse `fan_out`'s nullable-element output: outer `None` (field absent / /// scalar broadcast of nothing) and inner `None` (an explicit `null` list @@ -743,8 +958,9 @@ mod tests { } /// The native `bench_serving` payload (a `GenerateReqInput` superset) parses: - /// its `lora_path`/`return_routed_experts`/`image_data` are accepted-but-ignored, - /// so `split` succeeds and drops them while the real fields survive. + /// its `lora_path`/`return_routed_experts` are accepted-but-ignored and a + /// `null` `image_data` means "no multimodal input", so `split` succeeds + /// while the real fields survive. #[test] fn accepts_bench_serving_payload() { let (ps, is_batch) = requests( @@ -758,6 +974,111 @@ mod tests { assert_eq!(ps.len(), 1); assert_eq!(ps[0].text.as_deref(), Some("hi")); assert!(ps[0].stream); + assert!(!ps[0].has_multimodal()); + } + + /// Mm columns fan out per Python `_normalize_{image,video}_data`: a single + /// request passes the raw value through; a batch broadcasts a scalar image as + /// `[img]` per item, maps a list per item with matching lengths, and treats + /// `null`/`[]` as absent. + #[test] + fn split_mm_fanout_matches_python_normalize() { + let image_of = |p: &GenerateRequest| p.mm.as_ref().unwrap().image_data.clone().unwrap(); + + // Single request: raw value passes through untouched. + let (ps, _) = requests(r#"{"text": "a", "image_data": "http://x/i.jpg"}"#).unwrap(); + assert_eq!(image_of(&ps[0]).as_str(), Some("http://x/i.jpg")); + assert!(ps[0].has_multimodal()); + + // Batch + scalar image: broadcast, wrapped as a one-image list per item. + let (ps, _) = requests(r#"{"text": ["a", "b"], "image_data": "u"}"#).unwrap(); + for p in &ps { + assert_eq!(image_of(p).as_array().unwrap().len(), 1); + assert!(p.has_multimodal()); + } + + // Batch + per-item list: element i goes to item i. + let (ps, _) = requests(r#"{"text": ["a", "b"], "image_data": ["u1", "u2"]}"#).unwrap(); + assert_eq!(image_of(&ps[0]).as_str(), Some("u1")); + assert_eq!(image_of(&ps[1]).as_str(), Some("u2")); + + // Batch + wrong-length list is a 400. + assert!(requests(r#"{"text": ["a", "b"], "image_data": ["u1"]}"#).is_err()); + + // null / [] mean "no multimodal input". + let (ps, _) = requests(r#"{"text": "a", "image_data": null}"#).unwrap(); + assert!(!ps[0].has_multimodal()); + let (ps, _) = requests(r#"{"text": "a", "image_data": []}"#).unwrap(); + assert!(!ps[0].has_multimodal()); + + // Batch + scalar video: broadcast bare (not wrapped), per Python + // `_normalize_video_data`. + let (ps, _) = requests(r#"{"text": ["a", "b"], "video_data": "v"}"#).unwrap(); + let video = ps[1].mm.as_ref().unwrap().video_data.clone().unwrap(); + assert_eq!(video.as_str(), Some("v")); + assert!(ps[1].has_multimodal()); + } + + /// A scalar broadcast is budget-checked before the deep clones (16 MiB × + /// 4096 prompts would be 64 GiB and an abort); per-item lists clone nothing + /// and are never charged. + #[test] + fn oversized_mm_broadcast_rejected() { + let big = rmpv::Value::from("x".repeat(MAX_BROADCAST_CLONE_BYTES / 2 + 1)); + let err = split_mm_column(Some(big.clone()), 2, true, MmBroadcast::WrapInList) + .err() + .unwrap(); + assert!(err.contains("broadcast"), "{err}"); + // A per-item list of the same total size moves, not clones: accepted. + let list = rmpv::Value::Array(vec![big, rmpv::Value::from("y")]); + assert!(split_mm_column(Some(list), 2, true, MmBroadcast::WrapInList).is_ok()); + // Small scalars broadcast fine. + let small = rmpv::Value::from("u1"); + assert!(split_mm_column(Some(small), 2, true, MmBroadcast::AsIs).is_ok()); + } + + /// `mm_hashes` rides only on single requests (Python `__getitem__` + /// parity: batches drop it) and moves into the work item. + #[test] + fn mm_hashes_single_only() { + let (mut ps, _) = + requests(r#"{"text": "a", "image_data": "u", "mm_hashes": ["a1b2", "0xff"]}"#).unwrap(); + assert_eq!(ps[0].mm.as_ref().unwrap().mm_hashes, vec!["a1b2", "0xff"]); + assert_eq!(ps[0].take_mm_work().mm_hashes, vec!["a1b2", "0xff"]); + assert!(ps[0].mm.as_ref().unwrap().mm_hashes.is_empty()); + + // A batch cannot carry hashes (Python drops them), so it is rejected... + for body in [ + r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": [["x"], ["y"]]}"#, + r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": ["x", "y"]}"#, + ] { + let err = requests(body).err().unwrap(); + assert!(matches!(err, Error::Validation(_)), "{body}: {err:?}"); + } + // ...while an absent or empty field is not a payload and must still pass. + for body in [ + r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": null}"#, + r#"{"text": ["a", "b"], "image_data": ["u", "v"], "mm_hashes": []}"#, + ] { + assert!(requests(body).is_ok(), "{body}"); + } + } + + /// `take_mm_work` clones `text` (the scheduler header still needs it) and + /// moves everything the worker owns out of the request. + #[test] + fn mm_work_item_takes_owned_fields() { + let (mut ps, _) = + requests(r#"{"text": "hi", "image_data": ["u1", "u2"], "audio_data": "a"}"#).unwrap(); + let work = ps[0].take_mm_work(); + assert_eq!(work.text.as_deref(), Some("hi")); + assert!(work.input_ids.is_none()); + assert_eq!(work.image_data.unwrap().as_array().unwrap().len(), 2); + assert!(work.video_data.is_none()); + assert_eq!(work.audio_data.unwrap().as_str(), Some("a")); + // Moved out, not cloned; `text` survives for the header. + assert!(ps[0].mm.as_ref().unwrap().image_data.is_none()); + assert_eq!(ps[0].text.as_deref(), Some("hi")); } /// The body limit is disabled, so an unbounded batch turns a small body into an diff --git a/rust/sglang-server/src/mm.rs b/rust/sglang-server/src/mm.rs new file mode 100644 index 000000000..d08f6b9c2 --- /dev/null +++ b/rust/sglang-server/src/mm.rs @@ -0,0 +1,413 @@ +//! Multimodal worker pool. +//! +//! Rust threads drain requests parked in `Encoding` and run the `sglang-mm` +//! pipeline registered by `Server.start_mm_workers` (decode → preprocess → +//! placeholder expansion → M-RoPE, GIL-free). Each worker parks the result +//! buffers in the rid-keyed [`Sidecar`] and returns only the expanded ids; +//! Python attaches the buffers at drain time (`Server.take_mm`). Inputs the +//! pipeline cannot serve are rejected to the client — no Python fallback. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; + +use crate::message::MmRequest; +use crate::runtime::Runnable; +use crate::tokenizer::TextTokenizer; +use crate::tokenizer_manager::TmEvent; + +/// A named POSIX shared-memory segment owning its name: dropped → unlinked. +/// +/// Written by an MM worker so the TP broadcast carries a ~100-byte +/// `ShmPointerMMData` stub instead of the ~20 MB feature tensor, and every +/// rank maps it in parallel. Python's `materialize()` unlinks after cloning; +/// this `Drop` covers the paths where the buffers never reach Python (aborted +/// while parked, late result purged). +pub struct ShmSegment { + name: String, +} + +impl ShmSegment { + /// Create `/dev/shm/{name}` holding exactly `bytes`. No leading slash — + /// the name must suit Python's `SharedMemory(name=…)` (shm_open adds one). + pub fn create(name: String, bytes: &[u8]) -> Result { + let c_name = std::ffi::CString::new(format!("/{name}")) + .map_err(|_| "shm name contains NUL".to_string())?; + // SAFETY: plain POSIX calls on a name we own; every handle created + // below is closed/unmapped on all paths. + unsafe { + let fd = libc::shm_open( + c_name.as_ptr(), + libc::O_CREAT | libc::O_EXCL | libc::O_RDWR, + 0o600, + ); + if fd < 0 { + return Err(format!( + "shm_open({name}): {}", + std::io::Error::last_os_error() + )); + } + let segment = Self { name }; // unlink from here on any failure + if libc::ftruncate(fd, bytes.len() as libc::off_t) != 0 { + let e = std::io::Error::last_os_error(); + libc::close(fd); + return Err(format!("ftruncate({}): {e}", segment.name)); + } + let ptr = libc::mmap( + std::ptr::null_mut(), + bytes.len(), + libc::PROT_WRITE, + libc::MAP_SHARED, + fd, + 0, + ); + libc::close(fd); + if ptr == libc::MAP_FAILED { + return Err(format!( + "mmap({}): {}", + segment.name, + std::io::Error::last_os_error() + )); + } + std::ptr::copy_nonoverlapping(bytes.as_ptr(), ptr.cast::(), bytes.len()); + libc::munmap(ptr, bytes.len()); + Ok(segment) + } + } + + /// Hand the segment — and the duty to unlink — to the caller (Python, at + /// drain time). + pub fn into_name(self) -> String { + std::mem::take(&mut std::mem::ManuallyDrop::new(self).name) + } +} + +impl Drop for ShmSegment { + fn drop(&mut self) { + if let Ok(c_name) = std::ffi::CString::new(format!("/{}", self.name)) { + // SAFETY: unlinking a name we created; ENOENT (already unlinked + // by Python's materialize) is fine to ignore. + unsafe { libc::shm_unlink(c_name.as_ptr()) }; + } + } +} + +/// Unique segment names: the pid separates server restarts (a crash can leak +/// segments under the old pid), the counter separates results within one. +fn shm_name(item: usize) -> String { + static COUNTER: AtomicU64 = AtomicU64::new(0); + let n = COUNTER.fetch_add(1, Ordering::Relaxed); + format!("sglmm-{}-{n}-{item}", std::process::id()) +} + +/// Python parity: caller hashes override the computed ones so an external +/// router's keys align with the prefix cache. A length mismatch or malformed +/// entry warns and keeps the computed hash — never blocks the request. +fn apply_caller_hashes(hashes: &mut [u64], caller: &[String]) { + if caller.is_empty() { + return; + } + if caller.len() != hashes.len() { + tracing::warn!( + caller = caller.len(), + items = hashes.len(), + "mm_hashes length != mm item count; ignoring caller hashes" + ); + return; + } + for (hash, entry) in hashes.iter_mut().zip(caller) { + match parse_caller_hash(entry) { + Some(v) => *hash = v, + None => tracing::warn!(%entry, "malformed mm_hashes entry; keeping computed hash"), + } + } +} + +/// Hex of any width, as Python's `int(hex_hash, 16)` takes it (a full SHA-256 +/// being the common case), keeping the low 64 bits — only the low 30 are +/// observable, through `_compute_pad_value`. +fn parse_caller_hash(entry: &str) -> Option { + let hex = entry.strip_prefix("0x").unwrap_or(entry); + if hex.is_empty() || !hex.bytes().all(|b| b.is_ascii_hexdigit()) { + return None; + } + u64::from_str_radix(&hex[hex.len().saturating_sub(16)..], 16).ok() +} + +/// One parked result: the buffers the drain-time Python adapter needs (the +/// expanded `input_ids` travel separately, via `TmEvent::MmEncoded`). The qwen +/// drain shape (`sglang_mm::qwen_vl::pack_drain`); generalizes to a +/// named-tensor handoff once a family needs a different one. +pub struct MmSidecarEntry { + pub features: FeatureStore, + pub grids: Vec<[u32; 3]>, + pub hashes: Vec, + pub offsets: Vec<(u32, u32)>, + pub mrope: Vec, + pub mrope_delta: i64, +} + +/// Where a result's feature buffers live between worker and drain. +pub enum FeatureStore { + /// In-process; the drain wraps them zero-copy. Single-rank serving, or the + /// shm fallback. Under TP the whole buffer would ride `broadcast_pyobj`. + Inline(Vec), + /// One POSIX segment per item, written by the worker; only the names cross + /// ranks. See [`ShmSegment`]. + Shm(Vec), +} + +/// Results parked between a worker's `MmEncoded` and the scheduler drain, keyed +/// by rid. Owns the lifecycle so entries never leak: [`park`](Self::park) +/// strictly before `MmEncoded`, [`take`](Self::take) at the drain, +/// [`purge`](Self::purge) for requests that die while parked. +#[derive(Clone, Default)] +pub struct Sidecar(Arc>>); + +impl Sidecar { + pub fn park(&self, rid: String, entry: MmSidecarEntry) { + self.0.lock().unwrap().insert(rid, entry); + } + pub fn take(&self, rid: &str) -> Option { + self.0.lock().unwrap().remove(rid) + } + pub fn purge(&self, rid: &str) { + self.0.lock().unwrap().remove(rid); + } +} + +/// Shared state of the mm path, built once at `start_mm_workers`. +pub struct Context { + pub family: Box, + /// `None` under `skip_tokenizer_init` (requests must carry `input_ids`). + pub tokenizer: Option>, + pub sidecar: Sidecar, + /// Park feature buffers in POSIX shm. Set by the Python launcher + /// (`NativeMmHost._use_feature_shm`) exactly when the scheduler broadcasts + /// across TP ranks and will unwrap `ShmPointerMMData`. + pub feature_shm: bool, +} + +impl Context { + pub fn new( + spec_json: &str, + tokenizer: Option>, + sidecar: Sidecar, + ) -> Result { + let feature_shm = serde_json::from_str::(spec_json) + .ok() + .and_then(|v| v.get("feature_shm").and_then(|b| b.as_bool())) + .unwrap_or(false); + Ok(Self { + family: sglang_mm::registry::pipeline_from_spec(spec_json)?, + tokenizer, + sidecar, + feature_shm, + }) + } +} + +/// Run the pipeline for one request. `Ok` returns the final expanded ids, the +/// buffers already parked; `Err` rejects the request back to the client. +fn process( + ctx: &Context, + rid: &crate::ids::Rid, + mut work: crate::message::MmWorkItem, +) -> Result, String> { + let caller_hashes = std::mem::take(&mut work.mm_hashes); + let input = crate::message::mm_payload::to_mm_input(work)?; + let output = sglang_mm::driver::process(ctx.family.as_ref(), input, |text| { + let tokenizer = ctx.tokenizer.as_ref().ok_or_else(|| { + "skip_tokenizer_init is set: multimodal text prompts require input_ids".to_string() + })?; + tokenizer.encode(text).map_err(|error| error.to_string()) + })?; + let mut drain = sglang_mm::qwen_vl::pack_drain(output)?; + apply_caller_hashes(&mut drain.hashes, &caller_hashes); + let features = if ctx.feature_shm { + park_features_in_shm(&drain.features, &drain.grids) + } else { + FeatureStore::Inline(drain.features) + }; + ctx.sidecar.park( + rid.as_str().to_owned(), + MmSidecarEntry { + features, + grids: drain.grids, + hashes: drain.hashes, + offsets: drain.offsets, + mrope: drain.mrope, + mrope_delta: drain.mrope_delta, + }, + ); + Ok(drain.input_ids) +} + +/// Split the flat feature buffer per item (`t*h*w` rows per grid) and park each +/// slice in its own segment. Any shm failure (`/dev/shm` full, odd shape) falls +/// back to inline, as Python's `_wrap_shm_or_inline` does: degrade to the slow +/// path, never fail the request. +fn park_features_in_shm(features: &[f32], grids: &[[u32; 3]]) -> FeatureStore { + let total_rows: usize = grids + .iter() + .map(|g| g[0] as usize * g[1] as usize * g[2] as usize) + .sum(); + if total_rows == 0 || !features.len().is_multiple_of(total_rows) { + return FeatureStore::Inline(features.to_vec()); + } + let dim = features.len() / total_rows; + let mut segments = Vec::with_capacity(grids.len()); + let mut row = 0usize; + for (item, grid) in grids.iter().enumerate() { + let rows = grid[0] as usize * grid[1] as usize * grid[2] as usize; + let slice = &features[row * dim..(row + rows) * dim]; + row += rows; + match ShmSegment::create(shm_name(item), bytemuck::cast_slice(slice)) { + Ok(segment) => segments.push(segment), + Err(error) => { + tracing::warn!(%error, "mm: shm feature transport failed; falling back to inline"); + return FeatureStore::Inline(features.to_vec()); + } + } + } + FeatureStore::Shm(segments) +} + +/// One MM worker, spawned via `Runtime::spawn_mm_pool` (which owns the +/// pinning policy for this pool — see its docs). +pub struct MmWorker { + rx: flume::Receiver, + tm: flume::Sender, + ctx: Arc, +} + +impl MmWorker { + pub fn new( + rx: flume::Receiver, + tm: flume::Sender, + ctx: Arc, + ) -> Self { + Self { rx, tm, ctx } + } +} + +impl Runnable for MmWorker { + /// Drain until the mm channel closes (tm-ingress drops its sender on + /// shutdown). One request at a time, so the pool size bounds MM + /// concurrency; an error rejects the request back to the client. + fn run(self) { + while let Ok(req) = self.rx.recv() { + let rid = req.rid; + let event = match process(&self.ctx, &rid, req.work) { + Ok(input_ids) => { + tracing::debug!(%rid, tokens = input_ids.len(), "mm: processed"); + TmEvent::MmEncoded { rid, input_ids } + } + Err(message) => { + tracing::warn!(%rid, %message, "mm processing rejected"); + TmEvent::MmFailed { rid, message } + } + }; + if self.tm.send(event).is_err() { + return; // tm-ingress gone: shutdown + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Caller hashes override computed ones; mismatched lengths and malformed + /// entries fall back per item, never reject (Python parity). + #[test] + fn caller_hashes_override_with_fallback() { + let mut hashes = vec![1, 2, 3]; + apply_caller_hashes(&mut hashes, &[]); + assert_eq!(hashes, [1, 2, 3]); + + apply_caller_hashes(&mut hashes, &["ff".into()]); // length mismatch + assert_eq!(hashes, [1, 2, 3]); + + apply_caller_hashes(&mut hashes, &["ff".into(), "not-hex".into(), "0x10".into()]); + assert_eq!(hashes, [0xff, 2, 0x10]); + } + + /// A full SHA-256 (what routers send) keeps its low 64 bits rather than + /// falling back, so the pad value matches Python's wide `int`. + #[test] + fn caller_hashes_accept_arbitrary_width() { + let sha256 = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"; + let mut hashes = vec![1]; + apply_caller_hashes(&mut hashes, &[sha256.into()]); + assert_eq!(hashes, [0xa495991b7852b855]); + assert_eq!(hashes[0] % (1 << 30), 944_945_237); // int(sha256, 16) % (1 << 30) + + // Width alone is never malformed; a non-hex digit still is. + assert_eq!(parse_caller_hash(&"f".repeat(64)), Some(u64::MAX)); + assert_eq!(parse_caller_hash("0x"), None); + assert_eq!(parse_caller_hash(""), None); + } + + fn shm_path(name: &str) -> std::path::PathBuf { + std::path::Path::new("/dev/shm").join(name) + } + + /// The segment holds exactly the written bytes and dropping it unlinks — + /// the leak guard for results purged before Python takes them. + #[test] + fn segment_roundtrip_and_drop_unlinks() { + let name = shm_name(0); + let payload: Vec = (0..255u8).collect(); + let segment = ShmSegment::create(name.clone(), &payload).unwrap(); + assert_eq!(std::fs::read(shm_path(&name)).unwrap(), payload); + drop(segment); + assert!(!shm_path(&name).exists(), "drop must unlink"); + } + + /// `into_name` transfers the unlink duty to the caller (Python's + /// `materialize()`), so the segment must survive the handoff. + #[test] + fn into_name_disarms_the_unlink() { + let segment = ShmSegment::create(shm_name(0), &[1, 2, 3]).unwrap(); + let name = segment.into_name(); + assert!(shm_path(&name).exists(), "handoff must not unlink"); + // manual cleanup for the test + let c = std::ffi::CString::new(format!("/{name}")).unwrap(); + unsafe { libc::shm_unlink(c.as_ptr()) }; + } + + /// Per-item slicing follows the grid row counts, so Python's + /// `(rows, feature_dim)` reshape of a segment sees only its own item. + #[test] + fn park_splits_features_by_grid() { + // Two items: grids (1,2,2)=4 rows and (1,1,2)=2 rows, dim=3. + let features: Vec = (0..18).map(|i| i as f32).collect(); + let grids = [[1, 2, 2], [1, 1, 2]]; + let FeatureStore::Shm(segments) = park_features_in_shm(&features, &grids) else { + panic!("expected shm store"); + }; + assert_eq!(segments.len(), 2); + let read = |seg: &ShmSegment| -> Vec { std::fs::read(shm_path(&seg.name)).unwrap() }; + assert_eq!( + read(&segments[0]), + bytemuck::cast_slice::(&features[..12]) + ); + assert_eq!( + read(&segments[1]), + bytemuck::cast_slice::(&features[12..]) + ); + } + + /// A degenerate shape must degrade to inline, never a shm-side panic. + #[test] + fn shape_surprise_falls_back_inline() { + let features = vec![0.0f32; 7]; // not divisible by 2 rows + let grids = [[1, 1, 2]]; + assert!(matches!( + park_features_in_shm(&features, &grids), + FeatureStore::Inline(_) + )); + } +} diff --git a/rust/sglang-server/src/runtime.rs b/rust/sglang-server/src/runtime.rs index df5f425d4..0cfb9d920 100644 --- a/rust/sglang-server/src/runtime.rs +++ b/rust/sglang-server/src/runtime.rs @@ -8,6 +8,8 @@ //! * Detokenizer — M pinned OS threads / shards (CPU bound), core set C //! * TM ingress — 1 thread driving the ingress FSM //! * TM egress — 1 thread draining the egress ring → detok shards +//! * MM workers — K unpinned OS threads, spawned late via +//! [`Runtime::spawn_mm_pool`] (multimodal models only) //! //! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling //! axum's worker threads. @@ -38,6 +40,18 @@ pub use runnable::Runnable; pub struct Runtime { pub ingress: IngressConsumer, pub egress: EgressProducer, + /// Requests parked in `Encoding`, drained by the MM worker pool + /// (`Server.start_mm_workers`). Stays empty for non-multimodal models — + /// ingress never routes to it. + pub mm: flume::Receiver, + /// Back-channel for the MM workers' `MmEncoded` / `MmFailed` into tm-ingress. + pub tm: flume::Sender, + /// The loaded tokenizer, shared with the MM worker path (`None` under + /// `skip_tokenizer_init`). + pub tokenizer: Option>, + /// MM results parked between a worker's `MmEncoded` and the scheduler drain + /// (`Server.take_mm`). + pub mm_sidecar: crate::mm::Sidecar, /// Worker join handles, joined by `request_shutdown` / `Drop`. threads: Mutex>>, /// The single shutdown sender. @@ -49,6 +63,21 @@ pub struct Runtime { const SHUTDOWN_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); impl Runtime { + /// Spawn `workers` `mm-worker-{i}` threads into the shutdown join set — + /// late, once Python has built the mm spec (`Server::start_mm_workers`). + /// + /// Deliberately unpinned: the threads inherit the launch thread's affinity, + /// already narrowed by `RustServer.launch` to the server cores, so bursty + /// MM preprocessing floats over that whole set (rather than owning cores + /// that idle between bursts) and never preempts the scheduler's reserved + /// cores. + pub fn spawn_mm_pool(&self, workers: usize, ctx: Arc) { + let mut threads = self.threads.lock().unwrap(); + spawn_pool("mm-worker", None, workers.max(1), &mut threads, |_| { + crate::mm::MmWorker::new(self.mm.clone(), self.tm.clone(), ctx.clone()) + }); + } + /// Stop the runtime and join every worker thread (with a bounded wait). /// /// Dropping `shutdown_tx` wakes the tm-ingress/tm-egress selectors (which @@ -107,6 +136,10 @@ pub fn start(cfg: RuntimeConfig) -> Result { let (tm_tx, tm_rx) = flume::bounded::(cfg.rust_server_args.channel_cap); let (tok_tx, tok_rx) = flume::bounded::(cfg.rust_server_args.channel_cap); + // Encoding → MM worker pool. Bounded like the other stage edges so a slow + // pool back-pressures instead of buffering unboundedly. + let (mm_tx, mm_rx) = + flume::bounded::(cfg.rust_server_args.channel_cap); let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num; let mut detok_tx = Vec::with_capacity(detokenizer_worker_num); let mut detok_rx = Vec::with_capacity(detokenizer_worker_num); @@ -140,6 +173,14 @@ pub fn start(cfg: RuntimeConfig) -> Result { cfg.server_args.revision.as_deref(), skip_tokenizer_init, )?; + // The `TextTokenizer` view of it, shared by the tokenizer pool and the MM + // worker path (which encodes the placeholder-expanded prompt itself). + let text_tokenizer: Option> = dyn_tokenizer + .as_ref() + .map(|t| Arc::new(tokenizer::DynamoTokenizer::new(t.clone())) as _); + + // Shared: MM workers park, the Python drain pops, tm-ingress purges. + let mm_sidecar: crate::mm::Sidecar = Default::default(); // --- Detokenizer shards (pinned, CPU bound) --- { @@ -168,10 +209,9 @@ pub fn start(cfg: RuntimeConfig) -> Result { // --- Tokenizer pool (pinned, CPU bound) --- // Only spawned when a real tokenizer is loaded; under `skip_tokenizer_init` // there is none and ingress never routes to the pool, so we skip it. - if let Some(t) = &dyn_tokenizer { + if let Some(tokenizer) = &text_tokenizer { // Reuse the single loaded tokenizer (shared with the detok shards). - let tokenizer: Arc = - Arc::new(tokenizer::DynamoTokenizer::new(t.clone())); + let tokenizer = tokenizer.clone(); let tok_cores = plan.as_ref().map(|p| p.tok.clone()); // Workers share the MPMC inbox (`tok_rx`) and the read-only backend, so // each gets a cheap clone of both. @@ -220,6 +260,11 @@ pub fn start(cfg: RuntimeConfig) -> Result { .map(|c| vec![c]); let limits = tokenizer_manager::Limits::try_from(&*cfg.server_args) .map_err(|e| format!("ingress limits: {e}"))?; + let mm = tokenizer_manager::Mm { + enabled: cfg.server_args.model_is_multimodal(), + tx: mm_tx, + sidecar: mm_sidecar.clone(), + }; let mut parts = Some((tm_rx, ingress_tx)); // moved into the single worker let shutdown_rx = shutdown_rx.clone(); spawn_pool("tm-ingress", cores, 1, &mut threads, |_| { @@ -230,6 +275,7 @@ pub fn start(cfg: RuntimeConfig) -> Result { senders.clone(), ingress_tx, limits.clone(), + mm.clone(), shutdown_rx.clone(), ) }); @@ -282,6 +328,10 @@ pub fn start(cfg: RuntimeConfig) -> Result { Ok(Runtime { ingress: ingress_rx, egress: egress_tx, + mm: mm_rx, + tm: tm_tx, + tokenizer: text_tokenizer, + mm_sidecar, threads: Mutex::new(threads), shutdown_tx: Mutex::new(Some(shutdown_tx)), }) diff --git a/rust/sglang-server/src/runtime/config.rs b/rust/sglang-server/src/runtime/config.rs index 722f56964..b3c91d9ca 100644 --- a/rust/sglang-server/src/runtime/config.rs +++ b/rust/sglang-server/src/runtime/config.rs @@ -158,6 +158,11 @@ pub struct ModelConfig { /// boot ([`ServerArgs::validate_mandatory`]). #[serde(default)] pub vocab_size: Option, + /// Whether the model accepts multimodal inputs. Gates the MM Encoding branch + /// in tm-ingress; `false` silently ignores mm fields, as the Python + /// `TokenizerManager` does with `mm_processor is None`. + #[serde(default)] + pub is_multimodal: bool, /// Resolved default sampling parameters, stamped by /// `RustServer._build_server_args` from Python's /// `ModelConfig.get_default_sampling_params()`. Already gated on @@ -259,6 +264,12 @@ impl ServerArgs { self.disaggregation_mode == "prefill" } + /// Whether the served model is multimodal, from the scheduler's dump. See + /// [`ModelConfig::is_multimodal`]. + pub fn model_is_multimodal(&self) -> bool { + self.model_config.is_multimodal + } + /// Bind address `host:port`. `host` is expected to be an IP — the result is /// parsed as a `SocketAddr`, so a bare IPv6 host gets bracketed. pub fn bind(&self) -> String { diff --git a/rust/sglang-server/src/tokenizer_manager.rs b/rust/sglang-server/src/tokenizer_manager.rs index 29bf73e91..e5b25cc4f 100644 --- a/rust/sglang-server/src/tokenizer_manager.rs +++ b/rust/sglang-server/src/tokenizer_manager.rs @@ -14,7 +14,7 @@ mod egress; mod ingress; pub use egress::{ActivityCounter, Egress}; -pub use ingress::{Ingress, Limits}; +pub use ingress::{Ingress, Limits, Mm}; use crate::ids::Rid; use crate::message::{DetokMsg, Request}; @@ -36,6 +36,13 @@ pub enum TmEvent { /// A request back from the tokenizer pool: `PreSendValidating` (ids filled) on success, /// or `Failed` on a tokenize error. `drive` handles both. Tokenized(Request), + /// An MM worker finished a request parked in `Encoding`: `input_ids` are the + /// final placeholder-expanded prompt ids. The buffers ride the rid-keyed + /// sidecar (`Server.take_mm`), not this event. + MmEncoded { rid: Rid, input_ids: Vec }, + /// An MM worker rejected a request parked in `Encoding` (bad media URL, + /// unsupported modality, preprocess error, …). + MmFailed { rid: Rid, message: String }, } /// Producer-side handles, cloned into every stage that needs to emit. diff --git a/rust/sglang-server/src/tokenizer_manager/ingress.rs b/rust/sglang-server/src/tokenizer_manager/ingress.rs index 325a64433..aef1abe4e 100644 --- a/rust/sglang-server/src/tokenizer_manager/ingress.rs +++ b/rust/sglang-server/src/tokenizer_manager/ingress.rs @@ -17,14 +17,17 @@ //! The egress edges (Streaming/Finalizing/Completed) are driven on the egress //! side (see `egress` + `detokenizer`). +use std::collections::HashMap; + use bytes::Bytes; use crate::error::Error; use crate::fsm::{Event, RequestState, ValidationOutcome}; +use crate::ids::Rid; use crate::message::{ - AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, Request, - RequestKind, + AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, MmRequest, + Request, RequestKind, }; use crate::ring::IngressProducer; use crate::runtime::{Runnable, ServerArgs}; @@ -41,9 +44,29 @@ pub struct Ingress { senders: Senders, ingress: IngressProducer, limits: Limits, + mm: Mm, + /// Requests parked in `Encoding` while an MM worker processes their media; + /// resumed by `MmEncoded` / `MmFailed`. Only this thread touches it, so no + /// lock. + pending_mm: HashMap, shutdown: flume::Receiver<()>, } +/// The ingress side of the MM path. +#[derive(Clone)] +pub struct Mm { + /// Whether the model is multimodal. When false, mm fields are silently + /// ignored, as the Python `TokenizerManager` does with `mm_processor is + /// None`. + pub enabled: bool, + /// → MM worker pool (spawned via `Server.start_mm_workers`). + pub tx: flume::Sender, + /// Results sidecar. Purged here when a late result arrives for a request + /// that is no longer parked; otherwise it would leak, since only the + /// scheduler drain pops entries. + pub sidecar: crate::mm::Sidecar, +} + /// Longest client-supplied rid accepted. It keys the detok table and travels on /// every chunk, so its length is a recurring cost; Python mints 32-byte uuid hex. const MAX_RID_LEN: usize = 128; @@ -104,6 +127,7 @@ impl Ingress { senders: Senders, ingress: IngressProducer, limits: Limits, + mm: Mm, shutdown: flume::Receiver<()>, ) -> Self { Self { @@ -112,6 +136,8 @@ impl Ingress { senders, ingress, limits, + mm, + pending_mm: HashMap::new(), shutdown, } } @@ -124,7 +150,7 @@ enum Lane { } impl Runnable for Ingress { - fn run(self) { + fn run(mut self) { loop { // Select, not a drain-then-block: an abort arriving while the inbox is // idle must still be handled at once. @@ -139,6 +165,12 @@ impl Runnable for Ingress { Some(Lane::Event(TmEvent::Ingress(req) | TmEvent::Tokenized(req))) => { self.drive(req) } + Some(Lane::Event(TmEvent::MmEncoded { rid, input_ids })) => { + self.on_mm_encoded(rid, input_ids) + } + Some(Lane::Event(TmEvent::MmFailed { rid, message })) => { + self.on_mm_failed(rid, message) + } None => { // Shutdown, or the inbox closed. Drain whatever is still queued // on the abort lane first: those requests are in flight on the @@ -168,6 +200,9 @@ impl Ingress { if err.http_status() == 500 { tracing::error!(rid = %req.rid, error = %err, "ingress rejected request"); } + // A rejected request never reaches the scheduler drain, so purge any + // parked MM result (no-op for the common non-mm request). + self.mm.sidecar.purge(req.rid.as_str()); let _ = req.state.apply(Event::Error(err.clone())); let _ = req.sink.try_send(EgressItem::Error(err)); // client may be gone if registered { @@ -178,11 +213,12 @@ impl Ingress { } /// Drive a request through its ingress states until it terminates (failed or - /// pushed to the ring) or is handed to the tokenizer pool (re-entering as a - /// `Tokenized` event). Each arm acts and advances the FSM; the loop - /// re-dispatches. The arms are the design table's states, `Failed` the single - /// reject path. - fn drive(&self, mut req: Request) { + /// pushed to the ring), is handed to the tokenizer pool (re-entering as a + /// `Tokenized` event), or is parked in `pending_mm` awaiting an MM worker + /// (re-entering via `MmEncoded` / `MmFailed`). Each arm acts and advances + /// the FSM; the loop re-dispatches. The arms are the design table's states, + /// `Failed` the single reject path. + fn drive(&mut self, mut req: Request) { // Flipped once `register_detok` succeeds; `fail` must not deregister before // that (see `fail`). A pool return re-enters `drive` already registered. let mut registered = !matches!(req.state, RequestState::Received); @@ -235,6 +271,13 @@ impl Ingress { .normalize(self.limits.skip_tokenizer_init, self.limits.vocab_size) { Err(e) => Err(e), + // The native pipeline produces the final input_ids, + // so it wins even over a pre-tokenized prompt (which + // still needs placeholder expansion) — the same + // precedence as the Python TokenizerManager. + Ok(()) if self.mm.enabled && g.has_multimodal() => { + Ok(ValidationOutcome::HasMultimodal) + } // Client ids skip the pool; text goes to the tokenizer. Ok(()) if g.already_tokenized() => { Ok(ValidationOutcome::AlreadyTokenized) @@ -252,6 +295,40 @@ impl Ingress { } } } + // Hand off to the MM worker pool and park the request; it + // re-enters via `MmEncoded` (→ PreSendValidating) or `MmFailed` + // (→ reject). Doesn't loop. + RequestState::Encoding => { + let work = { + let RequestKind::Generate(g) = &mut req.kind else { + self.fail( + &mut req, + Error::Internal("non-generate request in Encoding".into()), + registered, + ); + return; + }; + g.take_mm_work() + }; + let msg = MmRequest { + rid: req.rid.clone(), + work, + }; + // Full = the pool can't keep up, so back-pressure like a full + // ingress ring. Disconnected = pool gone. + if let Err(e) = self.mm.tx.try_send(msg) { + let err = match e { + flume::TrySendError::Full(_) => Error::QueueFull, + flume::TrySendError::Disconnected(_) => { + Error::Internal("mm worker pool gone".into()) + } + }; + self.fail(&mut req, err, registered); + return; + } + self.pending_mm.insert(req.rid.clone(), req); + return; + } // Hand off to the tokenizer pool; it returns the request as a // `Tokenized` event (PreSendValidating, or Failed on error). // Doesn't loop. @@ -392,6 +469,34 @@ impl Ingress { } } + /// An MM worker finished a parked request: fill in the final expanded + /// `input_ids`, advance `Encoding → PreSendValidating`, and resume driving + /// (pre-send checks → ring). No pending entry means the request was already + /// rejected or aborted, so the result is dropped. + fn on_mm_encoded(&mut self, rid: Rid, input_ids: Vec) { + let Some(mut req) = self.pending_mm.remove(&rid) else { + tracing::debug!(rid = %rid, "mm result for unknown/finished request; dropped"); + // It will never reach the scheduler drain, so purge or leak. + self.mm.sidecar.purge(rid.as_str()); + return; + }; + if let RequestKind::Generate(g) = &mut req.kind { + g.input_ids = Some(input_ids); + } + let _ = req.state.apply(Event::EncodeDone); // Encoding → PreSendValidating + self.drive(req); + } + + /// An MM worker failed a parked request (bad URL, processor error): reject it + /// back to the client, as Python turns a per-request exception into a 400. + fn on_mm_failed(&mut self, rid: Rid, message: String) { + let Some(mut req) = self.pending_mm.remove(&rid) else { + tracing::debug!(rid = %rid, "mm failure for unknown/finished request; dropped"); + return; + }; + self.fail(&mut req, Error::Encode(message), true); // parked ⇒ registered + } + /// Client disconnected (or a detok terminal): deregister the sink, then push an /// `AbortReq(rid)` so the scheduler stops generating for it. /// @@ -400,8 +505,15 @@ impl Ingress { /// That wastes GPU work until the request finishes on its own, but it cannot be /// misdelivered — the rid is unique to this request for the process's lifetime /// ([`Rid::from_client`]), so no later request can ever answer to it. - fn on_abort(&self, source: AbortSource) { + /// + /// A request parked in `pending_mm` is cancelled here, so the worker's late + /// result lands in `on_mm_encoded`'s no-entry branch and purges the sidecar — + /// no generation runs for output nobody will read. + fn on_abort(&mut self, source: AbortSource) { let rid = source.rid().clone(); + if self.pending_mm.remove(&rid).is_some() { + tracing::debug!(rid = %rid, "abort cancelled request parked for MM"); + } let _ = self .senders .detok_for(&rid) @@ -616,12 +728,14 @@ mod tests { use tokio::sync::mpsc; /// An `Ingress` plus its detok-shard receiver, ring consumer (keep alive — - /// dropping it closes the ring → false QueueFull), and tm inbox sender. + /// dropping it closes the ring → false QueueFull), tm inbox sender, and the + /// mm-pool receiver (keep alive — dropping it makes mm submits fail). fn make_ingress() -> ( Ingress, flume::Receiver, IngressConsumer, flume::Sender, + flume::Receiver, ) { make_ingress_with(test_limits()) } @@ -633,6 +747,7 @@ mod tests { flume::Receiver, IngressConsumer, flume::Sender, + flume::Receiver, ) { make_ingress_inner(test_limits(), abort_rx) } @@ -644,6 +759,7 @@ mod tests { flume::Receiver, IngressConsumer, flume::Sender, + flume::Receiver, ) { let (abort_tx, abort_rx) = flume::unbounded::(); std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx @@ -658,6 +774,7 @@ mod tests { flume::Receiver, IngressConsumer, flume::Sender, + flume::Receiver, ) { let (tok_tx, _tok_rx) = flume::unbounded(); let (detok_tx, detok_rx) = flume::unbounded(); @@ -669,12 +786,30 @@ mod tests { }; let (ingress_producer, consumer) = ingress_ring(16); let (tm_tx, tm_rx) = flume::unbounded(); + let (mm_tx, mm_rx) = flume::unbounded(); // Keep the shutdown sender alive (leak) so its branch never fires — tests // end `run` by dropping `tm_tx`, not by shutdown. let (sd_tx, sd_rx) = flume::unbounded::<()>(); std::mem::forget(sd_tx); - let ingress = Ingress::new(tm_rx, abort_rx, senders, ingress_producer, limits, sd_rx); - (ingress, detok_rx, consumer, tm_tx) + let ingress = Ingress::new( + tm_rx, + abort_rx, + senders, + ingress_producer, + limits, + test_mm(mm_tx, true), + sd_rx, + ); + (ingress, detok_rx, consumer, tm_tx, mm_rx) + } + + /// An [`Mm`] over `tx` with a fresh sidecar. + fn test_mm(tx: flume::Sender, enabled: bool) -> Mm { + Mm { + enabled, + tx, + sidecar: Default::default(), + } } /// Both abort sources do the same two things: drop the detok entry so no @@ -696,7 +831,7 @@ mod tests { let (ingress_producer, consumer) = ingress_ring(16); let (sd_tx, sd_rx) = flume::unbounded::<()>(); std::mem::forget(sd_tx); - let ingress = Ingress::new( + let mut ingress = Ingress::new( flume::unbounded().1, flume::unbounded().1, Senders { @@ -707,6 +842,7 @@ mod tests { }, ingress_producer, test_limits(), + test_mm(flume::unbounded().0, true), sd_rx, ); @@ -960,7 +1096,7 @@ mod tests { /// to the ring, after registration — so it must be deregistered, not leaked. #[test] fn over_context_request_deregisters_and_never_reaches_the_ring() { - let (ingress, detok_rx, consumer, _tm_tx) = make_ingress_with(Limits { + let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress_with(Limits { context_len: 4, ..test_limits() }); @@ -993,7 +1129,7 @@ mod tests { /// pins. Nothing may reach the scheduler ring. #[test] fn detokenize_flows_register_then_decode_and_skips_the_ring() { - let (ingress, detok_rx, consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress(); let (tx, mut rx) = mpsc::channel(8); ingress.drive(Request { rid: "41".into(), @@ -1028,7 +1164,7 @@ mod tests { /// leak and no decode job to drop). #[test] fn detokenize_negative_ids_reject_before_registration() { - let (ingress, detok_rx, consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress(); let (tx, mut rx) = mpsc::channel(8); ingress.drive(Request { rid: "43".into(), @@ -1071,7 +1207,15 @@ mod tests { let (_tm_tx, tm_rx) = flume::unbounded(); let (sd_tx, sd_rx) = flume::unbounded::<()>(); std::mem::forget(sd_tx); - let ingress = Ingress::new(tm_rx, abort_rx, senders, producer, test_limits(), sd_rx); + let mut ingress = Ingress::new( + tm_rx, + abort_rx, + senders, + producer, + test_limits(), + test_mm(flume::unbounded().0, true), + sd_rx, + ); ingress.on_abort(AbortSource::Guard("pushed".into())); ingress.on_abort(AbortSource::Guard("dropped".into())); @@ -1108,7 +1252,7 @@ mod tests { #[test] fn pre_registration_failure_does_not_deregister() { // Rejected inside `validate` (out-of-vocab id), which runs before registration. - let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress(); let mut req = generate_req(41, SamplingParams::default()); if let RequestKind::Generate(g) = &mut req.kind { g.input_ids = Some(vec![2_000_000_000]); @@ -1121,7 +1265,7 @@ mod tests { ); // A post-registration reject still deregisters (the leak fix stays fixed). - let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress(); ingress.drive(generate_req( 42, SamplingParams { @@ -1140,7 +1284,7 @@ mod tests { /// sees `Register` then `Deregister`. Regression for RSS growth on bad input. #[test] fn rejected_request_deregisters_from_shard() { - let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress(); // top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it. let bad = SamplingParams { top_p: 2.0, @@ -1167,7 +1311,7 @@ mod tests { /// and kills the scheduler process (`make_ingress` bounds vocab at 1000). #[test] fn out_of_vocab_input_ids_rejected() { - let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress(); let mut req = generate_req(21, SamplingParams::default()); if let RequestKind::Generate(g) = &mut req.kind { g.input_ids = Some(vec![1, 2_000_000_000]); @@ -1185,7 +1329,7 @@ mod tests { /// Same guard for negative ids and for `token_ids_logprob` entries. #[test] fn negative_and_logprob_token_ids_rejected() { - let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress(); let mut req = generate_req(22, SamplingParams::default()); if let RequestKind::Generate(g) = &mut req.kind { g.input_ids = Some(vec![-1]); @@ -1196,7 +1340,7 @@ mod tests { Ok(_) => panic!("negative token id must not be admitted"), } - let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress(); let mut req = generate_req(23, SamplingParams::default()); if let RequestKind::Generate(g) = &mut req.kind { g.token_ids_logprob = Some(vec![999_999]); @@ -1211,7 +1355,7 @@ mod tests { /// A valid request is registered and handed onward — never deregistered. #[test] fn admitted_request_keeps_registration() { - let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress(); // Empty map → all sampling defaults, passes normalization. ingress.drive(generate_req(9, SamplingParams::default())); @@ -1229,7 +1373,7 @@ mod tests { /// path and deregistered, not leaked. #[test] fn tokenize_failure_deregisters_via_ingress() { - let (ingress, detok_rx, _consumer, tm_tx) = make_ingress(); + let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress(); // The pool marks a failed encode as `Failed(err)` before returning it. let mut req = generate_req(11, SamplingParams::default()); let _ = req @@ -1253,7 +1397,7 @@ mod tests { fn abort_deregisters_from_shard() { // Aborts arrive on their own unbounded lane now, not the request inbox. let (abort_tx, abort_rx) = flume::unbounded::(); - let (ingress, detok_rx, _consumer, tm_tx) = make_ingress_with_abort(abort_rx); + let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress_with_abort(abort_rx); abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap(); drop(abort_tx); drop(tm_tx); @@ -1270,7 +1414,7 @@ mod tests { /// rejected; its registration is untouched. #[test] fn tokenized_return_pushes_without_deregister() { - let (ingress, detok_rx, _consumer, tm_tx) = make_ingress(); + let (ingress, detok_rx, _consumer, tm_tx, _mm_rx) = make_ingress(); let mut req = generate_req(15, SamplingParams::default()); // Simulate a successful pool return: ids filled, PreSendValidating. if let RequestKind::Generate(g) = &mut req.kind { @@ -1293,7 +1437,7 @@ mod tests { #[test] fn tokenize_pool_gone_deregisters() { // `make_ingress` drops the tok receiver, so `tok.send` fails. - let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let (mut ingress, detok_rx, _consumer, _tm_tx, _mm_rx) = make_ingress(); // No ids → NeedsTokenize → Tokenizing branch. let mut req = generate_req(21, SamplingParams::default()); if let RequestKind::Generate(g) = &mut req.kind { @@ -1311,4 +1455,157 @@ mod tests { ); assert!(detok_rx.try_recv().is_err(), "no further shard messages"); } + + /// Build a generate request carrying an image. The parked entry and the + /// `MmEncoded` resume path agree on identity via the rid string. + fn mm_generate_req(rid: &str) -> Request { + let (tx, _rx) = mpsc::channel(8); + Request { + rid: rid.to_string().into(), + state: RequestState::Received, + sink: EgressSink::Local(tx), + kind: RequestKind::Generate(Box::new(GenerateRequest { + rid: rid.to_string().into(), + text: Some(" hi".into()), + mm: Some(Box::new(crate::message::MmData { + image_data: Some(rmpv::Value::from("data:image/jpeg;base64,xxxx")), + ..Default::default() + })), + ..Default::default() + })), + } + } + + /// An abort while the request is parked for MM cancels it: the pending + /// entry is removed, the worker's late result is dropped, and its parked + /// sidecar entry is purged — no scheduler work runs for a dead client. + #[test] + fn abort_cancels_parked_mm_request() { + let (mut ingress, _detok_rx, consumer, _tm_tx, mm_rx) = make_ingress(); + ingress.drive(mm_generate_req("mm-gone")); + mm_rx.try_recv().expect("parked to mm pool"); + + // The worker parks its result, as it always does before MmEncoded. + ingress.mm.sidecar.park( + "mm-gone".into(), + crate::mm::MmSidecarEntry { + features: crate::mm::FeatureStore::Inline(vec![]), + grids: vec![], + hashes: vec![], + offsets: vec![], + mrope: vec![], + mrope_delta: 0, + }, + ); + ingress.on_abort(AbortSource::Guard("mm-gone".to_string().into())); + assert_eq!(consumer.drain(16).headers.len(), 1, "only the AbortReq"); + + // The late result must be dropped, not queued, and the sidecar purged. + ingress.on_mm_encoded("mm-gone".to_string().into(), vec![5, 6]); + assert!( + consumer.drain(16).headers.is_empty(), + "cancelled, not queued" + ); + assert!(ingress.mm.sidecar.take("mm-gone").is_none(), "entry purged"); + } + + /// A multimodal request parks in `Encoding` (submitted to the mm worker + /// pool, not the tokenizer pool, not the ring) until `MmEncoded` resumes + /// it → ring. + #[test] + fn mm_request_parks_then_mm_encoded_pushes_to_ring() { + let (mut ingress, _detok_rx, consumer, _tm_tx, mm_rx) = make_ingress(); + ingress.drive(mm_generate_req("mm-1")); + + // Submitted to the mm pool with the typed work item; nothing on the ring yet. + let sub = mm_rx.try_recv().expect("mm pool must receive the request"); + assert_eq!(sub.rid.as_str(), "mm-1"); + assert_eq!(sub.work.text.as_deref(), Some(" hi")); + assert!(sub.work.input_ids.is_none(), "no client input_ids"); + assert_eq!( + sub.work.image_data.as_ref().and_then(|v| v.as_str()), + Some("data:image/jpeg;base64,xxxx") + ); + assert!(consumer.drain(16).headers.is_empty(), "parked, not queued"); + + // The worker returns the final expanded ids → pushed to the ring. + ingress.on_mm_encoded("mm-1".to_string().into(), vec![5, 6, 7, 8]); + let batch = consumer.drain(16); + assert_eq!(batch.headers.len(), 1); + assert_eq!( + batch.lengths, + vec![4], + "expanded ids ride the columnar cell" + ); + } + + /// A worker failure rejects the parked request (deregister, no ring push). + #[test] + fn mm_failure_rejects_parked_request() { + let (mut ingress, detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress(); + ingress.drive(mm_generate_req("mm-2")); + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. })), + "registered before parking", + ); + + ingress.on_mm_failed("mm-2".to_string().into(), "bad image".into()); + assert!( + matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) + if rid.as_str() == "mm-2"), + "mm failure must deregister", + ); + assert!(consumer.drain(16).headers.is_empty(), "nothing queued"); + } + + /// On a non-multimodal model (`Mm::enabled == false`), image_data is silently + /// ignored and the request tokenizes as plain text — the Python + /// TokenizerManager behavior when `mm_processor is None`. + #[test] + fn mm_fields_ignored_when_disabled() { + let (tok_tx, tok_rx) = flume::unbounded(); + let (detok_tx, _detok_rx) = flume::unbounded(); + let senders = Senders { + tm: flume::unbounded().0, + abort: flume::unbounded().0, + tok: tok_tx, + detok: vec![detok_tx], + }; + let (ingress_producer, _consumer) = ingress_ring(16); + let (_tm_tx, tm_rx) = flume::unbounded(); + let (mm_tx, mm_rx) = flume::unbounded(); + let (abort_tx, abort_rx) = flume::unbounded::(); + std::mem::forget(abort_tx); + let (sd_tx, sd_rx) = flume::unbounded::<()>(); + std::mem::forget(sd_tx); + let mut ingress = Ingress::new( + tm_rx, + abort_rx, + senders, + ingress_producer, + test_limits(), + test_mm(mm_tx, false), + sd_rx, + ); + + ingress.drive(mm_generate_req("mm-3")); + assert!( + mm_rx.try_recv().is_err(), + "mm disabled: nothing submitted to the mm channel", + ); + assert!( + tok_rx.try_recv().is_ok(), + "request must fall through to plain tokenization", + ); + } + + /// A late mm result for a rid that is no longer parked is dropped without + /// panicking (e.g. hash-collision overwrite) — regression guard. + #[test] + fn late_mm_result_is_dropped() { + let (mut ingress, _detok_rx, consumer, _tm_tx, _mm_rx) = make_ingress(); + ingress.on_mm_encoded("ghost".to_string().into(), vec![1]); + ingress.on_mm_failed("ghost".to_string().into(), "boom".into()); + assert!(consumer.drain(16).headers.is_empty()); + } } diff --git a/test/registered/unit/multimodal/rust/_mm_rust_utils.py b/test/registered/unit/multimodal/rust/_mm_rust_utils.py index b8cc6c759..7fe1c9be3 100644 --- a/test/registered/unit/multimodal/rust/_mm_rust_utils.py +++ b/test/registered/unit/multimodal/rust/_mm_rust_utils.py @@ -88,5 +88,12 @@ def image_bytes(width, height, seed=0): 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}) +def spec_json(config, image_token_id=IMAGE_TOKEN_ID, resample="aten_u8"): + return json.dumps( + { + "family": "qwen_vl", + "image_token_id": image_token_id, + "resample": resample, + **config, + } + ) diff --git a/test/registered/unit/multimodal/rust/qwen/_fixtures.py b/test/registered/unit/multimodal/rust/qwen/_fixtures.py new file mode 100644 index 000000000..e1bfc6d31 --- /dev/null +++ b/test/registered/unit/multimodal/rust/qwen/_fixtures.py @@ -0,0 +1,102 @@ +from types import SimpleNamespace + +import numpy as np +from tokenizers import Tokenizer, decoders +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import WhitespaceSplit +from transformers import ( + PreTrainedTokenizerFast, + Qwen2VLProcessor, + Qwen2VLVideoProcessor, +) +from transformers.models.qwen2_vl.image_processing_qwen2_vl import ( + Qwen2VLImageProcessor as HfQwenImageProcessor, +) + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.multimodal.processors.qwen_vl import ( # noqa: E402 + QwenVLImageProcessor, +) + +register_cpu_ci(est_time=0, suite="base-a-test-cpu", disabled="Qwen test fixtures") + + +def make_processor(config, image_processor_cls=None): + """A ``QwenVLImageProcessor`` over a tiny hand-built tokenizer. + ``image_processor_cls`` picks the HF backend; they resample differently.""" + image_processor_cls = image_processor_cls or HfQwenImageProcessor + vocab = [ + "", + "<|vision_start|>", + "<|image_pad|>", + "<|vision_end|>", + "hello", + "<|video_pad|>", + "", + ] + backend = Tokenizer( + WordLevel( + {token: index for index, token in enumerate(vocab)}, unk_token=vocab[0] + ) + ) + backend.pre_tokenizer, backend.decoder = WhitespaceSplit(), decoders.Fuse() + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=backend, + unk_token=vocab[0], + pad_token=vocab[-1], + additional_special_tokens=vocab[1:4] + [vocab[5]], + ) + processor = Qwen2VLProcessor( + image_processor=image_processor_cls(**config), + video_processor=Qwen2VLVideoProcessor(), + tokenizer=tokenizer, + ) + hf_config = SimpleNamespace( + model_type="qwen2_5_vl", + architectures=["Qwen2_5_VLForConditionalGeneration"], + vision_start_token_id=1, + image_token_id=2, + vision_end_token_id=3, + video_token_id=5, + vision_config=SimpleNamespace(spatial_merge_size=2, tokens_per_second=2), + ) + server_args = SimpleNamespace( + # Non-auto: get_resolved_model_impl would choke on a SimpleNamespace. + model_impl="sglang", + keep_mm_feature_on_device=False, + mm_feature_transport="cpu", + disable_fast_image_processor=True, + skip_tokenizer_init=False, + # Read by NativeMmHost._use_feature_shm (single-rank fixture → the + # inline zero-copy transport, like the 1-GPU e2e). + tp_size=1, + dist_init_addr=None, + mm_process_config={}, + mm_io_worker_num=1, + mm_processor_worker_num=1, + tokenizer_worker_num=1, + base_gpu_id=0, + ) + return QwenVLImageProcessor( + hf_config, server_args, processor, None, skip_mm_pool=True + ) + + +def snapshot(input_ids, output): + return { + "input_ids": tuple(input_ids), + "grids": tuple( + tuple(item.image_grid_thw.flatten().tolist()) for item in output.mm_items + ), + "offsets": tuple(item.offsets[0] for item in output.mm_items), + "features": np.concatenate( + [item.feature.detach().cpu().numpy() for item in output.mm_items] + ), + "mrope": output.mrope_positions.detach().cpu().numpy(), + "delta": int(output.mrope_position_delta.item()), + "tokens": (output.im_start_id, output.im_token_id, output.im_end_id), + } diff --git a/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py b/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py new file mode 100644 index 000000000..1f2f542a5 --- /dev/null +++ b/test/registered/unit/multimodal/rust/qwen/test_e2e_parity.py @@ -0,0 +1,162 @@ +"""End-to-end parity at the scheduler-input boundary. + +`test_preprocess.py` pins the `preprocess` binding; this drives the whole native +path — the `process_native_mm` driver, then `NativeMmHost.build_native_mm` — and +compares every field the scheduler reads against the Python `mm_processor`. +Bitwise, for both HF backends: the Rust resize clones PIL's fixed-point bicubic +and ATen's uint8 antialias kernel, so whichever one a server is configured with +is reproduced exactly. +""" + +import asyncio +import base64 +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +import numpy as np + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.rust_server import NativeMmHost # noqa: E402 + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _fixtures import make_processor, snapshot # noqa: E402 +from _mm_rust_utils import PROCESSOR_CONFIGS, image_bytes, load_core # noqa: E402 + +register_cpu_ci(est_time=40, suite="base-a-test-cpu") + +CORE = load_core() +DRIVER = getattr(getattr(CORE, "qwen_vl", None), "process_native_mm", None) + +# The fixture tokenizer's vocab (see `_fixtures.make_processor`): +# 1 = <|vision_start|>, 2 = <|image_pad|>, 3 = <|vision_end|>, 4 = "hello". +PROMPT_PER_IMAGE = [1, 2, 3, 4] + + +@unittest.skipUnless(DRIVER, "sglang-mm native Qwen driver not built") +class TestQwenE2eParity(CustomTestCase): + """The torchvision backend a default server runs.""" + + image_processor = "Qwen2VLImageProcessor" + + def setUp(self): + import transformers.models.qwen2_vl as qwen2_vl + + self.processor = make_processor( + PROCESSOR_CONFIGS["qwen2_5_vl"], getattr(qwen2_vl, self.image_processor) + ) + + def tearDown(self): + self.processor.io_executor.shutdown() + self.processor.cpu_executor.shutdown() + + # --- the two paths under comparison --- + + def native_spec(self): + """Resolve the spec through the production gate, so a gate that stops + recognizing this image processor fails here too.""" + from sglang.srt.managers.multimodal_processor import import_processors + + import_processors("sglang.srt.multimodal.processors") + # Skip __init__: it would build a processor; reuse the fixture's. + host = NativeMmHost.__new__(NativeMmHost) + host.model_config = SimpleNamespace(hf_config=self.processor.hf_config) + host._processor = self.processor._processor + host.server_args = self.processor.server_args + spec = host.resolve_native_spec() + self.assertIsNotNone(spec, f"gate rejected {self.image_processor}") + return spec + + def run_native(self, spec, sources): + """The Rust path: the `process_native_mm` driver, then the drain + adapter — the same two steps `RustServer.drain` performs.""" + ids, features, grids, hashes, offsets, mrope, delta = DRIVER( + PROMPT_PER_IMAGE * len(sources), sources, spec.rust_json() + ) + # The shape of Rust's MmHandoff, inline transport (test_build_native_mm + # pins the shm shape). + handoff = SimpleNamespace( + features=features, + shm_names=None, + grids=grids, + hashes=hashes, + offsets=offsets, + mrope=mrope, + mrope_delta=delta, + ) + return snapshot(ids, NativeMmHost.build_native_mm(spec, handoff)) + + def run_python(self, sources): + """The reference path: the Python `mm_processor` the scheduler would use.""" + output = asyncio.run( + self.processor.process_mm_data_async( + image_data=sources, + input_text=PROMPT_PER_IMAGE * len(sources), + request_obj=SimpleNamespace( + video_data=None, audio_data=None, rid="parity" + ), + ) + ) + return snapshot(output.input_ids, output) + + def assert_parity(self, spec, sources): + rust, python = self.run_native(spec, sources), self.run_python(sources) + for field in ("input_ids", "grids", "offsets", "delta", "tokens"): + with self.subTest(field=field): + self.assertEqual(rust[field], python[field]) + with self.subTest(field="mrope"): + np.testing.assert_array_equal(rust["mrope"], python["mrope"]) + with self.subTest(field="features"): + # Bytes, not allclose: the scheduler gets these float32 buffers verbatim. + self.assertEqual(rust["features"].tobytes(), python["features"].tobytes()) + + # --- inputs --- + + def source_forms(self, directory): + """One image in each accepted transport form, plus a two-image batch.""" + first, second = image_bytes(96, 80), image_bytes(112, 88, 1) + path = Path(directory) / "image.png" + path.write_bytes(first) + return { + "raw_bytes": [first], + "data_url": ["data:image/png;base64," + base64.b64encode(first).decode()], + "file_uri": [path.as_uri()], + "two_image_batch": [first, second], + } + + def test_parity_across_source_forms(self): + spec = self.native_spec() + with tempfile.TemporaryDirectory() as directory: + for form, sources in self.source_forms(directory).items(): + with self.subTest(form=form): + self.assert_parity(spec, sources) + + +class TestQwenE2eParityPil(TestQwenE2eParity): + """The PIL backend (`--disable-fast-image-processor`): the same parity + suite, plus the transport-invariance check — a property of the native path + alone, so running it under one backend is enough.""" + + image_processor = "Qwen2VLImageProcessorPil" + + def test_source_form_is_transport_only(self): + """Bytes, a data: URL and a file:// path must yield one identical result.""" + spec = self.native_spec() + with tempfile.TemporaryDirectory() as directory: + forms = self.source_forms(directory) + features = { + form: self.run_native(spec, forms[form])["features"].tobytes() + for form in ("raw_bytes", "data_url", "file_uri") + } + self.assertEqual(len(set(features.values())), 1, sorted(features)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/rust/qwen/test_native_mm_host.py b/test/registered/unit/multimodal/rust/qwen/test_native_mm_host.py new file mode 100644 index 000000000..f53618b84 --- /dev/null +++ b/test/registered/unit/multimodal/rust/qwen/test_native_mm_host.py @@ -0,0 +1,105 @@ +"""The mm-item hash contract of the native Qwen path. + +The two paths hash different things, deliberately. Native hashes are `content_hash` +over the raw encoded source bytes, computed on an MM worker; the Python path hashes +the decoded feature tensor. Both feed `set_pad_value`, so the native drain can skip +`hash_feature` on the scheduler loop — the point of precomputing them. + +Field-by-field parity of everything else is `test_e2e_parity.py`. +""" + +import asyncio +import base64 +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.mm_utils import hash_feature # noqa: E402 +from sglang.srt.managers.rust_server import NativeMmHost # noqa: E402 + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from _fixtures import make_processor # noqa: E402 +from _mm_rust_utils import PROCESSOR_CONFIGS, image_bytes, load_core # noqa: E402 + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + +CORE = load_core() +DRIVER = getattr(getattr(CORE, "qwen_vl", None), "process_native_mm", None) + + +def raw_bytes(source): + """The encoded bytes behind any accepted source form.""" + if isinstance(source, bytes): + return source + if source.startswith("data:"): + return base64.b64decode(source.split(",", 1)[1]) + return Path(source.removeprefix("file://")).read_bytes() + + +@unittest.skipUnless(DRIVER, "sglang-mm native Qwen driver not built") +class TestQwenNativeMmHashes(CustomTestCase): + def setUp(self): + from sglang.srt.managers.multimodal_processor import import_processors + + import_processors("sglang.srt.multimodal.processors") + self.processor = make_processor(PROCESSOR_CONFIGS["qwen2_5_vl"]) + + def tearDown(self): + self.processor.io_executor.shutdown() + self.processor.cpu_executor.shutdown() + + def native_hashes(self, sources): + """Per-item hashes the Rust driver returns, via the production gate.""" + host = NativeMmHost.__new__(NativeMmHost) + host.model_config = SimpleNamespace(hf_config=self.processor.hf_config) + host._processor = self.processor._processor + host.server_args = self.processor.server_args + spec = host.resolve_native_spec() + self.assertIsNotNone(spec, "gate rejected the fixture processor") + input_ids = [t for _ in sources for t in (1, 2, 3, 4)] + return DRIVER(input_ids, sources, spec.rust_json())[3] + + def test_native_hashes_the_raw_source_bytes(self): + """Same image in any source form hashes identically, because the hash is + over the bytes and not over anything the transport changes.""" + first, second = image_bytes(96, 80), image_bytes(112, 88, 1) + data_url = "data:image/png;base64," + base64.b64encode(first).decode() + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "image.png" + path.write_bytes(first) + for sources in ([first], [data_url], [path.as_uri()], [first, second]): + with self.subTest(n=len(sources), form=type(sources[0]).__name__): + hashes = self.native_hashes(sources) + self.assertEqual( + list(hashes), + [CORE.common.content_hash(raw_bytes(s)) for s in sources], + ) + + def test_python_hashes_the_feature(self): + """The contrast: `set_pad_value` on the Python path derives its hash from + the feature tensor, which is why the native path must precompute one.""" + output = asyncio.run( + self.processor.process_mm_data_async( + image_data=[image_bytes(96, 80)], + input_text=[1, 2, 3, 4], + request_obj=SimpleNamespace( + video_data=None, audio_data=None, rid="hash" + ), + ) + ) + for item in output.mm_items: + expected = hash_feature(item.feature) + item.set_pad_value() + self.assertEqual(item.hash, expected) + + +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 index d9843a045..72efeb937 100644 --- a/test/registered/unit/multimodal/rust/qwen/test_preprocess.py +++ b/test/registered/unit/multimodal/rust/qwen/test_preprocess.py @@ -5,14 +5,13 @@ Covers ``QwenVlProcessor::process_item`` and ``smart_resize`` in 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. +Both HF processors are pinned, because they resample differently and either can +be what a server 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 +clones both kernels, so each is asserted with zero tolerance under the +``resample`` its spec selects. """ import sys @@ -43,48 +42,39 @@ 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): + def _assert_bit_exact(self, processor, resample): + """Zero tolerance, so any drift in any stage — smart_resize geometry, the + fixed-point kernel, rescale/normalize, HF patch order — shows up here + instead of being absorbed.""" 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) + image_bytes(*size, seed=index), + spec_json(config, resample=resample), ) 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() + np.testing.assert_array_equal( + 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) + self._assert_bit_exact(Qwen2VLImageProcessorPil, "pil") - 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).""" + def test_features_match_torchvision_processor_exactly(self): from transformers.models.qwen2_vl.image_processing_qwen2_vl import ( Qwen2VLImageProcessor, ) - self._assert_matches(Qwen2VLImageProcessor, max_diff=0.035, mean_diff=1e-3) + self._assert_bit_exact(Qwen2VLImageProcessor, "aten_u8") def test_smart_resize_matches_python(self): from sglang.srt.multimodal.processors.qwen_vl import smart_resize diff --git a/test/registered/unit/multimodal/rust/shared/test_build_native_mm.py b/test/registered/unit/multimodal/rust/shared/test_build_native_mm.py new file mode 100644 index 000000000..c29ad76cf --- /dev/null +++ b/test/registered/unit/multimodal/rust/shared/test_build_native_mm.py @@ -0,0 +1,166 @@ +"""``NativeMmHost.build_native_mm`` (managers/rust_server.py): the drain-time +wrapping contracts — tensors are zero-copy views over the Rust-owned buffers, and +pad values come from worker-precomputed hashes, since the scheduler loop must +never hash features. Synthetic buffers, so this needs no Rust extension.""" + +import os +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.rust_server import NativeMmHost, NativeMmSpec # noqa: E402 + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestBuildNativeMm(CustomTestCase): + def setUp(self): + # feature_dim == 3 * temporal_patch_size * patch_size**2 == 6. + self.spec = NativeMmSpec( + family="qwen_vl", + feature_shm=False, + image_token_id=10, + patch_size=1, + merge_size=1, + temporal_patch_size=2, + min_pixels=1, + max_pixels=1 << 30, + image_mean=(0.0, 0.0, 0.0), + image_std=(1.0, 1.0, 1.0), + resample="aten_u8", + vision_start_token_id=11, + vision_end_token_id=12, + video_token_id=13, + ) + + GRIDS = [(1, 2, 2), (1, 1, 1)] + HASHES = [101, 202] + OFFSETS = [(2, 5), (8, 8)] + + def transport(self, features): + """Inline: the features ride the numpy array itself.""" + return dict(features=features, shm_names=None) + + def build(self): + features = np.arange(30, dtype=np.float32) + output = NativeMmHost.build_native_mm( + self.spec, + SimpleNamespace( # the shape of Rust's MmHandoff + grids=self.GRIDS, + hashes=self.HASHES, + offsets=self.OFFSETS, + mrope=np.arange(30, dtype=np.int64), + mrope_delta=-3, + **self.transport(features), + ), + ) + return output, features + + def test_wraps_and_slices_native_buffers(self): + output, features = self.build() + self.assertEqual( + [tuple(item.feature.shape) for item in output.mm_items], [(4, 6), (1, 6)] + ) + self.assertEqual([item.hash for item in output.mm_items], [101, 202]) + self.assertEqual( + [item.offsets for item in output.mm_items], [[(2, 5)], [(8, 8)]] + ) + self.assertEqual(tuple(output.mrope_positions.shape), (3, 10)) + self.assertEqual(output.mrope_position_delta.item(), -3) + self.assertEqual( + (output.im_start_id, output.im_token_id, output.im_end_id), (11, 10, 12) + ) + features[0] = 99 + self.assertEqual(output.mm_items[0].feature[0, 0].item(), 99) + + def test_optional_pad_values_use_precomputed_hashes(self): + from sglang.srt.managers.schedule_batch import _compute_pad_value + + # The whole point of worker-precomputed hashes is that the scheduler + # loop never runs hash_feature — make any call a hard failure. + with ( + patch.dict(os.environ, {"SGLANG_MM_PRECOMPUTE_HASH": "1"}), + patch( + "sglang.srt.managers.mm_utils.hash_feature", + side_effect=AssertionError("scheduler loop must not hash features"), + ), + ): + output, _ = self.build() + self.assertEqual( + [item.pad_value for item in output.mm_items], + [_compute_pad_value(101), _compute_pad_value(202)], + ) + + +class TestBuildNativeMmShm(TestBuildNativeMm): + """The shm entry shape (TP>1): features arrive as named POSIX segments, and + each item becomes a ``ShmPointerMMData`` stub whose ``materialize()`` yields + that item's slice — and unlinks, taking the cleanup duty exactly once.""" + + def setUp(self): + super().setUp() + self._segments = [] + + def tearDown(self): + # Defensive: unlink anything a failing test left behind. + for shm in self._segments: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def _park(self, features): + from multiprocessing import shared_memory + + names, row = [], 0 + for t, h, w in self.GRIDS: + n = t * h * w + payload = features[row * 6 : (row + n) * 6].tobytes() + shm = shared_memory.SharedMemory(create=True, size=len(payload)) + shm.buf[:] = payload + self._segments.append(shm) + names.append(shm.name) + row += n + return names + + def transport(self, features): + """Shm: the worker parked each item's slice in its own segment.""" + return dict(features=None, shm_names=self._park(features)) + + def test_wraps_and_slices_native_buffers(self): + import torch + + from sglang.srt.managers.mm_utils import ShmPointerMMData + + output, features = self.build() + for item in output.mm_items: + self.assertIsInstance(item.feature, ShmPointerMMData) + # The stub is a zero-copy view over the segment until materialized. + self.assertEqual( + [tuple(item.feature.shape) for item in output.mm_items], [(4, 6), (1, 6)] + ) + self.assertEqual( + [item.feature.precomputed_hash for item in output.mm_items], self.HASHES + ) + tensors = [item.feature.materialize() for item in output.mm_items] + expected = torch.from_numpy(features).reshape(-1, 6) + self.assertTrue(torch.equal(tensors[0], expected[:4])) + self.assertTrue(torch.equal(tensors[1], expected[4:])) + # materialize() unlinked: the names must be gone. + from multiprocessing import shared_memory + + for item in output.mm_items: + with self.assertRaises(FileNotFoundError): + shared_memory.SharedMemory(name=item.feature.shm_name) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/rust/shared/test_native_mm_gate.py b/test/registered/unit/multimodal/rust/shared/test_native_mm_gate.py new file mode 100644 index 000000000..151cd5aeb --- /dev/null +++ b/test/registered/unit/multimodal/rust/shared/test_native_mm_gate.py @@ -0,0 +1,61 @@ +"""The native-MM launch gate's family selection (managers/rust_server.py). + +``NATIVE_MM_FAMILIES`` decides which models the Rust pipeline serves natively; +for everything else ``native_mm_family_for`` must return ``None``, which +``RustServer.launch`` turns into a hard launch error. Pins that non-Qwen +multimodal models — Inkling being the in-tree case — keep their Python +processor and never match a native family, so growing the registry cannot +silently reroute them. +""" + +import unittest +from types import SimpleNamespace + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.multimodal_processor import ( # noqa: E402 + get_mm_processor_cls, + import_processors, +) +from sglang.srt.managers.rust_server import native_mm_family_for # noqa: E402 + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +def processor_cls_for(architecture, model_type): + """Through the production selection, as `resolve_native_spec` calls it.""" + hf_config = SimpleNamespace(architectures=[architecture], model_type=model_type) + return get_mm_processor_cls(hf_config, SimpleNamespace(model_impl="sglang")) + + +class TestNativeMmGate(CustomTestCase): + @classmethod + def setUpClass(cls): + import_processors("sglang.srt.multimodal.processors") + + def test_qwen_vl_resolves_its_family(self): + cls = processor_cls_for("Qwen2_5_VLForConditionalGeneration", "qwen2_5_vl") + family = native_mm_family_for(cls, "qwen2_5_vl") + self.assertEqual(family and family.name, "qwen_vl") + + def test_inkling_keeps_its_python_processor(self): + from sglang.srt.multimodal.processors.inkling import InklingMultimodalProcessor + + cls = processor_cls_for("InklingForConditionalGeneration", "inkling_model") + self.assertIs(cls, InklingMultimodalProcessor) + self.assertIsNone(native_mm_family_for(cls, "inkling_model")) + + def test_family_requires_both_processor_and_model_type(self): + qwen = processor_cls_for("Qwen2_5_VLForConditionalGeneration", "qwen2_5_vl") + self.assertIsNone(native_mm_family_for(qwen, "inkling_model")) + # Identity, not name: an override class must not match (the + # SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE contract). + impostor = type("QwenVLImageProcessor", (), {}) + self.assertIsNone(native_mm_family_for(impostor, "qwen2_5_vl")) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/rust/shared/test_partition_cores.py b/test/registered/unit/multimodal/rust/shared/test_partition_cores.py new file mode 100644 index 000000000..b23cf8437 --- /dev/null +++ b/test/registered/unit/multimodal/rust/shared/test_partition_cores.py @@ -0,0 +1,54 @@ +"""``RustServer._partition_cores`` (managers/rust_server.py): the pool cores must +be a *bounded* slice of this rank's allowed cores, not the whole remainder — +sibling TP ranks share the NUMA node, so an unbounded mask lets MM preprocessing +bursts preempt a sibling's CUDA-launch thread (measured: ~20 ms of ViT wall time +per image request on TP4). Pure computation, so no Rust extension needed.""" + +import unittest +from unittest.mock import patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.rust_server import RustServer # noqa: E402 + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def partition(node_cores, **kwargs): + with patch("os.sched_getaffinity", return_value=set(node_cores), create=True): + return RustServer._partition_cores(**kwargs) + + +class TestPartitionCores(CustomTestCase): + def test_pool_is_bounded_not_the_node_remainder(self): + # A 120-core NUMA node shared with sibling TP ranks: the pools must + # NOT get cores 2..119. + launch, pool = partition(range(120), mm_workers=8) + self.assertEqual(launch, [0, 1]) + self.assertEqual(pool, list(range(2, 14))) # max(8, 8 + 4) after reserve + + def test_budget_scales_with_mm_workers_with_a_floor(self): + _, pool_text = partition(range(120), mm_workers=0) + _, pool_mm = partition(range(120), mm_workers=16) + self.assertEqual(len(pool_text), 8) # the floor covers the I/O threads + self.assertEqual(len(pool_mm), 20) + + def test_small_allowance_degrades_gracefully(self): + # Fewer allowed cores than the budget: take what exists after the + # launch reserve, never raise. + launch, pool = partition(range(6), mm_workers=8) + self.assertEqual(launch, [0]) # min(2, 6 // 4) == 1 + self.assertEqual(pool, list(range(1, 6))) + # Below the split threshold: unpinned. + self.assertEqual(partition(range(3), mm_workers=8), (None, None)) + + def test_launch_and_pool_cores_are_disjoint(self): + launch, pool = partition(range(32), mm_workers=8) + self.assertFalse(set(launch) & set(pool)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/vlm/test_rust_native_mm_e2e.py b/test/registered/vlm/test_rust_native_mm_e2e.py new file mode 100644 index 000000000..412bc8953 --- /dev/null +++ b/test/registered/vlm/test_rust_native_mm_e2e.py @@ -0,0 +1,129 @@ +"""E2E: Rust tokenizer-manager native multimodal path (``SGLANG_RUST_SERVER=1``). + +Covers what the CPU parity units structurally cannot: the sidecar handoff, the +drain ordering, Rust-side tokenization of multimodal prompts, and the rejection +of inputs outside the native pipeline's scope (there is no Python fallback). +""" + +import base64 +import importlib.util +import io +import os +import unittest + +import numpy as np +import requests +from PIL import Image + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large") + +IMAGE_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/man_ironing_on_back_of_suv.png" +VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>" + + +def chat_prompt(question, image_count=1): + return ( + f"<|im_start|>user\n{VISION_BLOCK * image_count}{question}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + + +def solid_image_data_url(fmt): + buffer = io.BytesIO() + Image.fromarray(np.full((64, 64, 3), (255, 0, 0), dtype=np.uint8)).save( + buffer, format=fmt + ) + encoded = base64.b64encode(buffer.getvalue()).decode() + return f"data:image/{fmt.lower()};base64,{encoded}" + + +@unittest.skipIf( + importlib.util.find_spec("sglang.srt.server._core") is None, + "sglang-server rust extension not installed (e.g. AMD suite)", +) +class TestRustServerNativeMm(CustomTestCase): + env = {"SGLANG_RUST_SERVER": "1"} + + @classmethod + def setUpClass(cls): + cls.process = popen_launch_server( + DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST, + DEFAULT_URL_FOR_TEST, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=["--enable-multimodal", "--mem-fraction-static", "0.8"], + env={**os.environ, **cls.env}, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def generate(self, prompt, image_data): + response = requests.post( + DEFAULT_URL_FOR_TEST + "/generate", + json={ + "text": prompt, + "image_data": image_data, + "sampling_params": {"temperature": 0, "max_new_tokens": 48}, + }, + ) + self.assertEqual(response.status_code, 200, response.text) + return response.json()["text"].lower() + + def test_single_image_url(self): + text = self.generate( + chat_prompt("Describe this image in one sentence."), [IMAGE_URL] + ) + keywords = ("iron", "man", "taxi", "cab", "car", "suv", "street") + self.assertTrue(any(w in text for w in keywords), text) + + def test_two_images(self): + red = solid_image_data_url("PNG") + text = self.generate( + chat_prompt("What color is the second image?", image_count=2), + [IMAGE_URL, red], + ) + self.assertIn("red", text) + + def test_unsupported_format_is_rejected(self): + # An undecodable format must be rejected, never silently answered, and + # must not crash the server. PCX is the probe because feature unification + # in the server binary (dynamo-parsers → openai-harmony) widens + # sglang-mm's jpeg/png/webp/gif/bmp set to every image-crate default, so + # the probe has to be a format the image crate does not know at all. + response = requests.post( + DEFAULT_URL_FOR_TEST + "/generate", + json={ + "text": chat_prompt("What color is this image?"), + "image_data": [solid_image_data_url("PCX")], + "sampling_params": {"max_new_tokens": 8}, + }, + ) + self.assertIn(response.status_code, (400, 500), response.text) + + def test_corrupt_image_is_rejected(self): + response = requests.post( + DEFAULT_URL_FOR_TEST + "/generate", + json={ + "text": chat_prompt("Describe this image."), + "image_data": ["data:image/png;base64,aW52YWxpZA=="], + "sampling_params": {"max_new_tokens": 8}, + }, + ) + # Surfaced as Error::Encode (500); rejected without killing the server. + self.assertIn(response.status_code, (400, 500), response.text) + + +if __name__ == "__main__": + unittest.main(verbosity=3) diff --git a/test/registered/vlm/test_rust_native_mm_mmmu.py b/test/registered/vlm/test_rust_native_mm_mmmu.py new file mode 100644 index 000000000..1f5f72eb2 --- /dev/null +++ b/test/registered/vlm/test_rust_native_mm_mmmu.py @@ -0,0 +1,168 @@ +"""MMMU accuracy gate for the Rust tokenizer manager's native multimodal path. + +``test_rust_native_mm_e2e.py`` checks that the output is *valid*; this checks that +native Rust preprocessing yields *equally good* model inputs. A systematic skew +(wrong resample filter, channel order, normalization, patch layout) still reads as +fluent text and passes a keyword smoke check, but drops MMMU below the gate. + +The rust server has no ``/v1/chat/completions`` route yet, so the eval drives +``/generate`` with hand-rendered Qwen chat prompts instead of lmms-eval's OpenAI +client (``MMMUMixin``). +""" + +import importlib.util +import os +import tempfile +import time +import unittest + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.simple_eval_common import MessageList, SamplerBase +from sglang.test.simple_eval_mmmu_vlm import MMMUVLMEval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + dump_metric, + popen_launch_server, +) + +register_cuda_ci(est_time=900, stage="base-b", runner_config="1-gpu-large") + +MODEL = "Qwen/Qwen3.5-0.8B" +VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>" + +NUM_EXAMPLES = 100 +# Calibrated 2026-07-24 on H200: the native path scores 0.37 on this fixed subset +# at temperature 0 (two runs), matching the Python reference (0.37, same sampler +# and samples). The gate leaves headroom for batching nondeterminism. +MMMU_ACCURACY_THRESHOLD = 0.30 + + +class QwenGenerateVisionSampler(SamplerBase): + """Drive ``/generate`` with Qwen chat prompts and ``image_data``. + + ``MMMUVLMEval`` emits OpenAI-style messages mixing ``text`` and ``image_url`` + parts. This sampler renders the Qwen chat format by hand — each image part + becomes a ``VISION_BLOCK`` at its original position — and ships the images + through ``image_data``. + """ + + def __init__(self, base_url: str, max_tokens: int = 1024): + self.generate_url = base_url + "/generate" + self.max_tokens = max_tokens + + def __call__(self, message_list: MessageList) -> str: + segments = [] + images = [] + for message in message_list: + content = message["content"] + parts = ( + [{"type": "text", "text": content}] + if isinstance(content, str) + else content + ) + for part in parts: + if part["type"] == "image_url": + images.append(part["image_url"]["url"]) + segments.append(VISION_BLOCK) + else: + segments.append(part["text"]) + prompt = ( + "<|im_start|>user\n" + + "".join(segments) + + "<|im_end|>\n<|im_start|>assistant\n" + ) + payload = { + "text": prompt, + "image_data": images, + "sampling_params": { + "temperature": 0, + "max_new_tokens": self.max_tokens, + }, + } + # Retry transient failures but fail loudly when they persist: returning "" + # would silently degrade the score and blur the gate. + for attempt in range(3): + try: + response = requests.post(self.generate_url, json=payload, timeout=600) + response.raise_for_status() + return response.json()["text"] + except requests.RequestException: + if attempt == 2: + raise + time.sleep(2**attempt) + + +@unittest.skipIf( + importlib.util.find_spec("sglang.srt.server._core") is None, + "sglang-server rust extension not installed (e.g. AMD suite)", +) +class TestRustNativeMmMMMU(CustomTestCase): + @classmethod + def setUpClass(cls): + # Capture the server log so the test can pin that the native MM + # pipeline is active. + cls.log_dir = tempfile.TemporaryDirectory() + cls.server_logs = tuple( + open(os.path.join(cls.log_dir.name, name), "w") + for name in ("stdout.log", "stderr.log") + ) + cls.process = popen_launch_server( + MODEL, + DEFAULT_URL_FOR_TEST, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=["--enable-multimodal", "--mem-fraction-static", "0.8"], + env={**os.environ, "SGLANG_RUST_SERVER": "1"}, + return_stdout_stderr=cls.server_logs, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + if hasattr(cls, "server_logs"): + for f in cls.server_logs: + f.close() + if hasattr(cls, "log_dir"): + cls.log_dir.cleanup() + + def _read_server_log(self): + text = [] + for f in self.server_logs: + with open(f.name) as reader: + text.append(reader.read()) + return "\n".join(text) + + def test_mmmu_accuracy(self): + # Guard the path under test: if the model ever drops off + # NATIVE_MM_FAMILIES, launch fails and this names why. + self.assertIn( + "native MM pipeline enabled", + self._read_server_log(), + "rust server did not enable the native MM pipeline for " + f"{MODEL}; this test must exercise the native path", + ) + + eval_obj = MMMUVLMEval(num_examples=NUM_EXAMPLES, num_threads=32) + sampler = QwenGenerateVisionSampler(base_url=DEFAULT_URL_FOR_TEST) + result = eval_obj(sampler) + print(f"MMMU metrics: {result.metrics}") + dump_metric( + "mmmu_score", + result.score, + labels={"model": MODEL, "eval": "mmmu", "api": "generate-rust-native-mm"}, + ) + self.assertGreaterEqual( + result.score, + MMMU_ACCURACY_THRESHOLD, + f"Rust native MM path scored {result.score:.4f} on MMMU, below the " + f"{MMMU_ACCURACY_THRESHOLD:.2f} gate", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=3)