minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add external multimodal processors to the Rust frontend (#39329)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-09-15 15:38:47 -07:00
committed by GitHub
parent 2929a39927
commit b803cfa0c4
20 changed files with 1027 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+4 -1
View File
@@ -2276,7 +2276,7 @@ class Scheduler(
self.rust_server = None
return
rust_server = RustServer.launch(self)
rust_server = self.get_rust_server_class().launch(self)
self.rust_server = rust_server
# The rust server *is* the ingress source: SchedulerRequestReceiver
# drains its request ring (rust_server_mode) instead of a zmq socket.
@@ -2284,6 +2284,9 @@ class Scheduler(
# Park the idle loop on the request ring within the rank-0 rust-server
self.idle_sleeper = RustServerIdleSleeper(rust_server)
def get_rust_server_class(self) -> type[RustServer]:
return RustServer
def rust_server_tokenizer_path(self) -> str:
return get_serving().tokenizer_path
python/sglang/srt/rust_server/config.py
+14 -6
View File
@@ -1,11 +1,12 @@
"""Configuration handoff and CPU placement for the embedded Rust server."""
"""Configuration 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 types import ModuleType
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
from sglang.srt.arg_groups.overrides import resolving_view
from sglang.srt.managers.utils import compute_num_reserved_tokens
@@ -25,8 +26,10 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _build_server_args(scheduler: Scheduler) -> ServerArgs:
"""The typed launch handoff for the scheduler's embedded Rust server:
def _build_server_args(
scheduler: Scheduler, *, extension: Optional[ModuleType] = None
) -> ServerArgs:
"""The typed launch configuration 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
@@ -35,7 +38,7 @@ def _build_server_args(scheduler: Scheduler) -> ServerArgs:
running on a silently-defaulted knob."""
from sglang.srt.rust_extensions import load_rust_extension
ext = load_rust_extension("sglang.srt.rust_extensions._server")
ext = extension or load_rust_extension("sglang.srt.rust_extensions._server")
sa = resolving_view(scheduler.server_args)
mc = scheduler.model_config
@@ -100,6 +103,7 @@ def _build_server_args(scheduler: Scheduler) -> ServerArgs:
def _partition_cores(
mm_workers: int = 0,
server_core_budget: Optional[Callable[[int, int], int]] = None,
) -> Tuple[Optional[List[int]], Optional[List[int]]]:
"""Split this rank's allowed cores into ``(launch_cores, server_cores)``.
@@ -138,7 +142,11 @@ def _partition_cores(
# 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)
pool_budget = (
server_core_budget(len(allowed), mm_workers)
if server_core_budget is not None
else max(8, mm_workers + 4)
)
server_cores = allowed[reserve : reserve + pool_budget]
logger.info(
"rust server cores=%s, scheduler launch cores=%s",
python/sglang/srt/rust_server/server.py
+69 -42
View File
@@ -3,7 +3,7 @@
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
CPU-core partitioning, the typed `server_args`, and control-response
routing — out of `scheduler.py`. The scheduler holds an `Optional[RustServer]`
and delegates to it.
"""
@@ -14,6 +14,7 @@ import logging
import os
from array import array
from itertools import chain
from types import ModuleType
from typing import TYPE_CHECKING, Any, List, Optional
import msgspec
@@ -39,8 +40,9 @@ from sglang.srt.utils.network import NetworkAddress
if TYPE_CHECKING:
from sglang.srt.managers.io_struct import BatchTokenIDOutput
from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.rust_extensions._server import MmSpec, Server
from sglang.srt.rust_extensions._server import MmEncodedResult, MmSpec, Server
logger = logging.getLogger(__name__)
@@ -62,8 +64,50 @@ class RustServer:
self.server = server
self.http_port = http_port
self.mm_spec = mm_spec
self._multimodal_enabled = mm_spec is not None
self._max_per_poll = max_per_poll
@classmethod
def _load_extension(cls) -> ModuleType:
from sglang.srt.rust_extensions import load_rust_extension
return load_rust_extension("sglang.srt.rust_extensions._server")
def _start_multimodal(self, scheduler: Scheduler) -> None:
"""Start the model's Rust workers and retain their scheduler-side state."""
mm_host = RustMmProcessor(
server_args=scheduler.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."
)
self.server.start_mm_workers(self._build_mm_spec(mm_spec), mm_host.mm_workers)
self.mm_spec = mm_spec
@staticmethod
def _server_core_budget(allowed_core_count: int, mm_workers: int) -> int:
"""Maximum cores available to the Rust frontend and MM workers."""
return max(8, mm_workers + 4)
@classmethod
def _partition_cores(
cls, mm_workers: int = 0
) -> tuple[Optional[List[int]], Optional[List[int]]]:
return _partition_cores(
mm_workers=mm_workers,
server_core_budget=cls._server_core_budget,
)
@classmethod
def launch(cls, scheduler: Scheduler) -> RustServer:
"""Start the embedded Rust server threads and bind the listen port.
@@ -71,14 +115,9 @@ class RustServer:
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
# Preserve the DP startup log; ports use node-local offsets.
dp_rank = scheduler.ps.attn_dp_rank if scheduler.ps.dp_size > 1 else None
if get_exec().moe.is_ep_scale_joiner:
@@ -94,7 +133,7 @@ class RustServer:
listen_port = get_serving().port + local_dp_rank
listen_addr = NetworkAddress(get_serving().host, listen_port).to_host_port_str()
launch_cores, server_cores = _partition_cores(
launch_cores, server_cores = cls._partition_cores(
mm_workers=(
(get_mm().mm_processor_worker_num or RustMmProcessor.AUTO_MM_WORKERS)
if scheduler.model_config.is_multimodal
@@ -102,16 +141,17 @@ class RustServer:
)
)
server = Server(
_build_server_args(scheduler),
# None -> run unpinned; the list carries the pinning decision.
extension = cls._load_extension()
server = extension.Server(
_build_server_args(scheduler, extension=extension),
# None runs unpinned; otherwise the list carries the pinning decision.
cores=server_cores,
port_offset=local_dp_rank,
)
instance = cls(server, http_port=listen_port)
# 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
@@ -125,23 +165,8 @@ class RustServer:
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)
instance._start_multimodal(scheduler)
instance._multimodal_enabled = True
# Narrow the scheduler thread only after the server threads are launched.
if launch_cores is not None:
@@ -162,7 +187,11 @@ class RustServer:
dp_note,
)
return cls(server, http_port=listen_port, mm_spec=mm_spec)
return instance
def _wrap_mm_result(self, entry: MmEncodedResult) -> MultimodalProcessorOutput:
assert self.mm_spec is not None
return RustMmProcessor.wrap_encoded(self.mm_spec, entry)
def wait_request(self, timeout_ms: int) -> None:
"""Block until a request is pushed into the in-process ring or the timeout
@@ -215,20 +244,20 @@ 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):
if self._multimodal_enabled and isinstance(obj, TokenizedGenerateReqInput):
# The buffers were parked in the Rust result store 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.
encoded = self.server.take_mm_result(obj.rid)
if encoded is not None:
obj.mm_inputs = RustMmProcessor.wrap_encoded(self.mm_spec, encoded)
mm_result = self.server.take_mm_result(obj.rid)
if mm_result is not None:
obj.mm_inputs = self._wrap_mm_result(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
request (routed by rid), encoded as **msgpack** (the ring's message
format).
A msgspec struct is converted to a *named map* (``structs.asdict``, since
@@ -407,16 +436,14 @@ class RustServer:
len(rids),
)
@staticmethod
def _build_mm_spec(spec: RustMmSpec) -> MmSpec:
"""The typed MM handoff for ``Server.start_mm_workers``: the
@classmethod
def _build_mm_spec(cls, spec: RustMmSpec) -> MmSpec:
"""The typed MM configuration 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
:func:`_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")
ext = cls._load_extension()
family = {"qwen_vl": ext.MmFamily.QwenVl}[spec.family]
resample = {"aten_u8": ext.MmResample.AtenU8, "pil": ext.MmResample.Pil}[
spec.resample