From 1da86b9801d5a5d2eedc2131a9960ce3ed069ece Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 31 Aug 2026 12:28:43 -0700 Subject: [PATCH] [Rust] Split and rename embedded server components (#37220) --- python/sglang/srt/managers/rust_server.py | 869 ------------ python/sglang/srt/managers/scheduler.py | 2 +- .../scheduler_components/idle_sleeper.py | 2 +- .../scheduler_components/output_streamer.py | 2 +- .../scheduler_components/request_receiver.py | 2 +- python/sglang/srt/rust_server/config.py | 146 ++ python/sglang/srt/rust_server/multimodal.py | 319 +++++ python/sglang/srt/rust_server/server.py | 438 ++++++ rust/sglang-mm/Cargo.toml | 2 +- rust/sglang-mm/src/qwen_vl/mod.rs | 8 +- rust/sglang-server/Cargo.toml | 2 +- rust/sglang-server/README.md | 11 + rust/sglang-server/src/api_server/openai.rs | 3 + .../src/api_server/openai/template.rs | 1230 +---------------- .../api_server/openai/template_builtins.rs | 284 ++++ .../src/api_server/openai/template_legacy.rs | 605 ++++++++ .../src/api_server/openai/template_loader.rs | 351 +++++ rust/sglang-server/src/lib.rs | 49 +- rust/sglang-server/src/message/config.rs | 18 +- rust/sglang-server/src/message/sampling.rs | 14 +- .../src/multi_modality/worker.rs | 2 +- rust/sglang-server/src/tokenizer_manager.rs | 2 + .../src/tokenizer_manager/to_scheduler.rs | 1105 +-------------- .../tokenizer_manager/to_scheduler_tests.rs | 892 ++++++++++++ .../tokenizer_manager/to_scheduler_types.rs | 50 + .../to_scheduler_validation.rs | 160 +++ .../src/tokenizer_manager/wiring.rs | 2 +- rust/sglang-server/src/utils/runtime.rs | 15 +- .../unit/multimodal/rust/_mm_rust_utils.py | 2 +- .../unit/multimodal/rust/qwen/_fixtures.py | 2 +- .../unit/multimodal/rust/qwen/test_driver.py | 8 +- .../multimodal/rust/qwen/test_e2e_parity.py | 16 +- ...e_mm_host.py => test_rust_mm_processor.py} | 10 +- .../rust/qwen/test_token_layout_mrope.py | 6 +- ...ive_mm.py => test_build_rust_mm_output.py} | 15 +- .../unit/multimodal/rust/shared/test_fetch.py | 2 +- .../rust/shared/test_image_decode.py | 2 +- .../rust/shared/test_partition_cores.py | 6 +- ...native_mm_gate.py => test_rust_mm_gate.py} | 22 +- .../registered/vlm/test_rust_native_mm_e2e.py | 4 +- .../vlm/test_rust_native_mm_mmmu.py | 22 +- 41 files changed, 3409 insertions(+), 3293 deletions(-) delete mode 100644 python/sglang/srt/managers/rust_server.py create mode 100644 python/sglang/srt/rust_server/config.py create mode 100644 python/sglang/srt/rust_server/multimodal.py create mode 100644 python/sglang/srt/rust_server/server.py create mode 100644 rust/sglang-server/README.md create mode 100644 rust/sglang-server/src/api_server/openai/template_builtins.rs create mode 100644 rust/sglang-server/src/api_server/openai/template_legacy.rs create mode 100644 rust/sglang-server/src/api_server/openai/template_loader.rs create mode 100644 rust/sglang-server/src/tokenizer_manager/to_scheduler_tests.rs create mode 100644 rust/sglang-server/src/tokenizer_manager/to_scheduler_types.rs create mode 100644 rust/sglang-server/src/tokenizer_manager/to_scheduler_validation.rs rename test/registered/unit/multimodal/rust/qwen/{test_native_mm_host.py => test_rust_mm_processor.py} (93%) rename test/registered/unit/multimodal/rust/shared/{test_build_native_mm.py => test_build_rust_mm_output.py} (94%) rename test/registered/unit/multimodal/rust/shared/{test_native_mm_gate.py => test_rust_mm_gate.py} (69%) diff --git a/python/sglang/srt/managers/rust_server.py b/python/sglang/srt/managers/rust_server.py deleted file mode 100644 index 0b1f20f25..000000000 --- a/python/sglang/srt/managers/rust_server.py +++ /dev/null @@ -1,869 +0,0 @@ -"""Embedded Rust server lifecycle for the scheduler. - -The Rust server replaces the Python api-server + `TokenizerManager` + -`DetokenizerManager` stack (hence this module sits beside them in `managers/`), -running them as Rust threads inside the scheduler process. This wrapper keeps -all `SGLANG_RUST_SERVER` plumbing — startup, CPU-core partitioning, the -typed `server_args` handoff, and control-response routing — out of `scheduler.py`. The -scheduler holds an `Optional[RustServer]` and delegates to it. -""" - -from __future__ import annotations - -import importlib -import json -import logging -import os -from array import array -from itertools import chain -from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Optional, Tuple - -import msgspec - -from sglang.srt.arg_groups.overrides import resolving_view -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, - msgpack_decode_explained, -) -from sglang.srt.runtime_context import ( - get_disagg, - get_mm, - get_model, - get_observability, - get_parallel, - get_serving, -) -from sglang.srt.utils.flatten import ( - FlatPairColumns, - NestedRowColumns, - RaggedPairColumns, -) -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.rust_extensions._server import MmSpec, Server, ServerArgs - 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 (as the typed extension ``MmSpec``, see - :meth:`RustServer._build_mm_spec`), the ``_multimodal`` parity API - (: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 — the - JSON form the ``_multimodal`` parity API takes; the server itself is - handed the typed ``MmSpec`` instead.""" - 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 = get_mm().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() - - 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((get_mm().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.multimodal.transport import ( - determine_tensor_transport_mode, - ) - - return ( - get_parallel().tp_size > 1 - and determine_tensor_transport_mode() != "default" - and not get_serving().skip_tokenizer_init - ) - - @staticmethod - def build_native_mm(spec: NativeMmSpec, entry): - """Drain-time adapter: wrap the Rust-produced buffers of one ``MmEncodeResult`` - 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``). - - The server owns the api-server, tokenizermanager, tokenizer, and detokenizer - all implemented as Rust threads in scheduler process. - """ - - 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 - def launch(cls, scheduler: Scheduler) -> RustServer: - """Start the embedded Rust server threads and bind the listen port. - - The caller gates this (``SGLANG_RUST_SERVER`` + rank 0); this always - creates. - """ - from sglang.srt.rust_extensions import load_rust_extension - - Server = load_rust_extension("sglang.srt.rust_extensions._server").Server - - # Force turn off HF tokenizers rayon's unpinned global thread pool. - os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") - - server_args = scheduler.server_args - # `TokenizerManager` merges these under each request's own sampling params - # (`{**preferred, **obj.sampling_params}`), and this server replaces that - # manager wholesale — so honouring the flag is not implemented here yet. - # Refuse rather than run: silently dropping it means generating with - # sampling the operator did not configure, and `/get_model_info` would go on - # advertising values no request ever receives. - if get_serving().preferred_sampling_params: - raise ValueError( - "SGLANG_RUST_SERVER does not yet apply --preferred-sampling-params " - "(the Python TokenizerManager merges it into every request; the rust " - "ingress has no equivalent). Launch without SGLANG_RUST_SERVER, or " - "drop --preferred-sampling-params and send those values per request." - ) - http_addr = f"{get_serving().host}:{get_serving().port}" - - # Per-DP-rank HTTP port with client load balancing. `None` when DP is off, - # so the rank is not conflated with rank 0 of a one-rank group. - dp_rank = scheduler.ps.attn_dp_rank if scheduler.ps.dp_size > 1 else None - if dp_rank is not None: - http_addr = f"{get_serving().host}:{get_serving().port + dp_rank}" - - launch_cores, server_cores = cls._partition_cores( - mm_workers=( - (get_mm().mm_processor_worker_num or NativeMmHost.AUTO_MM_WORKERS) - if scheduler.model_config.is_multimodal - else 0 - ) - ) - - server = Server( - cls._build_server_args(scheduler), - # None -> run unpinned; the list carries the pinning decision. - cores=server_cores, - http_addr=http_addr, - ) - - # 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(cls._build_mm_spec(mm_spec), mm_host.mm_workers) - - # Narrow the scheduler thread only after the server threads are launched. - if launch_cores is not None: - try: - # pid 0 == this thread (the scheduler event-loop / launch thread). - os.sched_setaffinity(0, set(launch_cores)) - except OSError as e: - logger.warning("rust server: cannot pin scheduler launch thread: %s", e) - - # Under DP every rank runs its own server on its own port, so the rank is - # what tells two otherwise identical startup lines apart. - dp_note = ( - "" if dp_rank is None else f" (DP rank {dp_rank}/{scheduler.ps.dp_size})" - ) - logger.info( - "SGLANG_RUST_SERVER enabled, Rust server listen on %s%s", - http_addr, - dp_note, - ) - - return cls(server, mm_spec=mm_spec) - - def wait_request(self, timeout_ms: int) -> None: - """Block until a request is pushed into the in-process ring or the timeout - elapses. - """ - self.server.wait_request(timeout_ms) - - def drain(self, max_recv: int) -> List[Any]: - """Ingress: non-blocking drain of the in-process ring → list of decoded - request objects. The scheduler's request receiver calls this instead of - polling the zmq socket when `rust_server_mode` is set. - - The transfer is **columnar**: `recv_requests` returns an `IngressBatch` - of scalar msgpack `headers` (with `input_ids` omitted) plus one - concatenated raw int64 `data` buffer and per-request `lengths`, so the - large `input_ids` lists never go through msgpack. Each header is `msgpack_decode`d (yielding - the same `TokenizedGenerateReqInput` / control objects the zmq path - produces, so the IPC schema is tracked automatically) and its `input_ids` - slice is wrapped as the `array("q")` the scheduler expects. `recv_requests` - never waits: the ring drain is `try_recv` (returns the instant the ring - is dry, capped at `max_recv`) and the rest is one memcpy per header - plus one for the concatenated ids — same contract as `zmq.NOBLOCK`. - Parking for work is :meth:`wait_request`, which does release the GIL. - """ - limit = max_recv if max_recv > 0 else self._max_per_poll - batch = self.server.recv_requests(limit) - # Bind once: each attribute access converts the rust vec to a fresh list. - headers, data, lengths = batch.headers, batch.data, batch.lengths - if not headers: - return [] - - ids_view = memoryview(data) - out = [] - pos = 0 # byte offset into ids_buf - for header, n in zip(headers, lengths): - nbytes = n * 8 - try: - obj = msgpack_decode_explained(header) - except MsgpackDecodeError as e: - # Return 400 for malformed request field (e.g. token_ids_logprob=[[0]]. - logger.warning( - "rust ingress: dropping undecodable request %s: %s", e.rid, e.reason - ) - if e.rid is not None: - self.server.push_error(e.rid, f"invalid request: {e.reason}") - pos += nbytes - continue - if n: # generate request: attach its int64 ids slice as array("q") - ids = array("q") - 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 - - def push_control_output(self, recv_req, output) -> None: - """Push a control-request response through the egress ring to the waiting - request (routed by rid), encoded as **msgpack** (the ring's native - format). - - A msgspec struct is converted to a *named map* (``structs.asdict``, since - the IPC structs are ``array_like`` and would otherwise lose field names) - so the Rust api_server can shape it per-endpoint (e.g. /server_info) - before rendering JSON to the client — keeping JSON formatting off the - scheduler's GIL. - """ - - # Invariant: control requests always carry a rust-minted rid; without - # one the response is unroutable, so fail loudly rather than drop it. - assert ( - recv_req.rid is not None - ), f"control response without rid: {type(output).__name__}" - # No local try/except: a failed push propagates to run_scheduler_process's - # outer handler, which logs the full traceback (scheduler-fatal either way). - payload = ( - msgspec.structs.asdict(output) - if isinstance(output, msgspec.Struct) - else output - ) - # enc_hook stringifies non-native types (paths, enums); JSON - # rendering happens in Rust. - encoded = msgspec.msgpack.encode(payload, enc_hook=str) - - self.server.push_control_result(recv_req.rid, encoded) - - def push_generation(self, payload: BatchTokenIDOutput) -> None: - """Egress redirect for generation output (replaces the zmq detokenizer). - - Push the WHOLE batch into the Rust egress ring as one frame (-> detokenizer - shards -> client streams), mirroring the ingress ``input_ids`` split so the - bulk numeric columns never go through msgpack: - - - ``header``: msgpack ``BatchHeader`` positional array — the per-request - scalar columns (``rids, finish_reasons, prompt_tokens, tok_lens``) plus - the shape metadata for the optional families (``*_lens`` element counts - for the flat logprob columns, ``*_reqlens``/``*_poslens`` for the ragged - and hidden ones). - - ``data``: the raw little-endian numeric buffer — every column is a - 4-byte element (``f32`` values, ``i32`` indices), concatenated in the - order the Rust ``for_each_chunk`` reads them. - - Logprobs are columnar: output families are per-step deltas, input - (prefill) families ride once on the first chunk. Ragged families (top-k, - token-ids) flatten a per-position ``list[list]`` into flat ``val``/``idx`` - buffers plus a per-position ``lens`` vector (0 = null position). Hidden - states flatten to rows of floats (one row per output position). - """ - output_ids = payload.output_ids or [] - prompt_tokens = payload.prompt_tokens or [] - - # Hot-path guard: almost no decode step wants logprobs / hidden states, - # so only then pay the per-request flatten + buffer packing below. - has_extra = bool( - payload.output_token_logprobs_val - or payload.input_token_logprobs_val - or payload.output_top_logprobs_val - or payload.input_top_logprobs_val - or payload.output_token_ids_logprobs_val - or payload.input_token_ids_logprobs_val - or payload.output_hidden_states - ) - - # Runs on the scheduler's CUDA-launch thread every decode step, so each - # Python-level pass over the batch costs inter-token latency: `rids` are - # the plain rid strings (hashed to a routing key on the Rust side with a - # per-process seed, off the GIL — not parsed; a rid is any string), - # `finished_reasons` already `dict | None`, and `output_ids` entries are - # always `array("i")` (never None) so `map(len)` and a bare - # `chain.from_iterable` stay in C. - rids = payload.rids - finish_reasons = payload.finished_reasons - tok_lens = list(map(len, output_ids)) - flat_ids = array("i", chain.from_iterable(output_ids)) - - # Column order here MUST match BatchHeader (header_cols) and - # for_each_chunk's read order (data_cols); the extras contribution - # is ordered by the `extras` tuple below. - header_cols = [rids, finish_reasons, prompt_tokens, tok_lens] - data_cols = [flat_ids.tobytes()] - - if has_extra: - # The `extras` tuple is the SINGLE source of the extras column - # order — it must match the Rust ``BatchHeader`` fields and - # ``for_each_chunk``'s read order. - # - # TODO(perf): the per-request flatten assumes the logprob/hidden - # columns are ragged, non-contiguous nested Python lists — which is - # only an assumption. The scheduler moves these off the GPU with - # `tensor.tolist()`, so revisit whether the upstream values are - # still contiguous tensors; if so, ship raw bytes + a shape - # descriptor and skip the flatten entirely. - batch_size = len(rids) - extras = ( - FlatPairColumns( - "output_token_logprobs", - payload.output_token_logprobs_val or [], - payload.output_token_logprobs_idx or [], - ), - FlatPairColumns( - "input_token_logprobs", - payload.input_token_logprobs_val or [], - payload.input_token_logprobs_idx or [], - first_none_to_nan=True, - ), - RaggedPairColumns( - "output_top_logprobs", - payload.output_top_logprobs_val or [], - payload.output_top_logprobs_idx or [], - ), - RaggedPairColumns( - "input_top_logprobs", - payload.input_top_logprobs_val or [], - payload.input_top_logprobs_idx or [], - ), - RaggedPairColumns( - "output_token_ids_logprobs", - payload.output_token_ids_logprobs_val or [], - payload.output_token_ids_logprobs_idx or [], - ), - RaggedPairColumns( - "input_token_ids_logprobs", - payload.input_token_ids_logprobs_val or [], - payload.input_token_ids_logprobs_idx or [], - ), - NestedRowColumns( - "output_hidden_states", payload.output_hidden_states or [] - ), - ) - - # Every column is all-or-nothing per payload — which is also what makes - # a family's emptiness a reliable "nobody asked for this" signal. - active = [] - for extra in extras: - populated = False - for name, col in extra.columns(): - assert len(col) in ( - 0, - batch_size, - ), f"extras column {name}: {len(col)} entries for a batch of {batch_size}" - populated |= len(col) > 0 - if populated: - active.append(extra) - - # Flatten only the families someone asked for. `has_extra` above is a - # per-FRAME guard, so one client enabling logprobs used to drag all - # seven families through the per-request loop: at B=4096 that is 28,672 - # bound-method calls per decode step, materializing 12 columns of 4096 - # zeros nobody reads. Measured 0.37 ms -> 7.90 ms GIL-held per step, - # i.e. 25-75% of a decode step added to the scheduler's critical path. - # - # Skipping `accept` leaves a family's buffers empty, which is exactly - # the wire form the Rust decoder already treats as absent (`per_req_ok` - # admits an empty column, `lens_i` reads 0 for every request). The - # `header_cols`/`data_cols` loops below still walk all seven, so column - # ORDER and arity are unchanged — an inactive family contributes empty - # columns in place rather than disappearing. - for extra in active: - accept = extra.accept # hoisted: this is the hottest loop here - for i in range(batch_size): - accept(i) - - for extra in extras: - header_cols += extra.header_cols() - data_cols += extra.data_cols() - - header = msgspec.msgpack.encode(header_cols) - # Pass the raw column list; the Rust side concatenates it into the frame - # with the GIL released. - if not self.server.push_decode_result_batch(header, data_cols): - logger.warning( - "Rust egress closed; dropped batch of %d requests during shutdown", - len(rids), - ) - - @staticmethod - def _build_mm_spec(spec: NativeMmSpec) -> MmSpec: - """The typed MM handoff for ``Server.start_mm_workers``: the - :class:`NativeMmSpec` fields the Rust pipeline consumes, as the Rust - extension's own ``MmSpec`` class (same required-keyword contract as - :meth:`_build_server_args`; ``family`` / ``resample`` become the - extension's ``MmFamily`` / ``MmResample`` enums).""" - from sglang.srt.rust_extensions import load_rust_extension - - ext = load_rust_extension("sglang.srt.rust_extensions._server") - family = {"qwen_vl": ext.MmFamily.QwenVl}[spec.family] - resample = {"aten_u8": ext.MmResample.AtenU8, "pil": ext.MmResample.Pil}[ - spec.resample - ] - return ext.MmSpec( - family=family, - feature_shm=spec.feature_shm, - image_token_id=spec.image_token_id, - patch_size=spec.patch_size, - merge_size=spec.merge_size, - temporal_patch_size=spec.temporal_patch_size, - min_pixels=spec.min_pixels, - max_pixels=spec.max_pixels, - image_mean=spec.image_mean, - image_std=spec.image_std, - resample=resample, - ) - - @staticmethod - def _build_server_args(scheduler: Scheduler) -> ServerArgs: - """The typed launch handoff for the scheduler's embedded Rust server: - the ``server_args`` fields it reads, the already-resolved - ``model_config``, and launch-time facts — as the Rust extension's own - ``ServerArgs`` class. Its constructor takes every field as a required - keyword (see ``rust/sglang-server/src/message/config.rs``), so a - missing, extra or mistyped field fails here at boot rather than - running on a silently-defaulted knob.""" - from sglang.srt.rust_extensions import load_rust_extension - - ext = load_rust_extension("sglang.srt.rust_extensions._server") - - sa = resolving_view(scheduler.server_args) - mc = scheduler.model_config - disaggregation_mode = { - "null": ext.DisaggregationMode.Null, - "prefill": ext.DisaggregationMode.Prefill, - "decode": ext.DisaggregationMode.Decode, - }[get_disagg().disaggregation_mode] - return ext.ServerArgs( - model_path=get_model().model_path, - served_model_name=get_serving().served_model_name, - tokenizer_path=get_serving().tokenizer_path, - revision=get_model().revision, - load_format=get_model().load_format, - weight_version=get_serving().weight_version, - host=get_serving().host, - port=get_serving().port, - log_level=get_observability().log_level, - log_level_http=get_observability().log_level_http, - chat_template=get_serving().chat_template, - tool_call_parser=get_serving().tool_call_parser, - reasoning_parser=get_serving().reasoning_parser, - stream_response_default_include_usage=get_serving().stream_response_default_include_usage, - tokenizer_worker_num=get_serving().tokenizer_worker_num, - detokenizer_worker_num=get_serving().detokenizer_worker_num, - skip_tokenizer_init=get_serving().skip_tokenizer_init, - incremental_streaming_output=get_serving().incremental_streaming_output, - disaggregation_mode=disaggregation_mode, - model_config=ext.ModelConfig( - context_len=mc.context_len, - vocab_size=mc.vocab_size, - is_multimodal=mc.is_multimodal, - # Resolved default sampling params (generation_config.json when - # `--sampling-defaults model`, {} otherwise). The rust server - # consumes these for omitted temperature/top_p in chat - # conversions instead of hard-coding the OpenAI terminal - # defaults. - default_sampling_params=ext.DefaultSamplingParams( - **mc.get_default_sampling_params() - ), - ), - # `preferred_sampling_params` is deliberately absent: `launch` - # refuses to start when it is set, so the Rust server never needs it. - preferred_sampling_params=( - json.dumps(get_serving().preferred_sampling_params) - if get_serving().preferred_sampling_params is not None - else None - ), - allow_auto_truncate=get_serving().allow_auto_truncate, - enable_return_hidden_states=sa.enable_return_hidden_states, - # Not a `server_args` field: `TokenizerManager` derives it, and the - # rust ingress needs the same number for its total-token check. - num_reserved_tokens=compute_num_reserved_tokens(), - # Launch-time facts Python's /server_info reports from - # scheduler_info / the package — stamped here so the rust endpoint - # can serve them statically (no scheduler round-trip). - version=__version__, - max_total_num_tokens=scheduler.max_total_num_tokens, - ) - - @staticmethod - 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 - of this rank's NUMA-local cores (when affinity/NUMA bind is on), so the - partition stays NUMA-local. Returns ``(None, None)`` (server runs - unpinned, confined only by the process affinity) when the platform has - no affinity API or too few cores to split. - """ - if not hasattr(os, "sched_getaffinity"): - return None, None - try: - allowed = sorted(os.sched_getaffinity(0)) - except OSError as e: - logger.warning("rust server: cannot read cpu affinity: %s", e) - return None, None - - # Need enough cores to reserve launch cores and still pin the pools. - if len(allowed) < 4: - logger.info( - "rust server: only %d cores allowed; running pools unpinned", - len(allowed), - ) - return None, None - - # Keep a small slice for the launch loop; cap at 2 (the event loop is - # effectively serial) and never take more than a quarter of the cores. - reserve = min(2, len(allowed) // 4) - launch_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, - launch_cores, - ) - return launch_cores, server_cores diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8d33ea577..058854e23 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -198,7 +198,6 @@ from sglang.srt.managers.prefill_delayer import ( PrefillDelayerSinglePassExecutor, RecentPrefillBatchSizeTracker, ) -from sglang.srt.managers.rust_server import RustServer from sglang.srt.managers.schedule_batch import ( FINISH_ABORT, MultimodalInputs, @@ -294,6 +293,7 @@ from sglang.srt.observability.trace import process_tracing_init, trace_set_threa from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.platforms import current_platform from sglang.srt.plugins import load_plugins +from sglang.srt.rust_server.server import RustServer from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_params import TOP_K_ALL from sglang.srt.server_args import PortArgs, ServerArgs, compute_world_size diff --git a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py index 5f85391ac..54099ae3a 100644 --- a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py +++ b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py @@ -9,7 +9,7 @@ from sglang.srt.observability.req_time_stats import real_time from sglang.srt.platforms import current_platform if TYPE_CHECKING: - from sglang.srt.managers.rust_server import RustServer + from sglang.srt.rust_server.server import RustServer class IdleSleeper: diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 370091753..15659cf2d 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -38,7 +38,7 @@ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils.weight_versions import compute_weight_version_spans if TYPE_CHECKING: - from sglang.srt.managers.rust_server import RustServer + from sglang.srt.rust_server.server import RustServer logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index bcdfff5a3..cbb6ebb18 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -37,7 +37,7 @@ from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed.parallel_state_wrapper import ParallelState - from sglang.srt.managers.rust_server import RustServer + from sglang.srt.rust_server.server import RustServer from sglang.srt.server_args import ServerArgs from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook from sglang.test.scripted_runtime.tokenizer_recv_proxy import ( diff --git a/python/sglang/srt/rust_server/config.py b/python/sglang/srt/rust_server/config.py new file mode 100644 index 000000000..4ffad333e --- /dev/null +++ b/python/sglang/srt/rust_server/config.py @@ -0,0 +1,146 @@ +"""Configuration handoff and CPU placement for the embedded Rust server.""" + +from __future__ import annotations + +import json +import logging +import os +from typing import TYPE_CHECKING, List, Optional, Tuple + +from sglang.srt.arg_groups.overrides import resolving_view +from sglang.srt.managers.utils import compute_num_reserved_tokens +from sglang.srt.runtime_context import ( + get_disagg, + get_model, + get_observability, + get_serving, +) +from sglang.version import __version__ + +if TYPE_CHECKING: + from sglang.srt.managers.scheduler import Scheduler + from sglang.srt.rust_extensions._server import ServerArgs + +logger = logging.getLogger(__name__) + + +def _build_server_args(scheduler: Scheduler) -> ServerArgs: + """The typed launch handoff for the scheduler's embedded Rust server: + the ``server_args`` fields it reads, the already-resolved + ``model_config``, and launch-time facts — as the Rust extension's own + ``ServerArgs`` class. Its constructor takes every field as a required + keyword (see ``rust/sglang-server/src/message/config.rs``), so a + missing, extra or mistyped field fails here at boot rather than + running on a silently-defaulted knob.""" + from sglang.srt.rust_extensions import load_rust_extension + + ext = load_rust_extension("sglang.srt.rust_extensions._server") + + sa = resolving_view(scheduler.server_args) + mc = scheduler.model_config + disaggregation_mode = { + "null": ext.DisaggregationMode.Null, + "prefill": ext.DisaggregationMode.Prefill, + "decode": ext.DisaggregationMode.Decode, + }[get_disagg().disaggregation_mode] + return ext.ServerArgs( + model_path=get_model().model_path, + served_model_name=get_serving().served_model_name, + tokenizer_path=get_serving().tokenizer_path, + revision=get_model().revision, + load_format=get_model().load_format, + weight_version=get_serving().weight_version, + host=get_serving().host, + port=get_serving().port, + log_level=get_observability().log_level, + log_level_http=get_observability().log_level_http, + chat_template=get_serving().chat_template, + tool_call_parser=get_serving().tool_call_parser, + reasoning_parser=get_serving().reasoning_parser, + stream_response_default_include_usage=get_serving().stream_response_default_include_usage, + tokenizer_worker_num=get_serving().tokenizer_worker_num, + detokenizer_worker_num=get_serving().detokenizer_worker_num, + skip_tokenizer_init=get_serving().skip_tokenizer_init, + incremental_streaming_output=get_serving().incremental_streaming_output, + disaggregation_mode=disaggregation_mode, + model_config=ext.ModelConfig( + context_len=mc.context_len, + vocab_size=mc.vocab_size, + is_multimodal=mc.is_multimodal, + # Resolved default sampling params (generation_config.json when + # `--sampling-defaults model`, {} otherwise). The rust server + # consumes these for omitted temperature/top_p in chat + # conversions instead of hard-coding the OpenAI terminal + # defaults. + default_sampling_params=ext.DefaultSamplingParams( + **mc.get_default_sampling_params() + ), + ), + # `preferred_sampling_params` is deliberately absent: `launch` + # refuses to start when it is set, so the Rust server never needs it. + preferred_sampling_params=( + json.dumps(get_serving().preferred_sampling_params) + if get_serving().preferred_sampling_params is not None + else None + ), + allow_auto_truncate=get_serving().allow_auto_truncate, + enable_return_hidden_states=sa.enable_return_hidden_states, + # Not a `server_args` field: `TokenizerManager` derives it, and the + # rust ingress needs the same number for its total-token check. + num_reserved_tokens=compute_num_reserved_tokens(), + # Launch-time facts Python's /server_info reports from + # scheduler_info / the package — stamped here so the rust endpoint + # can serve them statically (no scheduler round-trip). + version=__version__, + max_total_num_tokens=scheduler.max_total_num_tokens, + ) + + +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 + of this rank's NUMA-local cores (when affinity/NUMA bind is on), so the + partition stays NUMA-local. Returns ``(None, None)`` (server runs + unpinned, confined only by the process affinity) when the platform has + no affinity API or too few cores to split. + """ + if not hasattr(os, "sched_getaffinity"): + return None, None + try: + allowed = sorted(os.sched_getaffinity(0)) + except OSError as e: + logger.warning("rust server: cannot read cpu affinity: %s", e) + return None, None + + # Need enough cores to reserve launch cores and still pin the pools. + if len(allowed) < 4: + logger.info( + "rust server: only %d cores allowed; running pools unpinned", + len(allowed), + ) + return None, None + + # Keep a small slice for the launch loop; cap at 2 (the event loop is + # effectively serial) and never take more than a quarter of the cores. + reserve = min(2, len(allowed) // 4) + launch_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, + launch_cores, + ) + return launch_cores, server_cores diff --git a/python/sglang/srt/rust_server/multimodal.py b/python/sglang/srt/rust_server/multimodal.py new file mode 100644 index 000000000..e48b601d5 --- /dev/null +++ b/python/sglang/srt/rust_server/multimodal.py @@ -0,0 +1,319 @@ +"""Multimodal support for the embedded Rust server.""" + +from __future__ import annotations + +import importlib +import logging +from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Optional, Tuple + +import msgspec + +from sglang.srt.environ import envs +from sglang.srt.runtime_context import get_mm, get_parallel, get_serving + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class RustMmSpec(msgspec.Struct, frozen=True, kw_only=True): + """Resolved parameters of the Rust MM pipeline for one model, + consumed by the Rust worker pool (as the typed extension ``MmSpec``, see + :meth:`RustServer._build_mm_spec`), the ``_multimodal`` parity API + (:meth:`rust_json`) and the drain adapter + (:meth:`RustMmProcessor.build_output`).""" + + 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. + 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 — the + JSON form the ``_multimodal`` parity API takes; the server itself is + handed the typed ``MmSpec`` instead.""" + 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 RustMmFamily(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:`RUST_MM_FAMILIES` + plus its Rust arm — the launch gate is data-driven.""" + + name: str + # The registered Python MM processor the Rust pipeline replaces, as + # "module:Class". Compared by identity, so an + # SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE override still disables the Rust 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 Rust resize reproduces bit-exactly, each mapped + # to the `resample` the Rust pipeline must use (see `RustMmSpec.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 + + +RUST_MM_FAMILIES: Tuple[RustMmFamily, ...] = ( + RustMmFamily( + 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 rust_mm_family_for( + mm_processor_cls: Any, model_type: Optional[str] +) -> Optional[RustMmFamily]: + """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 RUST_MM_FAMILIES if f.serves(mm_processor_cls, model_type)), None + ) + + +class RustMmProcessor: + """Builds and validates the 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_spec` resolves the pipeline parameters from. At drain + time :meth:`build_output` wraps the Rust-produced buffers into the + scheduler's ``MultimodalProcessorOutput``. + + There is no Python fallback: a model without a Rust MM 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 = get_mm().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() + + def resolve_spec(self) -> Optional[RustMmSpec]: + """The :class:`RustMmSpec` for this model, or ``None`` when it has no + Rust 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 Rust 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 = rust_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 Rust 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 by Rust; anything else disables the pipeline. + image_overrides = dict((get_mm().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 = RustMmSpec( + 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: Rust 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.multimodal.transport import ( + determine_tensor_transport_mode, + ) + + return ( + get_parallel().tp_size > 1 + and determine_tensor_transport_mode() != "default" + and not get_serving().skip_tokenizer_init + ) + + @staticmethod + def build_output(spec: RustMmSpec, entry): + """Drain-time adapter: wrap the Rust-produced buffers of one ``MmEncodeResult`` + 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_result``'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_result`. + 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), + ) diff --git a/python/sglang/srt/rust_server/server.py b/python/sglang/srt/rust_server/server.py new file mode 100644 index 000000000..47cac4b40 --- /dev/null +++ b/python/sglang/srt/rust_server/server.py @@ -0,0 +1,438 @@ +"""Embedded Rust server lifecycle for the scheduler. + +The Rust server replaces the Python api-server + `TokenizerManager` + +`DetokenizerManager` stack, running them as Rust threads inside the scheduler +process. This wrapper keeps all `SGLANG_RUST_SERVER` plumbing — startup, +CPU-core partitioning, the typed `server_args` handoff, and control-response +routing — out of `scheduler.py`. The scheduler holds an `Optional[RustServer]` +and delegates to it. +""" + +from __future__ import annotations + +import logging +import os +from array import array +from itertools import chain +from typing import TYPE_CHECKING, Any, List, Optional + +import msgspec + +from sglang.srt.managers.io_struct import TokenizedGenerateReqInput +from sglang.srt.managers.utils import ( + MsgpackDecodeError, + msgpack_decode_explained, +) +from sglang.srt.runtime_context import get_mm, get_serving +from sglang.srt.rust_server.config import _build_server_args, _partition_cores +from sglang.srt.rust_server.multimodal import ( + RUST_MM_FAMILIES, + RustMmProcessor, + RustMmSpec, +) +from sglang.srt.utils.flatten import ( + FlatPairColumns, + NestedRowColumns, + RaggedPairColumns, +) + +if TYPE_CHECKING: + from sglang.srt.managers.io_struct import BatchTokenIDOutput + from sglang.srt.managers.scheduler import Scheduler + from sglang.srt.rust_extensions._server import MmSpec, Server + +logger = logging.getLogger(__name__) + + +class RustServer: + """Owns the embedded multi-threaded Rust server (``sglang_server.Server``). + + The server owns the api-server, tokenizermanager, tokenizer, and detokenizer + all implemented as Rust threads in scheduler process. + """ + + def __init__( + self, + server: Server, + mm_spec: Optional[RustMmSpec] = None, + max_per_poll: int = 256, + ): + self.server = server + self.mm_spec = mm_spec + self._max_per_poll = max_per_poll + + @classmethod + def launch(cls, scheduler: Scheduler) -> RustServer: + """Start the embedded Rust server threads and bind the listen port. + + The caller gates this (``SGLANG_RUST_SERVER`` + rank 0); this always + creates. + """ + from sglang.srt.rust_extensions import load_rust_extension + + Server = load_rust_extension("sglang.srt.rust_extensions._server").Server + + # Force turn off HF tokenizers rayon's unpinned global thread pool. + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + + server_args = scheduler.server_args + # `TokenizerManager` merges these under each request's own sampling params + # (`{**preferred, **obj.sampling_params}`), and this server replaces that + # manager wholesale — so honouring the flag is not implemented here yet. + # Refuse rather than run: silently dropping it means generating with + # sampling the operator did not configure, and `/get_model_info` would go on + # advertising values no request ever receives. + if get_serving().preferred_sampling_params: + raise ValueError( + "SGLANG_RUST_SERVER does not yet apply --preferred-sampling-params " + "(the Python TokenizerManager merges it into every request; the rust " + "ingress has no equivalent). Launch without SGLANG_RUST_SERVER, or " + "drop --preferred-sampling-params and send those values per request." + ) + http_addr = f"{get_serving().host}:{get_serving().port}" + + # Per-DP-rank HTTP port with client load balancing. `None` when DP is off, + # so the rank is not conflated with rank 0 of a one-rank group. + dp_rank = scheduler.ps.attn_dp_rank if scheduler.ps.dp_size > 1 else None + if dp_rank is not None: + http_addr = f"{get_serving().host}:{get_serving().port + dp_rank}" + + launch_cores, server_cores = _partition_cores( + mm_workers=( + (get_mm().mm_processor_worker_num or RustMmProcessor.AUTO_MM_WORKERS) + if scheduler.model_config.is_multimodal + else 0 + ) + ) + + server = Server( + _build_server_args(scheduler), + # None -> run unpinned; the list carries the pinning decision. + cores=server_cores, + http_addr=http_addr, + ) + + # Multimodal models must have a 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 = RustMmProcessor( + server_args=server_args, + model_config=scheduler.model_config, + processor=scheduler.processor, + ) + mm_spec = mm_host.resolve_spec() + if mm_spec is None: + supported = sorted( + set(chain.from_iterable(f.model_types for f in RUST_MM_FAMILIES)) + ) + raise RuntimeError( + "SGLANG_RUST_SERVER=1: no 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(cls._build_mm_spec(mm_spec), mm_host.mm_workers) + + # Narrow the scheduler thread only after the server threads are launched. + if launch_cores is not None: + try: + # pid 0 == this thread (the scheduler event-loop / launch thread). + os.sched_setaffinity(0, set(launch_cores)) + except OSError as e: + logger.warning("rust server: cannot pin scheduler launch thread: %s", e) + + # Under DP every rank runs its own server on its own port, so the rank is + # what tells two otherwise identical startup lines apart. + dp_note = ( + "" if dp_rank is None else f" (DP rank {dp_rank}/{scheduler.ps.dp_size})" + ) + logger.info( + "SGLANG_RUST_SERVER enabled, Rust server listen on %s%s", + http_addr, + dp_note, + ) + + return cls(server, mm_spec=mm_spec) + + def wait_request(self, timeout_ms: int) -> None: + """Block until a request is pushed into the in-process ring or the timeout + elapses. + """ + self.server.wait_request(timeout_ms) + + def drain(self, max_recv: int) -> List[Any]: + """Ingress: non-blocking drain of the in-process ring → list of decoded + request objects. The scheduler's request receiver calls this instead of + polling the zmq socket when `rust_server_mode` is set. + + The transfer is **columnar**: `recv_requests` returns an `IngressBatch` + of scalar msgpack `headers` (with `input_ids` omitted) plus one + concatenated raw int64 `data` buffer and per-request `lengths`, so the + large `input_ids` lists never go through msgpack. Each header is `msgpack_decode`d (yielding + the same `TokenizedGenerateReqInput` / control objects the zmq path + produces, so the IPC schema is tracked automatically) and its `input_ids` + slice is wrapped as the `array("q")` the scheduler expects. `recv_requests` + never waits: the ring drain is `try_recv` (returns the instant the ring + is dry, capped at `max_recv`) and the rest is one memcpy per header + plus one for the concatenated ids — same contract as `zmq.NOBLOCK`. + Parking for work is :meth:`wait_request`, which does release the GIL. + """ + limit = max_recv if max_recv > 0 else self._max_per_poll + batch = self.server.recv_requests(limit) + # Bind once: each attribute access converts the rust vec to a fresh list. + headers, data, lengths = batch.headers, batch.data, batch.lengths + if not headers: + return [] + + ids_view = memoryview(data) + out = [] + pos = 0 # byte offset into ids_buf + for header, n in zip(headers, lengths): + nbytes = n * 8 + try: + obj = msgpack_decode_explained(header) + except MsgpackDecodeError as e: + # Return 400 for malformed request field (e.g. token_ids_logprob=[[0]]. + logger.warning( + "rust ingress: dropping undecodable request %s: %s", e.rid, e.reason + ) + if e.rid is not None: + self.server.push_error(e.rid, f"invalid request: {e.reason}") + pos += nbytes + continue + if n: # generate request: attach its int64 ids slice as array("q") + ids = array("q") + 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 Rust + # path. `None` for a text-only request on a multimodal model. + mm_result = self.server.take_mm_result(obj.rid) + if mm_result is not None: + obj.mm_inputs = RustMmProcessor.build_output( + self.mm_spec, mm_result + ) + out.append(obj) + return out + + def push_control_output(self, recv_req, output) -> None: + """Push a control-request response through the egress ring to the waiting + request (routed by rid), encoded as **msgpack** (the ring's native + format). + + A msgspec struct is converted to a *named map* (``structs.asdict``, since + the IPC structs are ``array_like`` and would otherwise lose field names) + so the Rust api_server can shape it per-endpoint (e.g. /server_info) + before rendering JSON to the client — keeping JSON formatting off the + scheduler's GIL. + """ + + # Invariant: control requests always carry a rust-minted rid; without + # one the response is unroutable, so fail loudly rather than drop it. + assert ( + recv_req.rid is not None + ), f"control response without rid: {type(output).__name__}" + # No local try/except: a failed push propagates to run_scheduler_process's + # outer handler, which logs the full traceback (scheduler-fatal either way). + payload = ( + msgspec.structs.asdict(output) + if isinstance(output, msgspec.Struct) + else output + ) + # enc_hook stringifies non-native types (paths, enums); JSON + # rendering happens in Rust. + encoded = msgspec.msgpack.encode(payload, enc_hook=str) + + self.server.push_control_result(recv_req.rid, encoded) + + def push_generation(self, payload: BatchTokenIDOutput) -> None: + """Egress redirect for generation output (replaces the zmq detokenizer). + + Push the WHOLE batch into the Rust egress ring as one frame (-> detokenizer + shards -> client streams), mirroring the ingress ``input_ids`` split so the + bulk numeric columns never go through msgpack: + + - ``header``: msgpack ``BatchHeader`` positional array — the per-request + scalar columns (``rids, finish_reasons, prompt_tokens, tok_lens``) plus + the shape metadata for the optional families (``*_lens`` element counts + for the flat logprob columns, ``*_reqlens``/``*_poslens`` for the ragged + and hidden ones). + - ``data``: the raw little-endian numeric buffer — every column is a + 4-byte element (``f32`` values, ``i32`` indices), concatenated in the + order the Rust ``for_each_chunk`` reads them. + + Logprobs are columnar: output families are per-step deltas, input + (prefill) families ride once on the first chunk. Ragged families (top-k, + token-ids) flatten a per-position ``list[list]`` into flat ``val``/``idx`` + buffers plus a per-position ``lens`` vector (0 = null position). Hidden + states flatten to rows of floats (one row per output position). + """ + output_ids = payload.output_ids or [] + prompt_tokens = payload.prompt_tokens or [] + + # Hot-path guard: almost no decode step wants logprobs / hidden states, + # so only then pay the per-request flatten + buffer packing below. + has_extra = bool( + payload.output_token_logprobs_val + or payload.input_token_logprobs_val + or payload.output_top_logprobs_val + or payload.input_top_logprobs_val + or payload.output_token_ids_logprobs_val + or payload.input_token_ids_logprobs_val + or payload.output_hidden_states + ) + + # Runs on the scheduler's CUDA-launch thread every decode step, so each + # Python-level pass over the batch costs inter-token latency: `rids` are + # the plain rid strings (hashed to a routing key on the Rust side with a + # per-process seed, off the GIL — not parsed; a rid is any string), + # `finished_reasons` already `dict | None`, and `output_ids` entries are + # always `array("i")` (never None) so `map(len)` and a bare + # `chain.from_iterable` stay in C. + rids = payload.rids + finish_reasons = payload.finished_reasons + tok_lens = list(map(len, output_ids)) + flat_ids = array("i", chain.from_iterable(output_ids)) + + # Column order here MUST match BatchHeader (header_cols) and + # for_each_chunk's read order (data_cols); the extras contribution + # is ordered by the `extras` tuple below. + header_cols = [rids, finish_reasons, prompt_tokens, tok_lens] + data_cols = [flat_ids.tobytes()] + + if has_extra: + # The `extras` tuple is the SINGLE source of the extras column + # order — it must match the Rust ``BatchHeader`` fields and + # ``for_each_chunk``'s read order. + # + # TODO(perf): the per-request flatten assumes the logprob/hidden + # columns are ragged, non-contiguous nested Python lists — which is + # only an assumption. The scheduler moves these off the GPU with + # `tensor.tolist()`, so revisit whether the upstream values are + # still contiguous tensors; if so, ship raw bytes + a shape + # descriptor and skip the flatten entirely. + batch_size = len(rids) + extras = ( + FlatPairColumns( + "output_token_logprobs", + payload.output_token_logprobs_val or [], + payload.output_token_logprobs_idx or [], + ), + FlatPairColumns( + "input_token_logprobs", + payload.input_token_logprobs_val or [], + payload.input_token_logprobs_idx or [], + first_none_to_nan=True, + ), + RaggedPairColumns( + "output_top_logprobs", + payload.output_top_logprobs_val or [], + payload.output_top_logprobs_idx or [], + ), + RaggedPairColumns( + "input_top_logprobs", + payload.input_top_logprobs_val or [], + payload.input_top_logprobs_idx or [], + ), + RaggedPairColumns( + "output_token_ids_logprobs", + payload.output_token_ids_logprobs_val or [], + payload.output_token_ids_logprobs_idx or [], + ), + RaggedPairColumns( + "input_token_ids_logprobs", + payload.input_token_ids_logprobs_val or [], + payload.input_token_ids_logprobs_idx or [], + ), + NestedRowColumns( + "output_hidden_states", payload.output_hidden_states or [] + ), + ) + + # Every column is all-or-nothing per payload — which is also what makes + # a family's emptiness a reliable "nobody asked for this" signal. + active = [] + for extra in extras: + populated = False + for name, col in extra.columns(): + assert len(col) in ( + 0, + batch_size, + ), f"extras column {name}: {len(col)} entries for a batch of {batch_size}" + populated |= len(col) > 0 + if populated: + active.append(extra) + + # Flatten only the families someone asked for. `has_extra` above is a + # per-FRAME guard, so one client enabling logprobs used to drag all + # seven families through the per-request loop: at B=4096 that is 28,672 + # bound-method calls per decode step, materializing 12 columns of 4096 + # zeros nobody reads. Measured 0.37 ms -> 7.90 ms GIL-held per step, + # i.e. 25-75% of a decode step added to the scheduler's critical path. + # + # Skipping `accept` leaves a family's buffers empty, which is exactly + # the wire form the Rust decoder already treats as absent (`per_req_ok` + # admits an empty column, `lens_i` reads 0 for every request). The + # `header_cols`/`data_cols` loops below still walk all seven, so column + # ORDER and arity are unchanged — an inactive family contributes empty + # columns in place rather than disappearing. + for extra in active: + accept = extra.accept # hoisted: this is the hottest loop here + for i in range(batch_size): + accept(i) + + for extra in extras: + header_cols += extra.header_cols() + data_cols += extra.data_cols() + + header = msgspec.msgpack.encode(header_cols) + # Pass the raw column list; the Rust side concatenates it into the frame + # with the GIL released. + if not self.server.push_decode_result_batch(header, data_cols): + logger.warning( + "Rust egress closed; dropped batch of %d requests during shutdown", + len(rids), + ) + + @staticmethod + def _build_mm_spec(spec: RustMmSpec) -> MmSpec: + """The typed MM handoff for ``Server.start_mm_workers``: the + :class:`RustMmSpec` fields the Rust pipeline consumes, as the Rust + extension's own ``MmSpec`` class (same required-keyword contract as + :meth:`_build_server_args`; ``family`` / ``resample`` become the + extension's ``MmFamily`` / ``MmResample`` enums).""" + from sglang.srt.rust_extensions import load_rust_extension + + ext = load_rust_extension("sglang.srt.rust_extensions._server") + family = {"qwen_vl": ext.MmFamily.QwenVl}[spec.family] + resample = {"aten_u8": ext.MmResample.AtenU8, "pil": ext.MmResample.Pil}[ + spec.resample + ] + return ext.MmSpec( + family=family, + feature_shm=spec.feature_shm, + image_token_id=spec.image_token_id, + patch_size=spec.patch_size, + merge_size=spec.merge_size, + temporal_patch_size=spec.temporal_patch_size, + min_pixels=spec.min_pixels, + max_pixels=spec.max_pixels, + image_mean=spec.image_mean, + image_std=spec.image_std, + resample=resample, + ) diff --git a/rust/sglang-mm/Cargo.toml b/rust/sglang-mm/Cargo.toml index ab6ed7339..1a2320022 100644 --- a/rust/sglang-mm/Cargo.toml +++ b/rust/sglang-mm/Cargo.toml @@ -22,7 +22,7 @@ features = ["python", "parallel"] # which renames the built artifact. name = "sglang_mm_core" # cdylib: the PyO3 module (`sglang.srt.rust_extensions._multimodal`). -# rlib: pure-Rust core linked by sglang-server's native MM path. +# rlib: pure-Rust core linked by sglang-server's Rust MM path. crate-type = ["cdylib", "rlib"] [features] diff --git a/rust/sglang-mm/src/qwen_vl/mod.rs b/rust/sglang-mm/src/qwen_vl/mod.rs index 55933dce3..415848d9c 100644 --- a/rust/sglang-mm/src/qwen_vl/mod.rs +++ b/rust/sglang-mm/src/qwen_vl/mod.rs @@ -65,7 +65,7 @@ pub struct QwenVlProcessor { lut: [[f32; 256]; 3], } -/// `1 / rescale_factor`; `resolve_native_spec` rejects any other factor. +/// `1 / rescale_factor`; `resolve_spec` rejects any other factor. const INV_RESCALE: f32 = 255.0; /// u8 → normalized f32, rounded as the mirrored processor rounds. The slow one @@ -406,7 +406,7 @@ mod python { /// `(pixel_values flat f32, (t, h, w))` for one preprocessed image. type PyProcessedImage<'py> = (Bound<'py, PyArray1>, (u32, u32, u32)); - /// Full native pipeline output at the scheduler boundary: + /// Full Rust pipeline output at the scheduler boundary: /// `(input_ids, features, grids, hashes, offsets, mrope, mrope_delta)`. type PyNativeOutput<'py> = ( Vec, @@ -486,7 +486,7 @@ mod python { /// `sglang-server` (whose message layer owns the wire-payload parsing). #[pyfunction] #[pyo3(signature = (input_ids, images, spec_json))] - fn process_native_mm<'py>( + fn process_mm<'py>( py: Python<'py>, input_ids: Option>, images: Vec, @@ -529,7 +529,7 @@ mod python { m.add_function(wrap_pyfunction!(preprocess, &m)?)?; m.add_function(wrap_pyfunction!(smart_resize_py, &m)?)?; m.add_function(wrap_pyfunction!(mrope_image_only_py, &m)?)?; - m.add_function(wrap_pyfunction!(process_native_mm, &m)?)?; + m.add_function(wrap_pyfunction!(process_mm, &m)?)?; parent.add_submodule(&m)?; Ok(()) } diff --git a/rust/sglang-server/Cargo.toml b/rust/sglang-server/Cargo.toml index b5ccddb09..2a11b3ecc 100644 --- a/rust/sglang-server/Cargo.toml +++ b/rust/sglang-server/Cargo.toml @@ -50,7 +50,7 @@ hf-hub = { version = "0.4", default-features = false } itertools = "0.14" # 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 +# Same major as the workspace pyo3: the zero-copy MM drain (`take_mm_result`) moves # Rust vectors into numpy arrays. numpy = "0.29.0" # Pinned EXACTLY: this crate's accepted grammar defines the diff --git a/rust/sglang-server/README.md b/rust/sglang-server/README.md new file mode 100644 index 000000000..9c249ff70 --- /dev/null +++ b/rust/sglang-server/README.md @@ -0,0 +1,11 @@ +# sglang-server + +`sglang-server` is SGLang's Rust HTTP frontend and request-processing pipeline. It exchanges typed requests and responses with the Python scheduler while keeping latency-sensitive work outside Python. + +## Code review principles + +1. **Use strongly typed boundaries.** Model every supported protocol shape with structs, enums, and validated newtypes; avoid opaque values such as `serde_json::Value` and `rmpv::Value` in production paths. +2. **Keep one canonical schema.** Rust and Python must derive their wire contracts from one source of truth, aligned with `io_struct.py`, instead of independently duplicating field names, order, defaults, or validation. +3. **Make protocol declarations minimal and declarative.** A reviewer should be able to understand the wire format from its type declarations alone, without tracing fillers, conversion code, macros, or repeated field lists. +4. **Use one representation per semantic stage.** Separate external input, normalized domain data, and wire data, and convert between them once at explicit boundaries; do not keep multiple overlapping representations of the same state. +5. **Design compatibility and safety explicitly.** Version protocols, reject unsupported or malformed inputs clearly, validate lengths and resource bounds before allocation, preserve invariants in types, and test compatibility across the real Rust and Python codecs. diff --git a/rust/sglang-server/src/api_server/openai.rs b/rust/sglang-server/src/api_server/openai.rs index 9bd8b493e..f3b2b0a14 100644 --- a/rust/sglang-server/src/api_server/openai.rs +++ b/rust/sglang-server/src/api_server/openai.rs @@ -14,6 +14,9 @@ mod completions; mod models; mod reasoning; mod template; +mod template_builtins; +mod template_legacy; +mod template_loader; mod tools; pub(super) use template::ChatFormatter; diff --git a/rust/sglang-server/src/api_server/openai/template.rs b/rust/sglang-server/src/api_server/openai/template.rs index f6131f686..c9b2467b9 100644 --- a/rust/sglang-server/src/api_server/openai/template.rs +++ b/rust/sglang-server/src/api_server/openai/template.rs @@ -6,49 +6,22 @@ //! `Conversation.get_prompt()` so there is exactly one implementation of the //! per-style formatting logic (no Jinja translation to drift). -use std::path::{Path, PathBuf}; +use std::path::PathBuf; -use dynamo_protocols::types::{ - ChatCompletionRequestAssistantMessageContent, ChatCompletionRequestAssistantMessageContentPart, - ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageContent, - ChatCompletionRequestSystemMessageContentPart, ChatCompletionRequestUserMessageContent, - ChatCompletionRequestUserMessageContentPart, CreateChatCompletionRequest, -}; -use dynamo_renderer::{ChatTemplate, ContextMixins, PromptContextMixin, PromptFormatter}; -use serde_json::Value; +use dynamo_protocols::types::CreateChatCompletionRequest; +use dynamo_renderer::PromptFormatter; use thiserror::Error; use crate::message::types::OneOrMany; -const SUPPORTED_STYLES: &[&str] = &[ - "ADD_COLON_SINGLE", - "ADD_COLON_TWO", - "ADD_COLON_SPACE_SINGLE", - "NO_COLON_SINGLE", - "NO_COLON_TWO", - "ADD_NEW_LINE_SINGLE", - "LLAMA2", - "LLAMA3", - "LLAMA4", - "CHATGLM", - "CHATML", - "CHATINTERN", - "DOLLY", - "RWKV", - "PHOENIX", - "ROBIN", - "FALCON_CHAT", - "CHATGLM3", - "DEEPSEEK_CHAT", - "METAMATH", - "DeepSeekVL2", - "QWEN2_VL_EMBED", - "QWEN2_AUDIO", - "GEMMA3", - "MPT", - "PADDLE_OCR", - "UNLIMITED_OCR", -]; +#[cfg(test)] +pub(super) use super::template_builtins::builtin_template; +pub(super) use super::template_legacy::LegacyFormatter; +#[cfg(test)] +pub(super) use super::template_legacy::LegacySpec; +#[cfg(test)] +use super::template_loader::infer_legacy_template_from_model_path; +pub(super) use super::template_loader::load_chat_formatter; /// A chat prompt formatter: either the model's HuggingFace Jinja template or a /// legacy SGLang conversation template. @@ -89,599 +62,6 @@ impl ChatFormatter { } } -/// A legacy conversation template, mirroring Python's `Conversation` fields. -#[derive(Debug, Clone)] -pub(super) struct LegacySpec { - /// Python `Conversation.name` — drives the CHATGLM round-offset quirk. - pub(super) name: String, - pub(super) system_template: String, - pub(super) system_message: String, - /// `(user_role, assistant_role)` — Python `Conversation.roles`. - pub(super) roles: (String, String), - pub(super) style: String, - pub(super) sep: String, - /// `None` = Python's `Conversation.sep2` default. Styles that alternate - /// seps (`seps[i % 2]`) need it set; Python crashes on `None` there and we - /// error deliberately. - pub(super) sep2: Option, - /// Python `Conversation.stop_str` (`str | list[str] | None`). - pub(super) stop_str: Option>, - pub(super) image_token: String, - pub(super) audio_token: String, -} - -impl Default for LegacySpec { - fn default() -> Self { - Self { - name: String::new(), - system_template: String::new(), - system_message: String::new(), - roles: (String::new(), String::new()), - style: String::new(), - sep: String::new(), - sep2: None, - stop_str: None, - image_token: "".into(), - audio_token: "