add python foundations for the rust server (env vars, msgpack decode helpers, flatten utils)

python/sglang/srt/environ.py

@@ -1193,6 +1193,13 @@ class Envs:
     SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL = EnvInt(0)
     SGLANG_KV_CANARY_ENABLE_MHA_V = EnvBool(False)
 
+    # ===================================================================
+    # Rust Server specific envs.
+    # ===================================================================
+    SGLANG_RUST_SERVER = EnvBool(False)
+    # Most batched requests one /generate HTTP call may expand into.
+    SGLANG_MAX_BATCH_REQS_PER_HTTP_REQ = EnvInt(4096)
+
 
 envs = Envs()
 EnvField._allow_set_name = False

python/sglang/srt/managers/utils.py

@@ -2,20 +2,25 @@
 
 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 @@ def copy_to_cpu(self):
 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,
+    )

python/sglang/srt/server_args.py

@@ -7317,7 +7317,7 @@ def _handle_tokenizer_batching(self):
                 "Please choose one tokenizer batching approach."
             )
 
-        if self.skip_tokenizer_init:
+        if self.skip_tokenizer_init and not envs.SGLANG_RUST_SERVER.get():
             # Tokenizer workers still serve HTTP / state / output work, so
             # their fanout is preserved; detokenizer workers only decode.
             if self.detokenizer_worker_num != 1:
@@ -9108,14 +9108,13 @@ def init_new(
             dist_init_host = na.host
             dist_init_port = na.port
 
-            # We need 5 consecutive ports from port_base for:
-            # port_base, detokenizer, rpc, metrics, scheduler.
-            # In multi-node, all nodes derive ports independently from
-            # dist_init_port, so the derivation must be deterministic
-            # (no availability-based search). If incrementing would
-            # overflow the valid TCP range, decrement instead.
-            NUM_DERIVED_PORTS = 5
-            if server_args.is_ep_joiner:
+            # Reserve port_base+0..NUM_DERIVED_PORTS-1 (6 fixed ports + dp_size
+            # rust-path slots); derive from server_args only (never dp_rank) so
+            # every init_new call agrees, decrementing below dist_init_port on
+            # overflow.
+            is_rust_server = envs.SGLANG_RUST_SERVER.get()
+            NUM_DERIVED_PORTS = 6 if not is_rust_server else 6 + server_args.dp_size
+            if server_args.is_ep_scale_joiner:
                 port_base = server_args.port + ZMQ_TCP_PORT_DELTA
                 if port_base + NUM_DERIVED_PORTS > 65535:
                     port_base = server_args.port - ZMQ_TCP_PORT_DELTA

python/sglang/srt/utils/flatten.py

@@ -0,0 +1,186 @@
+"""Flatten ragged (nested, variable-length) structures into flat value buffers
+plus per-position length vectors — the columnar wire layout used by the embedded
+Rust server's egress path (see ``managers/rust_server.py``).
+
+The ``*Columns`` classes accumulate one batch column-family each: feed them one
+request cell at a time (``accept``), then read the header contribution (length
+vectors) and data contribution (raw ``array`` buffers).
+"""
+
+from __future__ import annotations
+
+from array import array
+from typing import List
+
+
+def flatten_ragged(per_pos_val, per_pos_idx):
+    """Flatten per-position ``list[Optional[list]]`` val/idx pairs into flat
+    buffers + a shared ``lens`` vector (falsy position -> len 0 -> Rust ``null``).
+    idx must mirror val exactly (asserted): the wire pairs both buffers by the
+    one ``lens`` vector, so divergence shifts or drops token ids downstream."""
+
+    flat_val: List[float] = []
+    flat_idx: List[int] = []
+    lens: List[int] = []
+    if not per_pos_val:
+        assert not per_pos_idx, (
+            f"ragged idx column has {len(per_pos_idx)} positions but the val "
+            "column is empty"
+        )
+        return flat_val, flat_idx, lens
+    assert per_pos_idx is not None and len(per_pos_idx) == len(per_pos_val), (
+        f"ragged idx column has {len(per_pos_idx) if per_pos_idx else 0} "
+        f"positions, val column has {len(per_pos_val)}"
+    )
+    for p, pv in enumerate(per_pos_val):
+        pi = per_pos_idx[p]
+        if pv:
+            # A truthy position holds only real logprobs; a `None`/empty position
+            # is the falsy branch below (len 0), so no per-value None check here.
+            assert pi is not None and len(pi) == len(pv), (
+                f"position {p}: idx len "
+                f"{len(pi) if pi is not None else None} != val len {len(pv)}"
+            )
+            flat_val.extend(pv)
+            flat_idx.extend(pi)
+            lens.append(len(pv))
+        else:
+            assert not pi, f"position {p}: idx has {len(pi)} entries but val is empty"
+            lens.append(0)
+    return flat_val, flat_idx, lens
+
+
+def flatten_hidden(hs):
+    """Flatten one request's hidden states into a flat ``val`` buffer plus a
+    per-row ``lens`` vector (one row per output position). Each top-level element
+    becomes a single row; the Rust side reshapes back to ``list[list[float]]``,
+    matching ``meta_info["hidden_states"]``'s common per-position-vector shape.
+    """
+    vals: List[float] = []
+    lens: List[int] = []
+    if not hs:
+        return vals, lens
+    for row in hs:
+        flat = _flatten_floats(row)
+        vals.extend(flat)
+        lens.append(len(flat))
+    return vals, lens
+
+
+def _flatten_floats(x):
+    """Recursively flatten a (possibly nested) float structure into a flat list
+    of floats — handles the ``float | list[float]`` union inside a hidden-state
+    chunk."""
+    if isinstance(x, (int, float)):
+        return [float(x)]
+    out: List[float] = []
+    for e in x:
+        out.extend(_flatten_floats(e))
+    return out
+
+
+class FlatPairColumns:
+    """A flat val/idx column pair (e.g. per-token logprob values + token ids):
+    per-request element counts in the header, concatenated f32 + i32 buffers in
+    the data. ``first_none_to_nan`` maps a leading ``None`` cell element to NaN
+    (the input-logprob first-prompt-token sentinel)."""
+
+    def __init__(self, name, vals, idxs, first_none_to_nan=False):
+        self.name = name
+        self.vals = vals
+        self.idxs = idxs
+        self.first_none_to_nan = first_none_to_nan
+        self.v = array("f")
+        self.i = array("i")
+        self.lens = []
+
+    def columns(self):
+        return ((f"{self.name}_val", self.vals), (f"{self.name}_idx", self.idxs))
+
+    def accept(self, j):
+        vv = (self.vals[j] if self.vals else None) or []
+        ii = (self.idxs[j] if self.idxs else None) or []
+        # Parity assert, the flat twin of `flatten_ragged`'s: only `len(vv)` is
+        # recorded in `lens`, but both buffers are extended, so a longer idx column
+        # silently pushes every LATER column's offset out by the difference. The
+        # decoder cannot catch it — the data buffer only grows, so the receiver's
+        # bounds check still passes and it hands the client another column's bytes
+        # reinterpreted as logprobs, with a 200.
+        assert len(ii) == len(
+            vv
+        ), f"{self.name}: request {j} has {len(ii)} idx entries but {len(vv)} vals"
+        if self.first_none_to_nan and vv and vv[0] is None:
+            self.v.append(float("nan"))
+            self.v.extend(vv[1:])
+        else:
+            self.v.extend(vv)
+        self.i.extend(ii)
+        self.lens.append(len(vv))
+
+    def header_cols(self):
+        return [self.lens]
+
+    def data_cols(self):
+        return [self.v.tobytes(), self.i.tobytes()]
+
+
+class RaggedPairColumns:
+    """A per-position ragged val/idx column pair (e.g. top-k / token-ids
+    logprobs): per-request position counts + a flat per-position length stream
+    in the header, concatenated f32/i32 buffers in the data."""
+
+    def __init__(self, name, vals, idxs):
+        self.name = name
+        self.vals = vals
+        self.idxs = idxs
+        self.v = array("f")
+        self.i = array("i")
+        self.pos = []
+        self.req = []
+
+    def columns(self):
+        return ((f"{self.name}_val", self.vals), (f"{self.name}_idx", self.idxs))
+
+    def accept(self, j):
+        fv, fi, lens = flatten_ragged(
+            self.vals[j] if self.vals else None,
+            self.idxs[j] if self.idxs else None,
+        )
+        self.v.extend(fv)
+        self.i.extend(fi)
+        self.pos.extend(lens)
+        self.req.append(len(lens))
+
+    def header_cols(self):
+        return [self.req, self.pos]
+
+    def data_cols(self):
+        return [self.v.tobytes(), self.i.tobytes()]
+
+
+class NestedRowColumns:
+    """A nested-rows float column (e.g. hidden states): per-request row counts +
+    per-row length stream in the header, one concatenated f32 buffer in the
+    data."""
+
+    def __init__(self, name, rows):
+        self.name = name
+        self.rows = rows
+        self.v = array("f")
+        self.pos = []
+        self.req = []
+
+    def columns(self):
+        return ((self.name, self.rows),)
+
+    def accept(self, j):
+        hv, hlens = flatten_hidden(self.rows[j] if self.rows else None)
+        self.v.extend(hv)
+        self.pos.extend(hlens)
+        self.req.append(len(hlens))
+
+    def header_cols(self):
+        return [self.req, self.pos]
+
+    def data_cols(self):
+        return [self.v.tobytes()]