support rust sglang server (#29799)

python/sglang/srt/managers/rust_server.py

@@ -0,0 +1,392 @@
+"""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
+`server_args` blob, 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, Tuple
+
+import msgspec
+
+from sglang.srt.managers.utils import (
+    MsgpackDecodeError,
+    compute_num_reserved_tokens,
+    msgpack_decode_explained,
+)
+from sglang.srt.utils.flatten import (
+    FlatPairColumns,
+    NestedRowColumns,
+    RaggedPairColumns,
+)
+from sglang.version import __version__
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.io_struct import BatchTokenIDOutput
+    from sglang.srt.managers.scheduler import Scheduler
+    from sglang.srt.server._core import 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, max_per_poll: int = 256):
+        self.server = server
+        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.server._core import 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 server_args.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"{server_args.host}:{server_args.port}"
+        launch_cores, server_cores = cls._partition_cores()
+
+        server = Server(
+            # None -> run unpinned; the list carries the pinning decision.
+            cores=server_cores,
+            http_addr=http_addr,
+            server_args_json=cls._build_server_args(scheduler),
+        )
+
+        # 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)
+
+        logger.info(
+            "SGLANG_RUST_SERVER enabled, Rust server listen on %s",
+            http_addr,
+        )
+
+        return cls(server)
+
+    def wait_ingress(self, timeout_ms: int) -> None:
+        """Block until a request is pushed into the in-process ring or the timeout
+        elapses.
+        """
+        self.server.wait_ingress(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`
+        releases the GIL for the drain + concat, so this never holds the GIL
+        across a wait — same contract as `zmq.NOBLOCK`.
+        """
+        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
+            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_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_batch(header, data_cols):
+            logger.warning(
+                "Rust egress closed; dropped batch of %d requests during shutdown",
+                len(rids),
+            )
+
+    @staticmethod
+    def _build_server_args(scheduler: Scheduler) -> str:
+        """JSON blob of the scheduler's ``server_args`` for its embedded Rust
+        server (carries the already-resolved ``model_config``)."""
+
+        server_args = dict(vars(scheduler.server_args))
+        model_config = dict(vars(scheduler.model_config))
+        model_config["hf_config"] = None  # HF config is not JSON-serializable
+        server_args["model_config"] = model_config
+        # 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).
+        server_args["version"] = __version__
+        # Not a `server_args` field: `TokenizerManager` derives it, and the rust
+        # ingress needs the same number for its total-token check.
+        server_args["num_reserved_tokens"] = compute_num_reserved_tokens(
+            scheduler.server_args
+        )
+        server_args["max_total_num_tokens"] = scheduler.max_total_num_tokens
+
+        return msgspec.json.encode(server_args, enc_hook=str).decode("utf-8")
+
+    @staticmethod
+    def _partition_cores() -> 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]
+        server_cores = allowed[reserve:]
+        logger.info(
+            "rust server cores=%s, scheduler launch cores=%s",
+            server_cores,
+            launch_cores,
+        )
+        return launch_cores, server_cores

python/sglang/srt/managers/schedule_batch.py

@@ -1450,7 +1450,28 @@ class Req(ReqDllmMixin):
             # Check stop regex
             if len(self.sampling_params.stop_regex_strs) > 0:
                 for stop_regex_str in self.sampling_params.stop_regex_strs:
-                    if re.search(stop_regex_str, tail_str):
+                    # Seatbelt, not validation: patterns are checked at ingress
+                    # (Python's `normalize`, or the rust server's stricter
+                    # `stop_regex_bound`). This runs per decode step on the hot
+                    # path, so an `re.error` escaping here would take the whole
+                    # scheduler down over one malformed request. Fail that request
+                    # instead.
+                    try:
+                        matched = re.search(stop_regex_str, tail_str)
+                    except (re.error, RecursionError) as e:
+                        logger.warning(
+                            "req %s: invalid stop_regex %r (%s); aborting the request",
+                            self.rid,
+                            stop_regex_str,
+                            e,
+                        )
+                        self.finished_reason = FINISH_ABORT(
+                            f"invalid stop_regex {stop_regex_str!r}: {e}",
+                            HTTPStatus.BAD_REQUEST,
+                            "BadRequestError",
+                        )
+                        break
+                    if matched:
                         self.finished_reason = FINISHED_MATCHED_REGEX(
                             matched=stop_regex_str
                         )

python/sglang/srt/managers/scheduler.py

@@ -167,6 +167,7 @@ from sglang.srt.managers.prefill_delayer import (
     PrefillDelayer,
     PrefillDelayerSinglePassExecutor,
 )
+from sglang.srt.managers.rust_server import RustServer
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
     MultimodalInputs,
@@ -185,7 +186,10 @@ from sglang.srt.managers.scheduler_components.batch_result_processor import (
 )
 from sglang.srt.managers.scheduler_components.dp_attn import SchedulerDPAttnAdapter
 from sglang.srt.managers.scheduler_components.flush_wrapper import SchedulerFlushWrapper
-from sglang.srt.managers.scheduler_components.idle_sleeper import IdleSleeper
+from sglang.srt.managers.scheduler_components.idle_sleeper import (
+    IdleSleeper,
+    RustServerIdleSleeper,
+)
 from sglang.srt.managers.scheduler_components.invariant_checker import (
     SchedulerInvariantChecker,
     create_scheduler_watchdog,
@@ -597,6 +601,10 @@ class Scheduler(
 
         self.maybe_init_scripted_scheduler_hook()
 
+        # Start the embedded Rust frontend (rank 0) before the request receiver,
+        # which reads self.rust_ring_recv to pick its ingress transport.
+        self.maybe_init_rust_server()
+
         self.init_request_receiver()
 
         self.init_dp_attn_adapter()
@@ -672,9 +680,12 @@ class Scheduler(
         )
 
         self.load_snapshot_writer = None
+        self.recv_from_tokenizer = None
+
         if not is_rank_zero:
             return
 
+        self.recv_from_tokenizer = self.ipc_channels.recv_from_tokenizer
         dp_rank = self.ps.dp_rank if self.ps.dp_rank is not None else 0
         try:
             self.load_snapshot_writer = create_load_snapshot_writer(
@@ -1748,11 +1759,16 @@ class Scheduler(
 
             output = self._request_dispatcher(recv_req)
             if output is not None:
-                if not isinstance(output, RpcReqOutput):
-                    self.ipc_channels.send_to_tokenizer.send_output(output, recv_req)
-                else:
+                if self.rust_server is not None:
+                    # Embedded Rust server: every control-request response goes
+                    # back through the egress ring (the zmq tokenizer socket is
+                    # not consumed); the Rust api_server shapes it per-endpoint.
+                    self.rust_server.push_control_output(recv_req, output)
+                elif isinstance(output, RpcReqOutput):
                     if self.ipc_channels.recv_from_rpc is not None:
                         sock_send(self.ipc_channels.recv_from_rpc, output)
+                else:
+                    self.ipc_channels.send_to_tokenizer.send_output(output, recv_req)
 
         self.flush_wrapper.check_pending()
         if self.external_corpus_manager is not None:
@@ -1808,9 +1824,33 @@ class Scheduler(
         else:
             self.scripted_scheduler_hook = None
 
+    def maybe_init_rust_server(self) -> None:
+        """Start the embedded Rust server (rank 0) if ``SGLANG_RUST_SERVER`` is
+        set, and point the ingress receiver at it. All the plumbing lives in
+        ``RustServer`` (scheduler_components/rust_scheduler.py)."""
+
+        is_rank_zero = (
+            self.ps.pp_rank == 0
+            and self.ps.attn_tp_rank == 0
+            and self.ps.attn_cp_rank == 0
+        )
+        if not (envs.SGLANG_RUST_SERVER.get() and is_rank_zero):
+            # Always define the attribute: init_output_streamer and the
+            # process_input_requests hook read self.rust_server unconditionally.
+            self.rust_server = None
+            return
+
+        rust_server = RustServer.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.
+        self.recv_from_tokenizer = rust_server
+        # Park the idle loop on the request ring within the rank-0 rust-server
+        self.idle_sleeper = RustServerIdleSleeper(rust_server)
+
     def init_request_receiver(self) -> None:
         self.request_receiver = SchedulerRequestReceiver(
-            recv_from_tokenizer=self.ipc_channels.recv_from_tokenizer,
+            recv_from_tokenizer=self.recv_from_tokenizer,
             recv_from_rpc=self.ipc_channels.recv_from_rpc,
             recv_skipper=self.recv_skipper,
             input_blocker=self.input_blocker,
@@ -1936,6 +1976,7 @@ class Scheduler(
             spec_algorithm=self.spec_algorithm,
             disaggregation_mode=self.disaggregation_mode,
             enable_hicache_storage=lambda: self.enable_hicache_storage,
+            rust_server=self.rust_server,
         )
 
     def init_batch_result_processor(self) -> None:

python/sglang/srt/managers/scheduler_components/idle_sleeper.py

@@ -1,9 +1,16 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import zmq
 
 from sglang.srt.environ import envs
 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
+
 
 class IdleSleeper:
     """
@@ -33,3 +40,29 @@ class IdleSleeper:
         ):
             self.last_empty_time = real_time()
             current_platform.empty_cache()
+
+
+class RustServerIdleSleeper:
+    """Idle sleeper for the embedded Rust server.
+
+    The Rust ingress is an in-process request ring, not a zmq socket.
+    Instead park directly on the ring: ``wait_ingress`` blocks until
+    a request is pushed — the request ring wakes the parked thread
+    the instant a producer pushes, so there's no added latency for real
+    requests — or the timeout elapses.
+    """
+
+    def __init__(self, rust_server: RustServer, timeout_ms: int = 1000):
+        self.rust_server = rust_server
+        self.timeout_ms = timeout_ms
+        self.last_empty_time = real_time()
+        self.empty_cache_interval = envs.SGLANG_EMPTY_CACHE_INTERVAL.get()
+
+    def maybe_sleep(self):
+        self.rust_server.wait_ingress(self.timeout_ms)
+        if (
+            self.empty_cache_interval > 0
+            and real_time() - self.last_empty_time > self.empty_cache_interval
+        ):
+            self.last_empty_time = real_time()
+            current_platform.empty_cache()

python/sglang/srt/managers/scheduler_components/output_streamer.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import logging
 from dataclasses import dataclass, field
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     List,
@@ -29,6 +30,10 @@ from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
+if TYPE_CHECKING:
+    from sglang.srt.managers.rust_server import RustServer
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -45,6 +50,10 @@ class SchedulerOutputStreamer:
     spec_algorithm: SpeculativeAlgorithm
     disaggregation_mode: DisaggregationMode
     enable_hicache_storage: Callable[[], bool]
+    # When SGLANG_RUST_SERVER is on, generation output is pushed to the embedded
+    # Rust egress ring via `rust_server.push_generation` instead of the zmq
+    # detokenizer. None otherwise. (Rust-specific state lives in RustServer.)
+    rust_server: Optional[RustServer] = None
     _test_stream_output_count: int = 0
 
     def _get_storage_backend_type(self) -> str:
@@ -147,6 +156,7 @@ class SchedulerOutputStreamer:
             default_stream_interval=self.server_args.stream_interval,
             default_force_stream_interval=DEFAULT_FORCE_STREAM_INTERVAL,
             get_cached_tokens_details=self.get_cached_tokens_details,
+            rust_server_mode=self.rust_server is not None,
         )
         for req in reqs:
             if req is skip_req:
@@ -165,7 +175,10 @@ class SchedulerOutputStreamer:
             is_idle_batch=is_idle_batch,
         )
         if payload is not None:
-            self.send_to_detokenizer.send_output(payload)
+            if self.rust_server is not None:
+                self.rust_server.push_generation(payload)
+            else:
+                self.send_to_detokenizer.send_output(payload)
 
     def _maybe_log_time_stats(self, *, req: Req) -> None:
         if (
@@ -260,7 +273,6 @@ class _GenerationStreamAccumulator:
     default_stream_interval: int
     default_force_stream_interval: int
     get_cached_tokens_details: Callable[[Req], Optional[CachedTokensDetails]]
-
     rids: list = field(default_factory=list)
     http_worker_ipcs: list = field(default_factory=list)
     finished_reasons: list = field(default_factory=list)
@@ -307,6 +319,10 @@ class _GenerationStreamAccumulator:
     output_token_ids_logprobs_idx: Optional[list] = None
     output_token_sampling_mask: Optional[list] = None
     output_token_sampling_logprobs: Optional[list] = None
+    # Rust server mode: the Rust detokenizer reconstructs text/ids from the raw
+    # output tokens itself and never consumes the scheduler's incremental-detok
+    # offsets (decode_ids / read_offset), so that per-step bookkeeping is skipped.
+    rust_server_mode: bool = False
 
     def __post_init__(self) -> None:
         if self.return_hidden_states:
@@ -367,34 +383,40 @@ class _GenerationStreamAccumulator:
         send_token_offset = req.send_token_offset
         send_output_token_logprobs_offset = req.send_output_token_logprobs_offset
         self.rids.append(req.rid)
-        self.http_worker_ipcs.append(req.http_worker_ipc)
         self.finished_reasons.append(
             req.finished_reason.to_json() if req.finished_reason else None
         )
-        self.decoded_texts.append(req.decoded_text)
-        decode_ids, read_offset = req.init_incremental_detokenize()
-
-        self.decode_ids_list.append(decode_ids[req.send_decode_id_offset :])
-
         # Exclude the tokens after stop condition
         output_ids_ = req.output_ids_through_stop
-
-        req.send_decode_id_offset = len(decode_ids)
-        self.read_offsets.append(read_offset)
         self.output_ids.append(output_ids_[send_token_offset:])
         req.send_token_offset = len(output_ids_)
-        self.skip_special_tokens.append(req.sampling_params.skip_special_tokens)
-        self.spaces_between_special_tokens.append(
-            req.sampling_params.spaces_between_special_tokens
-        )
-        self.no_stop_trim.append(req.sampling_params.no_stop_trim)
         self.prompt_tokens.append(len(req.origin_input_ids))
-        self.reasoning_tokens.append(req.reasoning_tokens)
-        self.completion_tokens.append(len(output_ids_))
-        self.cached_tokens.append(req.cached_tokens)
 
-        # Collect detailed cache breakdown if available
-        self.cached_tokens_details.append(self.get_cached_tokens_details(req))
+        if not self.rust_server_mode:
+            # Everything below feeds the Python DetokenizerManager /
+            # TokenizerManager (incremental detok, meta_info, per-request metrics)
+            # or gets pickled into the payload (time_stats). The Rust server
+            # replaces those stages and builds its own metadata from the
+            # ChunkEvent, so `push_generation` never reads these — skip the whole
+            # block. The parallel lists stay empty; the payload goes straight to
+            # `push_generation`, which only indexes the fields appended above.
+            self.http_worker_ipcs.append(req.http_worker_ipc)
+            self.decoded_texts.append(req.decoded_text)
+            decode_ids, read_offset = req.init_incremental_detokenize()
+            self.decode_ids_list.append(decode_ids[req.send_decode_id_offset :])
+            req.send_decode_id_offset = len(decode_ids)
+            self.read_offsets.append(read_offset)
+            self.skip_special_tokens.append(req.sampling_params.skip_special_tokens)
+            self.spaces_between_special_tokens.append(
+                req.sampling_params.spaces_between_special_tokens
+            )
+            self.no_stop_trim.append(req.sampling_params.no_stop_trim)
+            self.reasoning_tokens.append(req.reasoning_tokens)
+            self.completion_tokens.append(len(output_ids_))
+            self.cached_tokens.append(req.cached_tokens)
+
+            # Collect detailed cache breakdown if available
+            self.cached_tokens_details.append(self.get_cached_tokens_details(req))
 
         # Multimodal prompt token counts. In disagg decode mode the prefill node
         # already computed these and transferred them via the metadata buffer

python/sglang/srt/managers/scheduler_components/request_receiver.py

@@ -15,6 +15,7 @@ import zmq
 from torch.distributed import barrier
 
 from sglang.srt.disaggregation.utils import prepare_abort
+from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import (
     BatchTokenizedEmbeddingReqInput,
     BatchTokenizedGenerateReqInput,
@@ -35,6 +36,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.server_args import ServerArgs
     from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook
     from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
@@ -44,7 +46,7 @@ if TYPE_CHECKING:
 
 @dataclass(kw_only=True, slots=True, frozen=True)
 class SchedulerRequestReceiver:
-    recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy]
+    recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy, RustServer]
     recv_from_rpc: Optional[zmq.Socket]
     recv_skipper: Any
     input_blocker: Any
@@ -103,6 +105,15 @@ class SchedulerRequestReceiver:
             if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
                 recv_reqs = []
 
+                # Rust ringbuffer backend: drain the in-process ring fed by the
+                # embedded Rust TokenizerManager instead of a zmq socket. Same
+                # non-blocking, msgpack-decoded contract as the zmq path below.
+                if envs.SGLANG_RUST_SERVER.get():
+                    recv_reqs.extend(
+                        self.recv_from_tokenizer.drain(self.max_recv_per_poll)
+                    )
+                    return recv_reqs
+
                 while True:
                     try:
                         if self.recv_limit_reached(len(recv_reqs)):

python/sglang/srt/managers/tokenizer_manager.py

@@ -94,7 +94,10 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem
 from sglang.srt.managers.scheduler_input_blocker import input_blocker_guard_region
 from sglang.srt.managers.tokenizer_control_mixin import TokenizerControlMixin
 from sglang.srt.managers.tokenizer_manager_score_mixin import TokenizerManagerScoreMixin
-from sglang.srt.managers.utils import is_health_check_generate_req
+from sglang.srt.managers.utils import (
+    compute_num_reserved_tokens,
+    is_health_check_generate_req,
+)
 from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
 from sglang.srt.observability.metrics_collector import (
     STAT_LOGGER_ROLE_TOKENIZER,
@@ -117,7 +120,6 @@ from sglang.srt.server_args import (
     ServerArgs,
     set_global_server_args_for_tokenizer,
 )
-from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
     configure_gc_warning,
     freeze_gc,
@@ -400,18 +402,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
         self.max_req_input_len = None  # Will be set later in engine.py
         self.enable_priority_scheduling = server_args.enable_priority_scheduling
         self.default_priority_value = server_args.default_priority_value
-        speculative_algorithm = SpeculativeAlgorithm.from_string(
-            server_args.speculative_algorithm
-        )
-        if speculative_algorithm.is_eagle():
-            # In the current eagle implementation, we store the draft tokens in the output token slots,
-            # so we need to reserve the space for the draft tokens.
-            self.num_reserved_tokens = max(
-                server_args.speculative_eagle_topk * server_args.speculative_num_steps,
-                server_args.max_speculative_num_draft_tokens,
-            )
-        else:
-            self.num_reserved_tokens = 0
+        self.num_reserved_tokens = compute_num_reserved_tokens(server_args)
         self.validate_total_tokens = True
 
     def init_tokenizer_and_processor(self):

python/sglang/srt/managers/utils.py

@@ -2,20 +2,25 @@ from __future__ import annotations
 
 import dataclasses
 import logging
+import re
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, List, Optional, Union
 
+import msgspec
 import torch
 
 from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
 from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.managers import io_struct
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.state_capturer.base import TopkCaptureOutput
 
 if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import GenerationBatchResult
+    from sglang.srt.server_args import ServerArgs
     from sglang.srt.speculative.eagle_info import EagleDraftInput
 
 
@@ -318,3 +323,73 @@ class EmbeddingBatchResult:
 def is_health_check_generate_req(recv_req):
     rid = getattr(recv_req, "rid", None)
     return rid is not None and rid.startswith(HEALTH_CHECK_RID_PREFIX)
+
+
+class MsgpackDecodeError(ValueError):
+    """A msgpack frame the typed decoder rejected, with the failure explained:
+    ``rid`` (when recoverable from the raw tagged array) and a human-readable
+    ``reason`` whose leading ``$[<n>]`` array index is resolved to the struct
+    field name.
+    """
+
+    def __init__(self, rid: Optional[str], reason: str):
+        super().__init__(reason)
+        self.rid = rid
+        self.reason = reason
+
+
+def msgpack_decode_explained(data: bytes) -> Any:
+    """`io_struct.msgpack_decode`, but a rejected frame raises
+    `MsgpackDecodeError` carrying the rid (recovered via an untyped re-decode of
+    the tagged array) and a reason with the failing field named — for callers
+    that must report the failure back to a client (e.g. the rust ingress)
+    instead of just crashing."""
+    # TODO: the hook_custom_types() currently only apply for unit tests, once it
+    # esclate to the main code, we can provide a function to access the _all_types
+
+    try:
+        return io_struct.msgpack_decode(data)
+    except Exception as e:
+        msg = str(e)
+        try:
+            arr = msgspec.msgpack.decode(data)
+        except Exception:
+            arr = None
+        if not (isinstance(arr, (list, tuple)) and arr):
+            raise MsgpackDecodeError(None, msg) from e
+        # Tagged array_like layout is [tag, *fields]; rid is the first field of
+        # every BaseReq struct.
+        rid = str(arr[1]) if len(arr) > 1 and arr[1] is not None else None
+        tag_to_fields = {
+            cls.__struct_config__.tag: cls.__struct_fields__
+            for cls in io_struct._all_types
+            if isinstance(cls, type) and issubclass(cls, msgspec.Struct)
+        }
+        fields = tag_to_fields.get(arr[0])
+        if fields is not None:
+            # Leading ``$[<n>]`` in a msgspec ValidationError path, e.g.
+            # ``$[12][0]``.
+            m = re.search(r"\$\[(\d+)\]", msg)
+            if m is not None:
+                idx = int(m.group(1))
+                if 1 <= idx <= len(fields):
+                    msg = f"{msg[:m.start()]}$.{fields[idx - 1]}{msg[m.end():]}"
+        raise MsgpackDecodeError(rid, msg) from e
+
+
+def compute_num_reserved_tokens(server_args: ServerArgs) -> int:
+    """Output token slots reserved per request, on top of its input.
+
+    The current eagle implementation stores draft tokens in the output token
+    slots, so the context budget has to account for them; every other algorithm
+    reserves nothing. Shared by `TokenizerManager` and the rust server's
+    `server_args` blob (`RustServer._build_server_args`), which needs the same
+    number to run the total-token check in Rust.
+    """
+    algorithm = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)
+    if not algorithm.is_eagle():
+        return 0
+    return max(
+        server_args.speculative_eagle_topk * server_args.speculative_num_steps,
+        server_args.max_speculative_num_draft_tokens,
+    )