[Rust] Split and rename embedded server components (#37220)
This commit is contained in:
Lianmin Zheng
@@ -1,869 +0,0 @@
|
||||
"""Embedded Rust server lifecycle for the scheduler.
|
||||
|
||||
The Rust server replaces the Python api-server + `TokenizerManager` +
|
||||
`DetokenizerManager` stack (hence this module sits beside them in `managers/`),
|
||||
running them as Rust threads inside the scheduler process. This wrapper keeps
|
||||
all `SGLANG_RUST_SERVER` plumbing — startup, CPU-core partitioning, the
|
||||
typed `server_args` handoff, and control-response routing — out of `scheduler.py`. The
|
||||
scheduler holds an `Optional[RustServer]` and delegates to it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from array import array
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Optional, Tuple
|
||||
|
||||
import msgspec
|
||||
|
||||
from sglang.srt.arg_groups.overrides import resolving_view
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.managers.io_struct import TokenizedGenerateReqInput
|
||||
from sglang.srt.managers.utils import (
|
||||
MsgpackDecodeError,
|
||||
compute_num_reserved_tokens,
|
||||
msgpack_decode_explained,
|
||||
)
|
||||
from sglang.srt.runtime_context import (
|
||||
get_disagg,
|
||||
get_mm,
|
||||
get_model,
|
||||
get_observability,
|
||||
get_parallel,
|
||||
get_serving,
|
||||
)
|
||||
from sglang.srt.utils.flatten import (
|
||||
FlatPairColumns,
|
||||
NestedRowColumns,
|
||||
RaggedPairColumns,
|
||||
)
|
||||
from sglang.version import __version__
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.managers.io_struct import BatchTokenIDOutput
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.srt.rust_extensions._server import MmSpec, Server, ServerArgs
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NativeMmSpec(msgspec.Struct, frozen=True, kw_only=True):
|
||||
"""Resolved parameters of the native Rust MM pipeline for one model,
|
||||
consumed by the Rust worker pool (as the typed extension ``MmSpec``, see
|
||||
:meth:`RustServer._build_mm_spec`), the ``_multimodal`` parity API
|
||||
(:meth:`rust_json`) and the drain adapter
|
||||
(:meth:`NativeMmHost.build_native_mm`)."""
|
||||
|
||||
family: str
|
||||
feature_shm: bool
|
||||
image_token_id: int
|
||||
patch_size: int
|
||||
merge_size: int
|
||||
temporal_patch_size: int
|
||||
min_pixels: int
|
||||
max_pixels: int
|
||||
image_mean: Tuple[float, ...]
|
||||
image_std: Tuple[float, ...]
|
||||
# Which HF processor the Rust resize must reproduce bit-exactly, from
|
||||
# `NativeMmHost.NATIVE_IMAGE_PROCESSORS`.
|
||||
resample: str
|
||||
vision_start_token_id: Optional[int]
|
||||
vision_end_token_id: Optional[int]
|
||||
video_token_id: Optional[int]
|
||||
|
||||
# Used by the drain adapter only; every other field goes to Rust.
|
||||
DRAIN_ONLY = ("vision_start_token_id", "vision_end_token_id", "video_token_id")
|
||||
|
||||
@property
|
||||
def feature_dim(self) -> int:
|
||||
return 3 * self.temporal_patch_size * self.patch_size * self.patch_size
|
||||
|
||||
def rust_json(self) -> str:
|
||||
"""The subset `sglang_mm::registry::pipeline_from_spec` parses — the
|
||||
JSON form the ``_multimodal`` parity API takes; the server itself is
|
||||
handed the typed ``MmSpec`` instead."""
|
||||
fields = (f for f in self.__struct_fields__ if f not in self.DRAIN_ONLY)
|
||||
return msgspec.json.encode({f: getattr(self, f) for f in fields}).decode()
|
||||
|
||||
|
||||
class NativeMmFamily(msgspec.Struct, frozen=True, kw_only=True):
|
||||
"""The Python half of one Rust MM family (an arm of
|
||||
`sglang_mm::registry::pipeline_from_spec`): which models it serves.
|
||||
Supporting a new model family = one entry in :data:`NATIVE_MM_FAMILIES`
|
||||
plus its Rust arm — the launch gate is data-driven."""
|
||||
|
||||
name: str
|
||||
# The registered Python mm-processor the native pipeline replaces, as
|
||||
# "module:Class". Compared by identity, so an
|
||||
# SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE override still disables the native path.
|
||||
mm_processor: str
|
||||
# Model types whose image-only M-RoPE matches the family's fast path.
|
||||
model_types: FrozenSet[str]
|
||||
# HF image processors the native resize reproduces bit-exactly, each mapped
|
||||
# to the `resample` the Rust pipeline must use (see `NativeMmSpec.resample`).
|
||||
image_processors: Dict[str, str]
|
||||
|
||||
def serves(self, mm_processor_cls: Any, model_type: Optional[str]) -> bool:
|
||||
module_name, _, class_name = self.mm_processor.partition(":")
|
||||
cls = getattr(importlib.import_module(module_name), class_name)
|
||||
return mm_processor_cls is cls and model_type in self.model_types
|
||||
|
||||
|
||||
NATIVE_MM_FAMILIES: Tuple[NativeMmFamily, ...] = (
|
||||
NativeMmFamily(
|
||||
name="qwen_vl",
|
||||
mm_processor="sglang.srt.multimodal.processors.qwen_vl:QwenVLImageProcessor",
|
||||
model_types=frozenset(
|
||||
(
|
||||
"qwen2_vl",
|
||||
"qwen2_5_vl",
|
||||
"qwen3_vl",
|
||||
"qwen3_vl_moe",
|
||||
"qwen3_5",
|
||||
"qwen3_5_moe",
|
||||
)
|
||||
),
|
||||
image_processors={
|
||||
"Qwen2VLImageProcessor": "aten_u8",
|
||||
"Qwen2VLImageProcessorFast": "aten_u8",
|
||||
"Qwen2VLImageProcessorPil": "pil",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def native_mm_family_for(
|
||||
mm_processor_cls: Any, model_type: Optional[str]
|
||||
) -> Optional[NativeMmFamily]:
|
||||
"""The declared family serving this model, or ``None`` — which
|
||||
:meth:`RustServer.launch` turns into a hard error (no Python fallback)."""
|
||||
return next(
|
||||
(f for f in NATIVE_MM_FAMILIES if f.serves(mm_processor_cls, model_type)), None
|
||||
)
|
||||
|
||||
|
||||
class NativeMmHost:
|
||||
"""Builds and validates the native Rust MM pipeline for one model.
|
||||
|
||||
Construction registers the same ``mm_processor`` mapping the Python
|
||||
TokenizerManager would build — not to process requests (the Rust worker pool
|
||||
does that, GIL-free) but as the source of truth
|
||||
:meth:`resolve_native_spec` resolves the pipeline parameters from. At drain
|
||||
time :meth:`build_native_mm` wraps the Rust-produced buffers into the
|
||||
scheduler's ``MultimodalProcessorOutput``.
|
||||
|
||||
There is no Python fallback: a model without a native spec fails at launch,
|
||||
and inputs outside the pipeline's scope are rejected per request.
|
||||
"""
|
||||
|
||||
# Rust mm-worker threads when --mm-processor-worker-num is 0. They are
|
||||
# GIL-free, so unlike the Python processor pool more than one always helps.
|
||||
AUTO_MM_WORKERS = 8
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
server_args: ServerArgs,
|
||||
model_config: ModelConfig,
|
||||
processor: Any = None,
|
||||
):
|
||||
# Lazy: this class exists only for multimodal models under
|
||||
# SGLANG_RUST_SERVER.
|
||||
from sglang.srt.managers.multimodal_processor import import_processors
|
||||
from sglang.srt.managers.tokenizer_manager import get_processor_wrapper
|
||||
|
||||
self.server_args = server_args
|
||||
self.model_config = model_config
|
||||
# Worker threads == max concurrently-processed mm requests.
|
||||
self.mm_workers = get_mm().mm_processor_worker_num or self.AUTO_MM_WORKERS
|
||||
|
||||
# The mapping the Python TokenizerManager builds in
|
||||
# init_tokenizer_and_processor. The caller's already-loaded HF
|
||||
# AutoProcessor is reused when available (identical construction args).
|
||||
import_processors("sglang.srt.multimodal.processors")
|
||||
if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get():
|
||||
import_processors(mm_process_pkg, overwrite=True)
|
||||
self._processor = processor or get_processor_wrapper()
|
||||
|
||||
def resolve_native_spec(self) -> Optional[NativeMmSpec]:
|
||||
"""The :class:`NativeMmSpec` for this model, or ``None`` when it has no
|
||||
native pipeline (the launch gate turns that into a hard error).
|
||||
|
||||
Carries only resolved settings — patch geometry, pixel limits,
|
||||
normalization, token ids — never the HF config, and is conservative by
|
||||
design: an unrecognized knob disables the native path rather than being
|
||||
approximated."""
|
||||
from sglang.srt.managers.multimodal_processor import get_mm_processor_cls
|
||||
|
||||
hf_config = self.model_config.hf_config
|
||||
mm_processor_cls = get_mm_processor_cls(
|
||||
hf_config, self.server_args, model_config=self.model_config
|
||||
)
|
||||
family = native_mm_family_for(
|
||||
mm_processor_cls, getattr(hf_config, "model_type", None)
|
||||
)
|
||||
if family is None:
|
||||
return None
|
||||
ip = getattr(self._processor, "image_processor", None)
|
||||
resample = family.image_processors.get(type(ip).__name__)
|
||||
if resample is None:
|
||||
return None
|
||||
# The native pipeline always resizes, rescales by 1/255 and normalizes;
|
||||
# Rust's fused normalize constants assume that factor. Anything else
|
||||
# would silently produce different features.
|
||||
stages = ("do_resize", "do_rescale", "do_normalize")
|
||||
if not all(getattr(ip, stage, True) for stage in stages):
|
||||
return None
|
||||
if getattr(ip, "rescale_factor", None) != 1 / 255:
|
||||
return None
|
||||
|
||||
# `--mm-process-config {"image": {...}}`: only pixel-limit overrides are
|
||||
# mirrored natively, anything else disables the pipeline.
|
||||
image_overrides = dict((get_mm().mm_process_config or {}).get("image", {}))
|
||||
if not set(image_overrides) <= {"min_pixels", "max_pixels"}:
|
||||
return None
|
||||
|
||||
size = getattr(ip, "size", None) or {}
|
||||
min_pixels = image_overrides.get(
|
||||
"min_pixels", getattr(ip, "min_pixels", None) or size.get("shortest_edge")
|
||||
)
|
||||
max_pixels = image_overrides.get(
|
||||
"max_pixels", getattr(ip, "max_pixels", None) or size.get("longest_edge")
|
||||
)
|
||||
try:
|
||||
spec = NativeMmSpec(
|
||||
family=family.name,
|
||||
feature_shm=self._use_feature_shm(),
|
||||
image_token_id=hf_config.image_token_id,
|
||||
patch_size=ip.patch_size,
|
||||
merge_size=ip.merge_size,
|
||||
temporal_patch_size=ip.temporal_patch_size,
|
||||
min_pixels=int(min_pixels),
|
||||
max_pixels=int(max_pixels),
|
||||
image_mean=tuple(float(x) for x in ip.image_mean),
|
||||
image_std=tuple(float(x) for x in ip.image_std),
|
||||
resample=resample,
|
||||
vision_start_token_id=getattr(hf_config, "vision_start_token_id", None),
|
||||
vision_end_token_id=getattr(hf_config, "vision_end_token_id", None),
|
||||
video_token_id=getattr(hf_config, "video_token_id", None),
|
||||
)
|
||||
except (AttributeError, TypeError): # missing/odd processor attrs
|
||||
return None
|
||||
logger.info("rust server: native MM pipeline enabled (family=%s)", family.name)
|
||||
return spec
|
||||
|
||||
def _use_feature_shm(self) -> bool:
|
||||
"""Whether to park feature buffers in POSIX shm rather than inline.
|
||||
|
||||
On exactly when the drained request is broadcast across TP ranks *and*
|
||||
the receiver's ``unwrap_shm_features`` will materialize the stubs (its
|
||||
gates: non-default tensor transport, no ``skip_tokenizer_init``).
|
||||
|
||||
Inline, the whole ~20 MB/image buffer rides ``broadcast_pyobj`` serially
|
||||
on the scheduler loop, so ranks 1..n start the TP-sharded ViT ~30 ms
|
||||
after rank 0 and every rank then stalls that long at the first
|
||||
collective. With shm the broadcast carries a ~100-byte stub and all ranks
|
||||
map in parallel — the transport the Python TokenizerManager already uses.
|
||||
Single-rank serving stays inline, where shm would only add a copy.
|
||||
"""
|
||||
from sglang.srt.multimodal.transport import (
|
||||
determine_tensor_transport_mode,
|
||||
)
|
||||
|
||||
return (
|
||||
get_parallel().tp_size > 1
|
||||
and determine_tensor_transport_mode() != "default"
|
||||
and not get_serving().skip_tokenizer_init
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def build_native_mm(spec: NativeMmSpec, entry):
|
||||
"""Drain-time adapter: wrap the Rust-produced buffers of one ``MmEncodeResult``
|
||||
into the scheduler's ``MultimodalProcessorOutput``. Wrapping only — load,
|
||||
resize, patchify, token expansion and M-RoPE all ran in Rust.
|
||||
|
||||
Runs on the scheduler loop, so it must stay copy-free *and* hash-free:
|
||||
``take_mm``'s numpy arrays own the Rust buffers, ``torch.from_numpy`` just
|
||||
views them, and each item's ``hash`` is worker-precomputed so
|
||||
``set_pad_value`` skips ``hash_feature``. Any per-byte work here — memcpy,
|
||||
sha256, tens of MB per image-heavy request — measurably inflates every
|
||||
running request's inter-token latency."""
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.mm_utils import ShmPointerMMData
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalProcessorOutput,
|
||||
)
|
||||
|
||||
shm_names = entry.shm_names
|
||||
if shm_names is None:
|
||||
features = torch.from_numpy(entry.features.reshape(-1, spec.feature_dim))
|
||||
items = []
|
||||
row = 0
|
||||
for index, ((t, h, w), item_hash, offset) in enumerate(
|
||||
zip(entry.grids, entry.hashes, entry.offsets)
|
||||
):
|
||||
n = t * h * w
|
||||
if shm_names is None:
|
||||
feature = features[row : row + n]
|
||||
else:
|
||||
# The worker parked this item's buffer in a named POSIX
|
||||
# segment (see `_use_feature_shm`). Build the stub in its
|
||||
# post-`__setstate__` form: rank 0 never pickle-roundtrips its
|
||||
# own copy, and `materialize()` needs the mapped view.
|
||||
# Ownership of the unlink moved here with `take_mm`.
|
||||
feature = ShmPointerMMData.__new__(ShmPointerMMData)
|
||||
feature.__setstate__(
|
||||
{
|
||||
"shm_name": shm_names[index],
|
||||
"shape": (n, spec.feature_dim),
|
||||
"dtype": torch.float32,
|
||||
"precomputed_hash": item_hash,
|
||||
}
|
||||
)
|
||||
items.append(
|
||||
MultimodalDataItem(
|
||||
modality=Modality.IMAGE,
|
||||
feature=feature,
|
||||
hash=item_hash,
|
||||
offsets=[tuple(offset)],
|
||||
model_specific_data={
|
||||
"image_grid_thw": torch.tensor([[t, h, w]], dtype=torch.long)
|
||||
},
|
||||
)
|
||||
)
|
||||
row += n
|
||||
if envs.SGLANG_MM_PRECOMPUTE_HASH.get():
|
||||
for item in items:
|
||||
item.set_pad_value()
|
||||
return MultimodalProcessorOutput(
|
||||
mm_items=items,
|
||||
im_token_id=spec.image_token_id,
|
||||
im_start_id=spec.vision_start_token_id,
|
||||
im_end_id=spec.vision_end_token_id,
|
||||
video_token_id=spec.video_token_id,
|
||||
mrope_positions=torch.from_numpy(entry.mrope.reshape(3, -1)),
|
||||
mrope_position_delta=torch.tensor([[entry.mrope_delta]], dtype=torch.long),
|
||||
)
|
||||
|
||||
|
||||
class RustServer:
|
||||
"""Owns the embedded multi-threaded Rust server (``sglang_server.Server``).
|
||||
|
||||
The server owns the api-server, tokenizermanager, tokenizer, and detokenizer
|
||||
all implemented as Rust threads in scheduler process.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server: Server,
|
||||
mm_spec: Optional[NativeMmSpec] = None,
|
||||
max_per_poll: int = 256,
|
||||
):
|
||||
self.server = server
|
||||
self.mm_spec = mm_spec
|
||||
self._max_per_poll = max_per_poll
|
||||
|
||||
@classmethod
|
||||
def launch(cls, scheduler: Scheduler) -> RustServer:
|
||||
"""Start the embedded Rust server threads and bind the listen port.
|
||||
|
||||
The caller gates this (``SGLANG_RUST_SERVER`` + rank 0); this always
|
||||
creates.
|
||||
"""
|
||||
from sglang.srt.rust_extensions import load_rust_extension
|
||||
|
||||
Server = load_rust_extension("sglang.srt.rust_extensions._server").Server
|
||||
|
||||
# Force turn off HF tokenizers rayon's unpinned global thread pool.
|
||||
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
||||
|
||||
server_args = scheduler.server_args
|
||||
# `TokenizerManager` merges these under each request's own sampling params
|
||||
# (`{**preferred, **obj.sampling_params}`), and this server replaces that
|
||||
# manager wholesale — so honouring the flag is not implemented here yet.
|
||||
# Refuse rather than run: silently dropping it means generating with
|
||||
# sampling the operator did not configure, and `/get_model_info` would go on
|
||||
# advertising values no request ever receives.
|
||||
if get_serving().preferred_sampling_params:
|
||||
raise ValueError(
|
||||
"SGLANG_RUST_SERVER does not yet apply --preferred-sampling-params "
|
||||
"(the Python TokenizerManager merges it into every request; the rust "
|
||||
"ingress has no equivalent). Launch without SGLANG_RUST_SERVER, or "
|
||||
"drop --preferred-sampling-params and send those values per request."
|
||||
)
|
||||
http_addr = f"{get_serving().host}:{get_serving().port}"
|
||||
|
||||
# Per-DP-rank HTTP port with client load balancing. `None` when DP is off,
|
||||
# so the rank is not conflated with rank 0 of a one-rank group.
|
||||
dp_rank = scheduler.ps.attn_dp_rank if scheduler.ps.dp_size > 1 else None
|
||||
if dp_rank is not None:
|
||||
http_addr = f"{get_serving().host}:{get_serving().port + dp_rank}"
|
||||
|
||||
launch_cores, server_cores = cls._partition_cores(
|
||||
mm_workers=(
|
||||
(get_mm().mm_processor_worker_num or NativeMmHost.AUTO_MM_WORKERS)
|
||||
if scheduler.model_config.is_multimodal
|
||||
else 0
|
||||
)
|
||||
)
|
||||
|
||||
server = Server(
|
||||
cls._build_server_args(scheduler),
|
||||
# None -> run unpinned; the list carries the pinning decision.
|
||||
cores=server_cores,
|
||||
http_addr=http_addr,
|
||||
)
|
||||
|
||||
# Multimodal models must have a native Rust pipeline — there is no Python
|
||||
# fallback.
|
||||
mm_spec = None
|
||||
if scheduler.model_config.is_multimodal:
|
||||
# New threads inherit the spawning thread's affinity, and this launch
|
||||
# thread still holds the full mask. Narrow it first so every MM thread
|
||||
# created below (the processor's executors, the Rust MM workers) stays
|
||||
# off the scheduler's reserved cores, where MM preprocessing would
|
||||
# preempt the scheduler loop and inflate inter-token latency.
|
||||
if server_cores is not None:
|
||||
try:
|
||||
os.sched_setaffinity(0, set(server_cores))
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
"rust server: cannot confine mm threads to server cores: %s", e
|
||||
)
|
||||
mm_host = NativeMmHost(
|
||||
server_args=server_args,
|
||||
model_config=scheduler.model_config,
|
||||
processor=scheduler.processor,
|
||||
)
|
||||
mm_spec = mm_host.resolve_native_spec()
|
||||
if mm_spec is None:
|
||||
supported = sorted(
|
||||
set(chain.from_iterable(f.model_types for f in NATIVE_MM_FAMILIES))
|
||||
)
|
||||
raise RuntimeError(
|
||||
"SGLANG_RUST_SERVER=1: no native Rust MM pipeline for "
|
||||
f"model_type={scheduler.model_config.hf_config.model_type!r} "
|
||||
f"(supported: {', '.join(supported)}; "
|
||||
"images only). Unset SGLANG_RUST_SERVER to serve this model."
|
||||
)
|
||||
server.start_mm_workers(cls._build_mm_spec(mm_spec), mm_host.mm_workers)
|
||||
|
||||
# Narrow the scheduler thread only after the server threads are launched.
|
||||
if launch_cores is not None:
|
||||
try:
|
||||
# pid 0 == this thread (the scheduler event-loop / launch thread).
|
||||
os.sched_setaffinity(0, set(launch_cores))
|
||||
except OSError as e:
|
||||
logger.warning("rust server: cannot pin scheduler launch thread: %s", e)
|
||||
|
||||
# Under DP every rank runs its own server on its own port, so the rank is
|
||||
# what tells two otherwise identical startup lines apart.
|
||||
dp_note = (
|
||||
"" if dp_rank is None else f" (DP rank {dp_rank}/{scheduler.ps.dp_size})"
|
||||
)
|
||||
logger.info(
|
||||
"SGLANG_RUST_SERVER enabled, Rust server listen on %s%s",
|
||||
http_addr,
|
||||
dp_note,
|
||||
)
|
||||
|
||||
return cls(server, mm_spec=mm_spec)
|
||||
|
||||
def wait_request(self, timeout_ms: int) -> None:
|
||||
"""Block until a request is pushed into the in-process ring or the timeout
|
||||
elapses.
|
||||
"""
|
||||
self.server.wait_request(timeout_ms)
|
||||
|
||||
def drain(self, max_recv: int) -> List[Any]:
|
||||
"""Ingress: non-blocking drain of the in-process ring → list of decoded
|
||||
request objects. The scheduler's request receiver calls this instead of
|
||||
polling the zmq socket when `rust_server_mode` is set.
|
||||
|
||||
The transfer is **columnar**: `recv_requests` returns an `IngressBatch`
|
||||
of scalar msgpack `headers` (with `input_ids` omitted) plus one
|
||||
concatenated raw int64 `data` buffer and per-request `lengths`, so the
|
||||
large `input_ids` lists never go through msgpack. Each header is `msgpack_decode`d (yielding
|
||||
the same `TokenizedGenerateReqInput` / control objects the zmq path
|
||||
produces, so the IPC schema is tracked automatically) and its `input_ids`
|
||||
slice is wrapped as the `array("q")` the scheduler expects. `recv_requests`
|
||||
never waits: the ring drain is `try_recv` (returns the instant the ring
|
||||
is dry, capped at `max_recv`) and the rest is one memcpy per header
|
||||
plus one for the concatenated ids — same contract as `zmq.NOBLOCK`.
|
||||
Parking for work is :meth:`wait_request`, which does release the GIL.
|
||||
"""
|
||||
limit = max_recv if max_recv > 0 else self._max_per_poll
|
||||
batch = self.server.recv_requests(limit)
|
||||
# Bind once: each attribute access converts the rust vec to a fresh list.
|
||||
headers, data, lengths = batch.headers, batch.data, batch.lengths
|
||||
if not headers:
|
||||
return []
|
||||
|
||||
ids_view = memoryview(data)
|
||||
out = []
|
||||
pos = 0 # byte offset into ids_buf
|
||||
for header, n in zip(headers, lengths):
|
||||
nbytes = n * 8
|
||||
try:
|
||||
obj = msgpack_decode_explained(header)
|
||||
except MsgpackDecodeError as e:
|
||||
# Return 400 for malformed request field (e.g. token_ids_logprob=[[0]].
|
||||
logger.warning(
|
||||
"rust ingress: dropping undecodable request %s: %s", e.rid, e.reason
|
||||
)
|
||||
if e.rid is not None:
|
||||
self.server.push_error(e.rid, f"invalid request: {e.reason}")
|
||||
pos += nbytes
|
||||
continue
|
||||
if n: # generate request: attach its int64 ids slice as array("q")
|
||||
ids = array("q")
|
||||
ids.frombytes(ids_view[pos : pos + nbytes])
|
||||
obj.input_ids = ids
|
||||
pos += nbytes
|
||||
if self.mm_spec is not None and isinstance(obj, TokenizedGenerateReqInput):
|
||||
# The buffers rode the Rust sidecar, parked before the ring push;
|
||||
# wrapping them into tensors is the only Python step of the native
|
||||
# path. `None` for a text-only request on a multimodal model.
|
||||
native = self.server.take_mm(obj.rid)
|
||||
if native is not None:
|
||||
obj.mm_inputs = NativeMmHost.build_native_mm(self.mm_spec, native)
|
||||
out.append(obj)
|
||||
return out
|
||||
|
||||
def push_control_output(self, recv_req, output) -> None:
|
||||
"""Push a control-request response through the egress ring to the waiting
|
||||
request (routed by rid), encoded as **msgpack** (the ring's native
|
||||
format).
|
||||
|
||||
A msgspec struct is converted to a *named map* (``structs.asdict``, since
|
||||
the IPC structs are ``array_like`` and would otherwise lose field names)
|
||||
so the Rust api_server can shape it per-endpoint (e.g. /server_info)
|
||||
before rendering JSON to the client — keeping JSON formatting off the
|
||||
scheduler's GIL.
|
||||
"""
|
||||
|
||||
# Invariant: control requests always carry a rust-minted rid; without
|
||||
# one the response is unroutable, so fail loudly rather than drop it.
|
||||
assert (
|
||||
recv_req.rid is not None
|
||||
), f"control response without rid: {type(output).__name__}"
|
||||
# No local try/except: a failed push propagates to run_scheduler_process's
|
||||
# outer handler, which logs the full traceback (scheduler-fatal either way).
|
||||
payload = (
|
||||
msgspec.structs.asdict(output)
|
||||
if isinstance(output, msgspec.Struct)
|
||||
else output
|
||||
)
|
||||
# enc_hook stringifies non-native types (paths, enums); JSON
|
||||
# rendering happens in Rust.
|
||||
encoded = msgspec.msgpack.encode(payload, enc_hook=str)
|
||||
|
||||
self.server.push_control_result(recv_req.rid, encoded)
|
||||
|
||||
def push_generation(self, payload: BatchTokenIDOutput) -> None:
|
||||
"""Egress redirect for generation output (replaces the zmq detokenizer).
|
||||
|
||||
Push the WHOLE batch into the Rust egress ring as one frame (-> detokenizer
|
||||
shards -> client streams), mirroring the ingress ``input_ids`` split so the
|
||||
bulk numeric columns never go through msgpack:
|
||||
|
||||
- ``header``: msgpack ``BatchHeader`` positional array — the per-request
|
||||
scalar columns (``rids, finish_reasons, prompt_tokens, tok_lens``) plus
|
||||
the shape metadata for the optional families (``*_lens`` element counts
|
||||
for the flat logprob columns, ``*_reqlens``/``*_poslens`` for the ragged
|
||||
and hidden ones).
|
||||
- ``data``: the raw little-endian numeric buffer — every column is a
|
||||
4-byte element (``f32`` values, ``i32`` indices), concatenated in the
|
||||
order the Rust ``for_each_chunk`` reads them.
|
||||
|
||||
Logprobs are columnar: output families are per-step deltas, input
|
||||
(prefill) families ride once on the first chunk. Ragged families (top-k,
|
||||
token-ids) flatten a per-position ``list[list]`` into flat ``val``/``idx``
|
||||
buffers plus a per-position ``lens`` vector (0 = null position). Hidden
|
||||
states flatten to rows of floats (one row per output position).
|
||||
"""
|
||||
output_ids = payload.output_ids or []
|
||||
prompt_tokens = payload.prompt_tokens or []
|
||||
|
||||
# Hot-path guard: almost no decode step wants logprobs / hidden states,
|
||||
# so only then pay the per-request flatten + buffer packing below.
|
||||
has_extra = bool(
|
||||
payload.output_token_logprobs_val
|
||||
or payload.input_token_logprobs_val
|
||||
or payload.output_top_logprobs_val
|
||||
or payload.input_top_logprobs_val
|
||||
or payload.output_token_ids_logprobs_val
|
||||
or payload.input_token_ids_logprobs_val
|
||||
or payload.output_hidden_states
|
||||
)
|
||||
|
||||
# Runs on the scheduler's CUDA-launch thread every decode step, so each
|
||||
# Python-level pass over the batch costs inter-token latency: `rids` are
|
||||
# the plain rid strings (hashed to a routing key on the Rust side with a
|
||||
# per-process seed, off the GIL — not parsed; a rid is any string),
|
||||
# `finished_reasons` already `dict | None`, and `output_ids` entries are
|
||||
# always `array("i")` (never None) so `map(len)` and a bare
|
||||
# `chain.from_iterable` stay in C.
|
||||
rids = payload.rids
|
||||
finish_reasons = payload.finished_reasons
|
||||
tok_lens = list(map(len, output_ids))
|
||||
flat_ids = array("i", chain.from_iterable(output_ids))
|
||||
|
||||
# Column order here MUST match BatchHeader (header_cols) and
|
||||
# for_each_chunk's read order (data_cols); the extras contribution
|
||||
# is ordered by the `extras` tuple below.
|
||||
header_cols = [rids, finish_reasons, prompt_tokens, tok_lens]
|
||||
data_cols = [flat_ids.tobytes()]
|
||||
|
||||
if has_extra:
|
||||
# The `extras` tuple is the SINGLE source of the extras column
|
||||
# order — it must match the Rust ``BatchHeader`` fields and
|
||||
# ``for_each_chunk``'s read order.
|
||||
#
|
||||
# TODO(perf): the per-request flatten assumes the logprob/hidden
|
||||
# columns are ragged, non-contiguous nested Python lists — which is
|
||||
# only an assumption. The scheduler moves these off the GPU with
|
||||
# `tensor.tolist()`, so revisit whether the upstream values are
|
||||
# still contiguous tensors; if so, ship raw bytes + a shape
|
||||
# descriptor and skip the flatten entirely.
|
||||
batch_size = len(rids)
|
||||
extras = (
|
||||
FlatPairColumns(
|
||||
"output_token_logprobs",
|
||||
payload.output_token_logprobs_val or [],
|
||||
payload.output_token_logprobs_idx or [],
|
||||
),
|
||||
FlatPairColumns(
|
||||
"input_token_logprobs",
|
||||
payload.input_token_logprobs_val or [],
|
||||
payload.input_token_logprobs_idx or [],
|
||||
first_none_to_nan=True,
|
||||
),
|
||||
RaggedPairColumns(
|
||||
"output_top_logprobs",
|
||||
payload.output_top_logprobs_val or [],
|
||||
payload.output_top_logprobs_idx or [],
|
||||
),
|
||||
RaggedPairColumns(
|
||||
"input_top_logprobs",
|
||||
payload.input_top_logprobs_val or [],
|
||||
payload.input_top_logprobs_idx or [],
|
||||
),
|
||||
RaggedPairColumns(
|
||||
"output_token_ids_logprobs",
|
||||
payload.output_token_ids_logprobs_val or [],
|
||||
payload.output_token_ids_logprobs_idx or [],
|
||||
),
|
||||
RaggedPairColumns(
|
||||
"input_token_ids_logprobs",
|
||||
payload.input_token_ids_logprobs_val or [],
|
||||
payload.input_token_ids_logprobs_idx or [],
|
||||
),
|
||||
NestedRowColumns(
|
||||
"output_hidden_states", payload.output_hidden_states or []
|
||||
),
|
||||
)
|
||||
|
||||
# Every column is all-or-nothing per payload — which is also what makes
|
||||
# a family's emptiness a reliable "nobody asked for this" signal.
|
||||
active = []
|
||||
for extra in extras:
|
||||
populated = False
|
||||
for name, col in extra.columns():
|
||||
assert len(col) in (
|
||||
0,
|
||||
batch_size,
|
||||
), f"extras column {name}: {len(col)} entries for a batch of {batch_size}"
|
||||
populated |= len(col) > 0
|
||||
if populated:
|
||||
active.append(extra)
|
||||
|
||||
# Flatten only the families someone asked for. `has_extra` above is a
|
||||
# per-FRAME guard, so one client enabling logprobs used to drag all
|
||||
# seven families through the per-request loop: at B=4096 that is 28,672
|
||||
# bound-method calls per decode step, materializing 12 columns of 4096
|
||||
# zeros nobody reads. Measured 0.37 ms -> 7.90 ms GIL-held per step,
|
||||
# i.e. 25-75% of a decode step added to the scheduler's critical path.
|
||||
#
|
||||
# Skipping `accept` leaves a family's buffers empty, which is exactly
|
||||
# the wire form the Rust decoder already treats as absent (`per_req_ok`
|
||||
# admits an empty column, `lens_i` reads 0 for every request). The
|
||||
# `header_cols`/`data_cols` loops below still walk all seven, so column
|
||||
# ORDER and arity are unchanged — an inactive family contributes empty
|
||||
# columns in place rather than disappearing.
|
||||
for extra in active:
|
||||
accept = extra.accept # hoisted: this is the hottest loop here
|
||||
for i in range(batch_size):
|
||||
accept(i)
|
||||
|
||||
for extra in extras:
|
||||
header_cols += extra.header_cols()
|
||||
data_cols += extra.data_cols()
|
||||
|
||||
header = msgspec.msgpack.encode(header_cols)
|
||||
# Pass the raw column list; the Rust side concatenates it into the frame
|
||||
# with the GIL released.
|
||||
if not self.server.push_decode_result_batch(header, data_cols):
|
||||
logger.warning(
|
||||
"Rust egress closed; dropped batch of %d requests during shutdown",
|
||||
len(rids),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_mm_spec(spec: NativeMmSpec) -> MmSpec:
|
||||
"""The typed MM handoff for ``Server.start_mm_workers``: the
|
||||
:class:`NativeMmSpec` fields the Rust pipeline consumes, as the Rust
|
||||
extension's own ``MmSpec`` class (same required-keyword contract as
|
||||
:meth:`_build_server_args`; ``family`` / ``resample`` become the
|
||||
extension's ``MmFamily`` / ``MmResample`` enums)."""
|
||||
from sglang.srt.rust_extensions import load_rust_extension
|
||||
|
||||
ext = load_rust_extension("sglang.srt.rust_extensions._server")
|
||||
family = {"qwen_vl": ext.MmFamily.QwenVl}[spec.family]
|
||||
resample = {"aten_u8": ext.MmResample.AtenU8, "pil": ext.MmResample.Pil}[
|
||||
spec.resample
|
||||
]
|
||||
return ext.MmSpec(
|
||||
family=family,
|
||||
feature_shm=spec.feature_shm,
|
||||
image_token_id=spec.image_token_id,
|
||||
patch_size=spec.patch_size,
|
||||
merge_size=spec.merge_size,
|
||||
temporal_patch_size=spec.temporal_patch_size,
|
||||
min_pixels=spec.min_pixels,
|
||||
max_pixels=spec.max_pixels,
|
||||
image_mean=spec.image_mean,
|
||||
image_std=spec.image_std,
|
||||
resample=resample,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_server_args(scheduler: Scheduler) -> ServerArgs:
|
||||
"""The typed launch handoff for the scheduler's embedded Rust server:
|
||||
the ``server_args`` fields it reads, the already-resolved
|
||||
``model_config``, and launch-time facts — as the Rust extension's own
|
||||
``ServerArgs`` class. Its constructor takes every field as a required
|
||||
keyword (see ``rust/sglang-server/src/message/config.rs``), so a
|
||||
missing, extra or mistyped field fails here at boot rather than
|
||||
running on a silently-defaulted knob."""
|
||||
from sglang.srt.rust_extensions import load_rust_extension
|
||||
|
||||
ext = load_rust_extension("sglang.srt.rust_extensions._server")
|
||||
|
||||
sa = resolving_view(scheduler.server_args)
|
||||
mc = scheduler.model_config
|
||||
disaggregation_mode = {
|
||||
"null": ext.DisaggregationMode.Null,
|
||||
"prefill": ext.DisaggregationMode.Prefill,
|
||||
"decode": ext.DisaggregationMode.Decode,
|
||||
}[get_disagg().disaggregation_mode]
|
||||
return ext.ServerArgs(
|
||||
model_path=get_model().model_path,
|
||||
served_model_name=get_serving().served_model_name,
|
||||
tokenizer_path=get_serving().tokenizer_path,
|
||||
revision=get_model().revision,
|
||||
load_format=get_model().load_format,
|
||||
weight_version=get_serving().weight_version,
|
||||
host=get_serving().host,
|
||||
port=get_serving().port,
|
||||
log_level=get_observability().log_level,
|
||||
log_level_http=get_observability().log_level_http,
|
||||
chat_template=get_serving().chat_template,
|
||||
tool_call_parser=get_serving().tool_call_parser,
|
||||
reasoning_parser=get_serving().reasoning_parser,
|
||||
stream_response_default_include_usage=get_serving().stream_response_default_include_usage,
|
||||
tokenizer_worker_num=get_serving().tokenizer_worker_num,
|
||||
detokenizer_worker_num=get_serving().detokenizer_worker_num,
|
||||
skip_tokenizer_init=get_serving().skip_tokenizer_init,
|
||||
incremental_streaming_output=get_serving().incremental_streaming_output,
|
||||
disaggregation_mode=disaggregation_mode,
|
||||
model_config=ext.ModelConfig(
|
||||
context_len=mc.context_len,
|
||||
vocab_size=mc.vocab_size,
|
||||
is_multimodal=mc.is_multimodal,
|
||||
# Resolved default sampling params (generation_config.json when
|
||||
# `--sampling-defaults model`, {} otherwise). The rust server
|
||||
# consumes these for omitted temperature/top_p in chat
|
||||
# conversions instead of hard-coding the OpenAI terminal
|
||||
# defaults.
|
||||
default_sampling_params=ext.DefaultSamplingParams(
|
||||
**mc.get_default_sampling_params()
|
||||
),
|
||||
),
|
||||
# `preferred_sampling_params` is deliberately absent: `launch`
|
||||
# refuses to start when it is set, so the Rust server never needs it.
|
||||
preferred_sampling_params=(
|
||||
json.dumps(get_serving().preferred_sampling_params)
|
||||
if get_serving().preferred_sampling_params is not None
|
||||
else None
|
||||
),
|
||||
allow_auto_truncate=get_serving().allow_auto_truncate,
|
||||
enable_return_hidden_states=sa.enable_return_hidden_states,
|
||||
# Not a `server_args` field: `TokenizerManager` derives it, and the
|
||||
# rust ingress needs the same number for its total-token check.
|
||||
num_reserved_tokens=compute_num_reserved_tokens(),
|
||||
# Launch-time facts Python's /server_info reports from
|
||||
# scheduler_info / the package — stamped here so the rust endpoint
|
||||
# can serve them statically (no scheduler round-trip).
|
||||
version=__version__,
|
||||
max_total_num_tokens=scheduler.max_total_num_tokens,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _partition_cores(
|
||||
mm_workers: int = 0,
|
||||
) -> Tuple[Optional[List[int]], Optional[List[int]]]:
|
||||
"""Split this rank's allowed cores into ``(launch_cores, server_cores)``.
|
||||
|
||||
Pure computation — no affinity is changed here. Both sets are a subset
|
||||
of this rank's NUMA-local cores (when affinity/NUMA bind is on), so the
|
||||
partition stays NUMA-local. Returns ``(None, None)`` (server runs
|
||||
unpinned, confined only by the process affinity) when the platform has
|
||||
no affinity API or too few cores to split.
|
||||
"""
|
||||
if not hasattr(os, "sched_getaffinity"):
|
||||
return None, None
|
||||
try:
|
||||
allowed = sorted(os.sched_getaffinity(0))
|
||||
except OSError as e:
|
||||
logger.warning("rust server: cannot read cpu affinity: %s", e)
|
||||
return None, None
|
||||
|
||||
# Need enough cores to reserve launch cores and still pin the pools.
|
||||
if len(allowed) < 4:
|
||||
logger.info(
|
||||
"rust server: only %d cores allowed; running pools unpinned",
|
||||
len(allowed),
|
||||
)
|
||||
return None, None
|
||||
|
||||
# Keep a small slice for the launch loop; cap at 2 (the event loop is
|
||||
# effectively serial) and never take more than a quarter of the cores.
|
||||
reserve = min(2, len(allowed) // 4)
|
||||
launch_cores = allowed[:reserve]
|
||||
# Bound the pool instead of taking the whole remainder: this rank's
|
||||
# allowed cores are usually the entire NUMA node, shared with the sibling
|
||||
# TP ranks' processes, so an unbounded mask lets MM preprocessing bursts
|
||||
# preempt a sibling's CUDA-launch thread and inflate every rank's forward
|
||||
# through the TP collectives. Measured on Qwen3.5-35B TP4 at one 720p
|
||||
# image per request: ~20 ms of ViT wall time on the worst sibling, gone
|
||||
# once bounded. The budget covers the CPU-hot threads (MM workers, plus
|
||||
# the I/O-shaped tokenizer/ingress/egress/api ones that are rarely all hot
|
||||
# at once) and leaves the rest of the node to the scheduler ranks.
|
||||
pool_budget = max(8, mm_workers + 4)
|
||||
server_cores = allowed[reserve : reserve + pool_budget]
|
||||
logger.info(
|
||||
"rust server cores=%s, scheduler launch cores=%s",
|
||||
server_cores,
|
||||
launch_cores,
|
||||
)
|
||||
return launch_cores, server_cores
|
||||
@@ -198,7 +198,6 @@ from sglang.srt.managers.prefill_delayer import (
|
||||
PrefillDelayerSinglePassExecutor,
|
||||
RecentPrefillBatchSizeTracker,
|
||||
)
|
||||
from sglang.srt.managers.rust_server import RustServer
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
FINISH_ABORT,
|
||||
MultimodalInputs,
|
||||
@@ -294,6 +293,7 @@ from sglang.srt.observability.trace import process_tracing_init, trace_set_threa
|
||||
from sglang.srt.parser.reasoning_parser import ReasoningParser
|
||||
from sglang.srt.platforms import current_platform
|
||||
from sglang.srt.plugins import load_plugins
|
||||
from sglang.srt.rust_server.server import RustServer
|
||||
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
||||
from sglang.srt.sampling.sampling_params import TOP_K_ALL
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs, compute_world_size
|
||||
|
||||
@@ -9,7 +9,7 @@ from sglang.srt.observability.req_time_stats import real_time
|
||||
from sglang.srt.platforms import current_platform
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.rust_server import RustServer
|
||||
from sglang.srt.rust_server.server import RustServer
|
||||
|
||||
|
||||
class IdleSleeper:
|
||||
|
||||
@@ -38,7 +38,7 @@ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
from sglang.srt.utils.weight_versions import compute_weight_version_spans
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.rust_server import RustServer
|
||||
from sglang.srt.rust_server.server import RustServer
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -37,7 +37,7 @@ from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
|
||||
from sglang.srt.managers.rust_server import RustServer
|
||||
from sglang.srt.rust_server.server import RustServer
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook
|
||||
from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
"""Configuration handoff and CPU placement for the embedded Rust server."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import TYPE_CHECKING, List, Optional, Tuple
|
||||
|
||||
from sglang.srt.arg_groups.overrides import resolving_view
|
||||
from sglang.srt.managers.utils import compute_num_reserved_tokens
|
||||
from sglang.srt.runtime_context import (
|
||||
get_disagg,
|
||||
get_model,
|
||||
get_observability,
|
||||
get_serving,
|
||||
)
|
||||
from sglang.version import __version__
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.srt.rust_extensions._server import ServerArgs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _build_server_args(scheduler: Scheduler) -> ServerArgs:
|
||||
"""The typed launch handoff for the scheduler's embedded Rust server:
|
||||
the ``server_args`` fields it reads, the already-resolved
|
||||
``model_config``, and launch-time facts — as the Rust extension's own
|
||||
``ServerArgs`` class. Its constructor takes every field as a required
|
||||
keyword (see ``rust/sglang-server/src/message/config.rs``), so a
|
||||
missing, extra or mistyped field fails here at boot rather than
|
||||
running on a silently-defaulted knob."""
|
||||
from sglang.srt.rust_extensions import load_rust_extension
|
||||
|
||||
ext = load_rust_extension("sglang.srt.rust_extensions._server")
|
||||
|
||||
sa = resolving_view(scheduler.server_args)
|
||||
mc = scheduler.model_config
|
||||
disaggregation_mode = {
|
||||
"null": ext.DisaggregationMode.Null,
|
||||
"prefill": ext.DisaggregationMode.Prefill,
|
||||
"decode": ext.DisaggregationMode.Decode,
|
||||
}[get_disagg().disaggregation_mode]
|
||||
return ext.ServerArgs(
|
||||
model_path=get_model().model_path,
|
||||
served_model_name=get_serving().served_model_name,
|
||||
tokenizer_path=get_serving().tokenizer_path,
|
||||
revision=get_model().revision,
|
||||
load_format=get_model().load_format,
|
||||
weight_version=get_serving().weight_version,
|
||||
host=get_serving().host,
|
||||
port=get_serving().port,
|
||||
log_level=get_observability().log_level,
|
||||
log_level_http=get_observability().log_level_http,
|
||||
chat_template=get_serving().chat_template,
|
||||
tool_call_parser=get_serving().tool_call_parser,
|
||||
reasoning_parser=get_serving().reasoning_parser,
|
||||
stream_response_default_include_usage=get_serving().stream_response_default_include_usage,
|
||||
tokenizer_worker_num=get_serving().tokenizer_worker_num,
|
||||
detokenizer_worker_num=get_serving().detokenizer_worker_num,
|
||||
skip_tokenizer_init=get_serving().skip_tokenizer_init,
|
||||
incremental_streaming_output=get_serving().incremental_streaming_output,
|
||||
disaggregation_mode=disaggregation_mode,
|
||||
model_config=ext.ModelConfig(
|
||||
context_len=mc.context_len,
|
||||
vocab_size=mc.vocab_size,
|
||||
is_multimodal=mc.is_multimodal,
|
||||
# Resolved default sampling params (generation_config.json when
|
||||
# `--sampling-defaults model`, {} otherwise). The rust server
|
||||
# consumes these for omitted temperature/top_p in chat
|
||||
# conversions instead of hard-coding the OpenAI terminal
|
||||
# defaults.
|
||||
default_sampling_params=ext.DefaultSamplingParams(
|
||||
**mc.get_default_sampling_params()
|
||||
),
|
||||
),
|
||||
# `preferred_sampling_params` is deliberately absent: `launch`
|
||||
# refuses to start when it is set, so the Rust server never needs it.
|
||||
preferred_sampling_params=(
|
||||
json.dumps(get_serving().preferred_sampling_params)
|
||||
if get_serving().preferred_sampling_params is not None
|
||||
else None
|
||||
),
|
||||
allow_auto_truncate=get_serving().allow_auto_truncate,
|
||||
enable_return_hidden_states=sa.enable_return_hidden_states,
|
||||
# Not a `server_args` field: `TokenizerManager` derives it, and the
|
||||
# rust ingress needs the same number for its total-token check.
|
||||
num_reserved_tokens=compute_num_reserved_tokens(),
|
||||
# Launch-time facts Python's /server_info reports from
|
||||
# scheduler_info / the package — stamped here so the rust endpoint
|
||||
# can serve them statically (no scheduler round-trip).
|
||||
version=__version__,
|
||||
max_total_num_tokens=scheduler.max_total_num_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _partition_cores(
|
||||
mm_workers: int = 0,
|
||||
) -> Tuple[Optional[List[int]], Optional[List[int]]]:
|
||||
"""Split this rank's allowed cores into ``(launch_cores, server_cores)``.
|
||||
|
||||
Pure computation — no affinity is changed here. Both sets are a subset
|
||||
of this rank's NUMA-local cores (when affinity/NUMA bind is on), so the
|
||||
partition stays NUMA-local. Returns ``(None, None)`` (server runs
|
||||
unpinned, confined only by the process affinity) when the platform has
|
||||
no affinity API or too few cores to split.
|
||||
"""
|
||||
if not hasattr(os, "sched_getaffinity"):
|
||||
return None, None
|
||||
try:
|
||||
allowed = sorted(os.sched_getaffinity(0))
|
||||
except OSError as e:
|
||||
logger.warning("rust server: cannot read cpu affinity: %s", e)
|
||||
return None, None
|
||||
|
||||
# Need enough cores to reserve launch cores and still pin the pools.
|
||||
if len(allowed) < 4:
|
||||
logger.info(
|
||||
"rust server: only %d cores allowed; running pools unpinned",
|
||||
len(allowed),
|
||||
)
|
||||
return None, None
|
||||
|
||||
# Keep a small slice for the launch loop; cap at 2 (the event loop is
|
||||
# effectively serial) and never take more than a quarter of the cores.
|
||||
reserve = min(2, len(allowed) // 4)
|
||||
launch_cores = allowed[:reserve]
|
||||
# Bound the pool instead of taking the whole remainder: this rank's
|
||||
# allowed cores are usually the entire NUMA node, shared with the sibling
|
||||
# TP ranks' processes, so an unbounded mask lets MM preprocessing bursts
|
||||
# preempt a sibling's CUDA-launch thread and inflate every rank's forward
|
||||
# through the TP collectives. Measured on Qwen3.5-35B TP4 at one 720p
|
||||
# image per request: ~20 ms of ViT wall time on the worst sibling, gone
|
||||
# once bounded. The budget covers the CPU-hot threads (MM workers, plus
|
||||
# the I/O-shaped tokenizer/ingress/egress/api ones that are rarely all hot
|
||||
# at once) and leaves the rest of the node to the scheduler ranks.
|
||||
pool_budget = max(8, mm_workers + 4)
|
||||
server_cores = allowed[reserve : reserve + pool_budget]
|
||||
logger.info(
|
||||
"rust server cores=%s, scheduler launch cores=%s",
|
||||
server_cores,
|
||||
launch_cores,
|
||||
)
|
||||
return launch_cores, server_cores
|
||||
@@ -0,0 +1,319 @@
|
||||
"""Multimodal support for the embedded Rust server."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Optional, Tuple
|
||||
|
||||
import msgspec
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.runtime_context import get_mm, get_parallel, get_serving
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RustMmSpec(msgspec.Struct, frozen=True, kw_only=True):
|
||||
"""Resolved parameters of the Rust MM pipeline for one model,
|
||||
consumed by the Rust worker pool (as the typed extension ``MmSpec``, see
|
||||
:meth:`RustServer._build_mm_spec`), the ``_multimodal`` parity API
|
||||
(:meth:`rust_json`) and the drain adapter
|
||||
(:meth:`RustMmProcessor.build_output`)."""
|
||||
|
||||
family: str
|
||||
feature_shm: bool
|
||||
image_token_id: int
|
||||
patch_size: int
|
||||
merge_size: int
|
||||
temporal_patch_size: int
|
||||
min_pixels: int
|
||||
max_pixels: int
|
||||
image_mean: Tuple[float, ...]
|
||||
image_std: Tuple[float, ...]
|
||||
# Which HF processor the Rust resize must reproduce bit-exactly.
|
||||
resample: str
|
||||
vision_start_token_id: Optional[int]
|
||||
vision_end_token_id: Optional[int]
|
||||
video_token_id: Optional[int]
|
||||
|
||||
# Used by the drain adapter only; every other field goes to Rust.
|
||||
DRAIN_ONLY = ("vision_start_token_id", "vision_end_token_id", "video_token_id")
|
||||
|
||||
@property
|
||||
def feature_dim(self) -> int:
|
||||
return 3 * self.temporal_patch_size * self.patch_size * self.patch_size
|
||||
|
||||
def rust_json(self) -> str:
|
||||
"""The subset `sglang_mm::registry::pipeline_from_spec` parses — the
|
||||
JSON form the ``_multimodal`` parity API takes; the server itself is
|
||||
handed the typed ``MmSpec`` instead."""
|
||||
fields = (f for f in self.__struct_fields__ if f not in self.DRAIN_ONLY)
|
||||
return msgspec.json.encode({f: getattr(self, f) for f in fields}).decode()
|
||||
|
||||
|
||||
class RustMmFamily(msgspec.Struct, frozen=True, kw_only=True):
|
||||
"""The Python half of one Rust MM family (an arm of
|
||||
`sglang_mm::registry::pipeline_from_spec`): which models it serves.
|
||||
Supporting a new model family = one entry in :data:`RUST_MM_FAMILIES`
|
||||
plus its Rust arm — the launch gate is data-driven."""
|
||||
|
||||
name: str
|
||||
# The registered Python MM processor the Rust pipeline replaces, as
|
||||
# "module:Class". Compared by identity, so an
|
||||
# SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE override still disables the Rust path.
|
||||
mm_processor: str
|
||||
# Model types whose image-only M-RoPE matches the family's fast path.
|
||||
model_types: FrozenSet[str]
|
||||
# HF image processors the Rust resize reproduces bit-exactly, each mapped
|
||||
# to the `resample` the Rust pipeline must use (see `RustMmSpec.resample`).
|
||||
image_processors: Dict[str, str]
|
||||
|
||||
def serves(self, mm_processor_cls: Any, model_type: Optional[str]) -> bool:
|
||||
module_name, _, class_name = self.mm_processor.partition(":")
|
||||
cls = getattr(importlib.import_module(module_name), class_name)
|
||||
return mm_processor_cls is cls and model_type in self.model_types
|
||||
|
||||
|
||||
RUST_MM_FAMILIES: Tuple[RustMmFamily, ...] = (
|
||||
RustMmFamily(
|
||||
name="qwen_vl",
|
||||
mm_processor="sglang.srt.multimodal.processors.qwen_vl:QwenVLImageProcessor",
|
||||
model_types=frozenset(
|
||||
(
|
||||
"qwen2_vl",
|
||||
"qwen2_5_vl",
|
||||
"qwen3_vl",
|
||||
"qwen3_vl_moe",
|
||||
"qwen3_5",
|
||||
"qwen3_5_moe",
|
||||
)
|
||||
),
|
||||
image_processors={
|
||||
"Qwen2VLImageProcessor": "aten_u8",
|
||||
"Qwen2VLImageProcessorFast": "aten_u8",
|
||||
"Qwen2VLImageProcessorPil": "pil",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def rust_mm_family_for(
|
||||
mm_processor_cls: Any, model_type: Optional[str]
|
||||
) -> Optional[RustMmFamily]:
|
||||
"""The declared family serving this model, or ``None`` — which
|
||||
:meth:`RustServer.launch` turns into a hard error (no Python fallback)."""
|
||||
return next(
|
||||
(f for f in RUST_MM_FAMILIES if f.serves(mm_processor_cls, model_type)), None
|
||||
)
|
||||
|
||||
|
||||
class RustMmProcessor:
|
||||
"""Builds and validates the Rust MM pipeline for one model.
|
||||
|
||||
Construction registers the same ``mm_processor`` mapping the Python
|
||||
TokenizerManager would build — not to process requests (the Rust worker pool
|
||||
does that, GIL-free) but as the source of truth
|
||||
:meth:`resolve_spec` resolves the pipeline parameters from. At drain
|
||||
time :meth:`build_output` wraps the Rust-produced buffers into the
|
||||
scheduler's ``MultimodalProcessorOutput``.
|
||||
|
||||
There is no Python fallback: a model without a Rust MM spec fails at launch,
|
||||
and inputs outside the pipeline's scope are rejected per request.
|
||||
"""
|
||||
|
||||
# Rust mm-worker threads when --mm-processor-worker-num is 0. They are
|
||||
# GIL-free, so unlike the Python processor pool more than one always helps.
|
||||
AUTO_MM_WORKERS = 8
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
server_args: ServerArgs,
|
||||
model_config: ModelConfig,
|
||||
processor: Any = None,
|
||||
):
|
||||
# Lazy: this class exists only for multimodal models under
|
||||
# SGLANG_RUST_SERVER.
|
||||
from sglang.srt.managers.multimodal_processor import import_processors
|
||||
from sglang.srt.managers.tokenizer_manager import get_processor_wrapper
|
||||
|
||||
self.server_args = server_args
|
||||
self.model_config = model_config
|
||||
# Worker threads == max concurrently-processed mm requests.
|
||||
self.mm_workers = get_mm().mm_processor_worker_num or self.AUTO_MM_WORKERS
|
||||
|
||||
# The mapping the Python TokenizerManager builds in
|
||||
# init_tokenizer_and_processor. The caller's already-loaded HF
|
||||
# AutoProcessor is reused when available (identical construction args).
|
||||
import_processors("sglang.srt.multimodal.processors")
|
||||
if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get():
|
||||
import_processors(mm_process_pkg, overwrite=True)
|
||||
self._processor = processor or get_processor_wrapper()
|
||||
|
||||
def resolve_spec(self) -> Optional[RustMmSpec]:
|
||||
"""The :class:`RustMmSpec` for this model, or ``None`` when it has no
|
||||
Rust pipeline (the launch gate turns that into a hard error).
|
||||
|
||||
Carries only resolved settings — patch geometry, pixel limits,
|
||||
normalization, token ids — never the HF config, and is conservative by
|
||||
design: an unrecognized knob disables the Rust path rather than being
|
||||
approximated."""
|
||||
from sglang.srt.managers.multimodal_processor import get_mm_processor_cls
|
||||
|
||||
hf_config = self.model_config.hf_config
|
||||
mm_processor_cls = get_mm_processor_cls(
|
||||
hf_config, self.server_args, model_config=self.model_config
|
||||
)
|
||||
family = rust_mm_family_for(
|
||||
mm_processor_cls, getattr(hf_config, "model_type", None)
|
||||
)
|
||||
if family is None:
|
||||
return None
|
||||
ip = getattr(self._processor, "image_processor", None)
|
||||
resample = family.image_processors.get(type(ip).__name__)
|
||||
if resample is None:
|
||||
return None
|
||||
# The Rust pipeline always resizes, rescales by 1/255 and normalizes;
|
||||
# Rust's fused normalize constants assume that factor. Anything else
|
||||
# would silently produce different features.
|
||||
stages = ("do_resize", "do_rescale", "do_normalize")
|
||||
if not all(getattr(ip, stage, True) for stage in stages):
|
||||
return None
|
||||
if getattr(ip, "rescale_factor", None) != 1 / 255:
|
||||
return None
|
||||
|
||||
# `--mm-process-config {"image": {...}}`: only pixel-limit overrides are
|
||||
# mirrored by Rust; anything else disables the pipeline.
|
||||
image_overrides = dict((get_mm().mm_process_config or {}).get("image", {}))
|
||||
if not set(image_overrides) <= {"min_pixels", "max_pixels"}:
|
||||
return None
|
||||
|
||||
size = getattr(ip, "size", None) or {}
|
||||
min_pixels = image_overrides.get(
|
||||
"min_pixels", getattr(ip, "min_pixels", None) or size.get("shortest_edge")
|
||||
)
|
||||
max_pixels = image_overrides.get(
|
||||
"max_pixels", getattr(ip, "max_pixels", None) or size.get("longest_edge")
|
||||
)
|
||||
try:
|
||||
spec = RustMmSpec(
|
||||
family=family.name,
|
||||
feature_shm=self._use_feature_shm(),
|
||||
image_token_id=hf_config.image_token_id,
|
||||
patch_size=ip.patch_size,
|
||||
merge_size=ip.merge_size,
|
||||
temporal_patch_size=ip.temporal_patch_size,
|
||||
min_pixels=int(min_pixels),
|
||||
max_pixels=int(max_pixels),
|
||||
image_mean=tuple(float(x) for x in ip.image_mean),
|
||||
image_std=tuple(float(x) for x in ip.image_std),
|
||||
resample=resample,
|
||||
vision_start_token_id=getattr(hf_config, "vision_start_token_id", None),
|
||||
vision_end_token_id=getattr(hf_config, "vision_end_token_id", None),
|
||||
video_token_id=getattr(hf_config, "video_token_id", None),
|
||||
)
|
||||
except (AttributeError, TypeError): # missing/odd processor attrs
|
||||
return None
|
||||
logger.info("rust server: Rust MM pipeline enabled (family=%s)", family.name)
|
||||
return spec
|
||||
|
||||
def _use_feature_shm(self) -> bool:
|
||||
"""Whether to park feature buffers in POSIX shm rather than inline.
|
||||
|
||||
On exactly when the drained request is broadcast across TP ranks *and*
|
||||
the receiver's ``unwrap_shm_features`` will materialize the stubs (its
|
||||
gates: non-default tensor transport, no ``skip_tokenizer_init``).
|
||||
|
||||
Inline, the whole ~20 MB/image buffer rides ``broadcast_pyobj`` serially
|
||||
on the scheduler loop, so ranks 1..n start the TP-sharded ViT ~30 ms
|
||||
after rank 0 and every rank then stalls that long at the first
|
||||
collective. With shm the broadcast carries a ~100-byte stub and all ranks
|
||||
map in parallel — the transport the Python TokenizerManager already uses.
|
||||
Single-rank serving stays inline, where shm would only add a copy.
|
||||
"""
|
||||
from sglang.srt.multimodal.transport import (
|
||||
determine_tensor_transport_mode,
|
||||
)
|
||||
|
||||
return (
|
||||
get_parallel().tp_size > 1
|
||||
and determine_tensor_transport_mode() != "default"
|
||||
and not get_serving().skip_tokenizer_init
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def build_output(spec: RustMmSpec, entry):
|
||||
"""Drain-time adapter: wrap the Rust-produced buffers of one ``MmEncodeResult``
|
||||
into the scheduler's ``MultimodalProcessorOutput``. Wrapping only — load,
|
||||
resize, patchify, token expansion and M-RoPE all ran in Rust.
|
||||
|
||||
Runs on the scheduler loop, so it must stay copy-free *and* hash-free:
|
||||
``take_mm_result``'s numpy arrays own the Rust buffers, ``torch.from_numpy`` just
|
||||
views them, and each item's ``hash`` is worker-precomputed so
|
||||
``set_pad_value`` skips ``hash_feature``. Any per-byte work here — memcpy,
|
||||
sha256, tens of MB per image-heavy request — measurably inflates every
|
||||
running request's inter-token latency."""
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.mm_utils import ShmPointerMMData
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalProcessorOutput,
|
||||
)
|
||||
|
||||
shm_names = entry.shm_names
|
||||
if shm_names is None:
|
||||
features = torch.from_numpy(entry.features.reshape(-1, spec.feature_dim))
|
||||
items = []
|
||||
row = 0
|
||||
for index, ((t, h, w), item_hash, offset) in enumerate(
|
||||
zip(entry.grids, entry.hashes, entry.offsets)
|
||||
):
|
||||
n = t * h * w
|
||||
if shm_names is None:
|
||||
feature = features[row : row + n]
|
||||
else:
|
||||
# The worker parked this item's buffer in a named POSIX
|
||||
# segment (see `_use_feature_shm`). Build the stub in its
|
||||
# post-`__setstate__` form: rank 0 never pickle-roundtrips its
|
||||
# own copy, and `materialize()` needs the mapped view.
|
||||
# Ownership of the unlink moved here with `take_mm_result`.
|
||||
feature = ShmPointerMMData.__new__(ShmPointerMMData)
|
||||
feature.__setstate__(
|
||||
{
|
||||
"shm_name": shm_names[index],
|
||||
"shape": (n, spec.feature_dim),
|
||||
"dtype": torch.float32,
|
||||
"precomputed_hash": item_hash,
|
||||
}
|
||||
)
|
||||
items.append(
|
||||
MultimodalDataItem(
|
||||
modality=Modality.IMAGE,
|
||||
feature=feature,
|
||||
hash=item_hash,
|
||||
offsets=[tuple(offset)],
|
||||
model_specific_data={
|
||||
"image_grid_thw": torch.tensor([[t, h, w]], dtype=torch.long)
|
||||
},
|
||||
)
|
||||
)
|
||||
row += n
|
||||
if envs.SGLANG_MM_PRECOMPUTE_HASH.get():
|
||||
for item in items:
|
||||
item.set_pad_value()
|
||||
return MultimodalProcessorOutput(
|
||||
mm_items=items,
|
||||
im_token_id=spec.image_token_id,
|
||||
im_start_id=spec.vision_start_token_id,
|
||||
im_end_id=spec.vision_end_token_id,
|
||||
video_token_id=spec.video_token_id,
|
||||
mrope_positions=torch.from_numpy(entry.mrope.reshape(3, -1)),
|
||||
mrope_position_delta=torch.tensor([[entry.mrope_delta]], dtype=torch.long),
|
||||
)
|
||||
@@ -0,0 +1,438 @@
|
||||
"""Embedded Rust server lifecycle for the scheduler.
|
||||
|
||||
The Rust server replaces the Python api-server + `TokenizerManager` +
|
||||
`DetokenizerManager` stack, running them as Rust threads inside the scheduler
|
||||
process. This wrapper keeps all `SGLANG_RUST_SERVER` plumbing — startup,
|
||||
CPU-core partitioning, the typed `server_args` handoff, and control-response
|
||||
routing — out of `scheduler.py`. The scheduler holds an `Optional[RustServer]`
|
||||
and delegates to it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from array import array
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
import msgspec
|
||||
|
||||
from sglang.srt.managers.io_struct import TokenizedGenerateReqInput
|
||||
from sglang.srt.managers.utils import (
|
||||
MsgpackDecodeError,
|
||||
msgpack_decode_explained,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_mm, get_serving
|
||||
from sglang.srt.rust_server.config import _build_server_args, _partition_cores
|
||||
from sglang.srt.rust_server.multimodal import (
|
||||
RUST_MM_FAMILIES,
|
||||
RustMmProcessor,
|
||||
RustMmSpec,
|
||||
)
|
||||
from sglang.srt.utils.flatten import (
|
||||
FlatPairColumns,
|
||||
NestedRowColumns,
|
||||
RaggedPairColumns,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.io_struct import BatchTokenIDOutput
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.srt.rust_extensions._server import MmSpec, Server
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RustServer:
|
||||
"""Owns the embedded multi-threaded Rust server (``sglang_server.Server``).
|
||||
|
||||
The server owns the api-server, tokenizermanager, tokenizer, and detokenizer
|
||||
all implemented as Rust threads in scheduler process.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server: Server,
|
||||
mm_spec: Optional[RustMmSpec] = None,
|
||||
max_per_poll: int = 256,
|
||||
):
|
||||
self.server = server
|
||||
self.mm_spec = mm_spec
|
||||
self._max_per_poll = max_per_poll
|
||||
|
||||
@classmethod
|
||||
def launch(cls, scheduler: Scheduler) -> RustServer:
|
||||
"""Start the embedded Rust server threads and bind the listen port.
|
||||
|
||||
The caller gates this (``SGLANG_RUST_SERVER`` + rank 0); this always
|
||||
creates.
|
||||
"""
|
||||
from sglang.srt.rust_extensions import load_rust_extension
|
||||
|
||||
Server = load_rust_extension("sglang.srt.rust_extensions._server").Server
|
||||
|
||||
# Force turn off HF tokenizers rayon's unpinned global thread pool.
|
||||
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
||||
|
||||
server_args = scheduler.server_args
|
||||
# `TokenizerManager` merges these under each request's own sampling params
|
||||
# (`{**preferred, **obj.sampling_params}`), and this server replaces that
|
||||
# manager wholesale — so honouring the flag is not implemented here yet.
|
||||
# Refuse rather than run: silently dropping it means generating with
|
||||
# sampling the operator did not configure, and `/get_model_info` would go on
|
||||
# advertising values no request ever receives.
|
||||
if get_serving().preferred_sampling_params:
|
||||
raise ValueError(
|
||||
"SGLANG_RUST_SERVER does not yet apply --preferred-sampling-params "
|
||||
"(the Python TokenizerManager merges it into every request; the rust "
|
||||
"ingress has no equivalent). Launch without SGLANG_RUST_SERVER, or "
|
||||
"drop --preferred-sampling-params and send those values per request."
|
||||
)
|
||||
http_addr = f"{get_serving().host}:{get_serving().port}"
|
||||
|
||||
# Per-DP-rank HTTP port with client load balancing. `None` when DP is off,
|
||||
# so the rank is not conflated with rank 0 of a one-rank group.
|
||||
dp_rank = scheduler.ps.attn_dp_rank if scheduler.ps.dp_size > 1 else None
|
||||
if dp_rank is not None:
|
||||
http_addr = f"{get_serving().host}:{get_serving().port + dp_rank}"
|
||||
|
||||
launch_cores, server_cores = _partition_cores(
|
||||
mm_workers=(
|
||||
(get_mm().mm_processor_worker_num or RustMmProcessor.AUTO_MM_WORKERS)
|
||||
if scheduler.model_config.is_multimodal
|
||||
else 0
|
||||
)
|
||||
)
|
||||
|
||||
server = Server(
|
||||
_build_server_args(scheduler),
|
||||
# None -> run unpinned; the list carries the pinning decision.
|
||||
cores=server_cores,
|
||||
http_addr=http_addr,
|
||||
)
|
||||
|
||||
# Multimodal models must have a Rust pipeline — there is no Python
|
||||
# fallback.
|
||||
mm_spec = None
|
||||
if scheduler.model_config.is_multimodal:
|
||||
# New threads inherit the spawning thread's affinity, and this launch
|
||||
# thread still holds the full mask. Narrow it first so every MM thread
|
||||
# created below (the processor's executors, the Rust MM workers) stays
|
||||
# off the scheduler's reserved cores, where MM preprocessing would
|
||||
# preempt the scheduler loop and inflate inter-token latency.
|
||||
if server_cores is not None:
|
||||
try:
|
||||
os.sched_setaffinity(0, set(server_cores))
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
"rust server: cannot confine mm threads to server cores: %s", e
|
||||
)
|
||||
mm_host = RustMmProcessor(
|
||||
server_args=server_args,
|
||||
model_config=scheduler.model_config,
|
||||
processor=scheduler.processor,
|
||||
)
|
||||
mm_spec = mm_host.resolve_spec()
|
||||
if mm_spec is None:
|
||||
supported = sorted(
|
||||
set(chain.from_iterable(f.model_types for f in RUST_MM_FAMILIES))
|
||||
)
|
||||
raise RuntimeError(
|
||||
"SGLANG_RUST_SERVER=1: no Rust MM pipeline for "
|
||||
f"model_type={scheduler.model_config.hf_config.model_type!r} "
|
||||
f"(supported: {', '.join(supported)}; "
|
||||
"images only). Unset SGLANG_RUST_SERVER to serve this model."
|
||||
)
|
||||
server.start_mm_workers(cls._build_mm_spec(mm_spec), mm_host.mm_workers)
|
||||
|
||||
# Narrow the scheduler thread only after the server threads are launched.
|
||||
if launch_cores is not None:
|
||||
try:
|
||||
# pid 0 == this thread (the scheduler event-loop / launch thread).
|
||||
os.sched_setaffinity(0, set(launch_cores))
|
||||
except OSError as e:
|
||||
logger.warning("rust server: cannot pin scheduler launch thread: %s", e)
|
||||
|
||||
# Under DP every rank runs its own server on its own port, so the rank is
|
||||
# what tells two otherwise identical startup lines apart.
|
||||
dp_note = (
|
||||
"" if dp_rank is None else f" (DP rank {dp_rank}/{scheduler.ps.dp_size})"
|
||||
)
|
||||
logger.info(
|
||||
"SGLANG_RUST_SERVER enabled, Rust server listen on %s%s",
|
||||
http_addr,
|
||||
dp_note,
|
||||
)
|
||||
|
||||
return cls(server, mm_spec=mm_spec)
|
||||
|
||||
def wait_request(self, timeout_ms: int) -> None:
|
||||
"""Block until a request is pushed into the in-process ring or the timeout
|
||||
elapses.
|
||||
"""
|
||||
self.server.wait_request(timeout_ms)
|
||||
|
||||
def drain(self, max_recv: int) -> List[Any]:
|
||||
"""Ingress: non-blocking drain of the in-process ring → list of decoded
|
||||
request objects. The scheduler's request receiver calls this instead of
|
||||
polling the zmq socket when `rust_server_mode` is set.
|
||||
|
||||
The transfer is **columnar**: `recv_requests` returns an `IngressBatch`
|
||||
of scalar msgpack `headers` (with `input_ids` omitted) plus one
|
||||
concatenated raw int64 `data` buffer and per-request `lengths`, so the
|
||||
large `input_ids` lists never go through msgpack. Each header is `msgpack_decode`d (yielding
|
||||
the same `TokenizedGenerateReqInput` / control objects the zmq path
|
||||
produces, so the IPC schema is tracked automatically) and its `input_ids`
|
||||
slice is wrapped as the `array("q")` the scheduler expects. `recv_requests`
|
||||
never waits: the ring drain is `try_recv` (returns the instant the ring
|
||||
is dry, capped at `max_recv`) and the rest is one memcpy per header
|
||||
plus one for the concatenated ids — same contract as `zmq.NOBLOCK`.
|
||||
Parking for work is :meth:`wait_request`, which does release the GIL.
|
||||
"""
|
||||
limit = max_recv if max_recv > 0 else self._max_per_poll
|
||||
batch = self.server.recv_requests(limit)
|
||||
# Bind once: each attribute access converts the rust vec to a fresh list.
|
||||
headers, data, lengths = batch.headers, batch.data, batch.lengths
|
||||
if not headers:
|
||||
return []
|
||||
|
||||
ids_view = memoryview(data)
|
||||
out = []
|
||||
pos = 0 # byte offset into ids_buf
|
||||
for header, n in zip(headers, lengths):
|
||||
nbytes = n * 8
|
||||
try:
|
||||
obj = msgpack_decode_explained(header)
|
||||
except MsgpackDecodeError as e:
|
||||
# Return 400 for malformed request field (e.g. token_ids_logprob=[[0]].
|
||||
logger.warning(
|
||||
"rust ingress: dropping undecodable request %s: %s", e.rid, e.reason
|
||||
)
|
||||
if e.rid is not None:
|
||||
self.server.push_error(e.rid, f"invalid request: {e.reason}")
|
||||
pos += nbytes
|
||||
continue
|
||||
if n: # generate request: attach its int64 ids slice as array("q")
|
||||
ids = array("q")
|
||||
ids.frombytes(ids_view[pos : pos + nbytes])
|
||||
obj.input_ids = ids
|
||||
pos += nbytes
|
||||
if self.mm_spec is not None and isinstance(obj, TokenizedGenerateReqInput):
|
||||
# The buffers rode the Rust sidecar, parked before the ring push;
|
||||
# wrapping them into tensors is the only Python step of the Rust
|
||||
# path. `None` for a text-only request on a multimodal model.
|
||||
mm_result = self.server.take_mm_result(obj.rid)
|
||||
if mm_result is not None:
|
||||
obj.mm_inputs = RustMmProcessor.build_output(
|
||||
self.mm_spec, mm_result
|
||||
)
|
||||
out.append(obj)
|
||||
return out
|
||||
|
||||
def push_control_output(self, recv_req, output) -> None:
|
||||
"""Push a control-request response through the egress ring to the waiting
|
||||
request (routed by rid), encoded as **msgpack** (the ring's native
|
||||
format).
|
||||
|
||||
A msgspec struct is converted to a *named map* (``structs.asdict``, since
|
||||
the IPC structs are ``array_like`` and would otherwise lose field names)
|
||||
so the Rust api_server can shape it per-endpoint (e.g. /server_info)
|
||||
before rendering JSON to the client — keeping JSON formatting off the
|
||||
scheduler's GIL.
|
||||
"""
|
||||
|
||||
# Invariant: control requests always carry a rust-minted rid; without
|
||||
# one the response is unroutable, so fail loudly rather than drop it.
|
||||
assert (
|
||||
recv_req.rid is not None
|
||||
), f"control response without rid: {type(output).__name__}"
|
||||
# No local try/except: a failed push propagates to run_scheduler_process's
|
||||
# outer handler, which logs the full traceback (scheduler-fatal either way).
|
||||
payload = (
|
||||
msgspec.structs.asdict(output)
|
||||
if isinstance(output, msgspec.Struct)
|
||||
else output
|
||||
)
|
||||
# enc_hook stringifies non-native types (paths, enums); JSON
|
||||
# rendering happens in Rust.
|
||||
encoded = msgspec.msgpack.encode(payload, enc_hook=str)
|
||||
|
||||
self.server.push_control_result(recv_req.rid, encoded)
|
||||
|
||||
def push_generation(self, payload: BatchTokenIDOutput) -> None:
|
||||
"""Egress redirect for generation output (replaces the zmq detokenizer).
|
||||
|
||||
Push the WHOLE batch into the Rust egress ring as one frame (-> detokenizer
|
||||
shards -> client streams), mirroring the ingress ``input_ids`` split so the
|
||||
bulk numeric columns never go through msgpack:
|
||||
|
||||
- ``header``: msgpack ``BatchHeader`` positional array — the per-request
|
||||
scalar columns (``rids, finish_reasons, prompt_tokens, tok_lens``) plus
|
||||
the shape metadata for the optional families (``*_lens`` element counts
|
||||
for the flat logprob columns, ``*_reqlens``/``*_poslens`` for the ragged
|
||||
and hidden ones).
|
||||
- ``data``: the raw little-endian numeric buffer — every column is a
|
||||
4-byte element (``f32`` values, ``i32`` indices), concatenated in the
|
||||
order the Rust ``for_each_chunk`` reads them.
|
||||
|
||||
Logprobs are columnar: output families are per-step deltas, input
|
||||
(prefill) families ride once on the first chunk. Ragged families (top-k,
|
||||
token-ids) flatten a per-position ``list[list]`` into flat ``val``/``idx``
|
||||
buffers plus a per-position ``lens`` vector (0 = null position). Hidden
|
||||
states flatten to rows of floats (one row per output position).
|
||||
"""
|
||||
output_ids = payload.output_ids or []
|
||||
prompt_tokens = payload.prompt_tokens or []
|
||||
|
||||
# Hot-path guard: almost no decode step wants logprobs / hidden states,
|
||||
# so only then pay the per-request flatten + buffer packing below.
|
||||
has_extra = bool(
|
||||
payload.output_token_logprobs_val
|
||||
or payload.input_token_logprobs_val
|
||||
or payload.output_top_logprobs_val
|
||||
or payload.input_top_logprobs_val
|
||||
or payload.output_token_ids_logprobs_val
|
||||
or payload.input_token_ids_logprobs_val
|
||||
or payload.output_hidden_states
|
||||
)
|
||||
|
||||
# Runs on the scheduler's CUDA-launch thread every decode step, so each
|
||||
# Python-level pass over the batch costs inter-token latency: `rids` are
|
||||
# the plain rid strings (hashed to a routing key on the Rust side with a
|
||||
# per-process seed, off the GIL — not parsed; a rid is any string),
|
||||
# `finished_reasons` already `dict | None`, and `output_ids` entries are
|
||||
# always `array("i")` (never None) so `map(len)` and a bare
|
||||
# `chain.from_iterable` stay in C.
|
||||
rids = payload.rids
|
||||
finish_reasons = payload.finished_reasons
|
||||
tok_lens = list(map(len, output_ids))
|
||||
flat_ids = array("i", chain.from_iterable(output_ids))
|
||||
|
||||
# Column order here MUST match BatchHeader (header_cols) and
|
||||
# for_each_chunk's read order (data_cols); the extras contribution
|
||||
# is ordered by the `extras` tuple below.
|
||||
header_cols = [rids, finish_reasons, prompt_tokens, tok_lens]
|
||||
data_cols = [flat_ids.tobytes()]
|
||||
|
||||
if has_extra:
|
||||
# The `extras` tuple is the SINGLE source of the extras column
|
||||
# order — it must match the Rust ``BatchHeader`` fields and
|
||||
# ``for_each_chunk``'s read order.
|
||||
#
|
||||
# TODO(perf): the per-request flatten assumes the logprob/hidden
|
||||
# columns are ragged, non-contiguous nested Python lists — which is
|
||||
# only an assumption. The scheduler moves these off the GPU with
|
||||
# `tensor.tolist()`, so revisit whether the upstream values are
|
||||
# still contiguous tensors; if so, ship raw bytes + a shape
|
||||
# descriptor and skip the flatten entirely.
|
||||
batch_size = len(rids)
|
||||
extras = (
|
||||
FlatPairColumns(
|
||||
"output_token_logprobs",
|
||||
payload.output_token_logprobs_val or [],
|
||||
payload.output_token_logprobs_idx or [],
|
||||
),
|
||||
FlatPairColumns(
|
||||
"input_token_logprobs",
|
||||
payload.input_token_logprobs_val or [],
|
||||
payload.input_token_logprobs_idx or [],
|
||||
first_none_to_nan=True,
|
||||
),
|
||||
RaggedPairColumns(
|
||||
"output_top_logprobs",
|
||||
payload.output_top_logprobs_val or [],
|
||||
payload.output_top_logprobs_idx or [],
|
||||
),
|
||||
RaggedPairColumns(
|
||||
"input_top_logprobs",
|
||||
payload.input_top_logprobs_val or [],
|
||||
payload.input_top_logprobs_idx or [],
|
||||
),
|
||||
RaggedPairColumns(
|
||||
"output_token_ids_logprobs",
|
||||
payload.output_token_ids_logprobs_val or [],
|
||||
payload.output_token_ids_logprobs_idx or [],
|
||||
),
|
||||
RaggedPairColumns(
|
||||
"input_token_ids_logprobs",
|
||||
payload.input_token_ids_logprobs_val or [],
|
||||
payload.input_token_ids_logprobs_idx or [],
|
||||
),
|
||||
NestedRowColumns(
|
||||
"output_hidden_states", payload.output_hidden_states or []
|
||||
),
|
||||
)
|
||||
|
||||
# Every column is all-or-nothing per payload — which is also what makes
|
||||
# a family's emptiness a reliable "nobody asked for this" signal.
|
||||
active = []
|
||||
for extra in extras:
|
||||
populated = False
|
||||
for name, col in extra.columns():
|
||||
assert len(col) in (
|
||||
0,
|
||||
batch_size,
|
||||
), f"extras column {name}: {len(col)} entries for a batch of {batch_size}"
|
||||
populated |= len(col) > 0
|
||||
if populated:
|
||||
active.append(extra)
|
||||
|
||||
# Flatten only the families someone asked for. `has_extra` above is a
|
||||
# per-FRAME guard, so one client enabling logprobs used to drag all
|
||||
# seven families through the per-request loop: at B=4096 that is 28,672
|
||||
# bound-method calls per decode step, materializing 12 columns of 4096
|
||||
# zeros nobody reads. Measured 0.37 ms -> 7.90 ms GIL-held per step,
|
||||
# i.e. 25-75% of a decode step added to the scheduler's critical path.
|
||||
#
|
||||
# Skipping `accept` leaves a family's buffers empty, which is exactly
|
||||
# the wire form the Rust decoder already treats as absent (`per_req_ok`
|
||||
# admits an empty column, `lens_i` reads 0 for every request). The
|
||||
# `header_cols`/`data_cols` loops below still walk all seven, so column
|
||||
# ORDER and arity are unchanged — an inactive family contributes empty
|
||||
# columns in place rather than disappearing.
|
||||
for extra in active:
|
||||
accept = extra.accept # hoisted: this is the hottest loop here
|
||||
for i in range(batch_size):
|
||||
accept(i)
|
||||
|
||||
for extra in extras:
|
||||
header_cols += extra.header_cols()
|
||||
data_cols += extra.data_cols()
|
||||
|
||||
header = msgspec.msgpack.encode(header_cols)
|
||||
# Pass the raw column list; the Rust side concatenates it into the frame
|
||||
# with the GIL released.
|
||||
if not self.server.push_decode_result_batch(header, data_cols):
|
||||
logger.warning(
|
||||
"Rust egress closed; dropped batch of %d requests during shutdown",
|
||||
len(rids),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_mm_spec(spec: RustMmSpec) -> MmSpec:
|
||||
"""The typed MM handoff for ``Server.start_mm_workers``: the
|
||||
:class:`RustMmSpec` fields the Rust pipeline consumes, as the Rust
|
||||
extension's own ``MmSpec`` class (same required-keyword contract as
|
||||
:meth:`_build_server_args`; ``family`` / ``resample`` become the
|
||||
extension's ``MmFamily`` / ``MmResample`` enums)."""
|
||||
from sglang.srt.rust_extensions import load_rust_extension
|
||||
|
||||
ext = load_rust_extension("sglang.srt.rust_extensions._server")
|
||||
family = {"qwen_vl": ext.MmFamily.QwenVl}[spec.family]
|
||||
resample = {"aten_u8": ext.MmResample.AtenU8, "pil": ext.MmResample.Pil}[
|
||||
spec.resample
|
||||
]
|
||||
return ext.MmSpec(
|
||||
family=family,
|
||||
feature_shm=spec.feature_shm,
|
||||
image_token_id=spec.image_token_id,
|
||||
patch_size=spec.patch_size,
|
||||
merge_size=spec.merge_size,
|
||||
temporal_patch_size=spec.temporal_patch_size,
|
||||
min_pixels=spec.min_pixels,
|
||||
max_pixels=spec.max_pixels,
|
||||
image_mean=spec.image_mean,
|
||||
image_std=spec.image_std,
|
||||
resample=resample,
|
||||
)
|
||||
@@ -22,7 +22,7 @@ features = ["python", "parallel"]
|
||||
# which renames the built artifact.
|
||||
name = "sglang_mm_core"
|
||||
# cdylib: the PyO3 module (`sglang.srt.rust_extensions._multimodal`).
|
||||
# rlib: pure-Rust core linked by sglang-server's native MM path.
|
||||
# rlib: pure-Rust core linked by sglang-server's Rust MM path.
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[features]
|
||||
|
||||
@@ -65,7 +65,7 @@ pub struct QwenVlProcessor {
|
||||
lut: [[f32; 256]; 3],
|
||||
}
|
||||
|
||||
/// `1 / rescale_factor`; `resolve_native_spec` rejects any other factor.
|
||||
/// `1 / rescale_factor`; `resolve_spec` rejects any other factor.
|
||||
const INV_RESCALE: f32 = 255.0;
|
||||
|
||||
/// u8 → normalized f32, rounded as the mirrored processor rounds. The slow one
|
||||
@@ -406,7 +406,7 @@ mod python {
|
||||
|
||||
/// `(pixel_values flat f32, (t, h, w))` for one preprocessed image.
|
||||
type PyProcessedImage<'py> = (Bound<'py, PyArray1<f32>>, (u32, u32, u32));
|
||||
/// Full native pipeline output at the scheduler boundary:
|
||||
/// Full Rust pipeline output at the scheduler boundary:
|
||||
/// `(input_ids, features, grids, hashes, offsets, mrope, mrope_delta)`.
|
||||
type PyNativeOutput<'py> = (
|
||||
Vec<i32>,
|
||||
@@ -486,7 +486,7 @@ mod python {
|
||||
/// `sglang-server` (whose message layer owns the wire-payload parsing).
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (input_ids, images, spec_json))]
|
||||
fn process_native_mm<'py>(
|
||||
fn process_mm<'py>(
|
||||
py: Python<'py>,
|
||||
input_ids: Option<Vec<i32>>,
|
||||
images: Vec<PyImageSource>,
|
||||
@@ -529,7 +529,7 @@ mod python {
|
||||
m.add_function(wrap_pyfunction!(preprocess, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(smart_resize_py, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(mrope_image_only_py, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(process_native_mm, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(process_mm, &m)?)?;
|
||||
parent.add_submodule(&m)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -50,7 +50,7 @@ hf-hub = { version = "0.4", default-features = false }
|
||||
itertools = "0.14"
|
||||
# POSIX shm for the MM feature fan-out (`mm::ShmSegment`).
|
||||
libc = "0.2"
|
||||
# Same major as the workspace pyo3: the zero-copy MM drain (`take_mm`) moves
|
||||
# Same major as the workspace pyo3: the zero-copy MM drain (`take_mm_result`) moves
|
||||
# Rust vectors into numpy arrays.
|
||||
numpy = "0.29.0"
|
||||
# Pinned EXACTLY: this crate's accepted grammar defines the
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
# sglang-server
|
||||
|
||||
`sglang-server` is SGLang's Rust HTTP frontend and request-processing pipeline. It exchanges typed requests and responses with the Python scheduler while keeping latency-sensitive work outside Python.
|
||||
|
||||
## Code review principles
|
||||
|
||||
1. **Use strongly typed boundaries.** Model every supported protocol shape with structs, enums, and validated newtypes; avoid opaque values such as `serde_json::Value` and `rmpv::Value` in production paths.
|
||||
2. **Keep one canonical schema.** Rust and Python must derive their wire contracts from one source of truth, aligned with `io_struct.py`, instead of independently duplicating field names, order, defaults, or validation.
|
||||
3. **Make protocol declarations minimal and declarative.** A reviewer should be able to understand the wire format from its type declarations alone, without tracing fillers, conversion code, macros, or repeated field lists.
|
||||
4. **Use one representation per semantic stage.** Separate external input, normalized domain data, and wire data, and convert between them once at explicit boundaries; do not keep multiple overlapping representations of the same state.
|
||||
5. **Design compatibility and safety explicitly.** Version protocols, reject unsupported or malformed inputs clearly, validate lengths and resource bounds before allocation, preserve invariants in types, and test compatibility across the real Rust and Python codecs.
|
||||
@@ -14,6 +14,9 @@ mod completions;
|
||||
mod models;
|
||||
mod reasoning;
|
||||
mod template;
|
||||
mod template_builtins;
|
||||
mod template_legacy;
|
||||
mod template_loader;
|
||||
mod tools;
|
||||
|
||||
pub(super) use template::ChatFormatter;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,284 @@
|
||||
//! Built-in legacy SGLang conversation templates.
|
||||
|
||||
use crate::message::types::OneOrMany;
|
||||
|
||||
use super::template_legacy::LegacySpec;
|
||||
|
||||
pub(super) fn builtin_template(name: &str) -> Option<LegacySpec> {
|
||||
let spec = match name {
|
||||
"llama-2" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n".into(),
|
||||
roles: ("[INST]".into(), "[/INST]".into()),
|
||||
style: "LLAMA2".into(),
|
||||
sep: " ".into(),
|
||||
sep2: Some(" </s><s>".into()),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"[INST]".into(),
|
||||
"[/INST]".into(),
|
||||
"<<SYS>>".into(),
|
||||
"<</SYS>>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"mistral" | "devstral" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "[SYSTEM_PROMPT]\n{system_message}\n[/SYSTEM_PROMPT]\n\n".into(),
|
||||
roles: ("[INST]".into(), "[/INST]".into()),
|
||||
style: "LLAMA2".into(),
|
||||
sep: " ".into(),
|
||||
sep2: Some(" </s><s>".into()),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"[INST]".into(),
|
||||
"[/INST]".into(),
|
||||
"[SYSTEM_PROMPT]".into(),
|
||||
"[/SYSTEM_PROMPT]".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"llama-4" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|header_start|>system<|header_end|>\n\n{system_message}<|eot|>"
|
||||
.into(),
|
||||
roles: ("user".into(), "assistant".into()),
|
||||
style: "LLAMA4".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|end_of_text|>".into(),
|
||||
"<|eot|>".into(),
|
||||
"<|eom|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"phi-4-mm" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}".into(),
|
||||
roles: ("<|user|>".into(), "<|assistant|>".into()),
|
||||
style: "NO_COLON_SINGLE".into(),
|
||||
sep: "<|end|>".into(),
|
||||
stop_str: Some(OneOrMany::One("<|end|>".into())),
|
||||
image_token: "<|endoftext10|>".into(),
|
||||
audio_token: "<|endoftext11|>".into(),
|
||||
..Default::default()
|
||||
},
|
||||
"chatml" | "chatml-llava" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "CHATML".into(),
|
||||
sep: "<|im_end|>".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|endoftext|>".into(),
|
||||
"<|im_end|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"vicuna_v1.1" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}".into(),
|
||||
system_message: "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.".into(),
|
||||
roles: ("USER".into(), "ASSISTANT".into()),
|
||||
style: "ADD_COLON_TWO".into(),
|
||||
sep: " ".into(),
|
||||
sep2: Some("</s>".into()),
|
||||
..Default::default()
|
||||
},
|
||||
"llama_3_vision" | "llava_llama_3" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>"
|
||||
.into(),
|
||||
system_message: "You are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.".into(),
|
||||
roles: ("user".into(), "assistant".into()),
|
||||
style: "LLAMA3".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|end_of_text|>".into(),
|
||||
"<|eot_id|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"internlm2-chat" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_COLON_SINGLE".into(),
|
||||
sep: "\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|im_end|>".into(),
|
||||
"<|action_end|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"internvl-2-5" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。".into(),
|
||||
roles: ("<|im_start|>user\n".into(), "<|im_start|>assistant\n".into()),
|
||||
style: "MPT".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|im_end|>".into(),
|
||||
"<|action_end|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"qwen2-vl" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|im_end|>".into()])),
|
||||
..Default::default()
|
||||
},
|
||||
"deepseek-ocr" => LegacySpec {
|
||||
name: name.into(),
|
||||
style: "NO_COLON_SINGLE".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<｜end▁of▁sentence｜>".into()])),
|
||||
..Default::default()
|
||||
},
|
||||
"unlimited-ocr" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}".into(),
|
||||
style: "UNLIMITED_OCR".into(),
|
||||
sep2: Some(String::new()),
|
||||
..Default::default()
|
||||
},
|
||||
"paddle-ocr" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|begin_of_sentence|>{system_message}".into(),
|
||||
roles: ("User".into(), "Assistant".into()),
|
||||
style: "PADDLE_OCR".into(),
|
||||
sep: "<|end_of_sentence|>".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|end_of_sentence|>".into()])),
|
||||
image_token: "<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>".into(),
|
||||
..Default::default()
|
||||
},
|
||||
"deepseek-vl2" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}".into(),
|
||||
roles: ("<|User|>".into(), "<|Assistant|>".into()),
|
||||
style: "DeepSeekVL2".into(),
|
||||
sep: "\n\n".into(),
|
||||
sep2: Some("<｜end▁of▁sentence｜>".into()),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"User:".into(),
|
||||
"<｜end▁of▁sentence｜>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"gemma-it" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<start_of_turn>user\n{system_message}\n\n".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<start_of_turn>user\n".into(), "<start_of_turn>model\n".into()),
|
||||
style: "GEMMA3".into(),
|
||||
sep: "<end_of_turn>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<end_of_turn>".into()])),
|
||||
image_token: "<start_of_image>".into(),
|
||||
audio_token: "<start_of_audio>".into(),
|
||||
..Default::default()
|
||||
},
|
||||
"gme-qwen2-vl" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "QWEN2_VL_EMBED".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::One("<|endoftext|>".into())),
|
||||
..Default::default()
|
||||
},
|
||||
"minicpmv" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}.".into(),
|
||||
system_message: "You are a helpful assistant".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|im_end|>".into(),
|
||||
"<|endoftext|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"janus-pro" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "{system_message}.".into(),
|
||||
system_message: "You are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language".into(),
|
||||
roles: ("User".into(), "Assistant".into()),
|
||||
style: "ADD_COLON_TWO".into(),
|
||||
sep: "\n\n".into(),
|
||||
sep2: Some("<｜end▁of▁sentence｜>".into()),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|User|>".into(),
|
||||
"<｜end▁of▁sentence｜>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"minicpmo" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
|
||||
.into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec![
|
||||
"<|im_end|>".into(),
|
||||
"<|endoftext|>".into(),
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"kimi-vl" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_system|>system<|im_middle|>{system_message}".into(),
|
||||
system_message: "You are a helpful assistant".into(),
|
||||
roles: (
|
||||
"<|im_user|>user<|im_middle|>".into(),
|
||||
"<|im_assistant|>assistant<|im_middle|>".into(),
|
||||
),
|
||||
style: "NO_COLON_SINGLE".into(),
|
||||
sep: "<|im_end|>".into(),
|
||||
stop_str: Some(OneOrMany::One("<|im_end|>".into())),
|
||||
..Default::default()
|
||||
},
|
||||
"qwen2-audio" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
system_message: "You are a helpful assistant.".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "QWEN2_AUDIO".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|im_end|>".into()])),
|
||||
audio_token: "Audio {idx}: <|audio_bos|><|AUDIO|><|audio_eos|>\n".into(),
|
||||
..Default::default()
|
||||
},
|
||||
"moss-vl" => LegacySpec {
|
||||
name: name.into(),
|
||||
system_template: "<|im_start|>system\n{system_message}".into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|im_end|>".into()])),
|
||||
..Default::default()
|
||||
},
|
||||
"points-v15-chat" => LegacySpec {
|
||||
name: name.into(),
|
||||
roles: ("<|im_start|>user".into(), "<|im_start|>assistant".into()),
|
||||
style: "ADD_NEW_LINE_SINGLE".into(),
|
||||
sep: "<|im_end|>\n".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|im_end|>".into()])),
|
||||
..Default::default()
|
||||
},
|
||||
"whisper" => LegacySpec {
|
||||
name: name.into(),
|
||||
style: "NO_COLON_SINGLE".into(),
|
||||
stop_str: Some(OneOrMany::Many(vec!["<|endoftext|>".into()])),
|
||||
audio_token: String::new(),
|
||||
..Default::default()
|
||||
},
|
||||
_ => return None,
|
||||
};
|
||||
Some(spec)
|
||||
}
|
||||
@@ -0,0 +1,605 @@
|
||||
//! Legacy SGLang conversation-template rendering.
|
||||
|
||||
use dynamo_protocols::types::{
|
||||
ChatCompletionRequestAssistantMessageContent, ChatCompletionRequestAssistantMessageContentPart,
|
||||
ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageContent,
|
||||
ChatCompletionRequestSystemMessageContentPart, ChatCompletionRequestUserMessageContent,
|
||||
ChatCompletionRequestUserMessageContentPart, CreateChatCompletionRequest,
|
||||
};
|
||||
|
||||
use crate::message::types::OneOrMany;
|
||||
|
||||
use super::template::TemplateError;
|
||||
|
||||
/// A legacy conversation template, mirroring Python's `Conversation` fields.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct LegacySpec {
|
||||
/// Python `Conversation.name` — drives the CHATGLM round-offset quirk.
|
||||
pub(super) name: String,
|
||||
pub(super) system_template: String,
|
||||
pub(super) system_message: String,
|
||||
/// `(user_role, assistant_role)` — Python `Conversation.roles`.
|
||||
pub(super) roles: (String, String),
|
||||
pub(super) style: String,
|
||||
pub(super) sep: String,
|
||||
/// `None` = Python's `Conversation.sep2` default. Styles that alternate
|
||||
/// seps (`seps[i % 2]`) need it set; Python crashes on `None` there and we
|
||||
/// error deliberately.
|
||||
pub(super) sep2: Option<String>,
|
||||
/// Python `Conversation.stop_str` (`str | list[str] | None`).
|
||||
pub(super) stop_str: Option<OneOrMany<String>>,
|
||||
pub(super) image_token: String,
|
||||
pub(super) audio_token: String,
|
||||
}
|
||||
|
||||
impl Default for LegacySpec {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
name: String::new(),
|
||||
system_template: String::new(),
|
||||
system_message: String::new(),
|
||||
roles: (String::new(), String::new()),
|
||||
style: String::new(),
|
||||
sep: String::new(),
|
||||
sep2: None,
|
||||
stop_str: None,
|
||||
image_token: "<image>".into(),
|
||||
audio_token: "<audio>".into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Native port of Python `generate_chat_conv` + `Conversation.get_prompt()`:
|
||||
/// fold system messages into the system prompt, keep user/assistant messages in
|
||||
/// order, always append the assistant opening, then render per `sep_style`.
|
||||
#[derive(Clone)]
|
||||
pub struct LegacyFormatter {
|
||||
pub(super) spec: LegacySpec,
|
||||
}
|
||||
|
||||
impl LegacyFormatter {
|
||||
pub(super) fn render(
|
||||
&self,
|
||||
request: &CreateChatCompletionRequest,
|
||||
) -> Result<String, TemplateError> {
|
||||
let mut system_message = self.spec.system_message.clone();
|
||||
let mut messages: Vec<(String, String)> = Vec::new();
|
||||
for message in &request.messages {
|
||||
match message {
|
||||
ChatCompletionRequestMessage::System(message) => {
|
||||
system_message = extract_system_text(&message.content)?;
|
||||
}
|
||||
ChatCompletionRequestMessage::User(message) => {
|
||||
let content = match &message.content {
|
||||
ChatCompletionRequestUserMessageContent::Text(text) => text.clone(),
|
||||
ChatCompletionRequestUserMessageContent::Array(parts) => {
|
||||
let mut text = String::new();
|
||||
for part in parts {
|
||||
match part {
|
||||
ChatCompletionRequestUserMessageContentPart::Text(part) => {
|
||||
text.push_str(&part.text);
|
||||
}
|
||||
// Python would splice media tokens in here;
|
||||
// the OpenAI adapter rejects media content
|
||||
// upstream, so this is unreachable — error
|
||||
// rather than silently drop.
|
||||
_ => {
|
||||
return Err(TemplateError::MediaContent { role: "user" });
|
||||
}
|
||||
}
|
||||
}
|
||||
text
|
||||
}
|
||||
};
|
||||
messages.push((self.spec.roles.0.clone(), content));
|
||||
}
|
||||
ChatCompletionRequestMessage::Assistant(message) => {
|
||||
let content = message
|
||||
.content
|
||||
.as_ref()
|
||||
.map(extract_assistant_text)
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
messages.push((self.spec.roles.1.clone(), content));
|
||||
}
|
||||
other => {
|
||||
return Err(TemplateError::UnsupportedRole {
|
||||
role: match other {
|
||||
ChatCompletionRequestMessage::Developer(_) => "developer",
|
||||
ChatCompletionRequestMessage::Tool(_) => "tool",
|
||||
ChatCompletionRequestMessage::Function(_) => "function",
|
||||
_ => unreachable!(),
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
// Python's `generate_chat_conv` appends the assistant opening.
|
||||
messages.push((self.spec.roles.1.clone(), String::new()));
|
||||
self.render_prompt(&system_message, &messages)
|
||||
}
|
||||
|
||||
fn render_prompt(
|
||||
&self,
|
||||
system_message: &str,
|
||||
messages: &[(String, String)],
|
||||
) -> Result<String, TemplateError> {
|
||||
let spec = &self.spec;
|
||||
// Python: `self.system_template.format(system_message=self.system_message)`.
|
||||
let system_prompt = spec
|
||||
.system_template
|
||||
.replace("{system_message}", system_message);
|
||||
let user_role = &spec.roles.0;
|
||||
let assistant_role = &spec.roles.1;
|
||||
let mut ret = String::new();
|
||||
|
||||
match spec.style.as_str() {
|
||||
"ADD_COLON_SINGLE" => {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"ADD_COLON_TWO" => {
|
||||
let sep2 = sep2(spec, "ADD_COLON_TWO")?;
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"ADD_COLON_SPACE_SINGLE" => {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}: "));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"ADD_NEW_LINE_SINGLE" => {
|
||||
if !system_message.is_empty() && !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}\n{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"QWEN2_VL_EMBED" => {
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}\n{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
match &spec.stop_str {
|
||||
Some(OneOrMany::One(stop)) => ret.push_str(stop),
|
||||
// Python `ret += self.stop_str` raises TypeError for
|
||||
// `None` / list; error deliberately instead.
|
||||
_ => {
|
||||
return Err(TemplateError::InvalidStopString {
|
||||
style: "QWEN2_VL_EMBED".into(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
"NO_COLON_SINGLE" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"NO_COLON_TWO" => {
|
||||
let sep2 = sep2(spec, "NO_COLON_TWO")?;
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"RWKV" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"{role}: {}",
|
||||
content.replace("\r\n", "\n").replace("\n\n", "\n")
|
||||
));
|
||||
ret.push_str("\n\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
"LLAMA4" => {
|
||||
if !system_message.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("<|header_start|>{role}<|header_end|>\n\n"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"<|header_start|>{role}<|header_end|>\n\n{}<|eot|>",
|
||||
content.trim()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
"LLAMA3" => {
|
||||
if !system_message.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("<|start_header_id|>{role}<|end_header_id|>\n\n"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"<|start_header_id|>{role}<|end_header_id|>\n\n{}<|eot_id|>",
|
||||
content.trim()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
"LLAMA2" => {
|
||||
let sep2 = sep2(spec, "LLAMA2")?;
|
||||
if system_message.is_empty() {
|
||||
ret.push_str("[INST] ");
|
||||
} else {
|
||||
ret.push_str(&system_prompt);
|
||||
}
|
||||
for (i, (_, content)) in messages.iter().enumerate() {
|
||||
// Python: `tag = self.roles[i % 2]` — parity, not the
|
||||
// stored role, and the first message has no tag.
|
||||
let tag = if i % 2 == 0 {
|
||||
user_role
|
||||
} else {
|
||||
assistant_role
|
||||
};
|
||||
if content.is_empty() {
|
||||
ret.push_str(tag);
|
||||
} else if i == 0 {
|
||||
ret.push_str(&format!("{content} "));
|
||||
} else {
|
||||
ret.push_str(&format!("{tag} {content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"CHATGLM" => {
|
||||
// Python: `round_add_n = 1 if self.name == "chatglm2" else 0`.
|
||||
let round_add_n = if spec.name == "chatglm2" { 1 } else { 0 };
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if i % 2 == 0 {
|
||||
ret.push_str(&format!("[Round {}]{}", i / 2 + round_add_n, spec.sep));
|
||||
}
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}："));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}：{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"CHATML" => {
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
ret.push('\n');
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}\n{content}{}\n", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"CHATGLM3" => {
|
||||
if !system_message.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}\n{content}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
"CHATINTERN" => {
|
||||
let sep2 = sep2(spec, "CHATINTERN")?;
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if i % 2 == 0 {
|
||||
ret.push_str("<s>");
|
||||
}
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"{role}:{content}{}\n",
|
||||
sep_even_odd(spec, sep2, i)
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
"DOLLY" => {
|
||||
let sep2 = sep2(spec, "DOLLY")?;
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:\n"));
|
||||
} else {
|
||||
ret.push_str(&format!(
|
||||
"{role}:\n{content}{}",
|
||||
sep_even_odd(spec, sep2, i)
|
||||
));
|
||||
if i % 2 == 1 {
|
||||
ret.push_str("\n\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"PHOENIX" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}: <s>"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: <s>{content}</s>"));
|
||||
}
|
||||
}
|
||||
}
|
||||
"ROBIN" => {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}:\n{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"FALCON_CHAT" => {
|
||||
if !system_message.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"METAMATH" => {
|
||||
let sep2 = sep2(spec, "METAMATH")?;
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
// Python: sep2 prefixes odd messages; sep ends even ones.
|
||||
if content.is_empty() {
|
||||
if i % 2 == 0 {
|
||||
ret.push_str(&format!("{role}:\n"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {sep2}"));
|
||||
}
|
||||
} else if i % 2 == 0 {
|
||||
ret.push_str(&format!("{role}:\n{content}{}", spec.sep));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {sep2}{content}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
"DEEPSEEK_CHAT" => {
|
||||
let sep2 = sep2(spec, "DEEPSEEK_CHAT")?;
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"DeepSeekVL2" => {
|
||||
let sep2 = sep2(spec, "DeepSeekVL2")?;
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}:"));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
"GEMMA3" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else if i == 0 {
|
||||
ret.push_str(&format!("{content}{}", spec.sep));
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"MPT" => {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"QWEN2_AUDIO" => {
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
let mut counter = 1usize;
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}\n"));
|
||||
} else {
|
||||
let mut message = content.clone();
|
||||
while message.contains(&spec.audio_token) {
|
||||
// Python: `audio_token.format(idx=counter)`. A
|
||||
// token without `{idx}` makes the replace a no-op
|
||||
// and Python's loop infinite; bail out instead of
|
||||
// hanging the server.
|
||||
let indexed = spec.audio_token.replace("{idx}", &counter.to_string());
|
||||
if indexed == spec.audio_token {
|
||||
break;
|
||||
}
|
||||
message = message.replacen(&spec.audio_token, &indexed, 1);
|
||||
counter += 1;
|
||||
}
|
||||
ret.push_str(&format!("{role}\n{message}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"PADDLE_OCR" => {
|
||||
ret.push_str(&system_prompt);
|
||||
for (role, content) in messages {
|
||||
if content.is_empty() {
|
||||
ret.push_str(&format!("{role}: "));
|
||||
} else if role == user_role {
|
||||
ret.push_str(&format!("{role}: "));
|
||||
if content.contains(&spec.image_token) {
|
||||
ret.push_str(
|
||||
&content
|
||||
.replace(&format!("{}\n", spec.image_token), &spec.image_token),
|
||||
);
|
||||
} else {
|
||||
ret.push_str(content);
|
||||
}
|
||||
ret.push('\n');
|
||||
} else {
|
||||
ret.push_str(&format!("{role}: {content}{}", spec.sep));
|
||||
}
|
||||
}
|
||||
}
|
||||
"UNLIMITED_OCR" => {
|
||||
let sep2 = sep2(spec, "UNLIMITED_OCR")?;
|
||||
if !system_prompt.is_empty() {
|
||||
ret.push_str(&system_prompt);
|
||||
ret.push_str(&spec.sep);
|
||||
}
|
||||
for (i, (role, content)) in messages.iter().enumerate() {
|
||||
if content.is_empty() {
|
||||
ret.push_str(role);
|
||||
} else {
|
||||
ret.push_str(&format!("{role}{content}{}", sep_even_odd(spec, sep2, i)));
|
||||
}
|
||||
}
|
||||
}
|
||||
other => {
|
||||
return Err(TemplateError::InvalidStyle {
|
||||
style: other.to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(ret)
|
||||
}
|
||||
}
|
||||
|
||||
/// Python `seps = [self.sep, self.sep2]` indexed by message parity.
|
||||
fn sep_even_odd<'a>(spec: &'a LegacySpec, sep2: &'a str, index: usize) -> &'a str {
|
||||
if index.is_multiple_of(2) {
|
||||
&spec.sep
|
||||
} else {
|
||||
sep2
|
||||
}
|
||||
}
|
||||
|
||||
fn sep2<'a>(spec: &'a LegacySpec, style: &str) -> Result<&'a str, TemplateError> {
|
||||
spec.sep2
|
||||
.as_deref()
|
||||
.ok_or_else(|| TemplateError::MissingSep2 {
|
||||
style: style.to_owned(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Python `generate_chat_conv` system extraction: a plain string, or an array
|
||||
/// with exactly one `text` part.
|
||||
fn extract_system_text(
|
||||
content: &ChatCompletionRequestSystemMessageContent,
|
||||
) -> Result<String, TemplateError> {
|
||||
match content {
|
||||
ChatCompletionRequestSystemMessageContent::Text(text) => Ok(text.clone()),
|
||||
ChatCompletionRequestSystemMessageContent::Array(parts) => {
|
||||
let mut texts = parts.iter().map(|part| match part {
|
||||
ChatCompletionRequestSystemMessageContentPart::Text(part) => part.text.as_str(),
|
||||
});
|
||||
match (texts.next(), texts.next()) {
|
||||
(Some(text), None) => Ok(text.to_owned()),
|
||||
_ => Err(TemplateError::NonTextContent { role: "system" }),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Python `generate_chat_conv` assistant extraction: a plain string, or an
|
||||
/// array with exactly one `text` part (`refusal` parts are rejected).
|
||||
fn extract_assistant_text(
|
||||
content: &ChatCompletionRequestAssistantMessageContent,
|
||||
) -> Result<String, TemplateError> {
|
||||
match content {
|
||||
ChatCompletionRequestAssistantMessageContent::Text(text) => Ok(text.clone()),
|
||||
ChatCompletionRequestAssistantMessageContent::Array(parts) => {
|
||||
let mut texts = parts.iter().filter_map(|part| match part {
|
||||
ChatCompletionRequestAssistantMessageContentPart::Text(part) => {
|
||||
Some(part.text.as_str())
|
||||
}
|
||||
ChatCompletionRequestAssistantMessageContentPart::Refusal(_) => None,
|
||||
});
|
||||
match (texts.next(), texts.next()) {
|
||||
(Some(text), None) => Ok(text.to_owned()),
|
||||
_ => Err(TemplateError::NonTextContent { role: "assistant" }),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,351 @@
|
||||
//! Chat-template loading and model-path inference.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use dynamo_renderer::{ChatTemplate, ContextMixins, PromptContextMixin, PromptFormatter};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::message::types::OneOrMany;
|
||||
|
||||
use super::template::{ChatFormatter, TemplateError};
|
||||
use super::template_builtins::builtin_template;
|
||||
use super::template_legacy::{LegacyFormatter, LegacySpec};
|
||||
|
||||
const SUPPORTED_STYLES: &[&str] = &[
|
||||
"ADD_COLON_SINGLE",
|
||||
"ADD_COLON_TWO",
|
||||
"ADD_COLON_SPACE_SINGLE",
|
||||
"NO_COLON_SINGLE",
|
||||
"NO_COLON_TWO",
|
||||
"ADD_NEW_LINE_SINGLE",
|
||||
"LLAMA2",
|
||||
"LLAMA3",
|
||||
"LLAMA4",
|
||||
"CHATGLM",
|
||||
"CHATML",
|
||||
"CHATINTERN",
|
||||
"DOLLY",
|
||||
"RWKV",
|
||||
"PHOENIX",
|
||||
"ROBIN",
|
||||
"FALCON_CHAT",
|
||||
"CHATGLM3",
|
||||
"DEEPSEEK_CHAT",
|
||||
"METAMATH",
|
||||
"DeepSeekVL2",
|
||||
"QWEN2_VL_EMBED",
|
||||
"QWEN2_AUDIO",
|
||||
"GEMMA3",
|
||||
"MPT",
|
||||
"PADDLE_OCR",
|
||||
"UNLIMITED_OCR",
|
||||
];
|
||||
|
||||
pub(super) fn load_chat_formatter(
|
||||
config_file: Option<&str>,
|
||||
model_path: Option<&str>,
|
||||
chat_template_arg: Option<&str>,
|
||||
) -> Result<ChatFormatter, TemplateError> {
|
||||
// Python resolves registry names before looking at the filesystem — and
|
||||
// before touching the tokenizer config, so a built-in name works even when
|
||||
// `tokenizer_config.json` is absent.
|
||||
if let Some(argument) = chat_template_arg
|
||||
&& let Some(spec) = builtin_template(argument)
|
||||
{
|
||||
return Ok(ChatFormatter::Legacy(Box::new(LegacyFormatter { spec })));
|
||||
}
|
||||
|
||||
// Python `load_chat_template` (no `--chat-template`): infer a legacy
|
||||
// template from the model path before falling back to the HF template, so
|
||||
// a legacy model whose config has no `chat_template` still gets one.
|
||||
if chat_template_arg.is_none()
|
||||
&& let Some(model_path) = model_path
|
||||
&& let Some(spec) = infer_legacy_template_from_model_path(model_path)
|
||||
{
|
||||
tracing::info!(%model_path, "inferred legacy chat template from model path");
|
||||
return Ok(ChatFormatter::Legacy(Box::new(LegacyFormatter { spec })));
|
||||
}
|
||||
|
||||
// Every remaining source builds the HF renderer around the tokenizer
|
||||
// config (the template itself, or the argument injected into it).
|
||||
let Some(config_file) = config_file else {
|
||||
return Err(TemplateError::MissingConfig);
|
||||
};
|
||||
|
||||
let config_path = Path::new(config_file);
|
||||
let config_text = read_to_string(config_path, "tokenizer config")?;
|
||||
let mut config = parse_json(&config_text, config_path, "tokenizer config")?;
|
||||
|
||||
let Some(argument) = chat_template_arg else {
|
||||
return formatter_from_config(&config);
|
||||
};
|
||||
|
||||
let path = Path::new(argument);
|
||||
if !path.exists() {
|
||||
return Err(TemplateError::NotFound {
|
||||
path: path.to_path_buf(),
|
||||
});
|
||||
}
|
||||
if !path.is_file() {
|
||||
return Err(TemplateError::NotFile {
|
||||
path: path.to_path_buf(),
|
||||
});
|
||||
}
|
||||
|
||||
if path.extension().and_then(|extension| extension.to_str()) == Some("jinja") {
|
||||
let template = read_to_string(path, "chat template")?;
|
||||
set_chat_template(
|
||||
&mut config,
|
||||
Value::String(template.trim_matches('\n').replace("\\n", "\n")),
|
||||
)?;
|
||||
return formatter_from_config(&config);
|
||||
}
|
||||
|
||||
let template_text = read_to_string(path, "chat template")?;
|
||||
let template = parse_json(&template_text, path, "chat template")?;
|
||||
|
||||
// HF-style JSON files may carry chat_template directly. Legacy SGLang
|
||||
// files carry Conversation fields and are translated below.
|
||||
if let Some(chat_template) = template.get("chat_template") {
|
||||
set_chat_template(&mut config, chat_template.clone())?;
|
||||
formatter_from_config(&config)
|
||||
} else {
|
||||
Ok(ChatFormatter::Legacy(Box::new(LegacyFormatter {
|
||||
spec: parse_legacy_template(&template, path)?,
|
||||
})))
|
||||
}
|
||||
}
|
||||
|
||||
/// Port of Python `get_conv_template_by_model_path` (conversation.py
|
||||
/// `matching_function_registry`, run in registration order): infer a legacy
|
||||
/// built-in template from the model path, optionally consulting the model's
|
||||
/// `config.json` `model_type`. `None` when nothing matches — the HF template
|
||||
/// is the fallback then, as in Python.
|
||||
pub(super) fn infer_legacy_template_from_model_path(model_path: &str) -> Option<LegacySpec> {
|
||||
let lower = model_path.to_lowercase();
|
||||
// Regexes without regex: every Python pattern here is a plain substring or
|
||||
// a `prefix.*suffix` pair, both on a lowercased path.
|
||||
let contains = |needle: &str| lower.contains(needle);
|
||||
let precedes = |prefix: &str, suffix: &str| {
|
||||
lower
|
||||
.find(prefix)
|
||||
.is_some_and(|start| lower[start + prefix.len()..].contains(suffix))
|
||||
};
|
||||
|
||||
if lower
|
||||
.split(|c: char| !c.is_alphanumeric())
|
||||
.any(|word| word == "points")
|
||||
{
|
||||
return builtin_template("points-v15-chat");
|
||||
}
|
||||
if precedes("moss", "vl") {
|
||||
return builtin_template("moss-vl");
|
||||
}
|
||||
if contains("internvl") {
|
||||
return builtin_template("internvl-2-5");
|
||||
}
|
||||
if contains("janus") {
|
||||
return builtin_template("janus-pro");
|
||||
}
|
||||
if contains("vicuna") || contains("llava-v1.5") || contains("llava-next-video-7b") {
|
||||
return builtin_template("vicuna_v1.1");
|
||||
}
|
||||
if precedes("deepseek", "vl2") {
|
||||
return builtin_template("deepseek-vl2");
|
||||
}
|
||||
if contains("llava-v1.6-34b")
|
||||
|| contains("llava-v1.6-yi-34b")
|
||||
|| contains("llava-next-video-34b")
|
||||
|| contains("llava-onevision-qwen2")
|
||||
{
|
||||
return builtin_template("chatml-llava");
|
||||
}
|
||||
// MiniCPM: 4.6+ uses its own template and must not fall back to the
|
||||
// legacy conv template.
|
||||
if contains("minicpm-v-4.6")
|
||||
|| contains("minicpm-v-4_6")
|
||||
|| contains("minicpm-o-4.6")
|
||||
|| contains("minicpm-o-4_6")
|
||||
{
|
||||
return None;
|
||||
}
|
||||
if contains("minicpm-v") {
|
||||
return builtin_template("minicpmv");
|
||||
}
|
||||
if contains("minicpm-o") {
|
||||
return builtin_template("minicpmo");
|
||||
}
|
||||
if contains("phi-4-multimodal") {
|
||||
return builtin_template("phi-4-mm");
|
||||
}
|
||||
if contains("deepseek-ocr") {
|
||||
return builtin_template("deepseek-ocr");
|
||||
}
|
||||
if contains("unlimited") {
|
||||
return builtin_template("unlimited-ocr");
|
||||
}
|
||||
if contains("paddleocr") {
|
||||
return builtin_template("paddle-ocr");
|
||||
}
|
||||
if contains("whisper") {
|
||||
return builtin_template("whisper");
|
||||
}
|
||||
|
||||
// Model-type matchers read `<model_path>/config.json` (local dirs only —
|
||||
// Python's `get_model_type` cannot resolve HF repo ids either).
|
||||
let model_type = read_model_type(model_path)?;
|
||||
// Python `MODEL_TYPE_TO_TEMPLATE`; minicpmv4_6 is deliberately absent.
|
||||
let name = match model_type.as_str() {
|
||||
"moss_vl" => "moss-vl",
|
||||
"internvl_chat" => "internvl-2-5",
|
||||
"multi_modality" => "janus-pro",
|
||||
"deepseek_vl_v2" => "deepseek-vl2",
|
||||
"minicpmv" => "minicpmv",
|
||||
"minicpmo" => "minicpmo",
|
||||
"phi4mm" => "phi-4-mm",
|
||||
"deepseek-ocr" => "deepseek-ocr",
|
||||
"unlimited-ocr" => "unlimited-ocr",
|
||||
"paddleocr_vl" => "paddle-ocr",
|
||||
_ => return None,
|
||||
};
|
||||
builtin_template(name)
|
||||
}
|
||||
|
||||
/// Python `get_model_type`: the `model_type` field of the model's `config.json`.
|
||||
fn read_model_type(model_path: &str) -> Option<String> {
|
||||
let config_path = Path::new(model_path).join("config.json");
|
||||
if !config_path.is_file() {
|
||||
return None;
|
||||
}
|
||||
let config: Value = parse_json(
|
||||
&read_to_string(&config_path, "model config").ok()?,
|
||||
&config_path,
|
||||
"model config",
|
||||
)
|
||||
.ok()?;
|
||||
config.get("model_type")?.as_str().map(str::to_owned)
|
||||
}
|
||||
|
||||
fn read_to_string(path: &Path, kind: &'static str) -> Result<String, TemplateError> {
|
||||
std::fs::read_to_string(path).map_err(|source| TemplateError::Read {
|
||||
kind,
|
||||
path: path.to_path_buf(),
|
||||
source,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_json(text: &str, path: &Path, kind: &'static str) -> Result<Value, TemplateError> {
|
||||
serde_json::from_str(text).map_err(|source| TemplateError::Parse {
|
||||
kind,
|
||||
path: path.to_path_buf(),
|
||||
source,
|
||||
})
|
||||
}
|
||||
|
||||
fn set_chat_template(config: &mut Value, chat_template: Value) -> Result<(), TemplateError> {
|
||||
let Some(config) = config.as_object_mut() else {
|
||||
return Err(TemplateError::ConfigNotObject);
|
||||
};
|
||||
config.insert("chat_template".to_string(), chat_template);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn formatter_from_config(config: &Value) -> Result<ChatFormatter, TemplateError> {
|
||||
let template: ChatTemplate = serde_json::from_value(config.clone())
|
||||
.map_err(|source| TemplateError::Config { source })?;
|
||||
if template.chat_template.is_none() {
|
||||
return Err(TemplateError::Missing);
|
||||
}
|
||||
let formatter = PromptFormatter::from_parts(
|
||||
template,
|
||||
ContextMixins::new(&[PromptContextMixin::OaiChat]),
|
||||
true,
|
||||
)
|
||||
.map_err(|error| TemplateError::Renderer {
|
||||
message: error.to_string(),
|
||||
})?;
|
||||
Ok(ChatFormatter::HuggingFace(formatter))
|
||||
}
|
||||
|
||||
/// Port of Python `_load_json_chat_template`: fields mirror `Conversation`
|
||||
/// exactly (missing `sep2`/`image_token`/`audio_token` stay at Python defaults).
|
||||
fn parse_legacy_template(value: &Value, path: &Path) -> Result<LegacySpec, TemplateError> {
|
||||
let object = value
|
||||
.as_object()
|
||||
.ok_or_else(|| TemplateError::LegacyNotObject {
|
||||
path: path.to_path_buf(),
|
||||
})?;
|
||||
|
||||
let required_string = |name: &str| -> Result<String, TemplateError> {
|
||||
object
|
||||
.get(name)
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.ok_or_else(|| TemplateError::LegacyMissingField {
|
||||
path: path.to_path_buf(),
|
||||
field: name.to_string(),
|
||||
})
|
||||
};
|
||||
|
||||
let style = required_string("sep_style")?;
|
||||
if !SUPPORTED_STYLES.contains(&style.as_str()) {
|
||||
return Err(TemplateError::UnknownStyle {
|
||||
path: path.to_path_buf(),
|
||||
style,
|
||||
});
|
||||
}
|
||||
|
||||
// Python `Conversation.stop_str: str | list[str] | None` — the key is
|
||||
// required (`template["stop_str"]` raises KeyError when missing), but an
|
||||
// explicit `null` value means `None`.
|
||||
let stop_str = match object.get("stop_str") {
|
||||
Some(Value::String(value)) => Some(OneOrMany::One(value.clone())),
|
||||
Some(Value::Array(values)) => {
|
||||
let strings = values
|
||||
.iter()
|
||||
.map(Value::as_str)
|
||||
.collect::<Option<Vec<_>>>()
|
||||
.ok_or_else(|| TemplateError::LegacyMissingField {
|
||||
path: path.to_path_buf(),
|
||||
field: "stop_str".to_string(),
|
||||
})?;
|
||||
Some(OneOrMany::Many(
|
||||
strings.into_iter().map(str::to_owned).collect(),
|
||||
))
|
||||
}
|
||||
Some(Value::Null) => None,
|
||||
Some(_) => {
|
||||
return Err(TemplateError::LegacyMissingField {
|
||||
path: path.to_path_buf(),
|
||||
field: "stop_str".to_string(),
|
||||
});
|
||||
}
|
||||
None => {
|
||||
return Err(TemplateError::LegacyMissingField {
|
||||
path: path.to_path_buf(),
|
||||
field: "stop_str".to_string(),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
Ok(LegacySpec {
|
||||
name: required_string("name")?,
|
||||
// Python: `system_template=template["system"] + "\n{system_message}"`.
|
||||
system_template: format!("{}\n{{system_message}}", required_string("system")?),
|
||||
system_message: object
|
||||
.get("system_message")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
roles: (required_string("user")?, required_string("assistant")?),
|
||||
style,
|
||||
sep: object
|
||||
.get("sep")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("\n")
|
||||
.to_string(),
|
||||
sep2: None,
|
||||
stop_str,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
@@ -34,31 +34,34 @@ fn value_error(context: &str, err: impl std::fmt::Display) -> PyErr {
|
||||
pyo3::exceptions::PyValueError::new_err(format!("{context}: {err}"))
|
||||
}
|
||||
|
||||
/// One drained MM result (see [`Server::take_mm`]), consumed by
|
||||
/// `RustServer.build_native_mm` to build the scheduler's
|
||||
/// One drained MM result (see [`Server::take_mm_result`]), consumed by
|
||||
/// `RustMmProcessor.build_output` to build the scheduler's
|
||||
/// `MultimodalProcessorOutput`.
|
||||
#[pyclass(frozen, get_all)]
|
||||
struct MmEncodeResult {
|
||||
/// *Generic.* All items' `pixel_values` concatenated, flat `f32` of logical
|
||||
/// shape `[sum(t*h*w), feature_dim]`; `Some` on the inline (single-rank) path.
|
||||
// General fields.
|
||||
/// All items' `pixel_values` concatenated as flat `f32` with logical shape
|
||||
/// `[sum(t*h*w), feature_dim]`; present on the inline (single-rank) path.
|
||||
features: Option<Py<numpy::PyArray1<f32>>>,
|
||||
/// *Generic.* Per-item POSIX shm segment name holding that item's features
|
||||
/// (`[t*h*w, feature_dim]` f32); `Some` on the TP-broadcast path.
|
||||
/// Per-item POSIX shared-memory segment holding `[t*h*w, feature_dim]` f32
|
||||
/// features; present on the TP-broadcast path.
|
||||
shm_names: Option<Vec<String>>,
|
||||
/// *Generic.* Per-item content hash of the raw source bytes (or the caller's
|
||||
/// `mm_hashes` override), precomputed so the drain never re-hashes.
|
||||
/// Per-item content hash of the raw source bytes, or the caller-provided
|
||||
/// `mm_hashes` override, precomputed so draining never re-hashes.
|
||||
hashes: Vec<u64>,
|
||||
/// *Generic.* Per-item inclusive `(start, end)` placeholder-token span in the
|
||||
/// expanded `input_ids`.
|
||||
/// Per-item inclusive `(start, end)` placeholder-token span in the expanded
|
||||
/// `input_ids`.
|
||||
offsets: Vec<(u32, u32)>,
|
||||
/// *Qwen-VL specific.* Per-item `image_grid_thw` `(t, h, w)` in patch units;
|
||||
/// `t*h*w` is also the item's row count in `features`.
|
||||
|
||||
// Qwen-VL-specific fields.
|
||||
/// Per-item `image_grid_thw` `(t, h, w)` in patch units; `t*h*w` is also the
|
||||
/// item's row count in `features`.
|
||||
grids: Vec<(u32, u32, u32)>,
|
||||
/// *Qwen-VL specific.* M-RoPE position ids, flat `i64` of row-major shape
|
||||
/// `[3, seq_len]` (temporal, height, width rows).
|
||||
/// M-RoPE position ids as flat `i64` with row-major shape `[3, seq_len]`
|
||||
/// (temporal, height, and width rows).
|
||||
mrope: Py<numpy::PyArray1<i64>>,
|
||||
/// *Qwen-VL specific.* M-RoPE delta, `max(mrope) + 1 - seq_len`, that decode
|
||||
/// adds to the plain sequence position.
|
||||
/// M-RoPE delta, `max(mrope) + 1 - seq_len`, added to the plain sequence
|
||||
/// position during decoding.
|
||||
mrope_delta: i64,
|
||||
}
|
||||
|
||||
@@ -93,7 +96,7 @@ impl Server {
|
||||
http_addr = None,
|
||||
to_scheduler_cap = 8192,
|
||||
from_scheduler_cap = 8192,
|
||||
channel_cap = 8192,
|
||||
stage_channel_cap = 8192,
|
||||
cores = None,
|
||||
))]
|
||||
// pyo3 `#[new]` constructor: the wide arg list is the Python-facing boot
|
||||
@@ -104,7 +107,7 @@ impl Server {
|
||||
http_addr: Option<String>,
|
||||
to_scheduler_cap: usize,
|
||||
from_scheduler_cap: usize,
|
||||
channel_cap: usize,
|
||||
stage_channel_cap: usize,
|
||||
cores: Option<Vec<usize>>,
|
||||
) -> PyResult<Self> {
|
||||
// `server_args` already arrived typed (pyo3 rejected any missing/extra/
|
||||
@@ -127,7 +130,7 @@ impl Server {
|
||||
http_api_worker_num: server_args.http_api_worker_num(),
|
||||
to_scheduler_cap,
|
||||
from_scheduler_cap,
|
||||
channel_cap,
|
||||
stage_channel_cap,
|
||||
cores,
|
||||
},
|
||||
server_args: std::sync::Arc::new(server_args),
|
||||
@@ -201,10 +204,10 @@ impl Server {
|
||||
}
|
||||
|
||||
/// Spawn the MM worker pool for the pipeline in `spec` (built from the
|
||||
/// resolved processor config; see `NativeMmHost.resolve_native_spec` and
|
||||
/// resolved processor config; see `RustMmProcessor.resolve_spec` and
|
||||
/// `RustServer._build_mm_spec`). Image-only requests are processed entirely
|
||||
/// in Rust and parked for [`Server::take_mm`]; anything the pipeline cannot
|
||||
/// serve is rejected back to the client — there is no Python fallback.
|
||||
/// in Rust and parked for [`Server::take_mm_result`]; anything the pipeline
|
||||
/// cannot serve is rejected back to the client — there is no Python fallback.
|
||||
fn start_mm_workers(&self, spec: MmSpec, workers: usize) -> PyResult<()> {
|
||||
let ctx = multi_modality::worker::Context::new(
|
||||
spec,
|
||||
@@ -224,7 +227,7 @@ impl Server {
|
||||
/// Runs on the scheduler loop between decode steps, so any per-byte work
|
||||
/// here — memcpy or hashing, tens of MB per image-heavy request — would
|
||||
/// stall every running request's ITL. Hence the worker-precomputed `hashes`.
|
||||
fn take_mm(&self, py: Python<'_>, rid: &str) -> Option<MmEncodeResult> {
|
||||
fn take_mm_result(&self, py: Python<'_>, rid: &str) -> Option<MmEncodeResult> {
|
||||
use numpy::IntoPyArray;
|
||||
|
||||
let res = self.rt.mm_sidecar.take(rid)?;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! Runtime configuration: the rust-server boot knobs
|
||||
//! ([`RustServerServerArgs`]), the scheduler's typed `server_args` handoff
|
||||
//! ([`ServerArgs`] / [`ModelConfig`]), the [`RuntimeConfig`] pairing them for
|
||||
//! `runtime::start`, and the native MM pipeline handoff ([`MmSpec`]).
|
||||
//! `runtime::start`, and the Rust MM pipeline handoff ([`MmSpec`]).
|
||||
//!
|
||||
//! [`ServerArgs`] / [`ModelConfig`] / [`DefaultSamplingParams`] /
|
||||
//! [`DisaggregationMode`] / [`MmSpec`] / [`MmFamily`] / [`MmResample`] are
|
||||
@@ -29,7 +29,7 @@ pub struct RustServerServerArgs {
|
||||
pub http_api_worker_num: usize,
|
||||
pub to_scheduler_cap: usize,
|
||||
pub from_scheduler_cap: usize,
|
||||
pub channel_cap: usize,
|
||||
pub stage_channel_cap: usize,
|
||||
/// CPU core ids the pools pin to (e.g. this rank's NUMA-local cores minus
|
||||
/// the scheduler's reserved launch cores). `None` → run unpinned.
|
||||
pub cores: Option<Vec<usize>>,
|
||||
@@ -42,7 +42,7 @@ impl Default for RustServerServerArgs {
|
||||
http_api_worker_num: 2,
|
||||
to_scheduler_cap: 8192,
|
||||
from_scheduler_cap: 8192,
|
||||
channel_cap: 8192,
|
||||
stage_channel_cap: 8192,
|
||||
cores: None,
|
||||
}
|
||||
}
|
||||
@@ -404,15 +404,15 @@ impl DefaultSamplingParams {
|
||||
}
|
||||
}
|
||||
|
||||
/// The native MM pipeline handoff, built by `RustServer._build_mm_spec` from
|
||||
/// the resolved `NativeMmSpec` and passed to `Server.start_mm_workers`. Same
|
||||
/// The Rust MM pipeline handoff, built by `RustServer._build_mm_spec` from
|
||||
/// the resolved `RustMmSpec` and passed to `Server.start_mm_workers`. Same
|
||||
/// contract as [`ServerArgs`]: every field is a required, typed constructor
|
||||
/// keyword, so a drifted Python caller fails at boot.
|
||||
#[pyo3::pyclass(frozen, from_py_object, module = "sglang.srt.rust_extensions._server")]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct MmSpec {
|
||||
/// Park feature buffers in POSIX shm rather than inline. Set by the Python
|
||||
/// launcher (`NativeMmHost._use_feature_shm`) exactly when the scheduler
|
||||
/// launcher (`RustMmProcessor._use_feature_shm`) exactly when the scheduler
|
||||
/// broadcasts across TP ranks and will unwrap `ShmPointerMMData`.
|
||||
pub feature_shm: bool,
|
||||
/// The family pipeline and its resolved processor parameters.
|
||||
@@ -475,7 +475,7 @@ impl MmSpec {
|
||||
|
||||
/// Which `sglang_mm` family pipeline serves the model — one variant per
|
||||
/// [`sglang_mm::registry::PipelineSpec`] arm. Exposed to Python as an enum
|
||||
/// (`MmFamily.QwenVl`); `NativeMmFamily.name` maps onto it at handoff.
|
||||
/// (`MmFamily.QwenVl`); `RustMmFamily.name` maps onto it at handoff.
|
||||
#[pyo3::pyclass(
|
||||
eq,
|
||||
frozen,
|
||||
@@ -487,9 +487,9 @@ pub enum MmFamily {
|
||||
QwenVl,
|
||||
}
|
||||
|
||||
/// The HF image processor the native resize must reproduce bit-exactly (see
|
||||
/// The HF image processor the Rust resize must reproduce bit-exactly (see
|
||||
/// [`sglang_mm::qwen_vl::Resampler`]). Exposed to Python as an enum
|
||||
/// (`MmResample.AtenU8` / `.Pil`); `NativeMmFamily.image_processors` maps each
|
||||
/// (`MmResample.AtenU8` / `.Pil`); `RustMmFamily.image_processors` maps each
|
||||
/// processor class onto it.
|
||||
#[pyo3::pyclass(
|
||||
eq,
|
||||
|
||||
@@ -72,7 +72,7 @@ fn max_new_tokens_default() -> Option<i64> {
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct SamplingParams {
|
||||
// --- API parameters (set by callers) ---
|
||||
// Output length and stopping.
|
||||
#[serde(default = "max_new_tokens_default")]
|
||||
pub max_new_tokens: Option<i64>,
|
||||
/// API input alias, copied to `stop_strs` then cleared by `normalize`.
|
||||
@@ -85,6 +85,8 @@ pub struct SamplingParams {
|
||||
/// API input alias, copied to `stop_regex_strs` then cleared by `normalize`.
|
||||
#[serde(default)]
|
||||
pub stop_regex: Option<OneOrMany<String>>,
|
||||
|
||||
// Sampling distribution and penalties.
|
||||
#[serde(
|
||||
default = "f64_one::default",
|
||||
deserialize_with = "f64_one::deserialize"
|
||||
@@ -125,6 +127,8 @@ pub struct SamplingParams {
|
||||
deserialize_with = "i64_zero::deserialize"
|
||||
)]
|
||||
pub min_new_tokens: i64,
|
||||
|
||||
// Sequence count and beam search.
|
||||
#[serde(
|
||||
default = "i64_one::default",
|
||||
deserialize_with = "i64_one::deserialize"
|
||||
@@ -134,6 +138,8 @@ pub struct SamplingParams {
|
||||
/// positional wire layout even though the rust path rejects it below.
|
||||
#[serde(default)]
|
||||
pub beam_width: Option<i64>,
|
||||
|
||||
// Structured-output constraints.
|
||||
#[serde(default)]
|
||||
pub json_schema: Option<String>,
|
||||
#[serde(default)]
|
||||
@@ -142,6 +148,8 @@ pub struct SamplingParams {
|
||||
pub ebnf: Option<String>,
|
||||
#[serde(default)]
|
||||
pub structural_tag: Option<String>,
|
||||
|
||||
// Output handling.
|
||||
#[serde(
|
||||
default = "bool_false::default",
|
||||
deserialize_with = "bool_false::deserialize"
|
||||
@@ -164,6 +172,8 @@ pub struct SamplingParams {
|
||||
pub no_stop_trim: bool,
|
||||
#[serde(default)]
|
||||
pub stream_interval: Option<i64>,
|
||||
|
||||
// Logit processing and reproducibility.
|
||||
/// Token id (as a string key, matching Python) → bias. Keys are vocab-bounded
|
||||
/// by [`verify`](Self::verify).
|
||||
#[serde(default)]
|
||||
@@ -175,7 +185,7 @@ pub struct SamplingParams {
|
||||
#[serde(default)]
|
||||
pub custom_params: Option<serde_json::Value>,
|
||||
|
||||
// --- Internal fields (populated by the pipeline below, not API-facing) ---
|
||||
// Normalized internal fields.
|
||||
//
|
||||
// All `skip_deserializing`: they are outputs of `normalize`, and a client that
|
||||
// could set them would be setting the pipeline's own state. `is_normalized` is
|
||||
|
||||
@@ -52,7 +52,7 @@ pub struct Context {
|
||||
pub tokenizer: Option<Arc<dyn TextTokenizer>>,
|
||||
pub sidecar: Sidecar,
|
||||
/// Park feature buffers in POSIX shm. Set by the Python launcher
|
||||
/// (`NativeMmHost._use_feature_shm`) exactly when the scheduler broadcasts
|
||||
/// (`RustMmProcessor._use_feature_shm`) exactly when the scheduler broadcasts
|
||||
/// across TP ranks and will unwrap `ShmPointerMMData`.
|
||||
pub feature_shm: bool,
|
||||
}
|
||||
|
||||
@@ -4,5 +4,7 @@ pub mod channel;
|
||||
pub mod detokenizer;
|
||||
pub mod from_scheduler;
|
||||
pub mod to_scheduler;
|
||||
mod to_scheduler_types;
|
||||
mod to_scheduler_validation;
|
||||
pub mod tokenizer;
|
||||
pub mod wiring;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,892 @@
|
||||
//! Tests for scheduler intake.
|
||||
|
||||
use super::*;
|
||||
use crate::message::request::GenerateRequest;
|
||||
use crate::message::response::ResponseSink;
|
||||
use crate::message::sampling::SamplingParams;
|
||||
use crate::tokenizer_manager::channel::{ToSchedulerRx, to_scheduler};
|
||||
use crate::utils::fsm::RequestState;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// An `Intake` plus its detok-shard receiver, to_scheduler channel consumer (keep alive —
|
||||
/// dropping it closes the channel → false QueueFull), tm inbox sender, and the
|
||||
/// mm-pool receiver (keep alive — dropping it makes mm submits fail).
|
||||
fn make_intake() -> (
|
||||
Intake,
|
||||
flume::Receiver<DetokMsg>,
|
||||
ToSchedulerRx,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
make_intake_with(test_limits())
|
||||
}
|
||||
|
||||
fn make_intake_with_abort(
|
||||
abort_rx: flume::Receiver<AbortSource>,
|
||||
) -> (
|
||||
Intake,
|
||||
flume::Receiver<DetokMsg>,
|
||||
ToSchedulerRx,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
make_intake_inner(test_limits(), abort_rx)
|
||||
}
|
||||
|
||||
fn make_intake_with(
|
||||
limits: Limits,
|
||||
) -> (
|
||||
Intake,
|
||||
flume::Receiver<DetokMsg>,
|
||||
ToSchedulerRx,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx
|
||||
make_intake_inner(limits, abort_rx)
|
||||
}
|
||||
|
||||
fn make_intake_inner(
|
||||
limits: Limits,
|
||||
abort_rx: flume::Receiver<AbortSource>,
|
||||
) -> (
|
||||
Intake,
|
||||
flume::Receiver<DetokMsg>,
|
||||
ToSchedulerRx,
|
||||
flume::Sender<TmEvent>,
|
||||
flume::Receiver<MmRequest>,
|
||||
) {
|
||||
let (tok_tx, _tok_rx) = flume::unbounded();
|
||||
let (detok_tx, detok_rx) = flume::unbounded();
|
||||
let senders = Senders {
|
||||
tok_manager_tx: flume::unbounded().0,
|
||||
abort_tx: flume::unbounded().0,
|
||||
tokenizer_tx: tok_tx,
|
||||
detokenizer_tx: vec![detok_tx],
|
||||
};
|
||||
let (to_scheduler_tx, consumer) = to_scheduler(16);
|
||||
let (tm_tx, tm_rx) = flume::unbounded();
|
||||
let (mm_tx, mm_rx) = flume::unbounded();
|
||||
// Keep the shutdown sender alive (leak) so its branch never fires — tests
|
||||
// end `run` by dropping `tm_tx`, not by shutdown.
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let intake = Intake::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
to_scheduler_tx,
|
||||
limits,
|
||||
test_mm(mm_tx, true),
|
||||
sd_rx,
|
||||
);
|
||||
(intake, detok_rx, consumer, tm_tx, mm_rx)
|
||||
}
|
||||
|
||||
/// An [`Mm`] over `tx` with a fresh sidecar.
|
||||
fn test_mm(tx: flume::Sender<MmRequest>, enabled: bool) -> Mm {
|
||||
Mm {
|
||||
enabled,
|
||||
tx,
|
||||
sidecar: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Both abort sources do the same two things: drop the detok entry so no
|
||||
/// further chunk can be delivered, and tell the scheduler to stop generating.
|
||||
///
|
||||
/// Neither releases anything, and nothing needs them to. Release ordering used
|
||||
/// to be the delicate part here — `AbortGuard::drop` releasing a rid right
|
||||
/// after enqueuing the abort ordered the SEND, not the EFFECT, so a retry of
|
||||
/// the same rid could `Register` ahead of the stale abort and be torn down by
|
||||
/// it. `Rid::from_client` removes the premise: a retry carries a different
|
||||
/// `Rid`, so no abort in flight can name it.
|
||||
#[test]
|
||||
fn every_abort_source_deregisters_and_stops_the_scheduler() {
|
||||
for source in [
|
||||
AbortSource::Guard("x".into()),
|
||||
AbortSource::Detok("x".into()),
|
||||
] {
|
||||
let (detok_tx, detok_rx) = flume::unbounded::<DetokMsg>();
|
||||
let (to_scheduler_tx, consumer) = to_scheduler(16);
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let mut intake = Intake::new(
|
||||
flume::unbounded().1,
|
||||
flume::unbounded().1,
|
||||
Senders {
|
||||
tok_manager_tx: flume::unbounded().0,
|
||||
abort_tx: flume::unbounded().0,
|
||||
tokenizer_tx: flume::unbounded().0,
|
||||
detokenizer_tx: vec![detok_tx],
|
||||
},
|
||||
to_scheduler_tx,
|
||||
test_limits(),
|
||||
test_mm(flume::unbounded().0, true),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
intake.on_abort(source.clone());
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "x"),
|
||||
"{source:?} must drop the detok entry",
|
||||
);
|
||||
assert_eq!(
|
||||
consumer.drain(8).headers.len(),
|
||||
1,
|
||||
"{source:?} must push an AbortReq so the scheduler stops",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// A context ceiling high enough that only a test which sets one on purpose
|
||||
/// can reach it. `context_len` is mandatory now, so "no ceiling" has to be a
|
||||
/// large number rather than `None`; kept well below `u64::MAX` so the
|
||||
/// `as i64` in the auto-truncate clamp cannot go negative if a future test
|
||||
/// does reach this path.
|
||||
const NO_CONTEXT_CEILING: u64 = 1 << 40;
|
||||
|
||||
/// The default test limits: a real tokenizer, vocab 1000, no context ceiling.
|
||||
/// Spelled out rather than `..Default::default()` — `Limits` deliberately has
|
||||
/// no `Default`, because a zero `vocab_size`/`context_len` would reject every
|
||||
/// request instead of behaving like "unset".
|
||||
fn test_limits() -> Limits {
|
||||
Limits {
|
||||
skip_tokenizer_init: false,
|
||||
vocab_size: 1000,
|
||||
context_len: NO_CONTEXT_CEILING,
|
||||
num_reserved_tokens: 0,
|
||||
allow_auto_truncate: false,
|
||||
enable_return_hidden_states: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn generate_req(id: u64, sampling_params: SamplingParams) -> Request {
|
||||
let (tx, _rx) = mpsc::channel(8);
|
||||
Request {
|
||||
rid: id.to_string().into(),
|
||||
state: RequestState::Received,
|
||||
sink: ResponseSink::Local(tx),
|
||||
kind: RequestKind::Generate(Box::new(GenerateRequest {
|
||||
rid: id.to_string().into(),
|
||||
input_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params,
|
||||
..Default::default()
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// `input + max_new_tokens` past the context window is an actionable 400, not a
|
||||
/// silently truncated 200 (Python `TokenizerManager._validate_one_request`).
|
||||
/// The message names both halves so the client can fix the right one.
|
||||
#[test]
|
||||
fn total_tokens_over_context_is_rejected() {
|
||||
let limits = Limits {
|
||||
context_len: 10,
|
||||
..test_limits()
|
||||
};
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens: Some(100),
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let err = check_total_tokens(&mut g, &limits).unwrap_err();
|
||||
let msg = err.to_string();
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(msg.contains("total of 103 tokens"), "{msg}");
|
||||
assert!(msg.contains("3 tokens from the input"), "{msg}");
|
||||
assert!(msg.contains("100 tokens for the completion"), "{msg}");
|
||||
// Exactly filling the window is allowed (Python compares with `>`).
|
||||
g.sampling_params.max_new_tokens = Some(7);
|
||||
assert!(check_total_tokens(&mut g, &limits).is_ok());
|
||||
assert_eq!(g.sampling_params.max_new_tokens, Some(7), "left alone");
|
||||
}
|
||||
|
||||
/// The reserved slots (eagle draft tokens) count as input, so a request can be
|
||||
/// rejected for them even when the prompt alone would fit.
|
||||
#[test]
|
||||
fn reserved_tokens_count_toward_the_limit() {
|
||||
let limits = Limits {
|
||||
context_len: 10,
|
||||
num_reserved_tokens: 5,
|
||||
..test_limits()
|
||||
};
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens: Some(3), // 3 + 3 fits, but 3 + 5 + 3 does not
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let msg = check_total_tokens(&mut g, &limits).unwrap_err().to_string();
|
||||
assert!(msg.contains("8 tokens from the input"), "{msg}");
|
||||
}
|
||||
|
||||
/// `--allow-auto-truncate` opts into clamping instead of rejecting; with no
|
||||
/// context length, or no `max_new_tokens` cap, there is nothing to check.
|
||||
#[test]
|
||||
fn auto_truncate_clamps_and_unknowns_skip() {
|
||||
let sp = |max_new_tokens| SamplingParams {
|
||||
max_new_tokens,
|
||||
..Default::default()
|
||||
};
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: sp(Some(100)),
|
||||
..Default::default()
|
||||
};
|
||||
let truncating = Limits {
|
||||
context_len: 10,
|
||||
allow_auto_truncate: true,
|
||||
..test_limits()
|
||||
};
|
||||
assert!(check_total_tokens(&mut g, &truncating).is_ok());
|
||||
assert_eq!(g.sampling_params.max_new_tokens, Some(7), "clamped to fit");
|
||||
|
||||
// Unknown context length → no ceiling to enforce.
|
||||
g.sampling_params = sp(Some(100));
|
||||
assert!(check_total_tokens(&mut g, &test_limits()).is_ok());
|
||||
assert_eq!(g.sampling_params.max_new_tokens, Some(100), "untouched");
|
||||
|
||||
// No cap requested → nothing to add to the input length, but the input
|
||||
// itself is still checked (see `input_length_is_checked_unconditionally`).
|
||||
g.sampling_params = sp(None);
|
||||
let roomy = Limits {
|
||||
context_len: 100,
|
||||
..test_limits()
|
||||
};
|
||||
assert!(check_total_tokens(&mut g, &roomy).is_ok());
|
||||
}
|
||||
|
||||
/// `max_new_tokens: null` means "no cap", NOT "skip the checks" — the input
|
||||
/// alone must still fit. Gating the whole function on `max_new_tokens` let an
|
||||
/// over-long prompt through to the scheduler with no error at all.
|
||||
/// Python compares with `>=`: a prompt that exactly fills the window leaves no
|
||||
/// room to generate.
|
||||
#[test]
|
||||
fn input_length_is_checked_unconditionally() {
|
||||
let limits = Limits {
|
||||
context_len: 3,
|
||||
..test_limits()
|
||||
};
|
||||
let req = |max_new_tokens| GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]), // exactly fills a 3-token window
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens,
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
for max_new_tokens in [None, Some(1)] {
|
||||
let err = check_total_tokens(&mut req(max_new_tokens), &limits)
|
||||
.expect_err("input == context_len must be rejected (Python uses >=)");
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("longer than the model's context"));
|
||||
}
|
||||
// One token shorter fits, with or without a cap.
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2]),
|
||||
..Default::default()
|
||||
};
|
||||
g.sampling_params.max_new_tokens = None;
|
||||
assert!(check_total_tokens(&mut g, &limits).is_ok());
|
||||
|
||||
// Under auto-truncate the input is cut to fit instead of rejected.
|
||||
let truncating = Limits {
|
||||
allow_auto_truncate: true,
|
||||
..limits.clone()
|
||||
};
|
||||
let mut g = req(None);
|
||||
assert!(check_total_tokens(&mut g, &truncating).is_ok());
|
||||
assert_eq!(
|
||||
g.input_ids.as_deref(),
|
||||
Some(&[1, 2, 3][..]),
|
||||
"fits at the cap"
|
||||
);
|
||||
}
|
||||
|
||||
/// The clamp runs AFTER `verify` (which happens in `Normalizing`), so lowering
|
||||
/// `max_new_tokens` can leave `min_new_tokens > max_new_tokens`. Nothing
|
||||
/// downstream re-checks — `is_normalized: true` makes the scheduler's own
|
||||
/// verify early-return — so the clamp has to re-assert it here.
|
||||
#[test]
|
||||
fn auto_truncate_cannot_invert_min_and_max_new_tokens() {
|
||||
let limits = Limits {
|
||||
context_len: 10,
|
||||
allow_auto_truncate: true,
|
||||
..test_limits()
|
||||
};
|
||||
let mut g = GenerateRequest {
|
||||
input_ids: Some(vec![1, 2, 3]), // clamps max_new_tokens to 7
|
||||
sampling_params: SamplingParams {
|
||||
max_new_tokens: Some(100),
|
||||
min_new_tokens: 50, // …which is below min_new_tokens
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let err = check_total_tokens(&mut g, &limits)
|
||||
.expect_err("a clamp that inverts min/max must 400, not ride the wire");
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("min_new_tokens"), "{err}");
|
||||
|
||||
// A clamp that keeps the invariant still clamps.
|
||||
g.sampling_params.min_new_tokens = 2;
|
||||
g.sampling_params.max_new_tokens = Some(100);
|
||||
assert!(check_total_tokens(&mut g, &limits).is_ok());
|
||||
assert_eq!(g.sampling_params.max_new_tokens, Some(7));
|
||||
}
|
||||
|
||||
/// `return_hidden_states` on a server not launched for it is a 400: the
|
||||
/// scheduler never computes them, so the request would otherwise 200 with
|
||||
/// `meta_info.hidden_states` silently missing.
|
||||
#[test]
|
||||
fn hidden_states_gated_on_server_support() {
|
||||
let req = |want| {
|
||||
let mut r = generate_req(31, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut r.kind {
|
||||
g.return_hidden_states = want;
|
||||
}
|
||||
r
|
||||
};
|
||||
let disabled = test_limits();
|
||||
let err = validate(&mut req(true), &disabled).unwrap_err();
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(
|
||||
err.to_string().contains("--enable-return-hidden-states"),
|
||||
"message must name the flag: {err}"
|
||||
);
|
||||
// Not asking for them (the client sent `false`, or sent nothing and
|
||||
// `into_requests` resolved the default), or asking on a server that
|
||||
// supports them, is fine.
|
||||
assert!(validate(&mut req(false), &disabled).is_ok());
|
||||
let enabled = Limits {
|
||||
enable_return_hidden_states: true,
|
||||
..test_limits()
|
||||
};
|
||||
assert!(validate(&mut req(true), &enabled).is_ok());
|
||||
}
|
||||
|
||||
/// End-to-end through `drive`: an over-context request is rejected on the way
|
||||
/// to the ring, after registration — so it must be deregistered, not leaked.
|
||||
#[test]
|
||||
fn over_context_request_deregisters_and_never_reaches_the_ring() {
|
||||
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake_with(Limits {
|
||||
context_len: 4,
|
||||
..test_limits()
|
||||
});
|
||||
intake.drive(generate_req(
|
||||
33,
|
||||
SamplingParams {
|
||||
max_new_tokens: Some(64),
|
||||
..Default::default()
|
||||
},
|
||||
));
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "33"),
|
||||
"registered before the check",
|
||||
);
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "33"),
|
||||
"must deregister on reject",
|
||||
);
|
||||
assert!(
|
||||
consumer.drain(16).headers.is_empty(),
|
||||
"must not reach the scheduler"
|
||||
);
|
||||
}
|
||||
|
||||
/// A `Detokenize` request terminates at the detok stage, and the shard must
|
||||
/// see its `Register` BEFORE its `Decode` — the shard delivers the result
|
||||
/// through the sink registered under that rid, so a `Decode` that arrives
|
||||
/// unregistered is silently dropped and the caller waits forever. Both
|
||||
/// messages ride one channel from this one thread, which is the FIFO this
|
||||
/// pins. Nothing may reach the scheduler ring.
|
||||
#[test]
|
||||
fn detokenize_flows_register_then_decode_and_skips_the_ring() {
|
||||
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let (tx, mut rx) = mpsc::channel(8);
|
||||
intake.drive(Request {
|
||||
rid: "41".into(),
|
||||
state: RequestState::Received,
|
||||
sink: ResponseSink::Local(tx),
|
||||
kind: RequestKind::Detokenize {
|
||||
token_ids: vec![7, 8, 9],
|
||||
},
|
||||
});
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "41"),
|
||||
"the sink must be registered before the decode job",
|
||||
);
|
||||
assert!(
|
||||
matches!(
|
||||
detok_rx.try_recv(),
|
||||
Ok(DetokMsg::Decode { rid, token_ids })
|
||||
if rid.as_str() == "41" && token_ids == [7, 8, 9]
|
||||
),
|
||||
"the decode job follows, ids intact",
|
||||
);
|
||||
assert!(
|
||||
consumer.drain(16).headers.is_empty(),
|
||||
"must never reach the scheduler"
|
||||
);
|
||||
assert!(
|
||||
rx.try_recv().is_err(),
|
||||
"no response until the shard answers"
|
||||
);
|
||||
}
|
||||
|
||||
/// Negative ids cannot decode (the shard's domain is `&[u32]`): rejected by
|
||||
/// `validate` at `Received` — an `Error` to the sink, and the shard sees
|
||||
/// NOTHING (validation runs before registration, so there is no entry to
|
||||
/// leak and no decode job to drop).
|
||||
#[test]
|
||||
fn detokenize_negative_ids_reject_before_registration() {
|
||||
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let (tx, mut rx) = mpsc::channel(8);
|
||||
intake.drive(Request {
|
||||
rid: "43".into(),
|
||||
state: RequestState::Received,
|
||||
sink: ResponseSink::Local(tx),
|
||||
kind: RequestKind::Detokenize {
|
||||
token_ids: vec![1, -1],
|
||||
},
|
||||
});
|
||||
let Ok(ResponseItem::Error(err)) = rx.try_recv() else {
|
||||
panic!("sink must receive the validation error");
|
||||
};
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("out of range"), "{err}");
|
||||
assert!(detok_rx.try_recv().is_err(), "shard never hears of it");
|
||||
assert!(consumer.drain(16).headers.is_empty());
|
||||
}
|
||||
|
||||
/// A dropped ring push is survivable, and this pins WHY. The ring is bounded,
|
||||
/// so under load the scheduler never learns to stop and keeps generating; its
|
||||
/// chunks then arrive for a rid the detok table no longer holds and are
|
||||
/// dropped. That wastes GPU work but cannot MISDELIVER, because
|
||||
/// `Rid::from_client` guarantees no later request ever answers to that rid.
|
||||
/// The detok entry is dropped either way — that is the half that must not
|
||||
/// depend on the ring.
|
||||
///
|
||||
/// Ring capacity 1: the first abort pushes, the second finds it full.
|
||||
#[test]
|
||||
fn abort_deregisters_even_when_the_ring_push_is_dropped() {
|
||||
let (tok_tx, _tok_rx) = flume::unbounded();
|
||||
let (detok_tx, detok_rx) = flume::unbounded();
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
let senders = Senders {
|
||||
tok_manager_tx: flume::unbounded().0,
|
||||
abort_tx,
|
||||
tokenizer_tx: tok_tx,
|
||||
detokenizer_tx: vec![detok_tx],
|
||||
};
|
||||
let (producer, _consumer) = to_scheduler(1);
|
||||
let (_tm_tx, tm_rx) = flume::unbounded();
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let mut intake = Intake::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
producer,
|
||||
test_limits(),
|
||||
test_mm(flume::unbounded().0, true),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
intake.on_abort(AbortSource::Guard("pushed".into()));
|
||||
intake.on_abort(AbortSource::Guard("dropped".into()));
|
||||
|
||||
// Both deregisters land regardless of whether the ring accepted the push.
|
||||
for expected in ["pushed", "dropped"] {
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == expected),
|
||||
"{expected}: the detok entry must be dropped even when the ring is full",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// The rid keys the detok table and rides on every chunk of every decode step,
|
||||
/// so an unbounded client-supplied one is a recurring cost, not a one-off.
|
||||
#[test]
|
||||
fn oversized_rid_is_rejected() {
|
||||
let mut req = generate_req(51, SamplingParams::default());
|
||||
req.rid = "x".repeat(MAX_RID_LEN + 1).into();
|
||||
let err = validate(&mut req, &test_limits()).expect_err("must be rejected");
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("over the"), "{err}");
|
||||
|
||||
// A uuid-sized rid — what Python mints — is nowhere near the cap.
|
||||
let mut req = generate_req(52, SamplingParams::default());
|
||||
req.rid = "0123456789abcdef0123456789abcdef".into();
|
||||
assert!(validate(&mut req, &test_limits()).is_ok());
|
||||
}
|
||||
|
||||
/// A request rejected BEFORE `register_detok` must not send `Deregister`: the
|
||||
/// handler is a bare `table.remove(&rid)`, so it would evict whatever entry
|
||||
/// holds that key — a concurrent request's sink — leaving that client hung with
|
||||
/// no terminal frame. Python validates before it inserts, so it cannot hit this.
|
||||
#[test]
|
||||
fn pre_registration_failure_does_not_deregister() {
|
||||
// Rejected inside `validate` (out-of-vocab id), which runs before registration.
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(41, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![2_000_000_000]);
|
||||
}
|
||||
intake.drive(req);
|
||||
assert!(
|
||||
detok_rx.try_recv().is_err(),
|
||||
"a pre-registration reject must send NOTHING to the shard — a Deregister \
|
||||
here removes a live request's sink"
|
||||
);
|
||||
|
||||
// A post-registration reject still deregisters (the leak fix stays fixed).
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
intake.drive(generate_req(
|
||||
42,
|
||||
SamplingParams {
|
||||
top_p: 2.0, // rejected by `normalize`, after registration
|
||||
..Default::default()
|
||||
},
|
||||
));
|
||||
assert!(matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. })));
|
||||
assert!(matches!(
|
||||
detok_rx.try_recv(),
|
||||
Ok(DetokMsg::Deregister { .. })
|
||||
));
|
||||
}
|
||||
|
||||
/// A request rejected at normalization (post-register) must not leak: the shard
|
||||
/// sees `Register` then `Deregister`. Regression for RSS growth on bad input.
|
||||
#[test]
|
||||
fn rejected_request_deregisters_from_shard() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
// top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it.
|
||||
let bad = SamplingParams {
|
||||
top_p: 2.0,
|
||||
..Default::default()
|
||||
};
|
||||
intake.drive(generate_req(7, bad));
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "7"),
|
||||
"expected Register for rid 7",
|
||||
);
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "7"),
|
||||
"expected Deregister for rid 7 (leak fix)",
|
||||
);
|
||||
assert!(
|
||||
detok_rx.try_recv().is_err(),
|
||||
"no further shard messages — registration fully cleaned up",
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: an out-of-vocabulary client token id must be rejected at
|
||||
/// with a 400 — passed through, it reaches the embedding lookup
|
||||
/// and kills the scheduler process (`make_intake` bounds vocab at 1000).
|
||||
#[test]
|
||||
fn out_of_vocab_input_ids_rejected() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(21, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![1, 2_000_000_000]);
|
||||
}
|
||||
intake.drive(req);
|
||||
// Rejected before registration: the only shard message is nothing at
|
||||
// all, or a Deregister if registration happened first — never a push.
|
||||
match detok_rx.try_recv() {
|
||||
Err(_) => {}
|
||||
Ok(DetokMsg::Deregister { .. }) => {}
|
||||
Ok(_) => panic!("out-of-vocab request must not be admitted"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Same guard for negative ids and for `token_ids_logprob` entries.
|
||||
#[test]
|
||||
fn negative_and_logprob_token_ids_rejected() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(22, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![-1]);
|
||||
}
|
||||
intake.drive(req);
|
||||
match detok_rx.try_recv() {
|
||||
Err(_) | Ok(DetokMsg::Deregister { .. }) => {}
|
||||
Ok(_) => panic!("negative token id must not be admitted"),
|
||||
}
|
||||
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(23, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.token_ids_logprob = Some(vec![999_999]);
|
||||
}
|
||||
intake.drive(req);
|
||||
match detok_rx.try_recv() {
|
||||
Err(_) | Ok(DetokMsg::Deregister { .. }) => {}
|
||||
Ok(_) => panic!("out-of-vocab token_ids_logprob must not be admitted"),
|
||||
}
|
||||
}
|
||||
|
||||
/// A valid request is registered and handed onward — never deregistered.
|
||||
#[test]
|
||||
fn admitted_request_keeps_registration() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
// Empty map → all sampling defaults, passes normalization.
|
||||
intake.drive(generate_req(9, SamplingParams::default()));
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "9"),
|
||||
"expected Register for rid 9",
|
||||
);
|
||||
assert!(
|
||||
detok_rx.try_recv().is_err(),
|
||||
"admitted request must not be deregistered",
|
||||
);
|
||||
}
|
||||
|
||||
/// A pool return in `Failed` state (failed encode) is rejected via the same
|
||||
/// path and deregistered, not leaked.
|
||||
#[test]
|
||||
fn tokenize_failure_deregisters_via_intake() {
|
||||
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake();
|
||||
// The pool marks a failed encode as `Failed(err)` before returning it.
|
||||
let mut req = generate_req(11, SamplingParams::default());
|
||||
let _ = req
|
||||
.state
|
||||
.apply(Event::Error(Error::Tokenize("boom".into())));
|
||||
tm_tx.send(TmEvent::Tokenized(req)).unwrap();
|
||||
// Close the inbox so the run loop returns after draining the one event.
|
||||
drop(tm_tx);
|
||||
intake.run();
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "11"),
|
||||
"tokenize failure must deregister rid 11",
|
||||
);
|
||||
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
||||
}
|
||||
|
||||
/// An abort deregisters (by the id hashed from the rid string), so a request
|
||||
/// aborted before any terminal chunk can't leak.
|
||||
#[test]
|
||||
fn abort_deregisters_from_shard() {
|
||||
// Aborts arrive on their own unbounded lane now, not the request inbox.
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake_with_abort(abort_rx);
|
||||
abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap();
|
||||
drop(abort_tx);
|
||||
drop(tm_tx);
|
||||
intake.run();
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "rid-13"),
|
||||
"abort must deregister by rid",
|
||||
);
|
||||
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
||||
}
|
||||
|
||||
/// A successful pool return (Queued, ids filled) is pushed to the ring, not
|
||||
/// rejected; its registration is untouched.
|
||||
#[test]
|
||||
fn tokenized_return_pushes_without_deregister() {
|
||||
let (intake, detok_rx, _consumer, tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(15, SamplingParams::default());
|
||||
// Simulate a successful pool return: ids filled, PreSendValidating.
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = Some(vec![1, 2, 3]);
|
||||
}
|
||||
req.state = RequestState::PreSendValidating;
|
||||
tm_tx.send(TmEvent::Tokenized(req)).unwrap();
|
||||
drop(tm_tx);
|
||||
intake.run();
|
||||
|
||||
// Pushed to the ring; the shard sees nothing.
|
||||
assert!(
|
||||
detok_rx.try_recv().is_err(),
|
||||
"a queued pool-return must be pushed, not touch the shard",
|
||||
);
|
||||
}
|
||||
|
||||
/// If the pool is gone, a request needing tokenization is rejected +
|
||||
/// deregistered, not silently dropped.
|
||||
#[test]
|
||||
fn tokenize_pool_gone_deregisters() {
|
||||
let (mut intake, detok_rx, _consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
let mut req = generate_req(21, SamplingParams::default());
|
||||
if let RequestKind::Generate(g) = &mut req.kind {
|
||||
g.input_ids = None;
|
||||
}
|
||||
intake.drive(req);
|
||||
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "21"),
|
||||
"expected Register for rid 21",
|
||||
);
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid }) if rid.as_str() == "21"),
|
||||
"pool-gone hand-off must deregister rid 21",
|
||||
);
|
||||
assert!(detok_rx.try_recv().is_err(), "no further shard messages");
|
||||
}
|
||||
|
||||
/// Build a generate request carrying an image. The parked entry and the
|
||||
/// `MmEncoded` resume path agree on identity via the rid string.
|
||||
fn mm_generate_req(rid: &str) -> Request {
|
||||
let (tx, _rx) = mpsc::channel(8);
|
||||
Request {
|
||||
rid: rid.to_string().into(),
|
||||
state: RequestState::Received,
|
||||
sink: ResponseSink::Local(tx),
|
||||
kind: RequestKind::Generate(Box::new(GenerateRequest {
|
||||
rid: rid.to_string().into(),
|
||||
text: Some("<image> hi".into()),
|
||||
mm: Some(Box::new(crate::message::request::MmData {
|
||||
image_data: Some(rmpv::Value::from("data:image/jpeg;base64,xxxx")),
|
||||
..Default::default()
|
||||
})),
|
||||
..Default::default()
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// An abort while the request is parked for MM cancels it: the pending
|
||||
/// entry is removed, the worker's late result is dropped, and its parked
|
||||
/// sidecar entry is purged — no scheduler work runs for a dead client.
|
||||
#[test]
|
||||
fn abort_cancels_parked_mm_request() {
|
||||
let (mut intake, _detok_rx, consumer, _tm_tx, mm_rx) = make_intake();
|
||||
intake.drive(mm_generate_req("mm-gone"));
|
||||
mm_rx.try_recv().expect("parked to mm pool");
|
||||
|
||||
// The worker parks its result, as it always does before MmEncoded.
|
||||
intake.mm.sidecar.park(
|
||||
"mm-gone".into(),
|
||||
crate::multi_modality::sidecar::MmSidecarEntry {
|
||||
features: crate::multi_modality::sidecar::FeatureStore::Inline(vec![]),
|
||||
grids: vec![],
|
||||
hashes: vec![],
|
||||
offsets: vec![],
|
||||
mrope: vec![],
|
||||
mrope_delta: 0,
|
||||
},
|
||||
);
|
||||
intake.on_abort(AbortSource::Guard("mm-gone".to_string().into()));
|
||||
assert_eq!(consumer.drain(16).headers.len(), 1, "only the AbortReq");
|
||||
|
||||
// The late result must be dropped, not queued, and the sidecar purged.
|
||||
intake.on_mm_encoded("mm-gone".to_string().into(), vec![5, 6]);
|
||||
assert!(
|
||||
consumer.drain(16).headers.is_empty(),
|
||||
"cancelled, not queued"
|
||||
);
|
||||
assert!(intake.mm.sidecar.take("mm-gone").is_none(), "entry purged");
|
||||
}
|
||||
|
||||
/// A multimodal request parks in `Encoding` (submitted to the mm worker
|
||||
/// pool, not the tokenizer pool, not the ring) until `MmEncoded` resumes
|
||||
/// it → ring.
|
||||
#[test]
|
||||
fn mm_request_parks_then_mm_encoded_pushes_to_ring() {
|
||||
let (mut intake, _detok_rx, consumer, _tm_tx, mm_rx) = make_intake();
|
||||
intake.drive(mm_generate_req("mm-1"));
|
||||
|
||||
// Submitted to the mm pool with the typed work item; nothing on the ring yet.
|
||||
let sub = mm_rx.try_recv().expect("mm pool must receive the request");
|
||||
assert_eq!(sub.rid.as_str(), "mm-1");
|
||||
assert_eq!(sub.work.text.as_deref(), Some("<image> hi"));
|
||||
assert!(sub.work.input_ids.is_none(), "no client input_ids");
|
||||
assert_eq!(
|
||||
sub.work.image_data.as_ref().and_then(|v| v.as_str()),
|
||||
Some("data:image/jpeg;base64,xxxx")
|
||||
);
|
||||
assert!(consumer.drain(16).headers.is_empty(), "parked, not queued");
|
||||
|
||||
// The worker returns the final expanded ids → pushed to the ring.
|
||||
intake.on_mm_encoded("mm-1".to_string().into(), vec![5, 6, 7, 8]);
|
||||
let batch = consumer.drain(16);
|
||||
assert_eq!(batch.headers.len(), 1);
|
||||
assert_eq!(
|
||||
batch.lengths,
|
||||
vec![4],
|
||||
"expanded ids ride the columnar cell"
|
||||
);
|
||||
}
|
||||
|
||||
/// A worker failure rejects the parked request (deregister, no ring push).
|
||||
#[test]
|
||||
fn mm_failure_rejects_parked_request() {
|
||||
let (mut intake, detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
intake.drive(mm_generate_req("mm-2"));
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { .. })),
|
||||
"registered before parking",
|
||||
);
|
||||
|
||||
intake.on_mm_failed("mm-2".to_string().into(), "bad image".into());
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Deregister { rid })
|
||||
if rid.as_str() == "mm-2"),
|
||||
"mm failure must deregister",
|
||||
);
|
||||
assert!(consumer.drain(16).headers.is_empty(), "nothing queued");
|
||||
}
|
||||
|
||||
/// On a non-multimodal model (`Mm::enabled == false`), image_data is silently
|
||||
/// ignored and the request tokenizes as plain text — the Python
|
||||
/// TokenizerManager behavior when `mm_processor is None`.
|
||||
#[test]
|
||||
fn mm_fields_ignored_when_disabled() {
|
||||
let (tok_tx, tok_rx) = flume::unbounded();
|
||||
let (detok_tx, _detok_rx) = flume::unbounded();
|
||||
let senders = Senders {
|
||||
tok_manager_tx: flume::unbounded().0,
|
||||
abort_tx: flume::unbounded().0,
|
||||
tokenizer_tx: tok_tx,
|
||||
detokenizer_tx: vec![detok_tx],
|
||||
};
|
||||
let (to_scheduler_tx, _consumer) = to_scheduler(16);
|
||||
let (_tm_tx, tm_rx) = flume::unbounded();
|
||||
let (mm_tx, mm_rx) = flume::unbounded();
|
||||
let (abort_tx, abort_rx) = flume::unbounded::<AbortSource>();
|
||||
std::mem::forget(abort_tx);
|
||||
let (sd_tx, sd_rx) = flume::unbounded::<()>();
|
||||
std::mem::forget(sd_tx);
|
||||
let mut intake = Intake::new(
|
||||
tm_rx,
|
||||
abort_rx,
|
||||
senders,
|
||||
to_scheduler_tx,
|
||||
test_limits(),
|
||||
test_mm(mm_tx, false),
|
||||
sd_rx,
|
||||
);
|
||||
|
||||
intake.drive(mm_generate_req("mm-3"));
|
||||
assert!(
|
||||
mm_rx.try_recv().is_err(),
|
||||
"mm disabled: nothing submitted to the mm channel",
|
||||
);
|
||||
assert!(
|
||||
tok_rx.try_recv().is_ok(),
|
||||
"request must fall through to plain tokenization",
|
||||
);
|
||||
}
|
||||
|
||||
/// A late mm result for a rid that is no longer parked is dropped without
|
||||
/// panicking (e.g. hash-collision overwrite) — regression guard.
|
||||
#[test]
|
||||
fn late_mm_result_is_dropped() {
|
||||
let (mut intake, _detok_rx, consumer, _tm_tx, _mm_rx) = make_intake();
|
||||
intake.on_mm_encoded("ghost".to_string().into(), vec![1]);
|
||||
intake.on_mm_failed("ghost".to_string().into(), "boom".into());
|
||||
assert!(consumer.drain(16).headers.is_empty());
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Configuration and multimodal handles for scheduler intake.
|
||||
|
||||
use crate::message::config::ServerArgs;
|
||||
use crate::message::request::MmRequest;
|
||||
|
||||
/// The intake side of the MM path.
|
||||
#[derive(Clone)]
|
||||
pub struct Mm {
|
||||
/// Whether the model is multimodal. When false, mm fields are silently
|
||||
/// ignored, as the Python `TokenizerManager` does with `mm_processor is
|
||||
/// None`.
|
||||
pub enabled: bool,
|
||||
/// → MM worker pool (spawned via `Server.start_mm_workers`).
|
||||
pub tx: flume::Sender<MmRequest>,
|
||||
/// Results sidecar. Purged here when a late result arrives for a request
|
||||
/// that is no longer parked; otherwise it would leak, since only the
|
||||
/// scheduler drain pops entries.
|
||||
pub sidecar: crate::multi_modality::sidecar::Sidecar,
|
||||
}
|
||||
|
||||
/// Resolved once at boot from the scheduler's `server_args`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Limits {
|
||||
/// Token-ids-in mode: a generate request must arrive already tokenized.
|
||||
pub skip_tokenizer_init: bool,
|
||||
/// `model_config.vocab_size`; bounds client-supplied token ids. A required
|
||||
/// field of the `ServerArgs` schema, so intake can check unconditionally.
|
||||
pub vocab_size: u64,
|
||||
/// `model_config.context_len`, the ceiling for input + `max_new_tokens`.
|
||||
pub context_len: u64,
|
||||
/// Output slots reserved on top of the input (eagle draft tokens).
|
||||
pub num_reserved_tokens: u64,
|
||||
/// Clamp `max_new_tokens` to what fits instead of rejecting the request.
|
||||
pub allow_auto_truncate: bool,
|
||||
/// Whether the server can produce hidden states at all.
|
||||
pub enable_return_hidden_states: bool,
|
||||
}
|
||||
|
||||
impl From<&ServerArgs> for Limits {
|
||||
fn from(sa: &ServerArgs) -> Self {
|
||||
Self {
|
||||
skip_tokenizer_init: sa.skip_tokenizer_init,
|
||||
vocab_size: sa.model_config.vocab_size,
|
||||
context_len: sa.model_config.context_len,
|
||||
num_reserved_tokens: sa.num_reserved_tokens,
|
||||
allow_auto_truncate: sa.allow_auto_truncate,
|
||||
enable_return_hidden_states: sa.enable_return_hidden_states,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
//! Request validation for scheduler intake.
|
||||
|
||||
use crate::message::request::{GenerateRequest, Request, RequestKind};
|
||||
use crate::utils::{
|
||||
error::Error,
|
||||
fsm::{Event, ValidationOutcome},
|
||||
};
|
||||
|
||||
use super::to_scheduler::MAX_RID_LEN;
|
||||
use super::to_scheduler_types::Limits;
|
||||
|
||||
/// `Received → Validating` + admissibility check. Under `skip_tokenizer_init` a
|
||||
/// generate request must already carry token ids (no tokenizer to byte-encode
|
||||
/// text); control requests carry none and are exempt.
|
||||
pub(super) fn validate(req: &mut Request, limits: &Limits) -> Result<(), Error> {
|
||||
let (skip_tokenizer_init, vocab_size) = (limits.skip_tokenizer_init, limits.vocab_size);
|
||||
let _ = req
|
||||
.state
|
||||
.apply(Event::Validated(ValidationOutcome::NeedsTokenize));
|
||||
|
||||
// The rid is the request's identity everywhere downstream: it keys the detok
|
||||
// table, and it rides on EVERY chunk of EVERY decode step. An unbounded
|
||||
// client-supplied rid is therefore a per-step cost, not a one-off. Python's is
|
||||
// a 32-byte uuid hex, so this is generous.
|
||||
// Measured on the CLIENT-facing form: the uniquifier `Rid::from_client` appends
|
||||
// is this server's own overhead, and charging the client for bytes it did not
|
||||
// send would reject a rid exactly at the documented limit.
|
||||
let client_rid_len = req.rid.client_facing().len();
|
||||
if client_rid_len > MAX_RID_LEN {
|
||||
return Err(Error::Validation(format!(
|
||||
"rid is {client_rid_len} bytes, over the {MAX_RID_LEN}-byte limit"
|
||||
)));
|
||||
}
|
||||
if skip_tokenizer_init
|
||||
&& matches!(&req.kind, RequestKind::Generate(g) if !g.already_tokenized())
|
||||
{
|
||||
// `Validation` (400), not `Tokenize` (500): the client sent a request this
|
||||
// server cannot serve, which is their error to fix — Python 400s it too.
|
||||
return Err(Error::Validation(
|
||||
"skip_tokenizer_init is set: request must provide input_ids".into(),
|
||||
));
|
||||
}
|
||||
|
||||
// Client-supplied token ids must be in-vocabulary: an out-of-range id
|
||||
// reaches the embedding lookup and kills the scheduler process, so 400
|
||||
// here instead — mirroring the Python `TokenizerManager` validation.
|
||||
if let RequestKind::Generate(g) = &req.kind {
|
||||
if let Some(ids) = &g.input_ids {
|
||||
for &id in ids {
|
||||
if id < 0 || id as u64 >= vocab_size {
|
||||
return Err(Error::Validation(format!(
|
||||
"input_ids contains out-of-vocabulary token id {id}; \
|
||||
valid range is [0, {vocab_size})"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(ids) = &g.token_ids_logprob {
|
||||
for &id in ids {
|
||||
if id < 0 || id as u64 >= vocab_size {
|
||||
return Err(Error::Validation(format!(
|
||||
"token_ids_logprob contains out-of-vocabulary token id \
|
||||
{id}; valid range is [0, {vocab_size})"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Detokenize ids must fit the shard's `&[u32]` decode domain. No vocab
|
||||
// bound — parity with the retired direct decode service: an unknown id is
|
||||
// the tokenizer's error to report, and nothing here reaches the scheduler's
|
||||
// embedding lookup.
|
||||
if let RequestKind::Detokenize { token_ids } = &req.kind {
|
||||
for &id in token_ids {
|
||||
if u32::try_from(id).is_err() {
|
||||
return Err(Error::Validation(format!("Token ID {id} is out of range")));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The scheduler only computes hidden states when launched for it, so without
|
||||
// this the request would 200 with `meta_info.hidden_states` silently absent
|
||||
// (Python `TokenizerManager._validate_one_request`).
|
||||
if !limits.enable_return_hidden_states
|
||||
&& matches!(&req.kind, RequestKind::Generate(g) if g.return_hidden_states)
|
||||
{
|
||||
return Err(Error::Validation(
|
||||
"The server is not configured to return the hidden states. \
|
||||
Please set `--enable-return-hidden-states` to enable this feature."
|
||||
.into(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The context-window checks that need the tokenized length, mirroring Python
|
||||
/// `TokenizerManager._validate_one_request`: the input alone must fit, and then
|
||||
/// input + `max_new_tokens` must fit. Without them the scheduler silently clamps
|
||||
/// and the client gets a 200 with a truncated completion instead of an actionable
|
||||
/// 400.
|
||||
///
|
||||
/// Under `allow_auto_truncate` both clamp instead of rejecting — the launch flag
|
||||
/// opted into that.
|
||||
pub(super) fn check_total_tokens(g: &mut GenerateRequest, limits: &Limits) -> Result<(), Error> {
|
||||
let max_req_len = limits.context_len;
|
||||
// Python counts the reserved slots as part of the input, so a request can be
|
||||
// rejected for them even when the prompt alone fits.
|
||||
let input_len =
|
||||
g.input_ids.as_ref().map_or(0, |ids| ids.len()) as u64 + limits.num_reserved_tokens;
|
||||
|
||||
// Input length first, and unconditionally: `max_new_tokens: null` means "no
|
||||
// cap", which must not disable this. Python's comparison is `>=` — a prompt
|
||||
// that exactly fills the window leaves no room to generate.
|
||||
if input_len >= max_req_len {
|
||||
if !limits.allow_auto_truncate {
|
||||
return Err(Error::Validation(format!(
|
||||
"The input ({input_len} tokens) is longer than the model's context \
|
||||
length ({max_req_len} tokens)."
|
||||
)));
|
||||
}
|
||||
if let Some(ids) = &mut g.input_ids {
|
||||
ids.truncate(max_req_len as usize);
|
||||
}
|
||||
}
|
||||
let input_len =
|
||||
g.input_ids.as_ref().map_or(0, |ids| ids.len()) as u64 + limits.num_reserved_tokens;
|
||||
|
||||
let Some(max_new_tokens) = g.sampling_params.max_new_tokens else {
|
||||
return Ok(()); // no cap requested → nothing to add to the input length
|
||||
};
|
||||
let total = input_len.saturating_add(max_new_tokens.max(0) as u64);
|
||||
if total <= max_req_len {
|
||||
return Ok(());
|
||||
}
|
||||
if !limits.allow_auto_truncate {
|
||||
return Err(Error::Validation(format!(
|
||||
"Requested token count exceeds the model's maximum context length of \
|
||||
{max_req_len} tokens. You requested a total of {total} tokens: {input_len} \
|
||||
tokens from the input messages and {max_new_tokens} tokens for the \
|
||||
completion. Please reduce the number of tokens in the input messages or \
|
||||
the completion to fit within the limit."
|
||||
)));
|
||||
}
|
||||
let clamped = max_req_len.saturating_sub(input_len) as i64;
|
||||
// Re-check what the clamp can break. `verify` already ran (in Normalizing), so
|
||||
// lowering `max_new_tokens` here can leave `min_new_tokens > max_new_tokens` —
|
||||
// and `is_normalized: true` stops the scheduler from re-verifying, so nothing
|
||||
// downstream would catch it. Python validates before it verifies; we can't
|
||||
// reorder the FSM, so we re-assert the one invariant the clamp can violate.
|
||||
if g.sampling_params.min_new_tokens > clamped {
|
||||
return Err(Error::Validation(format!(
|
||||
"min_new_tokens must be in [0, max_new_tokens({clamped})], got {}",
|
||||
g.sampling_params.min_new_tokens
|
||||
)));
|
||||
}
|
||||
g.sampling_params.max_new_tokens = Some(clamped);
|
||||
Ok(())
|
||||
}
|
||||
@@ -25,7 +25,7 @@ pub enum TmEvent {
|
||||
Tokenized(Request),
|
||||
/// An MM worker finished a request parked in `Encoding`: `input_ids` are the
|
||||
/// final placeholder-expanded prompt ids. The buffers ride the rid-keyed
|
||||
/// sidecar (`Server.take_mm`), not this event.
|
||||
/// sidecar (`Server.take_mm_result`), not this event.
|
||||
MmEncoded { rid: Rid, input_ids: Vec<i32> },
|
||||
/// An MM worker rejected a request parked in `Encoding` (bad media URL,
|
||||
/// unsupported modality, preprocess error, …).
|
||||
|
||||
@@ -51,7 +51,7 @@ pub struct Runtime {
|
||||
/// `skip_tokenizer_init`).
|
||||
pub tokenizer: Option<Arc<dyn tokenizer::TextTokenizer>>,
|
||||
/// MM results parked between a worker's `MmEncoded` and the scheduler drain
|
||||
/// (`Server.take_mm`).
|
||||
/// (`Server.take_mm_result`).
|
||||
pub mm_sidecar: crate::multi_modality::sidecar::Sidecar,
|
||||
/// Worker join handles, joined by `request_shutdown` / `Drop`.
|
||||
threads: Mutex<Vec<JoinHandle<()>>>,
|
||||
@@ -117,18 +117,19 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
|
||||
// --- inter-stage channels ---
|
||||
let (tok_manager_tx, tok_manager_rx) =
|
||||
flume::bounded::<TmEvent>(cfg.rust_server_args.channel_cap);
|
||||
flume::bounded::<TmEvent>(cfg.rust_server_args.stage_channel_cap);
|
||||
let (tokenizer_tx, tokenizer_rx) =
|
||||
flume::bounded::<crate::message::request::Request>(cfg.rust_server_args.channel_cap);
|
||||
flume::bounded::<crate::message::request::Request>(cfg.rust_server_args.stage_channel_cap);
|
||||
// Encoding → MM worker pool. Bounded like the other stage edges so a slow
|
||||
// pool back-pressures instead of buffering unboundedly.
|
||||
let (mm_worker_tx, mm_worker_rx) =
|
||||
flume::bounded::<crate::message::request::MmRequest>(cfg.rust_server_args.channel_cap);
|
||||
let (mm_worker_tx, mm_worker_rx) = flume::bounded::<crate::message::request::MmRequest>(
|
||||
cfg.rust_server_args.stage_channel_cap,
|
||||
);
|
||||
let detokenizer_worker_num = cfg.server_args.detokenizer_worker_num;
|
||||
let mut detokenizer_tx = Vec::with_capacity(detokenizer_worker_num);
|
||||
let mut detokenizer_rx = Vec::with_capacity(detokenizer_worker_num);
|
||||
for _ in 0..detokenizer_worker_num {
|
||||
let (tx, rx) = flume::bounded::<DetokMsg>(cfg.rust_server_args.channel_cap);
|
||||
let (tx, rx) = flume::bounded::<DetokMsg>(cfg.rust_server_args.stage_channel_cap);
|
||||
detokenizer_tx.push(tx);
|
||||
detokenizer_rx.push(rx);
|
||||
}
|
||||
@@ -303,7 +304,7 @@ pub fn start(cfg: RuntimeConfig) -> Result<Runtime, String> {
|
||||
rt.block_on(api_server::app::serve(
|
||||
listener,
|
||||
senders,
|
||||
cfg.rust_server_args.channel_cap,
|
||||
cfg.rust_server_args.stage_channel_cap,
|
||||
cfg.server_args.clone(),
|
||||
// Response heartbeat watched by `/health_generate`.
|
||||
response_activity,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Shared fixtures for the native Rust multimodal suites.
|
||||
"""Shared fixtures for the Rust multimodal suites.
|
||||
|
||||
Imported via ``sys.path`` from the sibling suites (unittest runs these files by
|
||||
path, so a package-relative import would break ``python <file>``); the module
|
||||
|
||||
@@ -76,7 +76,7 @@ def make_processor(case, config, image_processor_cls=None):
|
||||
skip_tokenizer_init=False,
|
||||
mm_preprocess_cache_size_mb=0,
|
||||
trust_mm_content_hashes=False,
|
||||
# Read by NativeMmHost._use_feature_shm (single-rank fixture → the
|
||||
# Read by RustMmProcessor._use_feature_shm (single-rank fixture → the
|
||||
# inline zero-copy transport, like the 1-GPU e2e).
|
||||
tp_size=1,
|
||||
dist_init_addr=None,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Native driver error paths: out-of-scope and malformed inputs are rejected.
|
||||
|
||||
Covers ``process`` in ``rust/sglang-mm/src/driver.rs`` (via the
|
||||
``_core.qwen_vl.process_native_mm`` binding). The wire-payload parsing that
|
||||
``_core.qwen_vl.process_mm`` binding). The wire-payload parsing that
|
||||
feeds this driver (modality/shape rejection) lives in ``sglang-server``'s
|
||||
message layer and is tested with the integration PR.
|
||||
|
||||
@@ -47,13 +47,13 @@ def gif_bytes():
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
QWEN_CORE and hasattr(QWEN_CORE, "process_native_mm"),
|
||||
QWEN_CORE and hasattr(QWEN_CORE, "process_mm"),
|
||||
"sglang-mm native Qwen driver not built",
|
||||
)
|
||||
class TestNativeDriverErrorPaths(CustomTestCase):
|
||||
def assert_rejected(self, input_ids, images, pattern, spec=SPEC):
|
||||
with self.assertRaisesRegex(ValueError, pattern):
|
||||
QWEN_CORE.process_native_mm(input_ids, images, spec)
|
||||
QWEN_CORE.process_mm(input_ids, images, spec)
|
||||
|
||||
def test_degenerate_geometry_rejected_not_panicked(self):
|
||||
"""A thin image against a tight ``max_pixels`` floors a side of the
|
||||
@@ -88,7 +88,7 @@ class TestNativeDriverErrorPaths(CustomTestCase):
|
||||
"""GIF moved from rejected to served when the pure-Rust webp/gif/bmp
|
||||
decoders were enabled; this pins the accept side of that contract flip
|
||||
(the reject side used to be asserted here and broke in CI)."""
|
||||
_, _, grids, _, offsets, _, _ = QWEN_CORE.process_native_mm(
|
||||
_, _, grids, _, offsets, _, _ = QWEN_CORE.process_mm(
|
||||
IMAGE_IDS, [gif_bytes()], SPEC
|
||||
)
|
||||
self.assertEqual(len(grids), 1)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""End-to-end parity at the scheduler-input boundary.
|
||||
|
||||
`test_preprocess.py` pins the `preprocess` binding; this drives the whole native
|
||||
path — the `process_native_mm` driver, then `NativeMmHost.build_native_mm` — and
|
||||
path — the `process_mm` driver, then `RustMmProcessor.build_output` — and
|
||||
compares every field the scheduler reads against the Python `mm_processor`.
|
||||
Bitwise, for both HF backends: the Rust resize clones PIL's fixed-point bicubic
|
||||
and ATen's uint8 antialias kernel, so whichever one a server is configured with
|
||||
@@ -23,7 +23,7 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||
|
||||
maybe_stub_sgl_kernel()
|
||||
|
||||
from sglang.srt.managers.rust_server import NativeMmHost # noqa: E402
|
||||
from sglang.srt.rust_server.multimodal import RustMmProcessor # noqa: E402
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
@@ -33,7 +33,7 @@ from _mm_rust_utils import PROCESSOR_CONFIGS, image_bytes, load_core # noqa: E4
|
||||
register_cpu_ci(est_time=40, suite="base-a-test-cpu")
|
||||
|
||||
CORE = load_core()
|
||||
DRIVER = getattr(getattr(CORE, "qwen_vl", None), "process_native_mm", None)
|
||||
DRIVER = getattr(getattr(CORE, "qwen_vl", None), "process_mm", None)
|
||||
|
||||
# The fixture tokenizer's vocab (see `_fixtures.make_processor`):
|
||||
# 1 = <|vision_start|>, 2 = <|image_pad|>, 3 = <|vision_end|>, 4 = "hello".
|
||||
@@ -68,21 +68,21 @@ class TestQwenE2eParity(CustomTestCase):
|
||||
|
||||
import_processors("sglang.srt.multimodal.processors")
|
||||
# Skip __init__: it would build a processor; reuse the fixture's.
|
||||
host = NativeMmHost.__new__(NativeMmHost)
|
||||
host = RustMmProcessor.__new__(RustMmProcessor)
|
||||
host.model_config = SimpleNamespace(hf_config=self.processor.hf_config)
|
||||
host._processor = self.processor._processor
|
||||
host.server_args = self.processor.server_args
|
||||
spec = host.resolve_native_spec()
|
||||
spec = host.resolve_spec()
|
||||
self.assertIsNotNone(spec, f"gate rejected {self.image_processor}")
|
||||
return spec
|
||||
|
||||
def run_native(self, spec, sources):
|
||||
"""The Rust path: the `process_native_mm` driver, then the drain
|
||||
"""The Rust path: the `process_mm` driver, then the drain
|
||||
adapter — the same two steps `RustServer.drain` performs."""
|
||||
ids, features, grids, hashes, offsets, mrope, delta = DRIVER(
|
||||
PROMPT_PER_IMAGE * len(sources), sources, spec.rust_json()
|
||||
)
|
||||
# The shape of Rust's MmEncodeResult, inline transport (test_build_native_mm
|
||||
# The shape of Rust's MmEncodeResult, inline transport (test_build_output
|
||||
# pins the shm shape).
|
||||
handoff = SimpleNamespace(
|
||||
features=features,
|
||||
@@ -93,7 +93,7 @@ class TestQwenE2eParity(CustomTestCase):
|
||||
mrope=mrope,
|
||||
mrope_delta=delta,
|
||||
)
|
||||
return snapshot(ids, NativeMmHost.build_native_mm(spec, handoff))
|
||||
return snapshot(ids, RustMmProcessor.build_output(spec, handoff))
|
||||
|
||||
def run_python(self, sources):
|
||||
"""The reference path: the Python `mm_processor` the scheduler would use."""
|
||||
|
||||
+5
-5
@@ -22,7 +22,7 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||
maybe_stub_sgl_kernel()
|
||||
|
||||
from sglang.srt.managers.mm_utils import hash_feature # noqa: E402
|
||||
from sglang.srt.managers.rust_server import NativeMmHost # noqa: E402
|
||||
from sglang.srt.rust_server.multimodal import RustMmProcessor # noqa: E402
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
@@ -32,7 +32,7 @@ from _mm_rust_utils import PROCESSOR_CONFIGS, image_bytes, load_core # noqa: E4
|
||||
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
|
||||
|
||||
CORE = load_core()
|
||||
DRIVER = getattr(getattr(CORE, "qwen_vl", None), "process_native_mm", None)
|
||||
DRIVER = getattr(getattr(CORE, "qwen_vl", None), "process_mm", None)
|
||||
|
||||
|
||||
def raw_bytes(source):
|
||||
@@ -45,7 +45,7 @@ def raw_bytes(source):
|
||||
|
||||
|
||||
@unittest.skipUnless(DRIVER, "sglang-mm native Qwen driver not built")
|
||||
class TestQwenNativeMmHashes(CustomTestCase):
|
||||
class TestQwenRustMmHashes(CustomTestCase):
|
||||
def setUp(self):
|
||||
from sglang.srt.managers.multimodal_processor import import_processors
|
||||
|
||||
@@ -58,11 +58,11 @@ class TestQwenNativeMmHashes(CustomTestCase):
|
||||
|
||||
def native_hashes(self, sources):
|
||||
"""Per-item hashes the Rust driver returns, via the production gate."""
|
||||
host = NativeMmHost.__new__(NativeMmHost)
|
||||
host = RustMmProcessor.__new__(RustMmProcessor)
|
||||
host.model_config = SimpleNamespace(hf_config=self.processor.hf_config)
|
||||
host._processor = self.processor._processor
|
||||
host.server_args = self.processor.server_args
|
||||
spec = host.resolve_native_spec()
|
||||
spec = host.resolve_spec()
|
||||
self.assertIsNotNone(spec, "gate rejected the fixture processor")
|
||||
input_ids = [t for _ in sources for t in (1, 2, 3, 4)]
|
||||
return DRIVER(input_ids, sources, spec.rust_json())[3]
|
||||
@@ -3,7 +3,7 @@
|
||||
Covers ``layout_by_placeholder`` / ``apply_layout`` in
|
||||
``rust/sglang-mm/src/common/token_layout.rs`` and ``mrope_image_only`` in
|
||||
``rust/sglang-mm/src/qwen_vl/mod.rs`` (via the
|
||||
``_core.qwen_vl.process_native_mm`` and ``mrope_image_only_py``
|
||||
``_core.qwen_vl.process_mm`` and ``mrope_image_only_py``
|
||||
bindings), against ``BaseMultimodalProcessor`` expansion/offsets and
|
||||
``MRotaryEmbedding.get_rope_index``.
|
||||
"""
|
||||
@@ -36,7 +36,7 @@ QWEN_CORE = getattr(load_core(), "qwen_vl", None)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
QWEN_CORE and hasattr(QWEN_CORE, "process_native_mm"),
|
||||
QWEN_CORE and hasattr(QWEN_CORE, "process_mm"),
|
||||
"sglang-mm native Qwen driver not built",
|
||||
)
|
||||
class TestQwenPromptGeometry(CustomTestCase):
|
||||
@@ -52,7 +52,7 @@ class TestQwenPromptGeometry(CustomTestCase):
|
||||
ids.extend((VISION_START_ID, IMAGE_TOKEN_ID, VISION_END_ID, 8))
|
||||
images = [image_bytes(96 + 8 * i, 80, i) for i in range(image_count)]
|
||||
with self.subTest(image_count=image_count):
|
||||
actual_ids, _, grids, _, offsets, _, _ = QWEN_CORE.process_native_mm(
|
||||
actual_ids, _, grids, _, offsets, _, _ = QWEN_CORE.process_mm(
|
||||
ids, images, spec_json(config)
|
||||
)
|
||||
counts = [t * h * w // config["merge_size"] ** 2 for t, h, w in grids]
|
||||
|
||||
+9
-6
@@ -1,4 +1,4 @@
|
||||
"""``NativeMmHost.build_native_mm`` (managers/rust_server.py): the drain-time
|
||||
"""``RustMmProcessor.build_output``: the drain-time
|
||||
wrapping contracts — tensors are zero-copy views over the Rust-owned buffers, and
|
||||
pad values come from worker-precomputed hashes, since the scheduler loop must
|
||||
never hash features. Synthetic buffers, so this needs no Rust extension."""
|
||||
@@ -15,15 +15,18 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||
|
||||
maybe_stub_sgl_kernel()
|
||||
|
||||
from sglang.srt.managers.rust_server import NativeMmHost, NativeMmSpec # noqa: E402
|
||||
from sglang.srt.rust_server.multimodal import ( # noqa: E402
|
||||
RustMmProcessor,
|
||||
RustMmSpec,
|
||||
)
|
||||
|
||||
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestBuildNativeMm(CustomTestCase):
|
||||
class TestBuildRustMmOutput(CustomTestCase):
|
||||
def setUp(self):
|
||||
# feature_dim == 3 * temporal_patch_size * patch_size**2 == 6.
|
||||
self.spec = NativeMmSpec(
|
||||
self.spec = RustMmSpec(
|
||||
family="qwen_vl",
|
||||
feature_shm=False,
|
||||
image_token_id=10,
|
||||
@@ -50,7 +53,7 @@ class TestBuildNativeMm(CustomTestCase):
|
||||
|
||||
def build(self):
|
||||
features = np.arange(30, dtype=np.float32)
|
||||
output = NativeMmHost.build_native_mm(
|
||||
output = RustMmProcessor.build_output(
|
||||
self.spec,
|
||||
SimpleNamespace( # the shape of Rust's MmEncodeResult
|
||||
grids=self.GRIDS,
|
||||
@@ -99,7 +102,7 @@ class TestBuildNativeMm(CustomTestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestBuildNativeMmShm(TestBuildNativeMm):
|
||||
class TestBuildRustMmOutputShm(TestBuildRustMmOutput):
|
||||
"""The shm entry shape (TP>1): features arrive as named POSIX segments, and
|
||||
each item becomes a ``ShmPointerMMData`` stub whose ``materialize()`` yields
|
||||
that item's slice — and unlinks, taking the cleanup duty exactly once."""
|
||||
@@ -28,7 +28,7 @@ FETCH = CORE and CORE.common.fetch_bytes
|
||||
|
||||
@unittest.skipUnless(FETCH, "sglang-mm fetch binding not built")
|
||||
class TestRustMediaSourceLoading(CustomTestCase):
|
||||
DATA = b"native-mm-source"
|
||||
DATA = b"rust-mm-source"
|
||||
|
||||
def test_inline_sources(self):
|
||||
encoded = base64.b64encode(self.DATA).decode()
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Model-independent image decode parity for native Rust MM.
|
||||
"""Model-independent image decode parity for Rust MM.
|
||||
|
||||
Covers ``decode_rgb`` in ``rust/sglang-mm/src/common/mod.rs`` (via the
|
||||
``_core.common.image_decode_rgb`` binding), against PIL's
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""``RustServer._partition_cores`` (managers/rust_server.py): the pool cores must
|
||||
"""``rust_server.config._partition_cores``: the pool cores must
|
||||
be a *bounded* slice of this rank's allowed cores, not the whole remainder —
|
||||
sibling TP ranks share the NUMA node, so an unbounded mask lets MM preprocessing
|
||||
bursts preempt a sibling's CUDA-launch thread (measured: ~20 ms of ViT wall time
|
||||
@@ -12,14 +12,14 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||
|
||||
maybe_stub_sgl_kernel()
|
||||
|
||||
from sglang.srt.managers.rust_server import RustServer # noqa: E402
|
||||
from sglang.srt.rust_server.config import _partition_cores # noqa: E402
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def partition(node_cores, **kwargs):
|
||||
with patch("os.sched_getaffinity", return_value=set(node_cores), create=True):
|
||||
return RustServer._partition_cores(**kwargs)
|
||||
return _partition_cores(**kwargs)
|
||||
|
||||
|
||||
class TestPartitionCores(CustomTestCase):
|
||||
|
||||
+11
-11
@@ -1,10 +1,10 @@
|
||||
"""The native-MM launch gate's family selection (managers/rust_server.py).
|
||||
"""The Rust-MM launch gate's family selection.
|
||||
|
||||
``NATIVE_MM_FAMILIES`` decides which models the Rust pipeline serves natively;
|
||||
for everything else ``native_mm_family_for`` must return ``None``, which
|
||||
``RUST_MM_FAMILIES`` decides which models the Rust pipeline serves;
|
||||
for everything else ``rust_mm_family_for`` must return ``None``, which
|
||||
``RustServer.launch`` turns into a hard launch error. Pins that non-Qwen
|
||||
multimodal models — Inkling being the in-tree case — keep their Python
|
||||
processor and never match a native family, so growing the registry cannot
|
||||
processor and never match a Rust family, so growing the registry cannot
|
||||
silently reroute them.
|
||||
"""
|
||||
|
||||
@@ -20,25 +20,25 @@ from sglang.srt.managers.multimodal_processor import ( # noqa: E402
|
||||
get_mm_processor_cls,
|
||||
import_processors,
|
||||
)
|
||||
from sglang.srt.managers.rust_server import native_mm_family_for # noqa: E402
|
||||
from sglang.srt.rust_server.multimodal import rust_mm_family_for # noqa: E402
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def processor_cls_for(architecture, model_type):
|
||||
"""Through the production selection, as `resolve_native_spec` calls it."""
|
||||
"""Through the production selection, as `resolve_spec` calls it."""
|
||||
hf_config = SimpleNamespace(architectures=[architecture], model_type=model_type)
|
||||
return get_mm_processor_cls(hf_config, SimpleNamespace(model_impl="sglang"))
|
||||
|
||||
|
||||
class TestNativeMmGate(CustomTestCase):
|
||||
class TestRustMmGate(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
import_processors("sglang.srt.multimodal.processors")
|
||||
|
||||
def test_qwen_vl_resolves_its_family(self):
|
||||
cls = processor_cls_for("Qwen2_5_VLForConditionalGeneration", "qwen2_5_vl")
|
||||
family = native_mm_family_for(cls, "qwen2_5_vl")
|
||||
family = rust_mm_family_for(cls, "qwen2_5_vl")
|
||||
self.assertEqual(family and family.name, "qwen_vl")
|
||||
|
||||
def test_inkling_keeps_its_python_processor(self):
|
||||
@@ -46,15 +46,15 @@ class TestNativeMmGate(CustomTestCase):
|
||||
|
||||
cls = processor_cls_for("InklingForConditionalGeneration", "inkling_model")
|
||||
self.assertIs(cls, InklingMultimodalProcessor)
|
||||
self.assertIsNone(native_mm_family_for(cls, "inkling_model"))
|
||||
self.assertIsNone(rust_mm_family_for(cls, "inkling_model"))
|
||||
|
||||
def test_family_requires_both_processor_and_model_type(self):
|
||||
qwen = processor_cls_for("Qwen2_5_VLForConditionalGeneration", "qwen2_5_vl")
|
||||
self.assertIsNone(native_mm_family_for(qwen, "inkling_model"))
|
||||
self.assertIsNone(rust_mm_family_for(qwen, "inkling_model"))
|
||||
# Identity, not name: an override class must not match (the
|
||||
# SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE contract).
|
||||
impostor = type("QwenVLImageProcessor", (), {})
|
||||
self.assertIsNone(native_mm_family_for(impostor, "qwen2_5_vl"))
|
||||
self.assertIsNone(rust_mm_family_for(impostor, "qwen2_5_vl"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
Covers what the CPU parity units structurally cannot: the sidecar handoff, the
|
||||
drain ordering, Rust-side tokenization of multimodal prompts, and the rejection
|
||||
of inputs outside the native pipeline's scope (there is no Python fallback).
|
||||
of inputs outside the Rust pipeline's scope (there is no Python fallback).
|
||||
"""
|
||||
|
||||
import base64
|
||||
@@ -51,7 +51,7 @@ def solid_image_data_url(fmt):
|
||||
importlib.util.find_spec("sglang.srt.rust_extensions._server") is None,
|
||||
"sglang-server rust extension not installed (e.g. AMD suite)",
|
||||
)
|
||||
class TestRustServerNativeMm(CustomTestCase):
|
||||
class TestRustServerMm(CustomTestCase):
|
||||
env = {"SGLANG_RUST_SERVER": "1"}
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""MMMU accuracy gate for the Rust tokenizer manager's native multimodal path.
|
||||
"""MMMU accuracy gate for the Rust tokenizer manager's multimodal path.
|
||||
|
||||
``test_rust_native_mm_e2e.py`` checks that the output is *valid*; this checks that
|
||||
native Rust preprocessing yields *equally good* model inputs. A systematic skew
|
||||
Rust preprocessing yields *equally good* model inputs. A systematic skew
|
||||
(wrong resample filter, channel order, normalization, patch layout) still reads as
|
||||
fluent text and passes a keyword smoke check, but drops MMMU below the gate.
|
||||
|
||||
@@ -36,7 +36,7 @@ MODEL = "Qwen/Qwen3.5-0.8B"
|
||||
VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
|
||||
NUM_EXAMPLES = 100
|
||||
# Calibrated 2026-07-24 on H200: the native path scores 0.37 on this fixed subset
|
||||
# Calibrated 2026-07-24 on H200: the Rust path scores 0.37 on this fixed subset
|
||||
# at temperature 0 (two runs), matching the Python reference (0.37, same sampler
|
||||
# and samples). The gate leaves headroom for batching nondeterminism.
|
||||
MMMU_ACCURACY_THRESHOLD = 0.30
|
||||
@@ -101,10 +101,10 @@ class QwenGenerateVisionSampler(SamplerBase):
|
||||
importlib.util.find_spec("sglang.srt.rust_extensions._server") is None,
|
||||
"sglang-server rust extension not installed (e.g. AMD suite)",
|
||||
)
|
||||
class TestRustNativeMmMMMU(CustomTestCase):
|
||||
class TestRustMmMMMU(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# Capture the server log so the test can pin that the native MM
|
||||
# Capture the server log so the test can pin that the Rust MM
|
||||
# pipeline is active.
|
||||
cls.log_dir = tempfile.TemporaryDirectory()
|
||||
cls.server_logs = tuple(
|
||||
@@ -139,12 +139,12 @@ class TestRustNativeMmMMMU(CustomTestCase):
|
||||
|
||||
def test_mmmu_accuracy(self):
|
||||
# Guard the path under test: if the model ever drops off
|
||||
# NATIVE_MM_FAMILIES, launch fails and this names why.
|
||||
# RUST_MM_FAMILIES, launch fails and this names why.
|
||||
self.assertIn(
|
||||
"native MM pipeline enabled",
|
||||
"Rust MM pipeline enabled",
|
||||
self._read_server_log(),
|
||||
"rust server did not enable the native MM pipeline for "
|
||||
f"{MODEL}; this test must exercise the native path",
|
||||
"rust server did not enable the Rust MM pipeline for "
|
||||
f"{MODEL}; this test must exercise the Rust path",
|
||||
)
|
||||
|
||||
eval_obj = MMMUVLMEval(num_examples=NUM_EXAMPLES, num_threads=32)
|
||||
@@ -154,12 +154,12 @@ class TestRustNativeMmMMMU(CustomTestCase):
|
||||
dump_metric(
|
||||
"mmmu_score",
|
||||
result.score,
|
||||
labels={"model": MODEL, "eval": "mmmu", "api": "generate-rust-native-mm"},
|
||||
labels={"model": MODEL, "eval": "mmmu", "api": "generate-rust-mm"},
|
||||
)
|
||||
self.assertGreaterEqual(
|
||||
result.score,
|
||||
MMMU_ACCURACY_THRESHOLD,
|
||||
f"Rust native MM path scored {result.score:.4f} on MMMU, below the "
|
||||
f"Rust MM path scored {result.score:.4f} on MMMU, below the "
|
||||
f"{MMMU_ACCURACY_THRESHOLD:.2f} gate",
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user