support rust sglang server (#29799)
This commit is contained in:
@@ -35,6 +35,7 @@ from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, U
|
|||||||
|
|
||||||
import aiohttp
|
import aiohttp
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
import orjson
|
||||||
import requests
|
import requests
|
||||||
from tqdm.asyncio import tqdm
|
from tqdm.asyncio import tqdm
|
||||||
from transformers import AutoTokenizer, PreTrainedTokenizerBase
|
from transformers import AutoTokenizer, PreTrainedTokenizerBase
|
||||||
@@ -710,12 +711,18 @@ async def async_request_sglang_generate(
|
|||||||
if not chunk_bytes:
|
if not chunk_bytes:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
|
# Cumulative chunks make parsing O(n^2) per request on this
|
||||||
|
# single asyncio thread; orjson on raw bytes is ~2.2x cheaper.
|
||||||
|
sse_data = (
|
||||||
|
chunk_bytes[6:]
|
||||||
|
if chunk_bytes.startswith(b"data: ")
|
||||||
|
else chunk_bytes
|
||||||
|
)
|
||||||
latency = time.perf_counter() - st
|
latency = time.perf_counter() - st
|
||||||
if chunk == "[DONE]":
|
if sse_data == b"[DONE]":
|
||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
data = json.loads(chunk)
|
data = orjson.loads(sse_data)
|
||||||
|
|
||||||
_meta_info = data.get("meta_info") or {}
|
_meta_info = data.get("meta_info") or {}
|
||||||
if _meta_info.get("spec_accept_length") is not None:
|
if _meta_info.get("spec_accept_length") is not None:
|
||||||
|
|||||||
@@ -55,6 +55,7 @@ from sglang.srt.entrypoints.engine_info_bootstrap_server import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.entrypoints.engine_score_mixin import EngineScoreMixin
|
from sglang.srt.entrypoints.engine_score_mixin import EngineScoreMixin
|
||||||
from sglang.srt.entrypoints.EngineBase import EngineBase
|
from sglang.srt.entrypoints.EngineBase import EngineBase
|
||||||
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.managers.data_parallel_controller import (
|
from sglang.srt.managers.data_parallel_controller import (
|
||||||
SCHEDULER_PIDS_ARG,
|
SCHEDULER_PIDS_ARG,
|
||||||
run_data_parallel_controller_process,
|
run_data_parallel_controller_process,
|
||||||
@@ -135,7 +136,7 @@ class SchedulerInitResult:
|
|||||||
scheduler_infos: List[Dict[str, Any]]
|
scheduler_infos: List[Dict[str, Any]]
|
||||||
all_child_pids: List[int] = dataclasses.field(default_factory=list)
|
all_child_pids: List[int] = dataclasses.field(default_factory=list)
|
||||||
wait_for_ready: Callable[[], None] = lambda: None
|
wait_for_ready: Callable[[], None] = lambda: None
|
||||||
wait_for_completion: Callable[[], None] = lambda: None
|
block_until_scheduler_exits: Callable[[], None] = lambda: None
|
||||||
engine_info_bootstrap_server: Optional[Any] = None
|
engine_info_bootstrap_server: Optional[Any] = None
|
||||||
|
|
||||||
|
|
||||||
@@ -231,6 +232,14 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
self.server_args = server_args
|
self.server_args = server_args
|
||||||
logger.info(f"{server_args=}")
|
logger.info(f"{server_args=}")
|
||||||
|
|
||||||
|
# Rust Server is not supported with the offline Engine API
|
||||||
|
if envs.SGLANG_RUST_SERVER.get():
|
||||||
|
raise ValueError(
|
||||||
|
"SGLANG_RUST_SERVER is not supported with the offline Engine "
|
||||||
|
"API; it only replaces the HTTP server path (`sglang serve`). "
|
||||||
|
"Unset SGLANG_RUST_SERVER to use sgl.Engine."
|
||||||
|
)
|
||||||
|
|
||||||
# Pre-initialize tokenizer_manager so the atexit handler in
|
# Pre-initialize tokenizer_manager so the atexit handler in
|
||||||
# shutdown() won't hit AttributeError.
|
# shutdown() won't hit AttributeError.
|
||||||
self.tokenizer_manager = None
|
self.tokenizer_manager = None
|
||||||
@@ -900,7 +909,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
if SCHEDULER_PIDS_ARG in info:
|
if SCHEDULER_PIDS_ARG in info:
|
||||||
all_child_pids.extend(info[SCHEDULER_PIDS_ARG])
|
all_child_pids.extend(info[SCHEDULER_PIDS_ARG])
|
||||||
|
|
||||||
def wait_for_completion():
|
def block_until_scheduler_exits():
|
||||||
for proc in scheduler_procs:
|
for proc in scheduler_procs:
|
||||||
proc.join()
|
proc.join()
|
||||||
logger.error(
|
logger.error(
|
||||||
@@ -913,7 +922,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
scheduler_infos=scheduler_infos,
|
scheduler_infos=scheduler_infos,
|
||||||
all_child_pids=all_child_pids,
|
all_child_pids=all_child_pids,
|
||||||
wait_for_ready=wait_for_ready,
|
wait_for_ready=wait_for_ready,
|
||||||
wait_for_completion=wait_for_completion,
|
block_until_scheduler_exits=block_until_scheduler_exits,
|
||||||
),
|
),
|
||||||
scheduler_procs,
|
scheduler_procs,
|
||||||
)
|
)
|
||||||
@@ -1074,7 +1083,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
server_args.host, server_args.port, server_args.enable_metrics
|
server_args.host, server_args.port, server_args.enable_metrics
|
||||||
)
|
)
|
||||||
|
|
||||||
scheduler_init_result.wait_for_completion()
|
scheduler_init_result.block_until_scheduler_exits()
|
||||||
return (
|
return (
|
||||||
None,
|
None,
|
||||||
None,
|
None,
|
||||||
@@ -1084,6 +1093,29 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
weight_cache_daemon_procs,
|
weight_cache_daemon_procs,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# The embedded Rust server (started inside the rank-0 scheduler) owns
|
||||||
|
# the API server, tokenization, and detokenization. In that mode we do
|
||||||
|
# not start the Python detokenizer subprocess(es) or tokenizer manager.
|
||||||
|
# Do not use RayEngine with the Rust server, as it is not supported.
|
||||||
|
if envs.SGLANG_RUST_SERVER.get():
|
||||||
|
scheduler_init_result.wait_for_ready()
|
||||||
|
# Set up subprocess liveness watchdog to detect crashes
|
||||||
|
processes = list(scheduler_procs or [])
|
||||||
|
names = [f"scheduler_{i}" for i in range(len(processes))]
|
||||||
|
subprocess_watchdog = SubprocessWatchdog(
|
||||||
|
processes=processes, process_names=names
|
||||||
|
)
|
||||||
|
subprocess_watchdog.start()
|
||||||
|
|
||||||
|
return (
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
port_args,
|
||||||
|
scheduler_init_result,
|
||||||
|
subprocess_watchdog,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
|
||||||
# Launch detokenizer process(es) — optionally fronted by a router when
|
# Launch detokenizer process(es) — optionally fronted by a router when
|
||||||
# detokenizer_worker_num > 1.
|
# detokenizer_worker_num > 1.
|
||||||
detoken_procs, detoken_names = cls._launch_detokenizer_subprocesses(
|
detoken_procs, detoken_names = cls._launch_detokenizer_subprocesses(
|
||||||
|
|||||||
@@ -2184,10 +2184,20 @@ def _execute_server_warmup(server_args: ServerArgs):
|
|||||||
and server_args.disaggregation_mode == "null"
|
and server_args.disaggregation_mode == "null"
|
||||||
and model_info["is_generation"]
|
and model_info["is_generation"]
|
||||||
):
|
):
|
||||||
|
served_model_name = ""
|
||||||
|
if not envs.SGLANG_RUST_SERVER.get():
|
||||||
|
served_model_name = _global_state.tokenizer_manager.served_model_name
|
||||||
|
else:
|
||||||
|
# _global_state.tokenizer_manager is not initialized in the rust server,
|
||||||
|
# so we need to get the model name from the model_info
|
||||||
|
served_model_name = model_info.get(
|
||||||
|
"model_path", server_args.served_model_name
|
||||||
|
)
|
||||||
|
served_model_name = served_model_name or server_args.model_path
|
||||||
# TODO: ChatCompletionRequest does not have bootstrap info required by disaggregation mode, disable image-warmup for now
|
# TODO: ChatCompletionRequest does not have bootstrap info required by disaggregation mode, disable image-warmup for now
|
||||||
# Only use chat completions format for generation models, not embedding models
|
# Only use chat completions format for generation models, not embedding models
|
||||||
json_data = {
|
json_data = {
|
||||||
"model": _global_state.tokenizer_manager.served_model_name,
|
"model": served_model_name,
|
||||||
"messages": [
|
"messages": [
|
||||||
{
|
{
|
||||||
"role": "user",
|
"role": "user",
|
||||||
@@ -2236,9 +2246,15 @@ def _execute_server_warmup(server_args: ServerArgs):
|
|||||||
verify=ssl_verify,
|
verify=ssl_verify,
|
||||||
)
|
)
|
||||||
assert res.status_code == 200, f"{res.text}"
|
assert res.status_code == 200, f"{res.text}"
|
||||||
|
# Skip server_status update for Rust server
|
||||||
|
if not envs.SGLANG_RUST_SERVER.get():
|
||||||
_global_state.tokenizer_manager.server_status = ServerStatus.Up
|
_global_state.tokenizer_manager.server_status = ServerStatus.Up
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
# TODO: @rainj-me fix this when Rust server supports disaggregation
|
||||||
|
assert (
|
||||||
|
not envs.SGLANG_RUST_SERVER.get()
|
||||||
|
), "Rust server is not supported for disaggregation warmup for now"
|
||||||
logger.info(f"Start of pd disaggregation warmup ...")
|
logger.info(f"Start of pd disaggregation warmup ...")
|
||||||
status_codes = asyncio.run(
|
status_codes = asyncio.run(
|
||||||
_send_disaggregation_warmup_requests(
|
_send_disaggregation_warmup_requests(
|
||||||
@@ -2691,6 +2707,20 @@ def launch_server(
|
|||||||
run_detokenizer_process_func=run_detokenizer_process_func,
|
run_detokenizer_process_func=run_detokenizer_process_func,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if envs.SGLANG_RUST_SERVER.get():
|
||||||
|
# The Rust server serves api-server, tokenizer, and detokenizer, so the
|
||||||
|
# main process has no Python HTTP server / tokenizer manager to run.
|
||||||
|
# Run a warmup /generate before advertising readiness: the Rust /health
|
||||||
|
# and /get_model_info endpoints are static (200 as soon as the server
|
||||||
|
# binds, before any forward pass), so without this the first real request
|
||||||
|
# pays the cold-start cost (observed as a >60s first generation).
|
||||||
|
if not server_args.skip_server_warmup:
|
||||||
|
_execute_server_warmup(server_args)
|
||||||
|
logger.info("The server is fired up and ready to roll!")
|
||||||
|
if launch_callback is not None:
|
||||||
|
launch_callback()
|
||||||
|
scheduler_init_result.block_until_scheduler_exits()
|
||||||
|
else:
|
||||||
_setup_and_run_http_server(
|
_setup_and_run_http_server(
|
||||||
server_args,
|
server_args,
|
||||||
tokenizer_manager,
|
tokenizer_manager,
|
||||||
|
|||||||
@@ -1258,6 +1258,13 @@ class Envs:
|
|||||||
SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL = EnvInt(0)
|
SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL = EnvInt(0)
|
||||||
SGLANG_KV_CANARY_ENABLE_MHA_V = EnvBool(False)
|
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()
|
envs = Envs()
|
||||||
EnvField._allow_set_name = False
|
EnvField._allow_set_name = False
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -1450,7 +1450,28 @@ class Req(ReqDllmMixin):
|
|||||||
# Check stop regex
|
# Check stop regex
|
||||||
if len(self.sampling_params.stop_regex_strs) > 0:
|
if len(self.sampling_params.stop_regex_strs) > 0:
|
||||||
for stop_regex_str in self.sampling_params.stop_regex_strs:
|
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(
|
self.finished_reason = FINISHED_MATCHED_REGEX(
|
||||||
matched=stop_regex_str
|
matched=stop_regex_str
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -167,6 +167,7 @@ from sglang.srt.managers.prefill_delayer import (
|
|||||||
PrefillDelayer,
|
PrefillDelayer,
|
||||||
PrefillDelayerSinglePassExecutor,
|
PrefillDelayerSinglePassExecutor,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.managers.rust_server import RustServer
|
||||||
from sglang.srt.managers.schedule_batch import (
|
from sglang.srt.managers.schedule_batch import (
|
||||||
FINISH_ABORT,
|
FINISH_ABORT,
|
||||||
MultimodalInputs,
|
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.dp_attn import SchedulerDPAttnAdapter
|
||||||
from sglang.srt.managers.scheduler_components.flush_wrapper import SchedulerFlushWrapper
|
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 (
|
from sglang.srt.managers.scheduler_components.invariant_checker import (
|
||||||
SchedulerInvariantChecker,
|
SchedulerInvariantChecker,
|
||||||
create_scheduler_watchdog,
|
create_scheduler_watchdog,
|
||||||
@@ -597,6 +601,10 @@ class Scheduler(
|
|||||||
|
|
||||||
self.maybe_init_scripted_scheduler_hook()
|
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_request_receiver()
|
||||||
|
|
||||||
self.init_dp_attn_adapter()
|
self.init_dp_attn_adapter()
|
||||||
@@ -672,9 +680,12 @@ class Scheduler(
|
|||||||
)
|
)
|
||||||
|
|
||||||
self.load_snapshot_writer = None
|
self.load_snapshot_writer = None
|
||||||
|
self.recv_from_tokenizer = None
|
||||||
|
|
||||||
if not is_rank_zero:
|
if not is_rank_zero:
|
||||||
return
|
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
|
dp_rank = self.ps.dp_rank if self.ps.dp_rank is not None else 0
|
||||||
try:
|
try:
|
||||||
self.load_snapshot_writer = create_load_snapshot_writer(
|
self.load_snapshot_writer = create_load_snapshot_writer(
|
||||||
@@ -1748,11 +1759,16 @@ class Scheduler(
|
|||||||
|
|
||||||
output = self._request_dispatcher(recv_req)
|
output = self._request_dispatcher(recv_req)
|
||||||
if output is not None:
|
if output is not None:
|
||||||
if not isinstance(output, RpcReqOutput):
|
if self.rust_server is not None:
|
||||||
self.ipc_channels.send_to_tokenizer.send_output(output, recv_req)
|
# Embedded Rust server: every control-request response goes
|
||||||
else:
|
# 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:
|
if self.ipc_channels.recv_from_rpc is not None:
|
||||||
sock_send(self.ipc_channels.recv_from_rpc, output)
|
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()
|
self.flush_wrapper.check_pending()
|
||||||
if self.external_corpus_manager is not None:
|
if self.external_corpus_manager is not None:
|
||||||
@@ -1808,9 +1824,33 @@ class Scheduler(
|
|||||||
else:
|
else:
|
||||||
self.scripted_scheduler_hook = None
|
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:
|
def init_request_receiver(self) -> None:
|
||||||
self.request_receiver = SchedulerRequestReceiver(
|
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_from_rpc=self.ipc_channels.recv_from_rpc,
|
||||||
recv_skipper=self.recv_skipper,
|
recv_skipper=self.recv_skipper,
|
||||||
input_blocker=self.input_blocker,
|
input_blocker=self.input_blocker,
|
||||||
@@ -1936,6 +1976,7 @@ class Scheduler(
|
|||||||
spec_algorithm=self.spec_algorithm,
|
spec_algorithm=self.spec_algorithm,
|
||||||
disaggregation_mode=self.disaggregation_mode,
|
disaggregation_mode=self.disaggregation_mode,
|
||||||
enable_hicache_storage=lambda: self.enable_hicache_storage,
|
enable_hicache_storage=lambda: self.enable_hicache_storage,
|
||||||
|
rust_server=self.rust_server,
|
||||||
)
|
)
|
||||||
|
|
||||||
def init_batch_result_processor(self) -> None:
|
def init_batch_result_processor(self) -> None:
|
||||||
|
|||||||
@@ -1,9 +1,16 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
import zmq
|
import zmq
|
||||||
|
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.observability.req_time_stats import real_time
|
from sglang.srt.observability.req_time_stats import real_time
|
||||||
from sglang.srt.platforms import current_platform
|
from sglang.srt.platforms import current_platform
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from sglang.srt.managers.rust_server import RustServer
|
||||||
|
|
||||||
|
|
||||||
class IdleSleeper:
|
class IdleSleeper:
|
||||||
"""
|
"""
|
||||||
@@ -33,3 +40,29 @@ class IdleSleeper:
|
|||||||
):
|
):
|
||||||
self.last_empty_time = real_time()
|
self.last_empty_time = real_time()
|
||||||
current_platform.empty_cache()
|
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()
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
|||||||
import logging
|
import logging
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import (
|
from typing import (
|
||||||
|
TYPE_CHECKING,
|
||||||
Any,
|
Any,
|
||||||
Callable,
|
Callable,
|
||||||
List,
|
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.server_args import ServerArgs
|
||||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from sglang.srt.managers.rust_server import RustServer
|
||||||
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -45,6 +50,10 @@ class SchedulerOutputStreamer:
|
|||||||
spec_algorithm: SpeculativeAlgorithm
|
spec_algorithm: SpeculativeAlgorithm
|
||||||
disaggregation_mode: DisaggregationMode
|
disaggregation_mode: DisaggregationMode
|
||||||
enable_hicache_storage: Callable[[], bool]
|
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
|
_test_stream_output_count: int = 0
|
||||||
|
|
||||||
def _get_storage_backend_type(self) -> str:
|
def _get_storage_backend_type(self) -> str:
|
||||||
@@ -147,6 +156,7 @@ class SchedulerOutputStreamer:
|
|||||||
default_stream_interval=self.server_args.stream_interval,
|
default_stream_interval=self.server_args.stream_interval,
|
||||||
default_force_stream_interval=DEFAULT_FORCE_STREAM_INTERVAL,
|
default_force_stream_interval=DEFAULT_FORCE_STREAM_INTERVAL,
|
||||||
get_cached_tokens_details=self.get_cached_tokens_details,
|
get_cached_tokens_details=self.get_cached_tokens_details,
|
||||||
|
rust_server_mode=self.rust_server is not None,
|
||||||
)
|
)
|
||||||
for req in reqs:
|
for req in reqs:
|
||||||
if req is skip_req:
|
if req is skip_req:
|
||||||
@@ -165,6 +175,9 @@ class SchedulerOutputStreamer:
|
|||||||
is_idle_batch=is_idle_batch,
|
is_idle_batch=is_idle_batch,
|
||||||
)
|
)
|
||||||
if payload is not None:
|
if payload is not None:
|
||||||
|
if self.rust_server is not None:
|
||||||
|
self.rust_server.push_generation(payload)
|
||||||
|
else:
|
||||||
self.send_to_detokenizer.send_output(payload)
|
self.send_to_detokenizer.send_output(payload)
|
||||||
|
|
||||||
def _maybe_log_time_stats(self, *, req: Req) -> None:
|
def _maybe_log_time_stats(self, *, req: Req) -> None:
|
||||||
@@ -260,7 +273,6 @@ class _GenerationStreamAccumulator:
|
|||||||
default_stream_interval: int
|
default_stream_interval: int
|
||||||
default_force_stream_interval: int
|
default_force_stream_interval: int
|
||||||
get_cached_tokens_details: Callable[[Req], Optional[CachedTokensDetails]]
|
get_cached_tokens_details: Callable[[Req], Optional[CachedTokensDetails]]
|
||||||
|
|
||||||
rids: list = field(default_factory=list)
|
rids: list = field(default_factory=list)
|
||||||
http_worker_ipcs: list = field(default_factory=list)
|
http_worker_ipcs: list = field(default_factory=list)
|
||||||
finished_reasons: 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_ids_logprobs_idx: Optional[list] = None
|
||||||
output_token_sampling_mask: Optional[list] = None
|
output_token_sampling_mask: Optional[list] = None
|
||||||
output_token_sampling_logprobs: 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:
|
def __post_init__(self) -> None:
|
||||||
if self.return_hidden_states:
|
if self.return_hidden_states:
|
||||||
@@ -367,28 +383,34 @@ class _GenerationStreamAccumulator:
|
|||||||
send_token_offset = req.send_token_offset
|
send_token_offset = req.send_token_offset
|
||||||
send_output_token_logprobs_offset = req.send_output_token_logprobs_offset
|
send_output_token_logprobs_offset = req.send_output_token_logprobs_offset
|
||||||
self.rids.append(req.rid)
|
self.rids.append(req.rid)
|
||||||
self.http_worker_ipcs.append(req.http_worker_ipc)
|
|
||||||
self.finished_reasons.append(
|
self.finished_reasons.append(
|
||||||
req.finished_reason.to_json() if req.finished_reason else None
|
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
|
# Exclude the tokens after stop condition
|
||||||
output_ids_ = req.output_ids_through_stop
|
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:])
|
self.output_ids.append(output_ids_[send_token_offset:])
|
||||||
req.send_token_offset = len(output_ids_)
|
req.send_token_offset = len(output_ids_)
|
||||||
|
self.prompt_tokens.append(len(req.origin_input_ids))
|
||||||
|
|
||||||
|
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.skip_special_tokens.append(req.sampling_params.skip_special_tokens)
|
||||||
self.spaces_between_special_tokens.append(
|
self.spaces_between_special_tokens.append(
|
||||||
req.sampling_params.spaces_between_special_tokens
|
req.sampling_params.spaces_between_special_tokens
|
||||||
)
|
)
|
||||||
self.no_stop_trim.append(req.sampling_params.no_stop_trim)
|
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.reasoning_tokens.append(req.reasoning_tokens)
|
||||||
self.completion_tokens.append(len(output_ids_))
|
self.completion_tokens.append(len(output_ids_))
|
||||||
self.cached_tokens.append(req.cached_tokens)
|
self.cached_tokens.append(req.cached_tokens)
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ import zmq
|
|||||||
from torch.distributed import barrier
|
from torch.distributed import barrier
|
||||||
|
|
||||||
from sglang.srt.disaggregation.utils import prepare_abort
|
from sglang.srt.disaggregation.utils import prepare_abort
|
||||||
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.managers.io_struct import (
|
from sglang.srt.managers.io_struct import (
|
||||||
BatchTokenizedEmbeddingReqInput,
|
BatchTokenizedEmbeddingReqInput,
|
||||||
BatchTokenizedGenerateReqInput,
|
BatchTokenizedGenerateReqInput,
|
||||||
@@ -35,6 +36,7 @@ from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
|
|||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.configs.model_config import ModelConfig
|
from sglang.srt.configs.model_config import ModelConfig
|
||||||
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
|
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.srt.server_args import ServerArgs
|
||||||
from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook
|
from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook
|
||||||
from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
|
from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
|
||||||
@@ -44,7 +46,7 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
@dataclass(kw_only=True, slots=True, frozen=True)
|
@dataclass(kw_only=True, slots=True, frozen=True)
|
||||||
class SchedulerRequestReceiver:
|
class SchedulerRequestReceiver:
|
||||||
recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy]
|
recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy, RustServer]
|
||||||
recv_from_rpc: Optional[zmq.Socket]
|
recv_from_rpc: Optional[zmq.Socket]
|
||||||
recv_skipper: Any
|
recv_skipper: Any
|
||||||
input_blocker: Any
|
input_blocker: Any
|
||||||
@@ -103,6 +105,15 @@ class SchedulerRequestReceiver:
|
|||||||
if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
|
if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
|
||||||
recv_reqs = []
|
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:
|
while True:
|
||||||
try:
|
try:
|
||||||
if self.recv_limit_reached(len(recv_reqs)):
|
if self.recv_limit_reached(len(recv_reqs)):
|
||||||
|
|||||||
@@ -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.scheduler_input_blocker import input_blocker_guard_region
|
||||||
from sglang.srt.managers.tokenizer_control_mixin import TokenizerControlMixin
|
from sglang.srt.managers.tokenizer_control_mixin import TokenizerControlMixin
|
||||||
from sglang.srt.managers.tokenizer_manager_score_mixin import TokenizerManagerScoreMixin
|
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.cpu_monitor import start_cpu_monitor_thread
|
||||||
from sglang.srt.observability.metrics_collector import (
|
from sglang.srt.observability.metrics_collector import (
|
||||||
STAT_LOGGER_ROLE_TOKENIZER,
|
STAT_LOGGER_ROLE_TOKENIZER,
|
||||||
@@ -117,7 +120,6 @@ from sglang.srt.server_args import (
|
|||||||
ServerArgs,
|
ServerArgs,
|
||||||
set_global_server_args_for_tokenizer,
|
set_global_server_args_for_tokenizer,
|
||||||
)
|
)
|
||||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
|
||||||
from sglang.srt.utils import (
|
from sglang.srt.utils import (
|
||||||
configure_gc_warning,
|
configure_gc_warning,
|
||||||
freeze_gc,
|
freeze_gc,
|
||||||
@@ -400,18 +402,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
self.max_req_input_len = None # Will be set later in engine.py
|
self.max_req_input_len = None # Will be set later in engine.py
|
||||||
self.enable_priority_scheduling = server_args.enable_priority_scheduling
|
self.enable_priority_scheduling = server_args.enable_priority_scheduling
|
||||||
self.default_priority_value = server_args.default_priority_value
|
self.default_priority_value = server_args.default_priority_value
|
||||||
speculative_algorithm = SpeculativeAlgorithm.from_string(
|
self.num_reserved_tokens = compute_num_reserved_tokens(server_args)
|
||||||
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.validate_total_tokens = True
|
self.validate_total_tokens = True
|
||||||
|
|
||||||
def init_tokenizer_and_processor(self):
|
def init_tokenizer_and_processor(self):
|
||||||
|
|||||||
@@ -2,20 +2,25 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import dataclasses
|
import dataclasses
|
||||||
import logging
|
import logging
|
||||||
|
import re
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import TYPE_CHECKING, Any, List, Optional, Union
|
from typing import TYPE_CHECKING, Any, List, Optional, Union
|
||||||
|
|
||||||
|
import msgspec
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
|
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
|
||||||
from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics
|
from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics
|
||||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
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.managers.schedule_batch import Req
|
||||||
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
|
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
|
from sglang.srt.state_capturer.base import TopkCaptureOutput
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.managers.scheduler import GenerationBatchResult
|
from sglang.srt.managers.scheduler import GenerationBatchResult
|
||||||
|
from sglang.srt.server_args import ServerArgs
|
||||||
from sglang.srt.speculative.eagle_info import EagleDraftInput
|
from sglang.srt.speculative.eagle_info import EagleDraftInput
|
||||||
|
|
||||||
|
|
||||||
@@ -318,3 +323,73 @@ class EmbeddingBatchResult:
|
|||||||
def is_health_check_generate_req(recv_req):
|
def is_health_check_generate_req(recv_req):
|
||||||
rid = getattr(recv_req, "rid", None)
|
rid = getattr(recv_req, "rid", None)
|
||||||
return rid is not None and rid.startswith(HEALTH_CHECK_RID_PREFIX)
|
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,
|
||||||
|
)
|
||||||
|
|||||||
@@ -401,7 +401,7 @@ class RayEngine(Engine):
|
|||||||
actor.run_event_loop.remote() for actor in scheduler_actors
|
actor.run_event_loop.remote() for actor in scheduler_actors
|
||||||
]
|
]
|
||||||
|
|
||||||
def wait_for_completion():
|
def block_until_scheduler_exits():
|
||||||
try:
|
try:
|
||||||
ray.get(event_loop_refs)
|
ray.get(event_loop_refs)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -410,7 +410,7 @@ class RayEngine(Engine):
|
|||||||
return (
|
return (
|
||||||
RaySchedulerInitResult(
|
RaySchedulerInitResult(
|
||||||
scheduler_infos=scheduler_infos,
|
scheduler_infos=scheduler_infos,
|
||||||
wait_for_completion=wait_for_completion,
|
block_until_scheduler_exits=block_until_scheduler_exits,
|
||||||
scheduler_actors=scheduler_actors,
|
scheduler_actors=scheduler_actors,
|
||||||
),
|
),
|
||||||
None,
|
None,
|
||||||
@@ -489,7 +489,7 @@ class RayEngine(Engine):
|
|||||||
|
|
||||||
event_loop_refs = controller.event_loop_refs
|
event_loop_refs = controller.event_loop_refs
|
||||||
|
|
||||||
def wait_for_completion():
|
def block_until_scheduler_exits():
|
||||||
try:
|
try:
|
||||||
ray.get(event_loop_refs)
|
ray.get(event_loop_refs)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -497,6 +497,6 @@ class RayEngine(Engine):
|
|||||||
|
|
||||||
return RaySchedulerInitResult(
|
return RaySchedulerInitResult(
|
||||||
scheduler_infos=scheduler_infos,
|
scheduler_infos=scheduler_infos,
|
||||||
wait_for_completion=wait_for_completion,
|
block_until_scheduler_exits=block_until_scheduler_exits,
|
||||||
scheduler_actors=controller.scheduler_actors,
|
scheduler_actors=controller.scheduler_actors,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -7377,7 +7377,7 @@ class ServerArgs:
|
|||||||
"Please choose one tokenizer batching approach."
|
"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
|
# Tokenizer workers still serve HTTP / state / output work, so
|
||||||
# their fanout is preserved; detokenizer workers only decode.
|
# their fanout is preserved; detokenizer workers only decode.
|
||||||
if self.detokenizer_worker_num != 1:
|
if self.detokenizer_worker_num != 1:
|
||||||
@@ -9190,14 +9190,13 @@ class PortArgs:
|
|||||||
dist_init_host = na.host
|
dist_init_host = na.host
|
||||||
dist_init_port = na.port
|
dist_init_port = na.port
|
||||||
|
|
||||||
# We need 5 consecutive ports from port_base for:
|
# Reserve port_base+0..NUM_DERIVED_PORTS-1 (6 fixed ports + dp_size
|
||||||
# port_base, detokenizer, rpc, metrics, scheduler.
|
# rust-path slots); derive from server_args only (never dp_rank) so
|
||||||
# In multi-node, all nodes derive ports independently from
|
# every init_new call agrees, decrementing below dist_init_port on
|
||||||
# dist_init_port, so the derivation must be deterministic
|
# overflow.
|
||||||
# (no availability-based search). If incrementing would
|
is_rust_server = envs.SGLANG_RUST_SERVER.get()
|
||||||
# overflow the valid TCP range, decrement instead.
|
NUM_DERIVED_PORTS = 6 if not is_rust_server else 6 + server_args.dp_size
|
||||||
NUM_DERIVED_PORTS = 5
|
if server_args.is_ep_scale_joiner:
|
||||||
if server_args.is_ep_joiner:
|
|
||||||
port_base = server_args.port + ZMQ_TCP_PORT_DELTA
|
port_base = server_args.port + ZMQ_TCP_PORT_DELTA
|
||||||
if port_base + NUM_DERIVED_PORTS > 65535:
|
if port_base + NUM_DERIVED_PORTS > 65535:
|
||||||
port_base = server_args.port - ZMQ_TCP_PORT_DELTA
|
port_base = server_args.port - ZMQ_TCP_PORT_DELTA
|
||||||
|
|||||||
@@ -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()]
|
||||||
@@ -16,6 +16,7 @@ from sglang.test.simple_eval_common import (
|
|||||||
ChatCompletionSampler,
|
ChatCompletionSampler,
|
||||||
CompletionSampler,
|
CompletionSampler,
|
||||||
Eval,
|
Eval,
|
||||||
|
GenerateSampler,
|
||||||
make_report,
|
make_report,
|
||||||
set_ulimit,
|
set_ulimit,
|
||||||
)
|
)
|
||||||
@@ -81,6 +82,13 @@ def run_eval_once(args, base_url: str, eval_obj: Eval) -> dict:
|
|||||||
**common_kwargs,
|
**common_kwargs,
|
||||||
stop=stop,
|
stop=stop,
|
||||||
)
|
)
|
||||||
|
elif api_mode == "generate":
|
||||||
|
# SGLang-native `/generate` (raw text + sampling_params), same stop defaults.
|
||||||
|
stop = getattr(args, "stop", ["Question", "Assistant:", "<|separator|>"])
|
||||||
|
sampler = GenerateSampler(
|
||||||
|
**common_kwargs,
|
||||||
|
stop=stop,
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
sampler = ChatCompletionSampler(
|
sampler = ChatCompletionSampler(
|
||||||
**common_kwargs,
|
**common_kwargs,
|
||||||
@@ -454,8 +462,8 @@ if __name__ == "__main__":
|
|||||||
"--api",
|
"--api",
|
||||||
type=str,
|
type=str,
|
||||||
default="chat",
|
default="chat",
|
||||||
choices=["chat", "completion"],
|
choices=["chat", "completion", "generate"],
|
||||||
help="API mode: 'chat' for /v1/chat/completions, 'completion' for /v1/completions",
|
help="API mode: 'chat' for /v1/chat/completions, 'completion' for /v1/completions, 'generate' for SGLang-native /generate",
|
||||||
)
|
)
|
||||||
parser.add_argument("--num-examples", type=int)
|
parser.add_argument("--num-examples", type=int)
|
||||||
parser.add_argument("--num-threads", type=int, default=512)
|
parser.add_argument("--num-threads", type=int, default=512)
|
||||||
|
|||||||
@@ -254,6 +254,100 @@ class CompletionSampler(SamplerBase):
|
|||||||
return ""
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
class GenerateSampler(SamplerBase):
|
||||||
|
"""
|
||||||
|
Sample from SGLang's native ``/generate`` endpoint (not the OpenAI-compatible
|
||||||
|
API). Sends raw text prompts with `sampling_params`, so it exercises the same
|
||||||
|
path as `bench_serving` rather than the `/v1/completions` wrapper.
|
||||||
|
|
||||||
|
`base_url` is the OpenAI-style URL the eval harness builds (``.../v1``); the
|
||||||
|
trailing ``/v1`` is stripped to reach the server root's ``/generate``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
base_url: str = None,
|
||||||
|
model: Optional[str] = None,
|
||||||
|
temperature: float = 0.0,
|
||||||
|
top_p: float = 1.0,
|
||||||
|
max_tokens: int = 2048,
|
||||||
|
stop: Optional[List[str]] = None,
|
||||||
|
):
|
||||||
|
self.client = LargerHttpxClient()
|
||||||
|
|
||||||
|
# The harness passes the OpenAI base (`.../v1`); `/generate` lives at root.
|
||||||
|
root = (base_url or "http://127.0.0.1:30000/v1").rstrip("/")
|
||||||
|
if root.endswith("/v1"):
|
||||||
|
root = root[: -len("/v1")]
|
||||||
|
self.generate_url = f"{root}/generate"
|
||||||
|
|
||||||
|
# `/generate` serves the loaded model and ignores a model field, so `model`
|
||||||
|
# is informational only; fill it from `/get_model_info` when unset.
|
||||||
|
if model is None:
|
||||||
|
try:
|
||||||
|
info = self.client.get(f"{root}/get_model_info").json()
|
||||||
|
model = info.get("model_path")
|
||||||
|
except Exception:
|
||||||
|
model = None
|
||||||
|
|
||||||
|
self.model = model
|
||||||
|
self.temperature = temperature
|
||||||
|
self.top_p = top_p
|
||||||
|
self.max_tokens = max_tokens
|
||||||
|
self.stop = stop
|
||||||
|
self._completion_tokens: list[int] = []
|
||||||
|
print(
|
||||||
|
f"GenerateSampler initialized with {self.generate_url=} {self.model=} "
|
||||||
|
f"{self.temperature=} {self.max_tokens=} {self.stop=}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def _pack_message(self, role: str, content: Any):
|
||||||
|
return {"role": str(role), "content": content}
|
||||||
|
|
||||||
|
def __call__(self, message_list: MessageList) -> str:
|
||||||
|
# Extract raw text from message list (eval objects pack prompt as a single user message)
|
||||||
|
prompt = "\n".join(
|
||||||
|
msg["content"]
|
||||||
|
for msg in message_list
|
||||||
|
if isinstance(msg.get("content"), str)
|
||||||
|
)
|
||||||
|
payload = {
|
||||||
|
"text": prompt,
|
||||||
|
"sampling_params": {
|
||||||
|
"temperature": self.temperature,
|
||||||
|
"top_p": self.top_p,
|
||||||
|
"max_new_tokens": self.max_tokens,
|
||||||
|
"stop": self.stop,
|
||||||
|
},
|
||||||
|
"stream": False,
|
||||||
|
}
|
||||||
|
trial = 0
|
||||||
|
while trial < 6:
|
||||||
|
try:
|
||||||
|
response = self.client.post(self.generate_url, json=payload)
|
||||||
|
# A 400 is a malformed request, not a transient failure — don't retry.
|
||||||
|
if response.status_code == 400:
|
||||||
|
print("Bad Request Error", response.text)
|
||||||
|
return ""
|
||||||
|
response.raise_for_status()
|
||||||
|
data = response.json()
|
||||||
|
meta_info = data.get("meta_info") or {}
|
||||||
|
completion_tokens = meta_info.get("completion_tokens")
|
||||||
|
if completion_tokens is not None:
|
||||||
|
self._completion_tokens.append(completion_tokens)
|
||||||
|
return data.get("text") or ""
|
||||||
|
except Exception as e:
|
||||||
|
exception_backoff = 2**trial
|
||||||
|
print(
|
||||||
|
f"Rate limit exception so wait and retry {trial} after {exception_backoff} sec",
|
||||||
|
e,
|
||||||
|
)
|
||||||
|
time.sleep(exception_backoff)
|
||||||
|
trial += 1
|
||||||
|
print(f"All retry attempts exhausted for request. Returning empty response.")
|
||||||
|
return ""
|
||||||
|
|
||||||
|
|
||||||
QUERY_TEMPLATE_MULTICHOICE = """
|
QUERY_TEMPLATE_MULTICHOICE = """
|
||||||
Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering.
|
Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering.
|
||||||
|
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import argparse
|
|||||||
import asyncio
|
import asyncio
|
||||||
import copy
|
import copy
|
||||||
import doctest
|
import doctest
|
||||||
|
import importlib.util
|
||||||
import inspect
|
import inspect
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
@@ -202,6 +203,23 @@ def is_h200_system():
|
|||||||
return envs.IS_H200.get()
|
return envs.IS_H200.get()
|
||||||
|
|
||||||
|
|
||||||
|
def is_rust_server_built():
|
||||||
|
"""Return whether the embedded Rust server extension (``SGLANG_RUST_SERVER``)
|
||||||
|
is importable.
|
||||||
|
|
||||||
|
``sglang/srt/server/`` is not in the source tree — it is produced by
|
||||||
|
``setup.py build_rust --inplace``, so on a build without it ``find_spec``
|
||||||
|
raises ``ModuleNotFoundError`` for the missing *parent* package rather than
|
||||||
|
returning ``None`` for the missing leaf. Suites gate a rust-server subclass on
|
||||||
|
this at class-definition time, so letting that escape would fail the whole
|
||||||
|
module import instead of skipping the one class.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
return importlib.util.find_spec("sglang.srt.server._core") is not None
|
||||||
|
except ModuleNotFoundError:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def _use_cached_default_models(model_repo: str):
|
def _use_cached_default_models(model_repo: str):
|
||||||
cache_dir = os.getenv("DEFAULT_MODEL_CACHE_DIR")
|
cache_dir = os.getenv("DEFAULT_MODEL_CACHE_DIR")
|
||||||
if cache_dir and model_repo:
|
if cache_dir and model_repo:
|
||||||
|
|||||||
@@ -31,8 +31,9 @@ uuid = { workspace = true }
|
|||||||
|
|
||||||
axum = { version = "0.8.9", features = ["json", "tokio"] }
|
axum = { version = "0.8.9", features = ["json", "tokio"] }
|
||||||
core_affinity = "0.8"
|
core_affinity = "0.8"
|
||||||
# the dynamo-tokenizers is deps on hf-hub, should bump version together
|
# pin the dynamo-tokenizers for now, 1.7.0 enables`serde_json/preserve_order`
|
||||||
dynamo-tokenizers = "1.5.3"
|
# transitively, which reorders every JSON object.
|
||||||
|
dynamo-tokenizers = "=1.5.3"
|
||||||
flume = "0.12.0"
|
flume = "0.12.0"
|
||||||
itertools = "0.14"
|
itertools = "0.14"
|
||||||
hf-hub = { version = "0.4", default-features = false }
|
hf-hub = { version = "0.4", default-features = false }
|
||||||
|
|||||||
@@ -144,6 +144,20 @@ impl Runnable for TokenizerWorker {
|
|||||||
tracing::error!("tokenizer pool received a non-generate request");
|
tracing::error!("tokenizer pool received a non-generate request");
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
|
// Size the scheduler's stop-match window in TOKENS, as Python's
|
||||||
|
// `normalize(tokenizer)` does.
|
||||||
|
let stop_tokens = g
|
||||||
|
.sampling_params
|
||||||
|
.stop_strs
|
||||||
|
.iter()
|
||||||
|
// A stop that won't encode falls back to its byte length rather
|
||||||
|
// than failing the request: still an over-estimate, never an
|
||||||
|
// under-estimate, so the scheduler cannot miss that stop.
|
||||||
|
.map(|s| self.tokenizer.encode(s).map_or(s.len(), |ids| ids.len()))
|
||||||
|
.max();
|
||||||
|
if let Some(n) = stop_tokens {
|
||||||
|
g.sampling_params.stop_str_max_len = n;
|
||||||
|
}
|
||||||
match self.tokenizer.encode(g.text.as_deref().unwrap_or("")) {
|
match self.tokenizer.encode(g.text.as_deref().unwrap_or("")) {
|
||||||
Ok(ids) => {
|
Ok(ids) => {
|
||||||
g.input_ids = Some(ids);
|
g.input_ids = Some(ids);
|
||||||
@@ -160,3 +174,68 @@ impl Runnable for TokenizerWorker {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::fsm::RequestState;
|
||||||
|
use crate::message::{EgressSink, GenerateRequest, RequestKind, SamplingParams};
|
||||||
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
|
/// One token per whitespace-separated word, so a stop's token count differs
|
||||||
|
/// from its byte count and the two units cannot be confused.
|
||||||
|
struct WordTokenizer;
|
||||||
|
impl TextTokenizer for WordTokenizer {
|
||||||
|
fn encode(&self, text: &str) -> Result<TokenIds, Error> {
|
||||||
|
Ok(text.split_whitespace().map(|_| 1i32).collect())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The scheduler's stop-match window must reach the wire as a TOKEN count, as
|
||||||
|
/// Python's `normalize(tokenizer)` produces.
|
||||||
|
///
|
||||||
|
/// `Normalizing` leaves a UTF-8 BYTE count there — a safe over-estimate, but it
|
||||||
|
/// makes the scheduler decode a longer tail on EVERY decode step of EVERY
|
||||||
|
/// request (14 tokens vs 6 for a typical stop set). This stage owns the
|
||||||
|
/// tokenizer, so it is where the exact count is resolved.
|
||||||
|
#[test]
|
||||||
|
fn tokenizing_replaces_the_byte_window_with_a_token_count() {
|
||||||
|
let (req_tx, req_rx) = flume::unbounded::<Request>();
|
||||||
|
let (tm_tx, tm_rx) = flume::unbounded::<TmEvent>();
|
||||||
|
|
||||||
|
// 8 bytes vs 3 "tokens" under WordTokenizer — units are distinguishable.
|
||||||
|
let sp = SamplingParams {
|
||||||
|
stop_strs: vec!["a bb ccc".to_string(), "dd".to_string()],
|
||||||
|
stop_str_max_len: 8, // what `normalize_stops` left: max BYTE length
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let (sink_tx, _sink_rx) = mpsc::channel(4);
|
||||||
|
req_tx
|
||||||
|
.send(Request {
|
||||||
|
rid: "1".into(),
|
||||||
|
state: RequestState::Tokenizing,
|
||||||
|
sink: EgressSink::Local(sink_tx),
|
||||||
|
kind: RequestKind::Generate(Box::new(GenerateRequest {
|
||||||
|
rid: "1".into(),
|
||||||
|
text: Some("hello world".into()),
|
||||||
|
sampling_params: sp,
|
||||||
|
..Default::default()
|
||||||
|
})),
|
||||||
|
})
|
||||||
|
.expect("send");
|
||||||
|
drop(req_tx); // closes the loop after one request
|
||||||
|
|
||||||
|
TokenizerWorker::new(req_rx, tm_tx, Arc::new(WordTokenizer)).run();
|
||||||
|
|
||||||
|
let TmEvent::Tokenized(req) = tm_rx.try_recv().expect("returned") else {
|
||||||
|
panic!("expected Tokenized");
|
||||||
|
};
|
||||||
|
let RequestKind::Generate(g) = &req.kind else {
|
||||||
|
panic!("expected generate");
|
||||||
|
};
|
||||||
|
assert_eq!(
|
||||||
|
g.sampling_params.stop_str_max_len, 3,
|
||||||
|
"must be the max TOKEN count (3), not the byte count (8)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -2,6 +2,8 @@
|
|||||||
python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_simple_decode
|
python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_simple_decode
|
||||||
python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_logprob_with_chunked_prefill
|
python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_logprob_with_chunked_prefill
|
||||||
python3 -m unittest test_srt_endpoint.TestTokenizeDetokenize
|
python3 -m unittest test_srt_endpoint.TestTokenizeDetokenize
|
||||||
|
python3 -m unittest test_srt_endpoint.TestRustServerEndpoint
|
||||||
|
python3 -m unittest test_srt_endpoint.TestRustServerLogprob
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import json
|
import json
|
||||||
@@ -24,17 +26,22 @@ from sglang.test.test_utils import (
|
|||||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
DEFAULT_URL_FOR_TEST,
|
DEFAULT_URL_FOR_TEST,
|
||||||
CustomTestCase,
|
CustomTestCase,
|
||||||
|
is_rust_server_built,
|
||||||
popen_launch_server,
|
popen_launch_server,
|
||||||
run_logprob_check,
|
run_logprob_check,
|
||||||
)
|
)
|
||||||
|
|
||||||
register_cuda_ci(est_time=160, stage="base-b", runner_config="1-gpu-small")
|
register_cuda_ci(est_time=260, stage="base-b", runner_config="1-gpu-small")
|
||||||
register_amd_ci(est_time=160, suite="stage-b-test-1-gpu-small-amd")
|
register_amd_ci(est_time=260, suite="stage-b-test-1-gpu-small-amd")
|
||||||
|
|
||||||
SERVER_ENV = {"SGLANG_USE_PICKLE_IPC": "0"}
|
SERVER_ENV = {"SGLANG_USE_PICKLE_IPC": "0"}
|
||||||
|
|
||||||
|
|
||||||
class TestSRTEndpoint(CustomTestCase):
|
class TestSRTEndpoint(CustomTestCase):
|
||||||
|
# Extra server-launch env; subclasses override to run the same suite
|
||||||
|
# against a different server flavor (e.g. SGLANG_RUST_SERVER=1).
|
||||||
|
env = {}
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls):
|
def setUpClass(cls):
|
||||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||||
@@ -46,7 +53,7 @@ class TestSRTEndpoint(CustomTestCase):
|
|||||||
# The tiny logprob chunk size routes this file's logprob tests
|
# The tiny logprob chunk size routes this file's logprob tests
|
||||||
# through the multi-chunk stitching path (requests at or below 64
|
# through the multi-chunk stitching path (requests at or below 64
|
||||||
# rows still cover the non-chunked path).
|
# rows still cover the non-chunked path).
|
||||||
env={**SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"},
|
env={**cls.env, **SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"},
|
||||||
other_args=(
|
other_args=(
|
||||||
"--enable-custom-logit-processor",
|
"--enable-custom-logit-processor",
|
||||||
"--mem-fraction-static",
|
"--mem-fraction-static",
|
||||||
@@ -844,5 +851,74 @@ class TestTokenizeDetokenize(CustomTestCase):
|
|||||||
self.assertEqual(r2.status_code, 500)
|
self.assertEqual(r2.status_code, 500)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Embedded Rust server (SGLANG_RUST_SERVER=1): rerun the whole endpoint suite
|
||||||
|
# against the rust api-server/tokenizer/detokenizer stack — the logprob tests
|
||||||
|
# exercise the columnar egress wire (`push_generation` extras -> Rust
|
||||||
|
# `BatchHeader`/`for_each_chunk` -> detok reshape) end to end. Suite
|
||||||
|
# surface the rust server does not implement yet is skipped explicitly below.
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
@unittest.skipUnless(
|
||||||
|
is_rust_server_built(),
|
||||||
|
"embedded rust server extension not built (e.g. AMD suite)",
|
||||||
|
)
|
||||||
|
class TestRustServerEndpoint(TestSRTEndpoint):
|
||||||
|
env = {"SGLANG_RUST_SERVER": "1"}
|
||||||
|
|
||||||
|
_RUST_TODO = "not implemented by the embedded Rust server yet"
|
||||||
|
|
||||||
|
@unittest.skip(f"custom_logit_processor request field {_RUST_TODO}")
|
||||||
|
def test_custom_logit_processor(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@unittest.skip(f"custom_logit_processor request field {_RUST_TODO}")
|
||||||
|
def test_custom_logit_processor_batch_mixed(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@unittest.skip(f"custom_logit_processor request field {_RUST_TODO}")
|
||||||
|
def test_stateful_custom_logit_processor(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@unittest.skip(f"custom_logit_processor request field {_RUST_TODO}")
|
||||||
|
def test_stateful_custom_logit_processor_batch_mixed(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@unittest.skip(f"/flush_cache endpoint + cached_tokens meta {_RUST_TODO}")
|
||||||
|
def test_cache_tokens(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def test_greedy_token_equals_top1(self):
|
||||||
|
"""Cross-column alignment guard for the columnar logprob wire: at
|
||||||
|
temperature 0 the chosen token must BE the top-1 entry of its own
|
||||||
|
position. A column shifted across requests or positions (the failure
|
||||||
|
mode a truncation-tolerant reader would mask) breaks this instantly."""
|
||||||
|
response = requests.post(
|
||||||
|
self.base_url + "/generate",
|
||||||
|
json={
|
||||||
|
"text": ["The capital of France is", "I have a very good idea on"],
|
||||||
|
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
|
||||||
|
"return_logprob": True,
|
||||||
|
"top_logprobs_num": 5,
|
||||||
|
"logprob_start_len": 0,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
self.assertEqual(response.status_code, 200, response.text)
|
||||||
|
for res in response.json():
|
||||||
|
meta = res["meta_info"]
|
||||||
|
out_lp = meta["output_token_logprobs"]
|
||||||
|
top = meta["output_top_logprobs"]
|
||||||
|
self.assertEqual(len(out_lp), meta["completion_tokens"])
|
||||||
|
self.assertEqual(len(top), len(out_lp))
|
||||||
|
# First prompt token's logprob is the None sentinel; it must
|
||||||
|
# survive the NaN wire encoding and come back as null.
|
||||||
|
self.assertIsNone(meta["input_token_logprobs"][0][0])
|
||||||
|
for (lp, tid, _), pos_top in zip(out_lp, top):
|
||||||
|
self.assertEqual(len(pos_top), 5)
|
||||||
|
self.assertEqual(pos_top[0][1], tid)
|
||||||
|
self.assertAlmostEqual(pos_top[0][0], lp, places=4)
|
||||||
|
vals = [t[0] for t in pos_top]
|
||||||
|
self.assertEqual(vals, sorted(vals, reverse=True))
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -9,13 +9,21 @@ from sglang.test.test_utils import (
|
|||||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
DEFAULT_URL_FOR_TEST,
|
DEFAULT_URL_FOR_TEST,
|
||||||
CustomTestCase,
|
CustomTestCase,
|
||||||
|
is_rust_server_built,
|
||||||
popen_launch_server,
|
popen_launch_server,
|
||||||
)
|
)
|
||||||
|
|
||||||
register_cuda_ci(est_time=62, stage="base-b", runner_config="1-gpu-large")
|
# Two classes run from this file: the default server plus the Rust-frontend
|
||||||
|
# variant (when the embedded extension is built), each launches a server + eval.
|
||||||
|
register_cuda_ci(est_time=124, stage="base-b", runner_config="1-gpu-large")
|
||||||
|
|
||||||
|
|
||||||
class TestModeloptFP8(CustomTestCase):
|
class TestModeloptFP8(CustomTestCase):
|
||||||
|
# Extra server env; the Rust-frontend subclass sets SGLANG_RUST_SERVER here.
|
||||||
|
env = None
|
||||||
|
# Eval endpoint. The Rust server exposes only the native `/generate`, so its
|
||||||
|
# subclass overrides this to "generate".
|
||||||
|
api = "completion"
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls):
|
def setUpClass(cls):
|
||||||
@@ -25,7 +33,15 @@ class TestModeloptFP8(CustomTestCase):
|
|||||||
cls.model,
|
cls.model,
|
||||||
cls.base_url,
|
cls.base_url,
|
||||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||||
other_args=["--quantization", "modelopt_fp8"],
|
other_args=[
|
||||||
|
"--quantization",
|
||||||
|
"modelopt_fp8",
|
||||||
|
"--tokenizer-worker-num",
|
||||||
|
"2",
|
||||||
|
"--detokenizer-worker-num",
|
||||||
|
"2",
|
||||||
|
],
|
||||||
|
env=cls.env,
|
||||||
)
|
)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -38,7 +54,7 @@ class TestModeloptFP8(CustomTestCase):
|
|||||||
base_url=self.base_url,
|
base_url=self.base_url,
|
||||||
model=self.model,
|
model=self.model,
|
||||||
eval_name="gsm8k",
|
eval_name="gsm8k",
|
||||||
api="completion",
|
api=self.api,
|
||||||
max_tokens=512,
|
max_tokens=512,
|
||||||
num_examples=200,
|
num_examples=200,
|
||||||
num_threads=200,
|
num_threads=200,
|
||||||
@@ -48,5 +64,20 @@ class TestModeloptFP8(CustomTestCase):
|
|||||||
self.assertGreater(metrics["score"], 0.70)
|
self.assertGreater(metrics["score"], 0.70)
|
||||||
|
|
||||||
|
|
||||||
|
@unittest.skipUnless(
|
||||||
|
is_rust_server_built(),
|
||||||
|
"embedded rust server extension not built",
|
||||||
|
)
|
||||||
|
class TestModeloptFP8WithRustServer(TestModeloptFP8):
|
||||||
|
"""Same model + eval, but served through the embedded Rust frontend
|
||||||
|
(`SGLANG_RUST_SERVER`). Guards the Rust tokenizer/detokenizer/completions path
|
||||||
|
against accuracy regressions: a bug there drops gsm8k score below the same
|
||||||
|
0.70 bar the default frontend must clear. Uses the native `/generate` endpoint
|
||||||
|
(the only API the Rust server exposes)."""
|
||||||
|
|
||||||
|
env = {"SGLANG_RUST_SERVER": "1"}
|
||||||
|
api = "generate"
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -0,0 +1,76 @@
|
|||||||
|
"""Run the `rust/` Cargo workspace's unit tests from the CPU CI suite.
|
||||||
|
|
||||||
|
The `rust/` workspace (sglang-grpc, sglang-mm, sglang-server) is compiled into
|
||||||
|
the wheel by setuptools-rust, but until now nothing ran `cargo test` in CI --
|
||||||
|
`.github/workflows/pr-test-rust.yml` and `pr-benchmark-rust.yml` are both
|
||||||
|
path-scoped to `sgl-model-gateway/**`, a different workspace. `lint.yml` covers
|
||||||
|
rustfmt/clippy via the pre-commit hooks, so this file only adds the test run.
|
||||||
|
|
||||||
|
The debug profile is deliberate: these are pure-logic tests (no timing or
|
||||||
|
codegen assertions), and the release profile costs a full LTO build for the
|
||||||
|
same coverage.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import shutil
|
||||||
|
import subprocess
|
||||||
|
import unittest
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
# base-c-test-cpu is where this was asked for, and it matches the repo's
|
||||||
|
# base-a + base-c dual-registration convention -- but base-c-test-cpu currently
|
||||||
|
# has no runner job in any workflow (it was carved out of base-b in #28623 to
|
||||||
|
# *reduce* CPU CI scope), so base-a-test-cpu is what actually executes.
|
||||||
|
register_cpu_ci(est_time=300, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
# repo root: test/registered/rust/<this file>
|
||||||
|
RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust"
|
||||||
|
|
||||||
|
# Not `est_time`: that is a scheduling hint for partition balancing (a rough
|
||||||
|
# average), this is a hard ceiling for the worst case. The 136 tests run in ~1s;
|
||||||
|
# what varies is the build. Cache-warm the workspace crates recompile in ~15s,
|
||||||
|
# but a Swatinem/rust-cache miss rebuilds all ~370 dependencies -- measured at
|
||||||
|
# 48s on 4 fast cores, so several minutes on a hosted runner.
|
||||||
|
#
|
||||||
|
# Capped below the 600s `timeout-minutes` on the suite's "Run test" step so a
|
||||||
|
# hang fails here, with output, instead of being killed as an opaque job
|
||||||
|
# timeout. The harness `--timeout-per-file` (1200s) is looser still.
|
||||||
|
BUILD_AND_RUN_TIMEOUT_S = 300
|
||||||
|
|
||||||
|
|
||||||
|
class TestCargoWorkspace(CustomTestCase):
|
||||||
|
def test_cargo_test_workspace(self):
|
||||||
|
# Not skipUnless: cargo is a hard dependency of the editable install
|
||||||
|
# (setuptools-rust builds sglang-grpc), so a missing toolchain is a
|
||||||
|
# broken environment, and a silently-skipped CI test is worthless.
|
||||||
|
self.assertIsNotNone(
|
||||||
|
shutil.which("cargo"),
|
||||||
|
"cargo not found on PATH; install a Rust toolchain "
|
||||||
|
"(scripts/ci/utils/install_rust_protoc.sh)",
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
(RUST_WORKSPACE / "Cargo.toml").is_file(),
|
||||||
|
f"rust workspace manifest not found at {RUST_WORKSPACE}",
|
||||||
|
)
|
||||||
|
|
||||||
|
proc = subprocess.run(
|
||||||
|
["cargo", "test", "--workspace"],
|
||||||
|
cwd=RUST_WORKSPACE,
|
||||||
|
capture_output=True,
|
||||||
|
text=True,
|
||||||
|
timeout=BUILD_AND_RUN_TIMEOUT_S,
|
||||||
|
)
|
||||||
|
# Print unconditionally so a green run still shows which tests ran.
|
||||||
|
print(proc.stdout)
|
||||||
|
self.assertEqual(
|
||||||
|
proc.returncode,
|
||||||
|
0,
|
||||||
|
f"`cargo test --workspace` failed in {RUST_WORKSPACE}\n"
|
||||||
|
f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user