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

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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)
 
 

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
 

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)

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,

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: