diff --git a/rust/sglang-server/src/fsm.rs b/rust/sglang-server/src/fsm.rs index 9206fa708410..afb1f58a7649 100644 --- a/rust/sglang-server/src/fsm.rs +++ b/rust/sglang-server/src/fsm.rs @@ -7,10 +7,10 @@ //! //! Port of the design enum: //! ```text -//! Received, Validating, Normalizing, Encoding, Tokenizing, Queued, -//! Streaming { chunks_sent }, Finalizing, Completed, Failed(Error), Aborted +//! Received, Validating, Normalizing, Encoding, Tokenizing, PreSendValidating, +//! Queued, Streaming { chunks_sent }, Finalizing, Completed, Failed(Error), +//! Aborted //! ``` -#![allow(dead_code)] // TODO: remove when the consumer PR lands use crate::error::Error; @@ -25,6 +25,10 @@ pub enum RequestState { Normalizing, Encoding, Tokenizing, + /// Every branch converges here with its final `input_ids`, for the checks + /// that need the tokenized length (the input + `max_new_tokens` ceiling). + /// The last state before the request leaves Rust. + PreSendValidating, Queued, Streaming { chunks_sent: u64, @@ -43,7 +47,7 @@ pub enum ValidationOutcome { HasMultimodal, /// Plain text → Tokenizing. NeedsTokenize, - /// Caller already supplied token ids → straight to Queued. + /// Caller already supplied token ids → straight to the pre-send checks. AlreadyTokenized, } @@ -57,9 +61,13 @@ pub enum Event { NeedsNormalize, EncodeDone, TokenizeDone, + /// The pre-send checks passed; the request may be pushed to the ring. + PreSendValidated, SchedulerPicked, // --- egress --- - Chunk { finish: bool }, + Chunk { + finish: bool, + }, FinalFrameSent, // --- terminal (valid from any state) --- Error(Error), @@ -112,14 +120,17 @@ impl RequestState { // ingress (Received, Validated(_)) => Validating, // Generate requests pass through Normalizing (sampling-param - // normalize/verify); control requests skip straight to Queued. + // normalize/verify); control requests skip it, having none. (Validating, NeedsNormalize) => Normalizing, - (Validating, Validated(AlreadyTokenized)) => Queued, + (Validating, Validated(AlreadyTokenized)) => PreSendValidating, (Normalizing, Validated(HasMultimodal)) => Encoding, (Normalizing, Validated(NeedsTokenize)) => Tokenizing, - (Normalizing, Validated(AlreadyTokenized)) => Queued, + (Normalizing, Validated(AlreadyTokenized)) => PreSendValidating, (Encoding, EncodeDone) => Tokenizing, - (Tokenizing, TokenizeDone) => Queued, + // Every ingress branch funnels through the pre-send checks, so they + // run exactly once per request no matter how it got its ids. + (Tokenizing, TokenizeDone) => PreSendValidating, + (PreSendValidating, PreSendValidated) => Queued, (Queued, SchedulerPicked) => Streaming { chunks_sent: 0 }, // egress (Streaming { chunks_sent }, Chunk { finish: false }) => Streaming { @@ -133,3 +144,59 @@ impl RequestState { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn after(mut state: RequestState, event: Event) -> RequestState { + state.apply(event).expect("edge must exist"); + state + } + + /// Every ingress branch — control, client-supplied ids, and text through the + /// tokenizer pool — must land in `PreSendValidating`, because that is where + /// the checks needing the final `input_ids` run. A branch that reached + /// `Queued` directly would skip them silently. + #[test] + fn every_branch_reaches_the_ring_through_pre_send_validating() { + for from in [ + after( + RequestState::Validating, + Event::Validated(ValidationOutcome::AlreadyTokenized), + ), + after( + RequestState::Normalizing, + Event::Validated(ValidationOutcome::AlreadyTokenized), + ), + after(RequestState::Tokenizing, Event::TokenizeDone), + ] { + assert!( + matches!(from, RequestState::PreSendValidating), + "branch bypassed the pre-send checks: {from:?}" + ); + assert!(matches!( + after(from, Event::PreSendValidated), + RequestState::Queued + )); + } + } + + /// The converse: `Queued` has no other in-edge, so the checks can't be skipped + /// by emitting the wrong event, and can't run twice. + #[test] + fn queued_has_no_other_in_edge() { + for mut state in [ + RequestState::Validating, + RequestState::Normalizing, + RequestState::Tokenizing, + RequestState::Queued, + ] { + assert_eq!( + state.apply(Event::PreSendValidated), + Err(TransitionError::Illegal), + "only PreSendValidating may enter Queued" + ); + } + } +} diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index f733bb9ddfeb..1d4e6e8fbbe1 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -16,6 +16,9 @@ mod error; mod fsm; mod ids; mod message; +mod ring; +mod runtime; +mod tokenizer_manager; mod utils; use pyo3::prelude::*; diff --git a/rust/sglang-server/src/message.rs b/rust/sglang-server/src/message.rs index 9b6e2754fcd3..8e2bb921fb2c 100644 --- a/rust/sglang-server/src/message.rs +++ b/rust/sglang-server/src/message.rs @@ -8,11 +8,73 @@ //! `SamplingParams` port), [`io_struct`] (the scheduler wire structs), [`types`] //! (the shared wire-shape adapters both directions use). +mod egress; +mod finish_reason; mod io_struct; mod request; mod sampling; mod types; -pub(crate) use request::GenerateRequest; +pub use egress::{ChunkEvent, EgressItem, EgressSink}; +pub(crate) use io_struct::{AbortReq, ControlRequest}; +pub use request::{GenerateRequest, RequestKind}; pub(crate) use sampling::{SamplingParams, SamplingParamsInput}; pub(crate) use types::{OneOrMany, OneOrManyItem, TokenIds}; + +use bytes::Bytes; + +use crate::fsm::RequestState; +use crate::ids::Rid; + +/// The owned request as it travels ingress stages (single owner, so `state` is +/// mutated lock-free). Common fields here; variant data in [`RequestKind`]. +#[derive(Debug)] +pub struct Request { + /// Client-visible request id (uuid hex) — what the scheduler wire and + /// `meta_info.id` carry. + pub rid: Rid, + pub state: RequestState, + /// Back-channel to the client connection for egress frames. + pub sink: EgressSink, + /// Discriminant + variant body (generate vs control). + pub kind: RequestKind, +} + +/// One ingress-ring entry, split columnar: the scalar `header` (msgpack, `input_ids` +/// omitted) + the raw int64 `ids` cell, so the big tensor never goes through msgpack. +#[derive(Debug)] +pub struct IngressMsg { + pub header: Bytes, + pub ids: Bytes, +} + +/// Messages to a Detokenizer shard. `Register` carries the per-request sink for +/// the shard's local `rid -> sink` map. The rid STRING is the identity: `Rid::hash` +/// picks the shard (collisions there merely co-locate, which is harmless), but two +/// distinct rids that hash alike must not be the same map entry — that evicted one +/// client's sink and delivered their tokens to the other's connection. Equal rids +/// cannot reach here from different requests: `Rid::from_client` uniquifies every +/// client-supplied one. +pub enum DetokMsg { + Register { + /// Client-visible rid string — kept in `DetokState` so the shard can + /// emit `TmEvent::Abort(rid)` (the wire needs the string, not the hash). + rid: Rid, + sink: EgressSink, + /// Decode logprob token ids to text here (CPU-bound) not on the api threads. + decode_logprob_text: bool, + /// `SamplingParams.no_stop_trim`: keep the matched stop; default trims it. + no_stop_trim: bool, + }, + /// One decode step's chunks for *this shard*. Batched because `tm-egress` blocks + /// per send, so one message per request cost ~1.3 µs × batch (5.1x at 4096). + Chunks(Vec), + /// Control result: one already-serialized payload delivered to the sink verbatim. + Result { rid: Rid, payload: bytes::Bytes }, + /// Terminal per-request failure → an `Error` to the sink (a 400, not a crash). + Fail { rid: Rid, message: String }, + /// Drop the `rid -> sink` entry for a request rejected before the scheduler + /// (the rejecting stage already answered the client); else `Register` leaks one + /// entry. + Deregister { rid: Rid }, +} diff --git a/rust/sglang-server/src/message/egress.rs b/rust/sglang-server/src/message/egress.rs new file mode 100644 index 000000000000..24d5b6d0e245 --- /dev/null +++ b/rust/sglang-server/src/message/egress.rs @@ -0,0 +1,1190 @@ +//! The egress (response) direction: the per-request back-channel the API +//! handler drains ([`EgressSink`] / [`EgressItem`]), the egress-ring frame +//! encodings (batch / control result / error), and the columnar batch decode +//! into per-request [`ChunkEvent`]s. + +use bytes::Bytes; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; + +use super::TokenIds; +use super::finish_reason::FinishReason; +use crate::error::Error; +use crate::ids::Rid; + +/// Per-request back-channel the detok shard writes egress frames to and the API +/// handler drains for SSE; bounded, and receiver-drop (disconnect) = stream end. +#[derive(Clone, Debug)] +pub enum EgressSink { + Local(mpsc::Sender), +} + +/// Why an [`EgressSink::try_send`] failed: `Full` = client backpressure, `Closed` +/// = client gone. Both terminal for a stream; the caller distinguishes for logging. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SinkError { + Full, + Closed, +} + +impl EgressSink { + /// Non-blocking send. `Err(Full)` = backpressure, `Err(Closed)` = client gone. + pub fn try_send(&self, item: EgressItem) -> Result<(), SinkError> { + match self { + EgressSink::Local(tx) => tx.try_send(item).map_err(|e| match e { + mpsc::error::TrySendError::Full(_) => SinkError::Full, + mpsc::error::TrySendError::Closed(_) => SinkError::Closed, + }), + } + } +} + +#[allow(dead_code)] // the receiver half is created inline in api_server::submit. +pub type EgressSource = mpsc::Receiver; + +/// What the connection handler receives on the egress stream: a detok-decoded +/// [`ChunkEvent`] (handler formats it), a verbatim control payload, or an error. +#[derive(Debug)] +pub enum EgressItem { + /// An intermediate streamed generation step (only sent for streaming reqs). + Frame(ChunkEvent), + /// The final generation step. + Done(ChunkEvent), + /// A control-request result: one verbatim payload (e.g. `/server_info`), + /// delivered as-is with no per-protocol formatting. + Control(Bytes), + /// Terminal failure: handler emits an error frame (stream) or status (unary). + Error(Error), +} + +/// Egress-ring frame tag (first byte, prepended Rust-side; Python wire unchanged): +/// a single control-request result payload. +pub const EGRESS_TAG_RESULT: u8 = 1; +/// A whole decode batch: msgpack columnar header + one concatenated raw buffer; +/// tm-egress decodes it into per-request [`ChunkEvent`]s (no per-request FFI). +pub const EGRESS_TAG_BATCH: u8 = 2; +/// A per-request failure `[rid, message]`: the Python drain couldn't decode a +/// header, so it routes a 400 back to that request instead of crashing the loop. +pub const EGRESS_TAG_ERROR: u8 = 3; + +/// Read `n` little-endian f32s from `data` at `*off`, advancing `*off`. `None` when +/// the range runs past the buffer (a malformed / positional-ABI-drifted frame): the +/// caller rejects the whole frame. Bounds-checked via `data.get` — clamping only the +/// end is unsafe because a prior bad length can push `*off` past `len`, making the +/// range reversed (`start > end`) and the slice panic. +fn take_f32(data: &[u8], off: &mut usize, n: usize) -> Option> { + let start = *off; + let end = start.checked_add(n.checked_mul(4)?)?; + let out = data + .get(start..end)? + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + *off = end; + Some(out) +} + +/// Read `n` little-endian i32s from `data` at `*off`, advancing `*off`. `None` past +/// the buffer end (see [`take_f32`]). +fn take_i32(data: &[u8], off: &mut usize, n: usize) -> Option> { + let start = *off; + let end = start.checked_add(n.checked_mul(4)?)?; + let out = data + .get(start..end)? + .chunks_exact(4) + .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + *off = end; + Some(out) +} + +/// Frame a decode batch: `[BATCH tag][u32 header len][header][data cols…]`. The +/// caller's `data_cols` are concatenated straight into the frame (one copy, no +/// `b"".join`); `header` is the msgpack [`BatchHeader`]. Runs off the GIL. +pub fn frame_egress_batch_cols(header: &[u8], data_cols: &[&[u8]]) -> Bytes { + let data_len: usize = data_cols.iter().map(|c| c.len()).sum(); + let mut buf = Vec::with_capacity(1 + 4 + header.len() + data_len); + buf.push(EGRESS_TAG_BATCH); + buf.extend_from_slice(&(header.len() as u32).to_le_bytes()); + buf.extend_from_slice(header); + for col in data_cols { + buf.extend_from_slice(col); + } + Bytes::from(buf) +} + +/// Columnar scalar header for a whole decode batch. All numeric fields are +/// `#[serde(default)]`; the hot path (no extras) emits just the first four. +#[derive(Debug, Default, Serialize, Deserialize)] +pub struct BatchHeader { + /// Request ids, as the same strings Python holds (`Req.rid`, uuid hex) — + /// hashed back to the internal routing key in `decode_one` + /// (`Rid::shard`), mirroring the control path. The wire has no + /// rid-shape coupling; any string is a valid rid. + pub rids: Vec, + pub finish_reasons: Vec>, + pub prompt_tokens: Vec, + pub tok_lens: Vec, + #[serde(default)] + pub out_lp_lens: Vec, + #[serde(default)] + pub in_lp_lens: Vec, + #[serde(default)] + pub out_top_reqlens: Vec, + #[serde(default)] + pub out_top_poslens: Vec, + #[serde(default)] + pub in_top_reqlens: Vec, + #[serde(default)] + pub in_top_poslens: Vec, + #[serde(default)] + pub out_tid_reqlens: Vec, + #[serde(default)] + pub out_tid_poslens: Vec, + #[serde(default)] + pub in_tid_reqlens: Vec, + #[serde(default)] + pub in_tid_poslens: Vec, + #[serde(default)] + pub hidden_reqlens: Vec, + #[serde(default)] + pub hidden_poslens: Vec, +} + +/// Read a request's flat logprob column (`l` val/idx pairs) from `data` at cursors +/// `cv`/`ci`, advancing them. Only this request is copied — no whole-column buffer. +/// `None` if either read runs past the buffer (see [`take_f32`]). +fn take_flat( + data: &[u8], + cv: &mut usize, + ci: &mut usize, + l: usize, +) -> Option<(Vec, Vec)> { + Some((take_f32(data, cv, l)?, take_i32(data, ci, l)?)) +} + +/// Read `np` per-position lengths from `poslens` at `*pcur`, advancing it. `None` +/// when the range runs past the column — the header's `reqlens` promised more +/// positions than `poslens` carries, so this request's lengths are unknowable. +/// Clamping instead would return a short `lens`, and since `lens` also drives how +/// far the val/idx cursors advance, every later request in that column would read +/// from the wrong offset: silently wrong logprobs rather than a rejected frame. +fn take_poslens(poslens: &[u32], pcur: &mut usize, np: usize) -> Option> { + let start = *pcur; + let end = start.checked_add(np)?; + let lens = poslens.get(start..end)?.to_vec(); + *pcur = end; + Some(lens) +} + +/// Read a request's ragged logprob column (`np` positions): its per-position `lens` +/// from `poslens` (advancing `pcur`), then that many val/idx from `cv`/`ci`. `None` +/// if either read runs past its buffer (see [`take_poslens`], [`take_f32`]). +fn take_ragged( + data: &[u8], + cv: &mut usize, + ci: &mut usize, + poslens: &[u32], + pcur: &mut usize, + np: usize, +) -> Option<(Vec, Vec, Vec)> { + let lens = take_poslens(poslens, pcur, np)?; + let nv: usize = lens.iter().map(|&x| x as usize).sum(); + Some((take_f32(data, cv, nv)?, take_i32(data, ci, nv)?, lens)) +} + +/// Like [`take_ragged`] but for hidden states — a val column + row `poslens`, no +/// idx column. `None` if either read runs past its buffer (see [`take_poslens`], +/// [`take_f32`]). +fn take_hidden( + data: &[u8], + cv: &mut usize, + poslens: &[u32], + pcur: &mut usize, + nr: usize, +) -> Option<(Vec, Vec)> { + let lens = take_poslens(poslens, pcur, nr)?; + let nv: usize = lens.iter().map(|&x| x as usize).sum(); + Some((take_f32(data, cv, nv)?, lens)) +} + +/// Decode a batch egress frame (tag stripped), calling `route` with each request's +/// [`ChunkEvent`] as it's decoded — one pass, no intermediate `Vec`, peak memory +/// one request. Column order matches `push_generation`. +/// +/// `ok == false` means the frame was rejected. The caller discards everything it +/// routed and fails the frame's requests instead of forwarding a partial fan-out +/// (see `tokenizer_manager::egress`), so a rejected frame delivers nothing — +/// `rids` exists precisely so those requests can be failed rather than left +/// waiting for a `Done` that no longer exists. +pub fn for_each_chunk(body: &[u8], mut route: impl FnMut(ChunkEvent)) -> Decoded { + let mut decoded = Decoded::default(); + if body.len() < 4 { + return decoded; + } + let hlen = u32::from_le_bytes([body[0], body[1], body[2], body[3]]) as usize; + let Some(header) = body.get(4..4 + hlen) else { + return decoded; + }; + // Every rejection past the header slice names its requests by re-reading the + // rid column, rather than the whole frame paying a clone for the error path. + macro_rules! reject { + () => {{ + decoded.rids = recover_rids(header); + return decoded; + }}; + } + + let data = &body[4 + hlen..]; + let mut h = match rmp_serde::from_slice::(header) { + Ok(h) => h, + Err(_) => reject!(), + }; + let n = h.rids.len(); + // Every column must agree with `rids`. A SHORT column contributes nothing to + // `base`, so the `base != data.len()` check below cannot see it: a short + // `tok_lens` delivered a 200 with an empty completion, and a short + // `finish_reasons` dropped the terminal marker so the request never completed + // and its unary drain pended forever. The producer already asserts this for the + // extras columns; the four core ones were unchecked. + if h.finish_reasons.len() != n || h.prompt_tokens.len() != n || h.tok_lens.len() != n { + reject!() + } + // The per-request extras columns are either absent (no request asked) or one + // entry per request — never partial. + let per_req_ok = |c: &[u32]| c.is_empty() || c.len() == n; + if !per_req_ok(&h.out_lp_lens) + || !per_req_ok(&h.in_lp_lens) + || !per_req_ok(&h.out_top_reqlens) + || !per_req_ok(&h.in_top_reqlens) + || !per_req_ok(&h.out_tid_reqlens) + || !per_req_ok(&h.in_tid_reqlens) + || !per_req_ok(&h.hidden_reqlens) + { + reject!() + } + let sum = |v: &[u32]| v.iter().map(|&x| x as usize).sum::(); + // Each ragged family's per-request counts must consume its position column + // EXACTLY. Only the deficit was caught (`take_poslens` runs off the end); a + // surplus left positions unread and delivered a truncated row with a 200. + if sum(&h.out_top_reqlens) != h.out_top_poslens.len() + || sum(&h.in_top_reqlens) != h.in_top_poslens.len() + || sum(&h.out_tid_reqlens) != h.out_tid_poslens.len() + || sum(&h.in_tid_reqlens) != h.in_tid_poslens.len() + || sum(&h.hidden_reqlens) != h.hidden_poslens.len() + { + reject!() + } + + // Per-column byte cursors, advanced per request — no whole-column read. Columns + // are concatenated in exactly this order, every element 4 bytes. + let mut base = 0usize; + let mut col = |count: usize| -> usize { + let start = base; + base += count * 4; + start + }; + // Each val/idx column pair shares one element count — sum it once. + let n_ids = sum(&h.tok_lens); + let n_olp = sum(&h.out_lp_lens); + let n_ilp = sum(&h.in_lp_lens); + let n_ot = sum(&h.out_top_poslens); + let n_it = sum(&h.in_top_poslens); + let n_od = sum(&h.out_tid_poslens); + let n_id = sum(&h.in_tid_poslens); + let n_h = sum(&h.hidden_poslens); + let mut c_ids = col(n_ids); + let mut c_olp_v = col(n_olp); + let mut c_olp_i = col(n_olp); + let mut c_ilp_v = col(n_ilp); + let mut c_ilp_i = col(n_ilp); + let mut c_ot_v = col(n_ot); + let mut c_ot_i = col(n_ot); + let mut c_it_v = col(n_it); + let mut c_it_i = col(n_it); + let mut c_od_v = col(n_od); + let mut c_od_i = col(n_od); + let mut c_id_v = col(n_id); + let mut c_id_i = col(n_id); + let mut c_h_v = col(n_h); + + // `col` summed every column's span into `base`, so a truncated frame is caught + // here — the one rejection that is genuinely whole-frame, since it precedes the + // routing loop. Past this point a failure can only be partial. + // + // The check is EQUALITY, not `>`: every column length is header-determined, so + // a data buffer longer than `base` means the header and the buffer disagree. + // Accepting the surplus let a producer-side val/idx mismatch through with a + // 200, every later column reading off by the difference. + if base != data.len() { + reject!() + } + + // Mirror of Python's `has_extra` guard: checking once per frame lets the + // per-request loop skip the extras machinery entirely on a plain decode frame. + let has_extras = !(h.out_lp_lens.is_empty() + && h.in_lp_lens.is_empty() + && h.out_top_reqlens.is_empty() + && h.in_top_reqlens.is_empty() + && h.out_tid_reqlens.is_empty() + && h.in_tid_reqlens.is_empty() + && h.hidden_reqlens.is_empty()); + + // Position cursors into the header's per-request `poslens` (ragged + hidden). + let (mut p_ot, mut p_it, mut p_od, mut p_id, mut p_h) = + (0usize, 0usize, 0usize, 0usize, 0usize); + let lens_i = |v: &[u32], i: usize| v.get(i).copied().unwrap_or(0) as usize; + + // Decode one request's slice of every column, advancing the cursors. `None` if a + // read overruns `data` (belt-and-suspenders past the upfront check, which already + // covers every column's span — so this is unreachable today). It stops the loop + // before slicing out of bounds, but requests decoded earlier in the frame are + // already routed; making that abort atomic would mean buffering the whole frame, + // which is exactly what the streaming decode avoids. + let mut decode_one = |i: usize| -> Option { + let token_ids = take_i32(data, &mut c_ids, lens_i(&h.tok_lens, i))?; + + // Plain decode frame (no request in the batch asked for logprobs/hidden): + // the extras columns are all zero-width, so skip reading them entirely. + let extras = if !has_extras { + None + } else { + let (out_lp_val, out_lp_idx) = + take_flat(data, &mut c_olp_v, &mut c_olp_i, lens_i(&h.out_lp_lens, i))?; + let (in_lp_val, in_lp_idx) = + take_flat(data, &mut c_ilp_v, &mut c_ilp_i, lens_i(&h.in_lp_lens, i))?; + let (out_top_val, out_top_idx, out_top_lens) = take_ragged( + data, + &mut c_ot_v, + &mut c_ot_i, + &h.out_top_poslens, + &mut p_ot, + lens_i(&h.out_top_reqlens, i), + )?; + let (in_top_val, in_top_idx, in_top_lens) = take_ragged( + data, + &mut c_it_v, + &mut c_it_i, + &h.in_top_poslens, + &mut p_it, + lens_i(&h.in_top_reqlens, i), + )?; + let (out_tid_val, out_tid_idx, out_tid_lens) = take_ragged( + data, + &mut c_od_v, + &mut c_od_i, + &h.out_tid_poslens, + &mut p_od, + lens_i(&h.out_tid_reqlens, i), + )?; + let (in_tid_val, in_tid_idx, in_tid_lens) = take_ragged( + data, + &mut c_id_v, + &mut c_id_i, + &h.in_tid_poslens, + &mut p_id, + lens_i(&h.in_tid_reqlens, i), + )?; + let (hidden_val, hidden_lens) = take_hidden( + data, + &mut c_h_v, + &h.hidden_poslens, + &mut p_h, + lens_i(&h.hidden_reqlens, i), + )?; + + // Even in an extras batch, most requests carry none — box only if this + // one actually does, so its `ChunkEvent` stays the small common frame. + let ex = ChunkExtras { + out_lp_val, + out_lp_idx, + in_lp_val, + in_lp_idx, + out_top_val, + out_top_idx, + out_top_lens, + in_top_val, + in_top_idx, + in_top_lens, + out_tid_val, + out_tid_idx, + out_tid_lens, + in_tid_val, + in_tid_idx, + in_tid_lens, + hidden_val, + hidden_lens, + // Explicit, NOT `..Default::default()` — same reason as `ChunkEvent` + // below: a new column must fail to compile here until it is decoded. + out_lp_txt: Vec::new(), + in_lp_txt: Vec::new(), + out_top_txt: Vec::new(), + in_top_txt: Vec::new(), + out_tid_txt: Vec::new(), + in_tid_txt: Vec::new(), + }; + (!ex.is_empty()).then(|| Box::new(ex)) + }; + + Some(ChunkEvent { + // Any string is a valid rid; hash to the routing key. An unknown + // rid routes to a shard whose table has no entry → dropped there. + rid: std::mem::take(&mut h.rids[i]).into(), + token_ids, + finish_reason: h.finish_reasons.get(i).cloned().flatten(), + prompt_tokens: h.prompt_tokens.get(i).copied().unwrap_or(0), + extras, + // Listed explicitly, NOT `..Default::default()`: a new column added to + // `ChunkEvent` and wired into the response must fail to compile here + // until it is actually decoded. With the struct-update syntax it + // compiled clean and silently shipped zeros. + text: String::new(), + completion_tokens: 0, + }) + }; + + for i in 0..n { + let Some(ev) = decode_one(i) else { reject!() }; + route(ev); + } + decoded.ok = true; + decoded +} + +/// Recover the rid column straight from the header bytes. +/// +/// Read through `rmpv` rather than a serde tuple: `(Vec,)` decodes ONLY a +/// 1-element array, while real headers carry 4 or 16 columns, so it returned +/// `Err(LengthMismatch(1))` every time and named nobody. Arity independence is the +/// whole point — the trigger for the typed-decode-failure path is Python appending +/// a column an older Rust build does not know. +/// +/// Called only when a frame is rejected, which is why the accepted path no longer +/// clones the column up front. That clone was discarded unread on every good frame +/// and cost ~27% of the whole decode at batch 4096, plus a third of the crate's +/// steady-state allocations. It also has to be a re-read rather than a snapshot: +/// `decode_one` moves each rid out of the header as it goes, so by the time the +/// routing loop can fail, the decoded header no longer holds the earlier ones. +fn recover_rids(header: &[u8]) -> Vec { + rmpv::decode::read_value(&mut &header[..]) + .ok() + .and_then(|v| match v { + rmpv::Value::Array(cols) => cols.into_iter().next(), + _ => None, + }) + .and_then(|c| match c { + rmpv::Value::Array(rids) => Some( + rids.iter() + .filter_map(|r| r.as_str().map(Rid::from)) + .collect(), + ), + _ => None, + }) + .unwrap_or_default() +} + +/// Outcome of [`for_each_chunk`]: whether the frame was accepted, plus the rids it +/// named. `rids` is populated as soon as the header parses, so a caller can fail +/// every request in a rejected frame — including ones whose chunk never decoded, +/// which would otherwise wait forever for a `Done` that no longer exists. +#[derive(Debug, Default)] +pub struct Decoded { + pub ok: bool, + pub rids: Vec, +} + +/// Frame a control result `[rid, payload]` for the egress ring (tag prepended). +pub fn frame_egress_result(rid: &str, payload: &[u8]) -> Bytes { + use rmpv::Value; + let arr = Value::Array(vec![Value::from(rid), Value::Binary(payload.to_vec())]); + let mut buf = Vec::with_capacity(1 + payload.len() + rid.len() + 8); + buf.push(EGRESS_TAG_RESULT); + let _ = rmpv::encode::write_value(&mut buf, &arr); + Bytes::from(buf) +} + +/// Frame a per-request failure `[rid, message]` for the egress ring — routes a +/// terminal error back to the owning request (→ HTTP 400) instead of crashing. +pub fn frame_egress_error(rid: &str, message: &str) -> Bytes { + use rmpv::Value; + let arr = Value::Array(vec![Value::from(rid), Value::from(message)]); + let mut buf = Vec::with_capacity(1 + rid.len() + message.len() + 8); + buf.push(EGRESS_TAG_ERROR); + let _ = rmpv::encode::write_value(&mut buf, &arr); + Bytes::from(buf) +} + +/// One scheduler output increment — the common, always-present frame. `token_ids` +/// / `prompt_tokens` / `finish_reason` arrive from Python (pre-decode); the detok +/// shard fills `text` in place. Deltas — fold with `OutputAccumulator` for +/// cumulative. Logprobs + hidden states (rare, and large) live behind the boxed +/// [`ChunkExtras`] (`None` unless requested) so this frame stays small even when +/// the decoder builds an inline array at up to batch 4096 per step. +/// +/// Not a wire type — built by `for_each_chunk` from the columnar [`BatchHeader`] +/// frame and moved between stages in-process (never serialized), so no serde. +#[derive(Debug, Clone, Default)] +pub struct ChunkEvent { + /// Client-visible rid — the request's IDENTITY. Moved out of the frame header + /// (which owns it and drops it), so carrying it costs no allocation. The shard + /// is still chosen by `Rid::shard`, but a hash collision there now only + /// co-locates two requests instead of merging them. + pub rid: Rid, + /// New token ids for this step. Empty allowed (e.g. metadata-only frames). + pub token_ids: TokenIds, + /// `None` while streaming, the [`FinishReason`] on the final chunk. + pub finish_reason: Option, + /// Prompt token count for this request (constant across its chunks). + pub prompt_tokens: u32, + /// Decoded text **delta** for this chunk (empty in skip mode / on partial UTF-8), + /// filled by the detok shard. `token_ids` doubles as `output_ids`; + /// `completion_tokens` is this chunk's count. + pub text: String, + pub completion_tokens: u64, + /// Logprob + hidden-state columns — `None` unless the request asked for them. + /// Boxed to keep the common token/text/finish frame small at large decode + /// batches (the decoder allocates it only when a column is non-empty). + pub extras: Option>, +} + +/// Logprob + hidden-state columns for a [`ChunkEvent`], allocated only when the +/// request enabled logprobs / hidden states. Columnar `val`/`idx` (+ ragged `lens`) +/// buffers arrive pre-decode; the detok shard fills the parallel `*_txt` columns +/// when `return_text_in_logprobs` is set. In-process only — no serde (see +/// [`ChunkEvent`]). +#[derive(Debug, Clone, Default)] +pub struct ChunkExtras { + /// Output-token logprobs (parallel `val`/`idx`, one entry per new output token). + pub out_lp_val: Vec, + pub out_lp_idx: Vec, + /// Input (prefill) token logprobs, sent once on the first chunk. + pub in_lp_val: Vec, + pub in_lp_idx: Vec, + /// Top-k logprobs (2-level ragged): flat `val`/`idx` + per-position `lens` (0 = + /// null). Output = per-step delta, input = once on the first chunk. + pub out_top_val: Vec, + pub out_top_idx: Vec, + pub out_top_lens: Vec, + pub in_top_val: Vec, + pub in_top_idx: Vec, + pub in_top_lens: Vec, + /// Token-ids logprobs (same ragged layout); set only when `token_ids_logprob` was. + pub out_tid_val: Vec, + pub out_tid_idx: Vec, + pub out_tid_lens: Vec, + pub in_tid_val: Vec, + pub in_tid_idx: Vec, + pub in_tid_lens: Vec, + /// Hidden states (dense f32): flat buffer + per-row lengths. Last-writer-wins + /// across chunks (the final message has the full set). + pub hidden_val: Vec, + pub hidden_lens: Vec, + /// Decoded logprob token text (`return_text_in_logprobs`), parallel to the + /// `*_idx` buffers; empty when not requested (the tuple's text slot stays null). + pub out_lp_txt: Vec, + pub in_lp_txt: Vec, + pub out_top_txt: Vec, + pub in_top_txt: Vec, + pub out_tid_txt: Vec, + pub in_tid_txt: Vec, +} + +impl ChunkExtras { + /// True when no logprob / hidden column carries data — lets the decoder skip the + /// box allocation for the common (extras-free) frame. + fn is_empty(&self) -> bool { + self.out_lp_val.is_empty() + && self.in_lp_val.is_empty() + && self.out_top_lens.is_empty() + && self.in_top_lens.is_empty() + && self.out_tid_lens.is_empty() + && self.in_tid_lens.is_empty() + && self.hidden_lens.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::finish_reason::{FinishKind, Matched}; + + #[test] + fn batch_cols_match_single_joined_buffer() { + let header = [1u8, 2, 3]; + let a = [10u8, 11]; + let b = [12u8, 13, 14]; + let multi = frame_egress_batch_cols(&header, &[&a[..], &b[..]]); + let joined: Vec = a.iter().chain(&b).copied().collect(); + let single = frame_egress_batch_cols(&header, &[joined.as_slice()]); + assert_eq!(multi, single); + assert_eq!(multi[0], EGRESS_TAG_BATCH); + assert_eq!( + u32::from_le_bytes([multi[1], multi[2], multi[3], multi[4]]), + 3 + ); + assert_eq!(&multi[5..8], &header); // header + assert_eq!(&multi[8..], &[10, 11, 12, 13, 14]); // columns end-to-end + } + + /// A batch frame (the fast path) decodes into per-request ChunkEvents, with + /// token ids sliced from the single concatenated buffer by `tok_lens`. The + /// header is a msgspec-style positional array (what Python emits). + #[test] + fn decodes_batch_frame() { + use rmpv::Value; + // 3 requests: rids "1","2","3"; finish [nil, {type:stop,matched:5}, nil]; + // prompt_tokens [4,5,6]; tok_lens [2,0,1] -> ids [10,11 | (none) | 12]. + let stop = Value::Map(vec![ + (Value::from("type"), Value::from("stop")), + (Value::from("matched"), Value::from(5)), + ]); + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1"), Value::from("2"), Value::from("3")]), + Value::Array(vec![Value::Nil, stop, Value::Nil]), + Value::Array(vec![ + Value::from(4u32), + Value::from(5u32), + Value::from(6u32), + ]), + Value::Array(vec![ + Value::from(2u32), + Value::from(0u32), + Value::from(1u32), + ]), + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = [10i32, 11, 12] + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect(); + + let framed = frame_egress_batch_cols(&header, &[&data]); + assert_eq!(framed[0], EGRESS_TAG_BATCH); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + assert_eq!(events.len(), 3); + assert_eq!(events[0].rid, Rid::from("1")); + assert_eq!(events[0].token_ids, vec![10, 11]); + assert_eq!(events[0].prompt_tokens, 4); + assert!(events[0].finish_reason.is_none()); + assert_eq!(events[1].rid, Rid::from("2")); + assert!(events[1].token_ids.is_empty()); + // The whole reason survives msgpack (type + matched), not just the type. + assert_eq!( + events[1].finish_reason, + Some( + FinishKind::Stop { + matched: Some(Matched::Token(5)) + } + .into() + ) + ); + assert_eq!(events[2].rid, Rid::from("3")); + assert_eq!(events[2].token_ids, vec![12]); + assert_eq!(events[2].prompt_tokens, 6); + // A plain decode frame carries no extras columns at all, so the per-frame + // `has_extras` guard must skip the extras machinery entirely for every + // request (this is the tm-egress hot path — see `for_each_chunk`). + assert!(events.iter().all(|e| e.extras.is_none())); + } + + /// A header whose column lengths exceed the data buffer (a Python/Rust + /// positional-ABI drift, or a truncated frame) is rejected: `for_each_chunk` + /// returns false and routes nothing — it must NOT panic the sole egress thread + /// on an out-of-bounds slice. Built the way Python emits (positional msgpack + /// header + concatenated data columns). + #[test] + fn rejects_frame_with_lengths_past_data() { + use rmpv::Value; + // 1 request: tok_lens[0]=10 claims 40 bytes and out_lp_lens[0]=1 puts the + // logprob column's base past the 4-byte data buffer. The old clamp-only-`end` + // code advanced the cursor past `len`, then sliced `data[40..4]` (start > end) + // and panicked. + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1")]), // rids + Value::Array(vec![Value::Nil]), // finish_reasons + Value::Array(vec![Value::from(0u32)]), // prompt_tokens + Value::Array(vec![Value::from(10u32)]), // tok_lens (claims 40 bytes) + Value::Array(vec![Value::from(1u32)]), // out_lp_lens (base now past data) + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = [0i32].iter().flat_map(|x| x.to_le_bytes()).collect(); // 4 bytes + + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut routed = 0usize; + let decoded = for_each_chunk(&framed[1..], |_| routed += 1); + assert!(!decoded.ok, "malformed frame must be rejected, not decoded"); + assert_eq!(routed, 0, "no request may be routed from a rejected frame"); + // The rid is still reported, so the caller can fail the request that was + // waiting on this frame instead of letting it hang. + assert_eq!(decoded.rids, vec![Rid::from("1")]); + } + + /// A header whose `reqlens` claim more positions than `poslens` carries is + /// rejected, not truncated. This drift passes the upfront `base > data.len()` + /// check — the data columns are exactly as long as `poslens` says — so only the + /// per-column bound catches it. Clamping (the old behavior) handed req0 a short + /// `lens`, which also under-advanced the val/idx cursors, so req1 read from the + /// wrong offset: a frame that decodes "successfully" into wrong logprobs. + #[test] + fn rejects_poslens_shorter_than_reqlens_claims() { + use rmpv::Value; + let f = |xs: &[f32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let i = |xs: &[i32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let arr_u = |xs: &[u32]| Value::Array(xs.iter().map(|&x| Value::from(x)).collect()); + let rids = Value::Array(vec![Value::from("1"), Value::from("2")]); + let finish = Value::Array(vec![Value::Nil, Value::Nil]); + + // Ragged column: reqs claim 2 + 1 = 3 positions, `out_top_poslens` has 2. + let header_arr = Value::Array(vec![ + rids.clone(), + finish.clone(), + arr_u(&[0, 0]), // prompt + arr_u(&[1, 1]), // tok_lens + arr_u(&[0, 0]), // out_lp_lens + arr_u(&[0, 0]), // in_lp_lens + arr_u(&[2, 1]), // out_top_reqlens — 3 positions claimed + arr_u(&[2, 2]), // out_top_poslens — only 2 supplied + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let mut data = Vec::new(); + data.extend(i(&[10, 20])); // token_ids + data.extend(f(&[-0.1, -0.2, -0.3, -0.4])); // out_top_val (sum of poslens = 4) + data.extend(i(&[1, 2, 3, 4])); // out_top_idx + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut routed = 0usize; + assert!( + !for_each_chunk(&framed[1..], |_| routed += 1).ok, + "ragged poslens drift must be rejected, not truncated" + ); + + // Same drift in the hidden column, which has no idx pair: 2 rows claimed, 1 supplied. + let header_arr = Value::Array(vec![ + rids, + finish, + arr_u(&[0, 0]), // prompt + arr_u(&[1, 1]), // tok_lens + arr_u(&[0, 0]), // out_lp_lens + arr_u(&[0, 0]), // in_lp_lens + arr_u(&[0, 0]), // out_top_reqlens + arr_u(&[]), // out_top_poslens + arr_u(&[0, 0]), // in_top_reqlens + arr_u(&[]), // in_top_poslens + arr_u(&[0, 0]), // out_tid_reqlens + arr_u(&[]), // out_tid_poslens + arr_u(&[0, 0]), // in_tid_reqlens + arr_u(&[]), // in_tid_poslens + arr_u(&[2, 0]), // hidden_reqlens — 2 rows claimed + arr_u(&[3]), // hidden_poslens — only 1 supplied + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let mut data = Vec::new(); + data.extend(i(&[10, 20])); // token_ids + data.extend(f(&[0.1, 0.2, 0.3])); // hidden_val (sum of poslens = 3) + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut routed = 0usize; + assert!( + !for_each_chunk(&framed[1..], |_| routed += 1).ok, + "hidden poslens drift must be rejected, not truncated" + ); + } + + /// A frame that fails at request 0 buckets nothing, so bucket-driven cleanup + /// would leave every request in it hanging. The rids come from the header, + /// which is fully parsed before the decode loop. + #[test] + fn rejected_frame_still_reports_all_its_rids() { + use rmpv::Value; + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("r0"), Value::from("r1")]), // rids + Value::Array(vec![Value::Nil, Value::Nil]), + Value::Array(vec![Value::from(0u32), Value::from(0u32)]), + Value::Array(vec![Value::from(9u32), Value::from(9u32)]), // tok_lens past data + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let framed = frame_egress_batch_cols(&header, &[&[0u8; 4][..]]); + let mut routed = 0usize; + let decoded = for_each_chunk(&framed[1..], |_| routed += 1); + assert!(!decoded.ok); + assert_eq!(routed, 0, "fails at request 0 — nothing bucketed"); + assert_eq!( + decoded.rids, + vec![Rid::from("r0"), Rid::from("r1")], + "both requests must be nameable so the caller can fail them" + ); + } + + /// A data buffer LONGER than the header's columns means the two disagree — a + /// producer-side val/idx mismatch. Accepting the surplus delivered another + /// column's bytes as logprobs with a 200. + #[test] + fn frame_longer_than_its_columns_is_rejected() { + use rmpv::Value; + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1")]), + Value::Array(vec![Value::Nil]), + Value::Array(vec![Value::from(0u32)]), + Value::Array(vec![Value::from(1u32)]), // tok_lens: 1 id = 4 bytes + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = vec![0u8; 8]; // 4 bytes too many + let framed = frame_egress_batch_cols(&header, &[&data]); + let decoded = for_each_chunk(&framed[1..], |_| {}); + assert!(!decoded.ok, "header and data must agree exactly"); + } + + /// Ingress/egress rid agreement: the rid decoded from the frame must be the + /// one Python sent, AND both sides must derive the same shard from it. The + /// partition key is memoized inside `Rid`, so a per-conversion hasher seed + /// would send a request's chunks to a shard that never registered it. + #[test] + fn uuid_rid_decodes_to_the_same_rid_and_shard() { + use rmpv::Value; + let rid = Rid::from("9f86d081884c7d659a2feaa0c55ad015"); + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from(rid.to_string())]), // rids + Value::Array(vec![Value::Nil]), // finish_reasons + Value::Array(vec![Value::from(1u32)]), // prompt_tokens + Value::Array(vec![Value::from(1u32)]), // tok_lens + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = [0i32].iter().flat_map(|x| x.to_le_bytes()).collect(); + + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + assert_eq!(events.len(), 1); + // The equality below reuses `from_rid`, so on its own it would hold for + // any decoder that hashed *something*. These pin the two failure modes it + // is meant to catch: a decoder that parsed the rid as an integer (a uuid + // is not numeric, so it would fall back to 0), and one that keyed off the + // request's position in the batch rather than its rid. + assert_eq!( + events[0].rid, rid, + "the rid IS the identity — carried, not hashed" + ); + assert_ne!(events[0].rid, Rid::from("0")); + // Same string, independently constructed → same shard, on any shard count. + for shards in [1usize, 3, 8, 64] { + assert_eq!( + events[0].rid.shard(shards), + Rid::from("9f86d081884c7d659a2feaa0c55ad015").shard(shards), + "the partition key must not depend on WHERE the Rid was built" + ); + } + assert_ne!( + events[0].rid, + Rid::from("another-rid"), + "distinct rids stay distinct — identity is the string, not a digest" + ); + } + + /// A batch frame carrying the numeric columns (extras path): 2 requests, + /// req0 with output logprobs + top-k + hidden, req1 empty. Verifies the + /// column-major data split by the header's reqlens/poslens. + #[test] + fn decodes_batch_frame_with_extras() { + use rmpv::Value; + let f = |xs: &[f32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let i = |xs: &[i32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let arr_u = |xs: &[u32]| Value::Array(xs.iter().map(|&x| Value::from(x)).collect()); + // header: rids, finish, prompt, tok_lens, out_lp_lens, in_lp_lens, + // out_top_reqlens, out_top_poslens, in_top_*, out_tid_*, in_tid_*, + // hidden_reqlens, hidden_poslens + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1"), Value::from("2")]), // rids + Value::Array(vec![Value::Nil, Value::Nil]), // finish + arr_u(&[3, 4]), // prompt + arr_u(&[1, 1]), // tok_lens + arr_u(&[2, 0]), // out_lp_lens + arr_u(&[0, 0]), // in_lp_lens + arr_u(&[1, 0]), // out_top_reqlens (req0: 1 pos) + arr_u(&[2]), // out_top_poslens (that pos: k=2) + arr_u(&[0, 0]), // in_top_reqlens + arr_u(&[]), // in_top_poslens + arr_u(&[0, 0]), // out_tid_reqlens + arr_u(&[]), // out_tid_poslens + arr_u(&[0, 0]), // in_tid_reqlens + arr_u(&[]), // in_tid_poslens + arr_u(&[1, 0]), // hidden_reqlens (req0: 1 row) + arr_u(&[3]), // hidden_poslens (dim 3) + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let mut data = Vec::new(); + data.extend(i(&[10, 20])); // token_ids: req0=[10], req1=[20] + data.extend(f(&[-0.5, -0.6])); // out_lp_val (req0, 2) + data.extend(i(&[10, 99])); // out_lp_idx + data.extend(f(&[-0.1, -0.2])); // out_top_val (1 pos, k=2) + data.extend(i(&[10, 11])); // out_top_idx + data.extend(f(&[0.1, 0.2, 0.3])); // hidden_val (1 row, dim 3) + + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + assert_eq!(events.len(), 2); + assert_eq!(events[0].token_ids, vec![10]); + let ex0 = events[0] + .extras + .as_deref() + .expect("req0 has logprob/hidden extras"); + assert_eq!(ex0.out_lp_val, vec![-0.5, -0.6]); + assert_eq!(ex0.out_lp_idx, vec![10, 99]); + assert_eq!(ex0.out_top_val, vec![-0.1, -0.2]); + assert_eq!(ex0.out_top_lens, vec![2]); + assert_eq!(ex0.hidden_val, vec![0.1, 0.2, 0.3]); + assert_eq!(ex0.hidden_lens, vec![3]); + // req1 has a token id but no numeric columns → no extras box allocated. + assert_eq!(events[1].token_ids, vec![20]); + assert!(events[1].extras.is_none()); + } + + /// An inactive extras family may arrive as EMPTY columns or as `batch_size` + /// zeros, and both must decode to exactly the same events. + /// + /// This is the contract the Python producer's per-family gating relies on. It + /// used to run all seven families whenever any one of them was active, so an + /// inactive family shipped 4096 zeros per column; it now skips `accept`, which + /// leaves those columns empty. `per_req_ok` admits either and `lens_i` reads 0 + /// past the end — but nothing pinned that, so a later tightening of the column + /// validation (say, requiring every column to be `batch_size` long) would + /// silently start rejecting whole frames, hanging every request in them. + #[test] + fn empty_and_zero_filled_inactive_families_decode_alike() { + use rmpv::Value; + let f = |xs: &[f32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let i = |xs: &[i32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let arr_u = |xs: &[u32]| Value::Array(xs.iter().map(|&x| Value::from(x)).collect()); + + // Only `out_lp` is active; the other six families are inactive. `zeros` + // picks how they are spelled: the old producer's per-request zeros, or the + // gated producer's empty column. + let build = |zeros: bool| { + // Inactive: the old producer ran `accept` anyway and appended a 0 per + // request; the gated producer skips it and leaves the column empty. + let reqlens = if zeros { arr_u(&[0, 0]) } else { arr_u(&[]) }; + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1"), Value::from("2")]), // rids + Value::Array(vec![Value::Nil, Value::Nil]), // finish + arr_u(&[3, 4]), // prompt + arr_u(&[1, 1]), // tok_lens + arr_u(&[2, 0]), // out_lp_lens (ACTIVE) + reqlens.clone(), // in_lp_lens + reqlens.clone(), // out_top_reqlens + arr_u(&[]), // out_top_poslens + reqlens.clone(), // in_top_reqlens + arr_u(&[]), // in_top_poslens + reqlens.clone(), // out_tid_reqlens + arr_u(&[]), // out_tid_poslens + reqlens.clone(), // in_tid_reqlens + arr_u(&[]), // in_tid_poslens + reqlens, // hidden_reqlens + arr_u(&[]), // hidden_poslens + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let mut data = Vec::new(); + data.extend(i(&[10, 20])); // token_ids + data.extend(f(&[-0.5, -0.6])); // out_lp_val (req0) + data.extend(i(&[10, 99])); // out_lp_idx + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!( + for_each_chunk(&framed[1..], |ev| events.push(ev)).ok, + "frame rejected (zeros={zeros})" + ); + events + }; + + let old = build(true); + let new = build(false); + assert_eq!(old.len(), 2); + assert_eq!(old.len(), new.len()); + for (o, n) in old.iter().zip(&new) { + assert_eq!(o.rid, n.rid); + assert_eq!(o.token_ids, n.token_ids); + match (o.extras.as_deref(), n.extras.as_deref()) { + (Some(a), Some(b)) => { + assert_eq!(a.out_lp_val, b.out_lp_val); + assert_eq!(a.out_lp_idx, b.out_lp_idx); + // Every inactive family stays empty in BOTH spellings. + assert!(a.in_lp_val.is_empty() && b.in_lp_val.is_empty()); + assert!(a.out_top_lens.is_empty() && b.out_top_lens.is_empty()); + assert!(a.in_top_lens.is_empty() && b.in_top_lens.is_empty()); + assert!(a.out_tid_lens.is_empty() && b.out_tid_lens.is_empty()); + assert!(a.in_tid_lens.is_empty() && b.in_tid_lens.is_empty()); + assert!(a.hidden_lens.is_empty() && b.hidden_lens.is_empty()); + } + (None, None) => {} + _ => panic!("extras presence differs between the two spellings"), + } + } + // req0 really did carry its active family through both spellings. + assert_eq!( + new[0] + .extras + .as_deref() + .expect("req0 has out_lp") + .out_lp_val, + vec![-0.5, -0.6] + ); + } + + /// All SEVEN extras families in one frame, each with a distinct length AND + /// distinct values. The existing extras test exercises only `out_lp` / + /// `out_top` / `hidden`, so transposing a header pair — `in_top_*` with + /// `out_tid_*`, say — leaves every assertion passing while the client receives + /// another request's logprobs under the wrong key. Lengths differ per family + /// (2/1/2/1/2/1/3 elements) so a swap misaligns the cursors too. + #[test] + fn decodes_all_extras_families_without_transposition() { + use rmpv::Value; + let f = |xs: &[f32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let i = |xs: &[i32]| -> Vec { xs.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let arr_u = |xs: &[u32]| Value::Array(xs.iter().map(|&x| Value::from(x)).collect()); + + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("1")]), // rids + Value::Array(vec![Value::Nil]), // finish + arr_u(&[9]), // prompt + arr_u(&[1]), // tok_lens + arr_u(&[2]), // out_lp_lens (2 flat) + arr_u(&[1]), // in_lp_lens (1 flat) + arr_u(&[1]), // out_top_reqlens (1 position… + arr_u(&[2]), // out_top_poslens …k=2) + arr_u(&[1]), // in_top_reqlens (1 position… + arr_u(&[1]), // in_top_poslens …k=1) + arr_u(&[1]), // out_tid_reqlens (1 position… + arr_u(&[2]), // out_tid_poslens …2 ids) + arr_u(&[1]), // in_tid_reqlens (1 position… + arr_u(&[1]), // in_tid_poslens …1 id) + arr_u(&[1]), // hidden_reqlens (1 row… + arr_u(&[3]), // hidden_poslens …dim 3) + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + + // Concatenated in `for_each_chunk`'s column order. + let mut data = Vec::new(); + data.extend(i(&[100])); // token_ids + data.extend(f(&[-1.1, -1.2])); // out_lp_val + data.extend(i(&[11, 12])); // out_lp_idx + data.extend(f(&[-2.1])); // in_lp_val + data.extend(i(&[21])); // in_lp_idx + data.extend(f(&[-3.1, -3.2])); // out_top_val + data.extend(i(&[31, 32])); // out_top_idx + data.extend(f(&[-4.1])); // in_top_val + data.extend(i(&[41])); // in_top_idx + data.extend(f(&[-5.1, -5.2])); // out_tid_val + data.extend(i(&[51, 52])); // out_tid_idx + data.extend(f(&[-6.1])); // in_tid_val + data.extend(i(&[61])); // in_tid_idx + data.extend(f(&[7.1, 7.2, 7.3])); // hidden_val + + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + assert_eq!(events.len(), 1); + let ex = events[0].extras.as_deref().expect("extras present"); + + assert_eq!(ex.out_lp_val, vec![-1.1, -1.2]); + assert_eq!(ex.out_lp_idx, vec![11, 12]); + assert_eq!(ex.in_lp_val, vec![-2.1]); + assert_eq!(ex.in_lp_idx, vec![21]); + assert_eq!(ex.out_top_val, vec![-3.1, -3.2]); + assert_eq!(ex.out_top_idx, vec![31, 32]); + assert_eq!(ex.out_top_lens, vec![2]); + assert_eq!(ex.in_top_val, vec![-4.1]); + assert_eq!(ex.in_top_idx, vec![41]); + assert_eq!(ex.in_top_lens, vec![1]); + assert_eq!(ex.out_tid_val, vec![-5.1, -5.2]); + assert_eq!(ex.out_tid_idx, vec![51, 52]); + assert_eq!(ex.out_tid_lens, vec![2]); + assert_eq!(ex.in_tid_val, vec![-6.1]); + assert_eq!(ex.in_tid_idx, vec![61]); + assert_eq!(ex.in_tid_lens, vec![1]); + assert_eq!(ex.hidden_val, vec![7.1, 7.2, 7.3]); + assert_eq!(ex.hidden_lens, vec![3]); + // Every byte of the data buffer was consumed by exactly one family. + assert_eq!(events[0].token_ids, vec![100]); + } + + /// Two DISTINCT rids that hash to the same shard must stay separate requests. + /// Identity is the rid string now, so a shard-hash collision can only co-locate + /// them — it can no longer merge their `DetokState`, which used to evict one + /// client's sink and deliver their tokens to the other's connection. + #[test] + fn colliding_rids_stay_distinct_requests() { + use rmpv::Value; + let header_arr = Value::Array(vec![ + Value::Array(vec![Value::from("alice"), Value::from("bob")]), + Value::Array(vec![Value::Nil, Value::Nil]), + Value::Array(vec![Value::from(0u32), Value::from(0u32)]), + Value::Array(vec![Value::from(1u32), Value::from(1u32)]), + ]); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &header_arr).unwrap(); + let data: Vec = [7i32, 8].iter().flat_map(|x| x.to_le_bytes()).collect(); + let framed = frame_egress_batch_cols(&header, &[&data]); + let mut events = Vec::new(); + assert!(for_each_chunk(&framed[1..], |ev| events.push(ev)).ok); + // Each chunk carries its OWN rid — the value a shard keys its table on. + assert_eq!(events[0].rid, Rid::from("alice")); + assert_eq!(events[1].rid, Rid::from("bob")); + assert_ne!(events[0].rid, events[1].rid); + } + + /// The common frame must stay small: logprob/hidden columns are boxed behind + /// `ChunkExtras`, so the inline decode array is a few KiB — not MiB — even at + /// batch 4096. A regression that inlines a rare column would blow this up. + /// + /// The budget moved 128 → 144 when the rid gained its memoized partition key + /// (`String` + `u64`). That costs 8 bytes × batch per decode step and buys not + /// re-hashing the rid on every chunk in the egress bucketing loop — a + /// deliberate trade, not drift. + #[test] + fn chunk_event_frame_stays_small() { + let sz = std::mem::size_of::(); + assert!( + sz <= 144, + "ChunkEvent grew to {sz} bytes; keep rare columns behind ChunkExtras" + ); + } +} + +#[cfg(test)] +mod rid_recovery_tests { + use super::*; + + /// A header this build cannot type-decode (Python appended a column) must + /// still name its requests, or every one of them hangs. The previous + /// `(Vec,)` tuple decoded ONLY a 1-element array, so it failed on every + /// real header — arity independence is the entire point of this path. + #[test] + fn rid_recovery_works_at_every_header_arity() { + use rmpv::Value; + for extra_cols in [0usize, 3, 15, 16] { + let mut cols = vec![Value::Array(vec![Value::from("a"), Value::from("b")])]; + // Columns of a type this build would reject (strings where u32 is + // expected) — the "Python widened a column" case. + cols.extend((0..extra_cols).map(|_| Value::from("unexpected"))); + let mut header = Vec::new(); + rmpv::encode::write_value(&mut header, &Value::Array(cols)).unwrap(); + let framed = frame_egress_batch_cols(&header, &[]); + let decoded = for_each_chunk(&framed[1..], |_| {}); + assert!(!decoded.ok, "arity {extra_cols}: must reject"); + assert_eq!( + decoded.rids, + vec![Rid::from("a"), Rid::from("b")], + "arity {extra_cols}: rids must survive so the caller can fail them" + ); + } + } +} diff --git a/rust/sglang-server/src/message/finish_reason.rs b/rust/sglang-server/src/message/finish_reason.rs new file mode 100644 index 000000000000..6f9b585038fa --- /dev/null +++ b/rust/sglang-server/src/message/finish_reason.rs @@ -0,0 +1,217 @@ +//! The terminal finish reason: Python's `FinishReasonDict` — what +//! `BaseFinishReason.to_json()` (schedule_batch.py) puts on the egress wire, and +//! what the API echoes back as `meta_info.finish_reason`. Ingress has no +//! counterpart; it rides in the [`BatchHeader`](super::egress::BatchHeader) and on +//! each terminal [`ChunkEvent`](super::ChunkEvent). + +use serde::{Deserialize, Serialize}; + +/// The stop that ended a request — Python's `matched` key, typed +/// `Union[str, int, List[int]]`: a stop token id, a stop string +/// (`FINISH_MATCHED_STR` / `FINISHED_MATCHED_REGEX`), or a multi-token stop +/// sequence. Untagged because the three wire shapes are disjoint, so the shape +/// alone picks the arm. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum Matched { + Token(i64), + Str(String), + Tokens(Vec), +} + +/// The finish reasons this build knows, keyed by the `type` tag Python writes. +/// Every field is optional so a reason missing one still classifies (and keeps its +/// tag) instead of falling through to [`FinishReason::Unknown`]. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum FinishKind { + /// `FINISH_MATCHED_TOKEN` / `FINISH_MATCHED_STR` / `FINISHED_MATCHED_REGEX` — + /// all three report `type: "stop"`, to match the OpenAI API's value. + Stop { + #[serde(default)] + matched: Option, + }, + /// `FINISH_LENGTH` — hit `max_new_tokens` (or the context limit). + Length { + #[serde(default)] + length: Option, + }, + /// `FINISH_ABORT` — a scheduler-side termination. Boxed because it is by far + /// the widest variant and the rarest: unboxed it would set the size of every + /// [`ChunkEvent`], including the plain stop/length ones (see + /// `chunk_event_frame_stays_small`). + Abort(Box), +} + +/// The `FINISH_ABORT` payload. `status_code`/`err_type` are `None` for a plain +/// abort and set for a request error (e.g. over-context → 400); Python emits all +/// three keys, nulls included, so none of them is skipped on the way out. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct AbortReason { + #[serde(default)] + pub message: Option, + #[serde(default)] + pub status_code: Option, + #[serde(default)] + pub err_type: Option, +} + +/// A terminal finish reason exactly as `BaseFinishReason.to_json()` +/// (schedule_batch.py) puts it on the wire — Python's `FinishReasonDict`. It +/// round-trips: the API echoes it verbatim as `meta_info.finish_reason`, so the +/// serialized form must stay key-for-key what Python would have sent. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum FinishReason { + Known(FinishKind), + /// A `type` this build doesn't know, kept as its raw map and echoed unchanged. + /// This arm is why the outer enum is untagged: a finish reason added Python-side + /// must not fail the header decode, which rejects the whole frame — every + /// request in the batch, not just the one that carried it. + Unknown(serde_json::Map), +} + +impl From for FinishReason { + fn from(kind: FinishKind) -> Self { + FinishReason::Known(kind) + } +} + +impl FinishReason { + /// The stop this request matched, if it stopped on one. `None` for + /// length/abort and for an unknown type. + pub fn matched(&self) -> Option<&Matched> { + match self { + FinishReason::Known(FinishKind::Stop { matched }) => matched.as_ref(), + _ => None, + } + } + + /// `Some((status, message))` when this is an abort carrying a `status_code` — + /// a scheduler-side request error the API surfaces as that HTTP status instead + /// of as a normal completion. A plain abort (no code) reads as `None`. + pub fn abort_status(&self) -> Option<(u16, &str)> { + match self { + FinishReason::Known(FinishKind::Abort(a)) => Some(( + a.status_code?, + a.message.as_deref().unwrap_or("request aborted"), + )), + _ => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Parse a finish reason from the map Python puts on the wire, so these tests + /// exercise the real deserialization and not just the classifier. + fn fr(v: serde_json::Value) -> Option { + Some(serde_json::from_value(v).expect("finish reason must parse")) + } + + /// Classify through an `Option`, the shape both API paths actually hold. + /// Owned message so a caller can classify a temporary inline. + fn abort_status(f: &Option) -> Option<(u16, String)> { + let (code, message) = f.as_ref()?.abort_status()?; + Some((code, message.to_string())) + } + + /// The classifier both API paths use: a validation abort yields its + /// `(code, message)` (the streaming path turns this into an SSE error event + /// instead of a normal `Done` frame); anything else yields `None`. + #[test] + fn abort_status_extracts_code_and_message() { + let (code, msg) = abort_status(&fr(serde_json::json!({ + "type": "abort", "message": "over the limit", "status_code": 400 + }))) + .expect("validation abort → (code, message)"); + assert_eq!(code, 400); + assert_eq!(msg, "over the limit"); + // Normal finish, bare abort (no status), and no finish → not an error. + assert!(abort_status(&fr(serde_json::json!({"type": "stop"}))).is_none()); + assert!(abort_status(&fr(serde_json::json!({"type": "abort"}))).is_none()); + assert!(abort_status(&None).is_none()); + } + + /// A normal finish, a bare abort (no status), and no finish are not errors + /// (the unary path returns them as a 200 result frame). + #[test] + fn non_error_finishes_stay_ok() { + assert!(abort_status(&fr(serde_json::json!({"type": "stop", "matched": 5}))).is_none()); + assert!(abort_status(&fr(serde_json::json!({"type": "length", "length": 8}))).is_none()); + assert!( + abort_status(&fr( + serde_json::json!({"type": "abort", "message": "Aborted"}) + )) + .is_none() + ); + assert!(abort_status(&None).is_none()); + } + + /// The `matched` accessor reads only a stop, and reads every shape of one. + #[test] + fn matched_reads_stops_only() { + let m = |v| fr(v).unwrap().matched().cloned(); + assert_eq!( + m(serde_json::json!({"type": "stop", "matched": 9})), + Some(Matched::Token(9)) + ); + assert_eq!( + m(serde_json::json!({"type": "stop", "matched": ""})), + Some(Matched::Str("".into())) + ); + assert_eq!( + m(serde_json::json!({"type": "stop", "matched": [9, 10]})), + Some(Matched::Tokens(vec![9, 10])) + ); + // A stop with no `matched`, and the non-stop reasons. + assert_eq!(m(serde_json::json!({"type": "stop"})), None); + assert_eq!(m(serde_json::json!({"type": "length", "length": 8})), None); + assert_eq!(m(serde_json::json!({"type": "abort"})), None); + } + + /// Python's abort dict always carries `status_code`/`err_type`, null included, + /// and the reason is echoed verbatim into `meta_info.finish_reason` — so the + /// typed form must serialize back key-for-key, nulls and all. A dropped null + /// key is a silently changed response body. + #[test] + fn finish_reason_round_trips_python_shapes() { + for wire in [ + serde_json::json!({"type": "stop", "matched": 9}), + serde_json::json!({"type": "stop", "matched": ""}), + serde_json::json!({"type": "stop", "matched": [9, 10]}), + serde_json::json!({"type": "length", "length": 8}), + serde_json::json!({ + "type": "abort", "message": "Aborted", + "status_code": null, "err_type": null + }), + serde_json::json!({ + "type": "abort", "message": "over the limit", + "status_code": 400, "err_type": "BadRequestError" + }), + ] { + let parsed: FinishReason = serde_json::from_value(wire.clone()).unwrap(); + assert!( + matches!(parsed, FinishReason::Known(_)), + "must classify, not fall back to Unknown: {wire}" + ); + assert_eq!(serde_json::to_value(&parsed).unwrap(), wire); + } + } + + /// A `type` this build doesn't know must not fail the decode — the header holds + /// the whole batch, so a rejected frame drops every request in it, not just the + /// one carrying the new reason. It is kept verbatim and echoed unchanged. + #[test] + fn unknown_finish_type_is_preserved_not_rejected() { + let wire = serde_json::json!({"type": "tool_calls", "name": "search"}); + let parsed: FinishReason = serde_json::from_value(wire.clone()).unwrap(); + assert!(matches!(parsed, FinishReason::Unknown(_))); + assert_eq!(serde_json::to_value(&parsed).unwrap(), wire); + // Unknown reasons classify as neither an abort nor a matched stop. + assert!(parsed.abort_status().is_none()); + assert!(parsed.matched().is_none()); + } +} diff --git a/rust/sglang-server/src/message/request.rs b/rust/sglang-server/src/message/request.rs index 9b2c5e3d9eb7..00277a543455 100644 --- a/rust/sglang-server/src/message/request.rs +++ b/rust/sglang-server/src/message/request.rs @@ -486,6 +486,10 @@ fn fan_out( mod tests { use super::*; + /// Vocab size for tests that aren't about the vocab bound (see + /// `sampling::tests::TEST_VOCAB`). + const TEST_VOCAB: u64 = 1000; + fn requests(body: &str) -> Result<(Vec, bool), Error> { serde_json::from_str::(body) .unwrap() @@ -558,7 +562,7 @@ mod tests { // Parallel sampling is rejected where Python reads it — in the params, // at normalization (the ingress step), not here. let (mut ps, _) = requests(r#"{"text": "a", "sampling_params": {"n": 2}}"#).unwrap(); - assert!(ps[0].sampling_params.normalize(false, None).is_err()); + assert!(ps[0].sampling_params.normalize(false, TEST_VOCAB).is_err()); } /// Unported `GenerateReqInput` fields are IGNORED, not rejected. diff --git a/rust/sglang-server/src/message/sampling.rs b/rust/sglang-server/src/message/sampling.rs index 309dfbafb2e3..99a051eaaae1 100644 --- a/rust/sglang-server/src/message/sampling.rs +++ b/rust/sglang-server/src/message/sampling.rs @@ -26,41 +26,197 @@ //! * `n > 1` (parallel sampling) is rejected — the rust egress maps one rid to //! one response, so every sample past the first would be dropped. -use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; +use std::fmt; +use serde::de::value::{MapAccessDeserializer, SeqAccessDeserializer}; +use serde::de::{MapAccess, SeqAccess, Visitor}; +use serde::{Deserialize, Deserializer, Serialize}; + +use super::OneOrMany; use crate::error::Error; +use crate::utils::regex::RegexPattern; + +/// `_SAMPLING_EPS` — temperatures in `[0, eps)` mean greedy decoding. +const SAMPLING_EPS: f64 = 1e-6; +/// `TOP_K_ALL = 1 << 30` — `top_k` sentinel for "consider the whole vocabulary". +const TOP_K_ALL: i64 = 1 << 30; +/// Most stop STRINGS accepted per request. The scheduler scans the decoded text +/// once per stop per decode step, so this is a per-step multiplier: 50k stops +/// measured 20.4 ms/step from a 586 KB body. +const MAX_STOP_COUNT: usize = 32; +/// Longest `stop_regex` accepted. A 1 MB literal pattern takes ~677 ms just to +/// compile, and that cost lands on the scheduler. +const MAX_STOP_REGEX_LEN: usize = 256; +/// Most `stop_regex` patterns accepted per request. Python's `re` cache holds 512 +/// (`re._MAXCACHE`), so past that every pattern recompiles on every decode step. +const MAX_STOP_REGEX_COUNT: usize = 32; + +/// One module per field default, each exposing the two hooks serde needs under +/// one name: `default` (key absent) and `deserialize` (key present — including +/// an explicit `null`, which Python's `__post_init__` maps back to the default: +/// "callers can pass null without crashing verify"). They cannot be one function +/// — serde calls `default()` with no arguments and `deserialize(deserializer)` — +/// but `deserialize` defers to `default()`, so the value is written once. +macro_rules! defaulted { + ($($name:ident: $ty:ty = $value:expr;)*) => {$( + mod $name { + use serde::{Deserialize, Deserializer}; + + pub(super) fn default() -> $ty { $value } + + pub(super) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<$ty, D::Error> { + Ok(Option::<$ty>::deserialize(d)?.unwrap_or_else(default)) + } + } + )*}; +} + +defaulted! { + f64_one: f64 = 1.0; + f64_zero: f64 = 0.0; + i64_top_k_all: i64 = super::TOP_K_ALL; + i64_zero: i64 = 0; + i64_one: i64 = 1; + bool_false: bool = false; + bool_true: bool = true; +} + +/// `max_new_tokens` is `Optional[int] = 128`: an *absent* key means 128, but an +/// explicit `null` means None (no limit) — so it keeps its `Option` rather than +/// going through [`defaulted`]. +fn max_new_tokens_default() -> Option { + Some(128) +} /// The sampling parameters of one `/generate` request. Deserialized from the /// client's `sampling_params` object (unknown keys are a 400, mirroring Python's /// `SamplingParams(**kwargs)` TypeError) and serialized by field name into the /// scheduler header once [`normalize`](Self::normalize) has run. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct SamplingParams { + // --- API parameters (set by callers) --- + #[serde(default = "max_new_tokens_default")] pub max_new_tokens: Option, + /// API input alias, copied to `stop_strs` then cleared by `normalize`. + #[serde(default)] + pub stop: Option>, + /// Python `Optional[Set[int]]`. A `null` *element* is a 400 here where Python + /// filters it out — a typed list can't hold one, and it is malformed input. + #[serde(default)] + pub stop_token_ids: Option>, + /// API input alias, copied to `stop_regex_strs` then cleared by `normalize`. + #[serde(default)] + pub stop_regex: Option>, + #[serde( + default = "f64_one::default", + deserialize_with = "f64_one::deserialize" + )] pub temperature: f64, -} + #[serde( + default = "f64_one::default", + deserialize_with = "f64_one::deserialize" + )] + pub top_p: f64, + #[serde( + default = "i64_top_k_all::default", + deserialize_with = "i64_top_k_all::deserialize" + )] + pub top_k: i64, + #[serde( + default = "f64_zero::default", + deserialize_with = "f64_zero::deserialize" + )] + pub min_p: f64, + #[serde( + default = "f64_zero::default", + deserialize_with = "f64_zero::deserialize" + )] + pub frequency_penalty: f64, + #[serde( + default = "f64_zero::default", + deserialize_with = "f64_zero::deserialize" + )] + pub presence_penalty: f64, + #[serde( + default = "f64_one::default", + deserialize_with = "f64_one::deserialize" + )] + pub repetition_penalty: f64, + #[serde( + default = "i64_zero::default", + deserialize_with = "i64_zero::deserialize" + )] + pub min_new_tokens: i64, + #[serde( + default = "i64_one::default", + deserialize_with = "i64_one::deserialize" + )] + pub n: i64, + #[serde(default)] + pub json_schema: Option, + #[serde(default)] + pub regex: Option, + #[serde(default)] + pub ebnf: Option, + #[serde(default)] + pub structural_tag: Option, + #[serde( + default = "bool_false::default", + deserialize_with = "bool_false::deserialize" + )] + pub ignore_eos: bool, + #[serde( + default = "bool_true::default", + deserialize_with = "bool_true::deserialize" + )] + pub skip_special_tokens: bool, + #[serde( + default = "bool_true::default", + deserialize_with = "bool_true::deserialize" + )] + pub spaces_between_special_tokens: bool, + #[serde( + default = "bool_false::default", + deserialize_with = "bool_false::deserialize" + )] + pub no_stop_trim: bool, + #[serde(default)] + pub stream_interval: Option, + /// Token id (as a string key, matching Python) → bias. Keys are vocab-bounded + /// by [`verify`](Self::verify). + #[serde(default)] + pub logit_bias: Option>, + #[serde(default)] + pub sampling_seed: Option, + /// Opaque JSON object forwarded to a custom logit processor. Python types it + /// as `Dict[str, JsonScalar | list | dict]`; it is never inspected here. + #[serde(default)] + pub custom_params: Option, -impl SamplingParams { - /// Normalize the fields, applying defaults and coercing types. Mirrors - /// Python's `SamplingParams.__post_init__` → `normalize`. - pub fn normalize( - &mut self, - _skip_tokenizer_init: bool, - _vocab_size: Option, - ) -> Result<(), Error> { - todo!() - } - - /// Verify the normalized fields are in range. Mirrors Python's - /// `SamplingParams.verify`. - pub fn verify(&self) -> Result<(), String> { - todo!() - } - - pub fn max_tokens_len(&self) -> usize { - todo!() - } + // --- Internal fields (populated by the pipeline below, not API-facing) --- + // + // 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 + // the dangerous one — `{"is_normalized": true, "temperature": 0.0}` makes + // `post_init` early-return, so the greedy mapping never runs and temperature 0 + // reaches the scheduler's `logits.div_()`; `stop` would likewise be dropped + // without ever reaching `stop_strs`. They still SERIALIZE: the scheduler needs + // them on the wire. + /// From `stop`; a list after `normalize` (Python widens str → [str] there). + #[serde(skip_deserializing)] + pub stop_strs: Vec, + /// From `stop_regex`. + #[serde(skip_deserializing)] + pub stop_regex_strs: Vec, + #[serde(skip_deserializing)] + pub stop_str_max_len: usize, + #[serde(skip_deserializing)] + pub stop_regex_max_len: usize, + /// Set by `normalize`; tells the scheduler its own pass can early-return. + #[serde(skip_deserializing)] + pub is_normalized: bool, } /// The `/generate` body's `sampling_params`: one object (broadcast to every @@ -71,10 +227,772 @@ impl SamplingParams { /// losing the field-level message ("unknown field `temperature`, expected one of /// …") that makes a typo actionable. Object-vs-list is unambiguous here, so a /// single `deserialize_any` dispatch keeps the inner error verbatim. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq)] pub enum SamplingParamsInput { /// Boxed: `SamplingParams` is ~440 bytes, so an inline variant would make /// every `GenerateBody` that big regardless of which form arrived. One(Box), Many(Vec), } + +impl<'de> Deserialize<'de> for SamplingParamsInput { + fn deserialize>(deserializer: D) -> Result { + struct InputVisitor; + + impl<'de> Visitor<'de> for InputVisitor { + type Value = SamplingParamsInput; + + fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("a sampling_params object, or a list of them (one per prompt)") + } + + fn visit_map>(self, map: A) -> Result { + SamplingParams::deserialize(MapAccessDeserializer::new(map)) + .map(|p| SamplingParamsInput::One(Box::new(p))) + } + + fn visit_seq>(self, seq: A) -> Result { + Vec::deserialize(SeqAccessDeserializer::new(seq)).map(SamplingParamsInput::Many) + } + } + + deserializer.deserialize_any(InputVisitor) + } +} + +impl Default for SamplingParams { + fn default() -> Self { + // Each field reads the same `default()` the serde attribute above names, + // so the values still live in one place — without re-parsing `{}` on + // every call (once per prompt of every params-less `/generate`). The + // struct literal makes completeness a compile error. + Self { + max_new_tokens: max_new_tokens_default(), + stop: None, + stop_token_ids: None, + stop_regex: None, + temperature: f64_one::default(), + top_p: f64_one::default(), + top_k: i64_top_k_all::default(), + min_p: f64_zero::default(), + frequency_penalty: f64_zero::default(), + presence_penalty: f64_zero::default(), + repetition_penalty: f64_one::default(), + min_new_tokens: i64_zero::default(), + n: i64_one::default(), + json_schema: None, + regex: None, + ebnf: None, + structural_tag: None, + ignore_eos: bool_false::default(), + skip_special_tokens: bool_true::default(), + spaces_between_special_tokens: bool_true::default(), + no_stop_trim: bool_false::default(), + custom_params: None, + stream_interval: None, + logit_bias: None, + sampling_seed: None, + stop_strs: Vec::new(), + stop_regex_strs: Vec::new(), + stop_str_max_len: 0, + stop_regex_max_len: 0, + is_normalized: false, + } + } +} + +impl SamplingParams { + /// `__post_init__` → `normalize` → `verify`, the order + /// `TokenizerManager._create_tokenized_object` runs them in. `Err` is a + /// request-local 400. `skip_tokenizer_init` stands in for Python's + /// `tokenizer is None`; `vocab_size` bounds `logit_bias` keys. + pub fn normalize(&mut self, skip_tokenizer_init: bool, vocab_size: u64) -> Result<(), Error> { + self.post_init(); + self.normalize_stops(skip_tokenizer_init)?; + self.verify(vocab_size) + } + + /// Python `__post_init__` (minus the null-to-default coercions, which the + /// null-tolerant deserializers above already did): copy the API aliases into + /// the internal fields and apply the greedy / `top_k` special cases. + fn post_init(&mut self) { + // Python's `__post_init__` guard. Without it a second `normalize` reads + // the aliases `normalize_stops` already cleared and silently wipes + // `stop_strs`/`stop_regex_strs` to empty — the request would stop + // matching its stop strings. + if self.is_normalized { + return; + } + // Moved out, not cloned: `normalize_stops` clears both aliases anyway. + self.stop_strs = take_one_or_many(self.stop.take()); + self.stop_regex_strs = take_one_or_many(self.stop_regex.take()); + // Python drops null entries and maps an empty set to None. + if self.stop_token_ids.as_ref().is_some_and(|v| v.is_empty()) { + self.stop_token_ids = None; + } + if (0.0..SAMPLING_EPS).contains(&self.temperature) { + // Greedy: temperature ~0 → temperature=1.0, top_k=1. + self.temperature = 1.0; + self.top_k = 1; + } + if self.top_k == -1 { + self.top_k = TOP_K_ALL; // -1 disables top_k → whole vocabulary + } + } + + /// Python `normalize(tokenizer)`: size the stop match windows, reject + /// tokenizer-dependent features when there is no tokenizer, and clear the API + /// aliases so they don't ride the wire twice. + fn normalize_stops(&mut self, skip_tokenizer_init: bool) -> Result<(), Error> { + // Match window: UTF-8 byte length is a safe upper bound on the token count. + self.stop_str_max_len = self.stop_strs.iter().map(|s| s.len()).max().unwrap_or(0); + // Validate + bound every stop_regex here, before it can reach the + // scheduler's `re.search` (see `RegexPattern`). A rejected pattern is a + // 400 for this request; an accepted one carries a bound the scheduler uses + // to size its match window. + if self.stop_strs.len() > MAX_STOP_COUNT { + return Err(bad(format!( + "at most {MAX_STOP_COUNT} stop strings are allowed, got {}", + self.stop_strs.len() + ))); + } + if self.stop_regex_strs.len() > MAX_STOP_REGEX_COUNT { + return Err(bad(format!( + "at most {MAX_STOP_REGEX_COUNT} stop_regex patterns are allowed, got {}", + self.stop_regex_strs.len() + ))); + } + let mut stop_regex_max_len = 0; + for pattern in &self.stop_regex_strs { + if pattern.len() > MAX_STOP_REGEX_LEN { + return Err(bad(format!( + "stop_regex is {} bytes, over the {MAX_STOP_REGEX_LEN}-byte limit", + pattern.len() + ))); + } + let pattern = RegexPattern::try_from(pattern.as_str()) + .map_err(|e| bad(format!("stop_regex {pattern:?} is invalid: {e}")))?; + stop_regex_max_len = stop_regex_max_len.max(pattern.max_len()); + } + self.stop_regex_max_len = stop_regex_max_len; + + // Python `raise_if_tokenizer_required`: these need `tokenizer.decode` / + // `eos_token_id`, which `skip_tokenizer_init` does not have. + if skip_tokenizer_init { + if !self.stop_strs.is_empty() { + return Err(bad( + "stop is unavailable when skip_tokenizer_init=True (requires a \ + tokenizer to decode tokens to text for matching)" + .into(), + )); + } + if !self.stop_regex_strs.is_empty() { + return Err(bad( + "stop_regex is unavailable when skip_tokenizer_init=True (requires a \ + tokenizer to decode tokens to text for matching)" + .into(), + )); + } + if self.min_new_tokens > 0 { + return Err(bad(format!( + "min_new_tokens={} is unavailable when skip_tokenizer_init=True \ + (requires a tokenizer for eos_token_id)", + self.min_new_tokens + ))); + } + } + + self.stop = None; + self.stop_regex = None; + self.is_normalized = true; + Ok(()) + } + + /// Python `verify(vocab_size)` — the same ranges, messages and mutual + /// exclusions, plus the rust-server `n == 1` restriction. + fn verify(&self, vocab_size: u64) -> Result<(), Error> { + if !self.temperature.is_finite() || self.temperature < 0.0 { + return Err(bad(format!( + "temperature must be a non-negative finite number, got {}", + self.temperature + ))); + } + if !(self.top_p > 0.0 && self.top_p <= 1.0) { + return Err(bad(format!("top_p must be in (0, 1], got {}", self.top_p))); + } + if !(0.0..=1.0).contains(&self.min_p) { + return Err(bad(format!("min_p must be in [0, 1], got {}", self.min_p))); + } + if self.top_k < 1 { + return Err(bad(format!( + "top_k must be -1 (disable) or at least 1, got {}", + self.top_k + ))); + } + if !(-2.0..=2.0).contains(&self.frequency_penalty) { + return Err(bad(format!( + "frequency_penalty must be in [-2, 2], got {}", + self.frequency_penalty + ))); + } + if !(-2.0..=2.0).contains(&self.presence_penalty) { + return Err(bad(format!( + "presence_penalty must be in [-2, 2], got {}", + self.presence_penalty + ))); + } + if !(self.repetition_penalty > 0.0 && self.repetition_penalty <= 2.0) { + return Err(bad(format!( + "repetition_penalty must be in (0, 2], got {}", + self.repetition_penalty + ))); + } + if self.min_new_tokens < 0 { + return Err(bad(format!( + "min_new_tokens must be non-negative, got {}", + self.min_new_tokens + ))); + } + // `None` = no limit, so the max_new_tokens checks only apply when set. + if let Some(max_new_tokens) = self.max_new_tokens { + if max_new_tokens < 0 { + return Err(bad(format!( + "max_new_tokens must be at least 0, got {max_new_tokens}" + ))); + } + if self.min_new_tokens > max_new_tokens { + return Err(bad(format!( + "min_new_tokens must be in [0, max_new_tokens({max_new_tokens})], got {}", + self.min_new_tokens + ))); + } + } + // A non-numeric bias key raises in the scheduler's `int(key)`, and an + // out-of-vocabulary one would index past the logits row, so both are + // rejected here (Python `verify` does the same, in that order). Only the + // *range* check needs the vocab size (`None` = unknown, skip it); the key + // format is checked either way, since `int(key)` runs regardless. + if let Some(logit_bias) = &self.logit_bias { + for key in logit_bias.keys() { + let token_id: u64 = key + .parse() + .map_err(|_| bad(format!("logit_bias keys must be token ids, got {key:?}")))?; + if token_id >= vocab_size { + return Err(bad(format!( + "logit_bias must have keys in [0, {}], got {token_id}", + vocab_size - 1 + ))); + } + } + } + // Grammars are mutually exclusive. + let grammars = [&self.json_schema, &self.regex, &self.ebnf] + .iter() + .filter(|g| g.is_some()) + .count(); + if grammars > 1 { + return Err(bad( + "Only one of regex, json_schema, or ebnf can be set".into() + )); + } + // Not a Python restriction: the rust egress maps one rid to one response, + // so parallel sampling would drop all but the first sample. This is the + // only place it is rejected — `n` lives in `sampling_params`, where + // Python reads it, and the `/generate` body has no `n` of its own. + if self.n != 1 { + return Err(bad(format!( + "n must be 1 (parallel sampling is not supported), got {}", + self.n + ))); + } + Ok(()) + } +} + +fn bad(msg: String) -> Error { + Error::Validation(msg) +} + +/// Widen a `str | [str]` API alias into the list form the internal field holds. +fn take_one_or_many(v: Option>) -> Vec { + match v { + None => Vec::new(), + Some(OneOrMany::One(s)) => vec![s], + Some(OneOrMany::Many(v)) => v, + } +} + +#[cfg(test)] +mod tests { + + use super::*; + + /// Vocab size for tests that aren't about the vocab bound at all. It is + /// mandatory now (`ServerArgs::validate_mandatory` rejects a boot without + /// one), so there is no longer an "unknown vocab" case to pass instead — + /// this is just a value large enough to stay out of the way. + const TEST_VOCAB: u64 = 1000; + + /// End to end through the path `/generate` takes: a bounded `stop_regex` + /// reaches the wire with its real length, and a malformed one is a 400. + #[test] + fn stop_regex_bound_reaches_the_wire() { + assert_eq!(norm(r#"{"stop_regex": "\\d{6}"}"#).stop_regex_max_len, 6); + assert_eq!(norm(r#"{"temperature": 0.7}"#).stop_regex_max_len, 0); + let _ = norm_err(r#"{"stop_regex": "("}"#); + let _ = norm_err(r#"{"stop_regex": ["\\d{6}", "("]}"#); + } + + /// Parse client JSON exactly as `/generate` does, then run the full pipeline. + fn norm(json: &str) -> SamplingParams { + let mut sp: SamplingParams = serde_json::from_str(json).expect("parses"); + sp.normalize(false, TEST_VOCAB).expect("normalizes"); + sp + } + + fn norm_err(json: &str) -> Error { + let mut sp: SamplingParams = serde_json::from_str(json).expect("parses"); + sp.normalize(false, TEST_VOCAB).expect_err("must reject") + } + + /// The wire shape the scheduler decodes: a map of field names → values — the + /// same `serde_json::Value` the header encoder hands to msgpack. + fn wire(sp: &SamplingParams) -> serde_json::Value { + serde_json::to_value(sp).expect("serializes") + } + + fn get(v: &serde_json::Value, key: &str) -> Option { + v.get(key).cloned() + } + + #[test] + fn greedy_sets_temp_one_topk_one() { + let sp = norm(r#"{"temperature": 0.0}"#); + assert_eq!(sp.temperature, 1.0); + assert_eq!(sp.top_k, 1); + assert!(sp.is_normalized); + } + + #[test] + fn topk_minus_one_becomes_all() { + assert_eq!(norm(r#"{"temperature": 0.7}"#).top_k, TOP_K_ALL); + assert_eq!( + norm(r#"{"top_k": -1, "temperature": 0.7}"#).top_k, + TOP_K_ALL + ); + } + + #[test] + fn stop_list_and_max_len_by_bytes() { + let sp = norm(r#"{"stop": ["Question:", "\n\n"]}"#); + assert_eq!(sp.stop_strs.len(), 2); + assert_eq!(sp.stop_str_max_len, 9); // "Question:" (ASCII) + // The API alias is cleared, so it never rides the wire twice. + assert!(sp.stop.is_none()); + } + + /// A multi-byte stop char must use its byte length as the window bound: `𓀀` + /// is 1 char but 4 UTF-8 bytes (and 3 tokens on Qwen3). Char count (1) would + /// under-size the tail and miss the stop; byte count (4) ≥ the token span. + #[test] + fn stop_str_max_len_uses_bytes_not_chars() { + let sp = norm(r#"{"stop": "𓀀"}"#); + assert_eq!("𓀀".chars().count(), 1); + assert_eq!("𓀀".len(), 4); + assert_eq!(sp.stop_strs, vec!["𓀀".to_string()]); // scalar widened to a list + assert_eq!(sp.stop_str_max_len, 4); + } + + #[test] + fn no_stop_yields_empty_list_zero_len() { + let sp = norm(r#"{"temperature": 0.0}"#); + assert!(sp.stop_strs.is_empty()); + assert_eq!(sp.stop_str_max_len, 0); + } + + /// The wire map is what the scheduler's msgspec decoder reads by field name: + /// the normalized values must be present under the Python names, and the + /// `is_normalized` flag must be set so its own pass early-returns. + #[test] + fn wire_map_carries_python_field_names() { + let sp = norm(r#"{"temperature": 0.7, "max_new_tokens": 64, "ignore_eos": true}"#); + let w = wire(&sp); + assert_eq!(get(&w, "temperature").unwrap().as_f64(), Some(0.7)); + assert_eq!(get(&w, "max_new_tokens").unwrap().as_i64(), Some(64)); + assert_eq!(get(&w, "ignore_eos").unwrap().as_bool(), Some(true)); + assert_eq!(get(&w, "top_k").unwrap().as_i64(), Some(TOP_K_ALL)); + assert_eq!(get(&w, "is_normalized").unwrap().as_bool(), Some(true)); + assert_eq!(get(&w, "stop_str_max_len").unwrap().as_i64(), Some(0)); + // Unset optionals ride as null, NOT omitted: the msgpack wire is + // positional (`array_like=True`), so a skipped field would shift every + // later one. JSON keeps the names, which is what this test is about. + assert!(get(&w, "regex").unwrap().is_null()); + assert!(get(&w, "stop").unwrap().is_null()); + } + + /// `max_new_tokens` is the one field where absent and null differ: absent = + /// 128 (the Python field default), explicit null = None (no limit). + #[test] + fn max_new_tokens_null_is_unlimited_absent_is_default() { + assert_eq!(norm("{}").max_new_tokens, Some(128)); + assert_eq!(norm(r#"{"max_new_tokens": null}"#).max_new_tokens, None); + // None = no limit, so a large min_new_tokens is not a range error. + let sp = norm(r#"{"max_new_tokens": null, "min_new_tokens": 4096}"#); + assert_eq!(sp.min_new_tokens, 4096); + } + + /// The 30 wire slots, in Python's declaration order. + /// + /// `SamplingParams` is `msgspec.Struct(array_like=True)` on the Python side, so + /// the header carries an ARRAY and every field is identified by POSITION. Two + /// things follow, and both are asserted below: the order must match + /// `SamplingParams.__struct_fields__` exactly, and no field may be omitted — + /// a `skip_serializing_if` anywhere would shorten the array and shift every + /// later field onto the wrong scheduler slot. + /// + /// KEEP IN SYNC with `sampling_params.py`. This list is an external-source + /// literal: it is the Python declaration order, not this file's. + const WIRE_ORDER: &[&str] = &[ + "max_new_tokens", + "stop", + "stop_token_ids", + "stop_regex", + "temperature", + "top_p", + "top_k", + "min_p", + "frequency_penalty", + "presence_penalty", + "repetition_penalty", + "min_new_tokens", + "n", + "json_schema", + "regex", + "ebnf", + "structural_tag", + "ignore_eos", + "skip_special_tokens", + "spaces_between_special_tokens", + "no_stop_trim", + "stream_interval", + "logit_bias", + "sampling_seed", + "custom_params", + "stop_strs", + "stop_regex_strs", + "stop_str_max_len", + "stop_regex_max_len", + "is_normalized", + ]; + + /// Every field reaches the wire, at the position Python expects. + /// + /// Each slot is given a DISTINCT value so a swap of two same-typed neighbours + /// is caught by value, not just by arity — the failure mode a length check + /// alone would wave through. Regression for the map-vs-array break: this used + /// to serialize as a map, which `array_like=True` rejects outright + /// (`Expected array, got object`), so every generate request failed to decode. + #[test] + fn wire_is_positional_and_complete() { + let sp = SamplingParams { + max_new_tokens: Some(11), + stop_token_ids: Some(vec![12]), + temperature: 0.13, + top_p: 0.14, + top_k: 15, + min_p: 0.16, + frequency_penalty: 0.17, + presence_penalty: 0.18, + repetition_penalty: 0.19, + min_new_tokens: 20, + n: 1, + json_schema: Some("22".into()), + regex: Some("23".into()), + ebnf: Some("24".into()), + structural_tag: Some("25".into()), + ignore_eos: true, + skip_special_tokens: false, + spaces_between_special_tokens: false, + no_stop_trim: true, + stream_interval: Some(30), + sampling_seed: Some(31), + ..Default::default() + }; + let buf = rmp_serde::to_vec(&sp).expect("serializes"); + let v = rmpv::decode::read_value(&mut &buf[..]).expect("decodes"); + let arr = v + .as_array() + .expect("array_like=True means an ARRAY, not a map"); + + assert_eq!( + arr.len(), + WIRE_ORDER.len(), + "every field must be emitted: a shorter array shifts later fields onto \ + the wrong scheduler slot" + ); + // Spot-check the positions whose neighbours share a type, where a swap + // would otherwise be invisible. + let at = |name: &str| WIRE_ORDER.iter().position(|f| *f == name).unwrap(); + assert_eq!(arr[at("max_new_tokens")].as_i64(), Some(11)); + assert_eq!(arr[at("temperature")].as_f64(), Some(0.13)); + assert_eq!(arr[at("top_p")].as_f64(), Some(0.14)); + assert_eq!(arr[at("top_k")].as_i64(), Some(15)); + assert_eq!(arr[at("min_p")].as_f64(), Some(0.16)); + assert_eq!(arr[at("frequency_penalty")].as_f64(), Some(0.17)); + assert_eq!(arr[at("presence_penalty")].as_f64(), Some(0.18)); + assert_eq!(arr[at("repetition_penalty")].as_f64(), Some(0.19)); + assert_eq!(arr[at("json_schema")].as_str(), Some("22")); + assert_eq!(arr[at("regex")].as_str(), Some("23")); + assert_eq!(arr[at("ebnf")].as_str(), Some("24")); + assert_eq!(arr[at("structural_tag")].as_str(), Some("25")); + assert_eq!(arr[at("ignore_eos")].as_bool(), Some(true)); + assert_eq!(arr[at("skip_special_tokens")].as_bool(), Some(false)); + assert_eq!(arr[at("no_stop_trim")].as_bool(), Some(true)); + assert_eq!(arr[at("stream_interval")].as_i64(), Some(30)); + assert_eq!(arr[at("sampling_seed")].as_i64(), Some(31)); + // Unset optionals ride as nil rather than being skipped. + assert!(arr[at("stop")].is_nil()); + assert!(arr[at("logit_bias")].is_nil()); + assert!(arr[at("custom_params")].is_nil()); + // `normalize` outputs occupy the tail. + assert!(arr[at("stop_strs")].is_array()); + assert_eq!(arr[at("is_normalized")].as_bool(), Some(false)); + } + + #[test] + fn verify_rejects_out_of_range() { + for (json, want) in [ + (r#"{"top_p": 2.0}"#, "top_p"), + (r#"{"top_k": 0, "temperature": 0.7}"#, "top_k"), + (r#"{"min_p": 1.5}"#, "min_p"), + (r#"{"frequency_penalty": 3.0}"#, "frequency_penalty"), + (r#"{"presence_penalty": -3.0}"#, "presence_penalty"), + (r#"{"repetition_penalty": 0.0}"#, "repetition_penalty"), + ( + r#"{"max_new_tokens": 8, "min_new_tokens": 9}"#, + "min_new_tokens", + ), + (r#"{"temperature": -0.1}"#, "temperature"), + (r#"{"max_new_tokens": -1}"#, "max_new_tokens"), + (r#"{"regex": "a", "ebnf": "b"}"#, "Only one of"), + (r#"{"n": 2}"#, "n must be 1"), + ] { + let err = norm_err(json).to_string(); + assert!( + err.contains(want), + "{json} must be rejected for {want}: {err}" + ); + } + } + + /// The inclusive bounds must ACCEPT their endpoints. Only the rejecting side + /// was covered, and far from the edge (`frequency_penalty: 3.0`), so flipping + /// any `..=` to `..` — or `>= 1` to `> 1` — would 400 legitimate requests + /// without failing a single test. + #[test] + fn verify_accepts_inclusive_boundaries() { + for json in [ + r#"{"top_p": 1.0, "temperature": 0.7}"#, + r#"{"min_p": 0.0, "temperature": 0.7}"#, + r#"{"min_p": 1.0, "temperature": 0.7}"#, + r#"{"top_k": 1, "temperature": 0.7}"#, + r#"{"frequency_penalty": 2.0}"#, + r#"{"frequency_penalty": -2.0}"#, + r#"{"presence_penalty": 2.0}"#, + r#"{"presence_penalty": -2.0}"#, + r#"{"repetition_penalty": 2.0}"#, + r#"{"max_new_tokens": 0}"#, + r#"{"min_new_tokens": 0}"#, + // min == max is in range: `[0, max_new_tokens]` is inclusive. + r#"{"max_new_tokens": 8, "min_new_tokens": 8}"#, + // Greedy: temperature 0 is the documented sentinel, not an under-run. + r#"{"temperature": 0.0}"#, + r#"{"n": 1}"#, + ] { + let mut sp: SamplingParams = serde_json::from_str(json).expect("parses"); + sp.normalize(false, TEST_VOCAB) + .unwrap_or_else(|e| panic!("{json} is in range but was rejected: {e}")); + } + } + + /// And the first value past each endpoint is still rejected — the pair of + /// tests brackets the boundary instead of testing one side of it. + #[test] + fn verify_rejects_just_past_the_boundaries() { + for json in [ + r#"{"top_p": 0.0, "temperature": 0.7}"#, // exclusive lower bound + r#"{"repetition_penalty": 0.0}"#, // exclusive lower bound + r#"{"top_k": 0, "temperature": 0.7}"#, + r#"{"min_p": 1.0000001, "temperature": 0.7}"#, + r#"{"frequency_penalty": 2.0000001}"#, + r#"{"presence_penalty": -2.0000001}"#, + r#"{"repetition_penalty": 2.0000001}"#, + r#"{"max_new_tokens": -1}"#, + r#"{"min_new_tokens": -1}"#, + r#"{"max_new_tokens": 8, "min_new_tokens": 9}"#, + ] { + let mut sp: SamplingParams = serde_json::from_str(json).expect("parses"); + assert!( + sp.normalize(false, TEST_VOCAB).is_err(), + "{json} is out of range but was accepted" + ); + } + } + + /// A wrong JSON type for a numeric field is rejected at parse time — it must + /// NOT silently fall back to the default (`temperature: "bad"` has different + /// semantics than an unset temperature). + #[test] + fn wrong_typed_field_is_rejected() { + for json in [ + r#"{"temperature": "bad"}"#, + r#"{"top_k": "bad"}"#, + r#"{"max_new_tokens": "bad"}"#, + r#"{"stop": 3}"#, + ] { + assert!( + serde_json::from_str::(json).is_err(), + "{json} must not parse" + ); + } + } + + /// An unknown key is a 400, mirroring Python's `SamplingParams(**kwargs)` + /// TypeError — a typo must not be silently ignored. (The bogus key is + /// deliberately not a near-miss of a real field: an editor spell-checker + /// kept "correcting" a misspelling here into a valid name, which silently + /// turned this assertion into a tautology.) + #[test] + fn unknown_field_is_rejected() { + assert!(serde_json::from_str::(r#"{"zzz_not_a_field": 1}"#).is_err()); + // ...while every declared field still parses. + assert!(serde_json::from_str::(r#"{"temperature": 0.7}"#).is_ok()); + } + + /// A present-but-null non-optional field keeps the default (Python's + /// `x if x is not None`) — null is absent, not a wrong type. + #[test] + fn null_field_keeps_default() { + let sp = norm(r#"{"temperature": null, "top_k": null, "skip_special_tokens": null}"#); + assert_eq!(sp.temperature, 1.0); + assert_eq!(sp.top_k, TOP_K_ALL); + assert!(sp.skip_special_tokens); + } + + /// `normalize` must be idempotent: `post_init` reads the API aliases, which + /// `normalize_stops` clears, so without Python's `if self.is_normalized: + /// return` guard a second call wipes `stop_strs` and drops the stop bound to + /// zero — silently, leaving a request that never stops. + #[test] + fn normalize_is_idempotent() { + let mut once = norm(r#"{"stop": ["END", "STOP"], "stop_regex": "\\d{3}"}"#); + let twice = { + let mut p = once.clone(); + p.normalize(false, TEST_VOCAB).expect("second normalize"); + p + }; + assert_eq!(once, twice, "a second normalize must change nothing"); + assert_eq!(twice.stop_strs, vec!["END".to_string(), "STOP".to_string()]); + assert_eq!(twice.stop_str_max_len, 4); + assert_eq!(twice.stop_regex_max_len, 3); + + // Greedy handling must not re-fire either: temperature is 1.0 after the + // first pass, which is not in the greedy window. + once.normalize(false, TEST_VOCAB).unwrap(); + assert_eq!(once.top_k, twice.top_k); + } + + /// `skip_tokenizer_init` has no tokenizer, so the text-matching stop features + /// and `min_new_tokens` (needs eos_token_id) are 400s, not silent no-ops. + /// Mirrors Python `raise_if_tokenizer_required`. + #[test] + fn tokenizer_dependent_features_rejected_without_tokenizer() { + for json in [ + r#"{"stop": "END"}"#, + r#"{"stop_regex": "\\d+"}"#, + r#"{"min_new_tokens": 1}"#, + ] { + let mut sp: SamplingParams = serde_json::from_str(json).expect("parses"); + assert!( + sp.normalize(true, TEST_VOCAB).is_err(), + "{json} must be rejected under skip_tokenizer_init" + ); + } + // The same params are fine when a tokenizer is present. + let mut sp: SamplingParams = serde_json::from_str(r#"{"stop": "END"}"#).unwrap(); + assert!(sp.normalize(false, TEST_VOCAB).is_ok()); + } + + /// `logit_bias` keys index the logits row, so an out-of-vocab id is a 400 + /// (Python `verify`'s vocab bound). The bound is exclusive, and it always + /// applies — `vocab_size` is mandatory, so there is no "unknown vocab" path + /// that skips this. + #[test] + fn logit_bias_keys_are_vocab_bounded() { + let mut sp: SamplingParams = + serde_json::from_str(r#"{"logit_bias": {"1000": 1.0}}"#).unwrap(); + assert!(sp.clone().normalize(false, 1000).is_err()); + assert!(sp.normalize(false, 1001).is_ok()); + + let mut sp: SamplingParams = + serde_json::from_str(r#"{"logit_bias": {"999": -1.0}}"#).unwrap(); + assert!(sp.normalize(false, 1000).is_ok()); + } + + /// The key *format* check is separate from the vocab bound: the scheduler + /// does `logit_bias[i, int(key)]`, so a key that is not a parseable + /// non-negative integer has to be a 400 in its own right — a range check + /// alone would let `"abc"` or `"1.5"` through to that indexing. + #[test] + fn logit_bias_keys_must_be_parseable_token_ids() { + for json in [ + r#"{"logit_bias": {"abc": 1.0}}"#, + r#"{"logit_bias": {"-1": 1.0}}"#, + r#"{"logit_bias": {"1.5": 1.0}}"#, + r#"{"logit_bias": {"": 1.0}}"#, + ] { + // Every key here is well inside TEST_VOCAB's range (or unparsable), + // so only the format check can be what rejects it. + let _ = norm_err(json); + } + assert!(norm(r#"{"logit_bias": {"7": 1.0}}"#).logit_bias.is_some()); + } + + /// Both `stop_regex` caps, neither of which had a test: deleting either `if` + /// left the suite green. The count cap bounds per-step recompilation (Python's + /// `re` cache is 512 entries); the length cap bounds compile time (a 1 MB + /// literal pattern measured ~677 ms). + #[test] + fn stop_regex_count_and_length_are_capped() { + let over: Vec = (0..MAX_STOP_REGEX_COUNT + 1) + .map(|i| format!("a{i}")) + .collect(); + let json = serde_json::json!({ "stop_regex": over }).to_string(); + assert!(norm_err(&json).to_string().contains("at most")); + + let at_cap: Vec = (0..MAX_STOP_REGEX_COUNT).map(|i| format!("a{i}")).collect(); + let json = serde_json::json!({ "stop_regex": at_cap }).to_string(); + assert!( + serde_json::from_str::(&json) + .unwrap() + .normalize(false, TEST_VOCAB) + .is_ok(), + "the cap itself must be accepted" + ); + + let long = "a".repeat(MAX_STOP_REGEX_LEN + 1); + let json = serde_json::json!({ "stop_regex": long }).to_string(); + let err = norm_err(&json).to_string(); + assert!(err.contains("over the"), "{err}"); + } + + /// The commoner field had no limit at all: the scheduler scans the decoded text + /// once per stop per decode step. + #[test] + fn stop_string_count_is_capped() { + let stops: Vec = (0..MAX_STOP_COUNT + 1).map(|i| i.to_string()).collect(); + let json = serde_json::json!({ "stop": stops }).to_string(); + let err = norm_err(&json).to_string(); + assert!(err.contains("at most"), "{err}"); + } +} diff --git a/rust/sglang-server/src/ring.rs b/rust/sglang-server/src/ring.rs new file mode 100644 index 000000000000..f8d0fa1f2575 --- /dev/null +++ b/rust/sglang-server/src/ring.rs @@ -0,0 +1,251 @@ +//! The two Rust↔Python boundary queues. +//! +//! In embedded mode the Rust frontend threads and the Python scheduler loop +//! share one process, so these are in-process `flume` channels — literal +//! `mpsc`/`mpmc`, no shared memory, no serialization beyond the msgpack bytes +//! the payload already is. +//! +//! GIL note: the Python side only ever calls the *non-blocking* `drain` / +//! `try_push` methods while holding the GIL, and the Rust worker threads only +//! ever push/drain raw `Bytes` — neither side touches a `PyObject` off-thread, +//! so the producer threads never need the GIL. + +use std::sync::Mutex; +use std::time::Duration; + +use bytes::Bytes; + +use crate::message::IngressMsg; + +/// Ingress: TokenizerManager → scheduler `recv_requests`. +/// Producers are Rust TM workers; the single consumer is the Python thread. +/// Carries [`IngressMsg`] (columnar: scalar header + raw int64 ids cell), not a +/// single msgpack blob, so the large `input_ids` tensor bypasses msgpack. +#[derive(Clone)] +pub struct IngressProducer { + tx: flume::Sender, +} + +pub struct IngressConsumer { + rx: flume::Receiver, + /// One-slot buffer holding a message consumed by a blocking [`wait`] so the + /// scheduler can park on idle without losing it — the next [`drain`] returns + /// it first. Only ever touched by the single consumer (the Python thread), + /// so contention is nil; the `Mutex` is just for interior mutability across + /// the `&self` methods. + /// + /// [`wait`]: IngressConsumer::wait + /// [`drain`]: IngressConsumer::drain + stash: Mutex>, +} + +/// A drained ingress batch in **columnar** (struct-of-arrays) form. The `ids` +/// cells are kept *un-concatenated* so the pyo3 boundary can copy them straight +/// into one `PyBytes` (no intermediate buffer); `ids_total` is their summed +/// length, precomputed for that single allocation. +#[derive(Default)] +pub struct IngressColumns { + /// Per-request scalar msgpack header (`input_ids` omitted). + pub headers: Vec, + /// Per-request raw little-endian int64 ids cell (empty for control reqs). + pub ids: Vec, + /// Per-request token count (`ids` cell length / 8). + pub lengths: Vec, + /// Sum of all `ids` cell byte lengths. + pub ids_total: usize, +} + +impl IngressProducer { + /// Non-blocking push. Returns `false` on a full ring (backpressure) so the + /// caller can fail the request rather than block a worker thread. + #[inline] + pub fn try_push(&self, msg: IngressMsg) -> bool { + self.tx.try_send(msg).is_ok() + } +} + +impl IngressConsumer { + /// Drain up to `max` messages into a columnar [`IngressColumns`], returning + /// immediately when the ring runs dry — mirrors the scheduler's existing + /// `zmq.NOBLOCK` loop in `request_receiver._pull_raw_reqs`. Splitting headers + /// from ids here (off the GIL) leaves `recv_requests` a thin marshaling shim. + /// + /// Non-blocking by construction: `try_recv` returns `Err(TryRecvError::Empty)` + /// instantly when the ring is empty, and `Err(_) => break` exits the loop + /// right away. + pub fn drain(&self, max: usize) -> IngressColumns { + let mut batch = IngressColumns::default(); + // A message parked by a prior blocking `wait` is delivered first. + if let Some(m) = self.stash.lock().unwrap().take() { + push_msg(&mut batch, m); + } + while batch.headers.len() < max { + match self.rx.try_recv() { + Ok(m) => push_msg(&mut batch, m), + Err(_) => break, // Empty or Disconnected -> stop now + } + } + batch + } + + /// Park up to `timeout` for at least one incoming message, so the idle + /// scheduler loop sleeps instead of spinning at 100% CPU. The message is + /// **stashed, not returned** — the next [`drain`](Self::drain) yields it — + /// so this composes with the existing non-blocking drain flow. Returns + /// whether a message is now available. `flume` wakes the parked thread the + /// instant a producer pushes, so this adds no latency to real requests. + pub fn wait(&self, timeout: Duration) -> bool { + if self.stash.lock().unwrap().is_some() { + return true; + } + match self.rx.recv_timeout(timeout) { + Ok(m) => { + *self.stash.lock().unwrap() = Some(m); + true + } + Err(_) => false, // Timeout or Disconnected + } + } +} + +/// Append one drained message's columnar cells to the batch. +#[inline] +fn push_msg(batch: &mut IngressColumns, m: IngressMsg) { + batch.ids_total += m.ids.len(); + batch.lengths.push((m.ids.len() / 8) as u32); // int64 cell → tokens + batch.headers.push(m.header); + batch.ids.push(m.ids); +} + +/// Egress: scheduler output (`push_chunk`) → Rust egress dispatcher. +/// The single producer is the Python thread; the consumer is the dispatcher. +#[derive(Clone)] +pub struct EgressProducer { + tx: flume::Sender, +} + +pub struct EgressConsumer { + rx: flume::Receiver, +} + +impl EgressProducer { + /// Blocking push: parks until the ring has space, so a full ring applies + /// backpressure to the scheduler instead of dropping output the scheduler has + /// already committed (advanced `send_token_offset` for). The GIL is released + /// around the call, so parking here doesn't stall other Python threads. + /// `false` only when the consumer is gone (runtime shutdown), where the frame + /// is unavoidably lost. + pub fn push(&self, msg: Bytes) -> bool { + self.tx.send(msg).is_ok() + } + + /// Non-blocking push, so the pyo3 boundary can try to hand the frame over + /// while still holding the GIL and detach only when it would actually park. + /// Releasing the GIL is not free: reacquiring it waits out the interpreter's + /// switch interval (5 ms by default), which dwarfs the sub-microsecond push + /// it was protecting. + /// + /// Hands the frame BACK on a full ring (`Err(Some(msg))`) so the caller can + /// retry it under [`push`](Self::push) without rebuilding it. `Err(None)` is + /// the consumer being gone (shutdown), where the frame is unavoidably lost. + #[inline] + pub fn try_push(&self, msg: Bytes) -> Result<(), Option> { + match self.tx.try_send(msg) { + Ok(()) => Ok(()), + Err(flume::TrySendError::Full(msg)) => Err(Some(msg)), + Err(flume::TrySendError::Disconnected(_)) => Err(None), + } + } +} + +impl EgressConsumer { + /// The underlying receiver, so the dispatcher can drain it via + /// [`tokenizer_manager::recv`](crate::tokenizer_manager::recv) (data + shutdown select). + pub fn receiver(&self) -> &flume::Receiver { + &self.rx + } +} + +/// Build both halves of a bounded ring. +pub fn ingress_ring(cap: usize) -> (IngressProducer, IngressConsumer) { + let (tx, rx) = flume::bounded(cap); + ( + IngressProducer { tx }, + IngressConsumer { + rx, + stash: Mutex::new(None), + }, + ) +} + +pub fn egress_ring(cap: usize) -> (EgressProducer, EgressConsumer) { + let (tx, rx) = flume::bounded(cap); + (EgressProducer { tx }, EgressConsumer { rx }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn msg(h: &'static [u8]) -> IngressMsg { + IngressMsg { + header: Bytes::from_static(h), + ids: Bytes::new(), + } + } + + /// `wait` parks when empty (times out), stashes a pushed message + /// non-destructively, and the next `drain` returns it. + #[test] + fn wait_stashes_then_drain_returns_it() { + let (tx, rx) = ingress_ring(8); + // Empty ring → times out, nothing stashed. + assert!(!rx.wait(Duration::from_millis(1))); + // Push one, then wait stashes it (returns true). + assert!(tx.try_push(msg(b"a"))); + assert!(rx.wait(Duration::from_millis(200))); + // Idempotent: already stashed, returns true without touching the ring. + assert!(rx.wait(Duration::from_millis(1))); + // Drain yields the stashed message, then the ring is empty. + assert_eq!(rx.drain(16).headers.len(), 1); + assert!(rx.drain(16).headers.is_empty()); + } + + /// A blocked `wait` is woken the instant a producer pushes (no polling). + #[test] + fn wait_wakes_on_push() { + let (tx, rx) = ingress_ring(8); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(20)); + let _ = tx.try_push(msg(b"a")); + }); + // Generous timeout, but it should return well before it as soon as the + // push lands. + assert!(rx.wait(Duration::from_secs(5))); + assert_eq!(rx.drain(16).headers.len(), 1); + } + + /// A full egress ring parks the producer until the consumer drains — the + /// committed frame is delivered in order, never dropped. + #[test] + fn egress_push_blocks_until_drained() { + let (tx, rx) = egress_ring(1); + assert!(tx.push(Bytes::from_static(b"a"))); // fits; ring now full + let t = std::thread::spawn(move || tx.push(Bytes::from_static(b"b"))); + // The parked push can't have completed while the ring is full. + std::thread::sleep(Duration::from_millis(20)); + // Drain one → frees a slot → the parked push lands. + assert_eq!(rx.receiver().recv().unwrap(), Bytes::from_static(b"a")); + assert!(t.join().unwrap(), "push should succeed once space frees"); + assert_eq!(rx.receiver().recv().unwrap(), Bytes::from_static(b"b")); + } + + /// A closed ring (consumer gone → shutdown) returns `false` instead of + /// parking forever, so a scheduler blocked in `push` unblocks on teardown. + #[test] + fn egress_push_returns_false_when_closed() { + let (tx, rx) = egress_ring(1); + drop(rx); + assert!(!tx.push(Bytes::from_static(b"x"))); + } +} diff --git a/rust/sglang-server/src/runtime.rs b/rust/sglang-server/src/runtime.rs new file mode 100644 index 000000000000..e3b0f7bdcdc4 --- /dev/null +++ b/rust/sglang-server/src/runtime.rs @@ -0,0 +1,21 @@ +//! Runtime bootstrap: wires channels, pins CPU-bound pools, starts the tokio +//! API server, and returns a handle the Python boundary uses for +//! `recv_requests` (ingress drain) and `push_batch` (egress push). +//! +//! Thread layout: +//! * API server — tokio multi-thread runtime (I/O bound), pinned core set A +//! * Tokenizer — N pinned OS threads (CPU bound), core set B +//! * Detokenizer — M pinned OS threads / shards (CPU bound), core set C +//! * TM ingress — 1 thread driving the ingress FSM +//! * TM egress — 1 thread draining the egress ring → detok shards +//! +//! Keeping CPU-bound tokenize/detokenize off the async executor avoids stalling +//! axum's worker threads. + +mod config; +mod runnable; + +pub use config::ServerArgs; + +// Re-export so stages keep importing `crate::runtime::Runnable`. +pub use runnable::Runnable; diff --git a/rust/sglang-server/src/runtime/config.rs b/rust/sglang-server/src/runtime/config.rs new file mode 100644 index 000000000000..ca5a557a9f8f --- /dev/null +++ b/rust/sglang-server/src/runtime/config.rs @@ -0,0 +1,203 @@ +//! Runtime configuration: the rust-server boot knobs +//! ([`RustServerServerArgs`]), the typed view of the scheduler's `server_args` +//! dump ([`ServerArgs`] / [`ModelConfig`]), and the [`RuntimeConfig`] pairing +//! them for `runtime::start`. + +use std::net::SocketAddr; +use std::sync::Arc; + +/// Boot knobs specific to the embedded rust server — none of these exist in +/// the Python `server_args` dump (see [`ServerArgs`]); they arrive as explicit +/// `Server::start` parameters. +#[derive(Clone, Debug)] +pub struct RustServerServerArgs { + pub http_addr: SocketAddr, + pub api_worker_num: usize, + pub ingress_ring_cap: usize, + pub egress_ring_cap: usize, + pub 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>, +} + +impl Default for RustServerServerArgs { + fn default() -> Self { + Self { + http_addr: "127.0.0.1:30000".parse().unwrap(), + api_worker_num: 2, + ingress_ring_cap: 8192, + egress_ring_cap: 8192, + channel_cap: 8192, + cores: None, + } + } +} + +#[derive(Clone, Debug)] +pub struct RuntimeConfig { + /// Rust-server-only boot knobs (listen address, pool/ring sizes, pinning). + pub rust_server_args: RustServerServerArgs, + /// The scheduler's `server_args` dump (worker counts, tokenizer source, + /// config-endpoint metadata). `Arc` so cloning the config (and, downstream, + /// each `AppState`) is cheap; immutable after construction. + pub server_args: Arc, +} + +impl Default for RuntimeConfig { + fn default() -> Self { + Self { + rust_server_args: RustServerServerArgs::default(), + server_args: Arc::new( + ServerArgs::from_json("{}").expect("empty server_args blob parses"), + ), + } + } +} + +/// The scheduler's startup blob (`RustServer._build_server_args`) parsed once into +/// typed fields: values are post-`__post_init__`, unknown keys (e.g. `api_key`) are dropped. +#[derive(Debug, serde::Deserialize)] +pub struct ServerArgs { + /// HF repo id / local dir of the model, reported by `/get_model_info`. + #[serde(default)] + pub model_path: String, + /// Model name reported by `/v1/models` and `/server_info`. + #[serde(default)] + pub served_model_name: String, + /// Tokenizer source (model dir / `tokenizer.json` / HF repo id). Empty only + /// in minimal standalone blobs — then boot requires `skip_tokenizer_init`. + #[serde(default)] + pub tokenizer_path: String, + /// HF revision, used only when `tokenizer_path` is a repo id. `None` → main. + #[serde(default)] + pub revision: Option, + /// HTTP bind address (see [`Self::bind`]). + #[serde(default = "default_host")] + pub host: String, + #[serde(default = "default_port")] + pub port: u16, + /// Log levels driving the access log — uvicorn runs at + /// `log_level_http or log_level` (see [`Self::http_access_log_enabled`]). + #[serde(default = "default_log_level")] + pub log_level: String, + #[serde(default)] + pub log_level_http: Option, + /// Pinned tokenizer threads / detok shards (Python asserts both ≥ 1). + #[serde(default = "default_worker_num")] + pub tokenizer_worker_num: usize, + #[serde(default = "default_worker_num")] + pub detokenizer_worker_num: usize, + /// Token-ids-in / token-ids-out mode: no tokenizer load, raw `output_ids` + /// frames (drives the `Skip` detok backend and the ingress branch). + #[serde(default)] + pub skip_tokenizer_init: bool, + /// Streamed `/generate` frames carry per-step deltas instead of cumulative + /// text. Matches the Python `TokenizerManager`. + #[serde(default)] + pub incremental_streaming_output: bool, + /// The resolved Python `ModelConfig`, attached to the blob at dump time. + #[serde(default)] + pub model_config: ModelConfig, + /// Default sampling params advertised by `/get_model_info`, verbatim from + /// `server_args.preferred_sampling_params` (a JSON object or null). + #[serde(default)] + pub preferred_sampling_params: Option, + /// Over-long inputs are truncated to fit the context instead of 400ing, and + /// `max_new_tokens` is clamped rather than rejected (Python + /// `TokenizerManager._validate_one_request`). + #[serde(default)] + pub allow_auto_truncate: bool, + /// `return_hidden_states` is refused unless the server was launched with it: + /// the scheduler simply won't produce them, so the request would 200 with the + /// field silently missing. + #[serde(default)] + pub enable_return_hidden_states: bool, + /// Output slots reserved per request on top of its input (eagle stores draft + /// tokens there). Not a `server_args` field — `TokenizerManager` derives it and + /// `RustServer._build_server_args` stamps it in, so both sides count alike. + #[serde(default)] + pub num_reserved_tokens: u64, + /// Launch-time stamps (not `server_args` fields): sglang package version + /// and the scheduler-derived KV token capacity, reported by `/server_info`. + #[serde(default)] + pub version: Option, + #[serde(default)] + pub max_total_num_tokens: Option, +} + +/// The slice of the resolved Python `ModelConfig` the rust server reads. +#[derive(Debug, Default, serde::Deserialize)] +pub struct ModelConfig { + /// Resolved context length (`max_model_len` in `/v1/models`); mandatory at + /// boot ([`ServerArgs::validate_mandatory`]). + #[serde(default)] + pub context_len: Option, + /// Bounds client-supplied token ids — ingress 400s out-of-vocab ids before + /// they crash the scheduler's embedding lookup; mandatory at + /// boot ([`ServerArgs::validate_mandatory`]). + #[serde(default)] + pub vocab_size: Option, +} + +fn default_host() -> String { + "127.0.0.1".into() +} +fn default_port() -> u16 { + 30000 +} +fn default_log_level() -> String { + "info".into() +} +fn default_worker_num() -> usize { + 1 +} + +impl ServerArgs { + /// Parse the blob; errors on malformed JSON or a wrongly-typed field. + pub fn from_json(s: &str) -> Result { + serde_json::from_str(s).map_err(|e| e.to_string()) + } + + /// Fail fast at startup if a field an endpoint depends on is missing. + pub fn validate_mandatory(&self) -> Result<(), String> { + if self.served_model_name.is_empty() { + return Err("no 'served_model_name' in server_args".into()); + } + if self.model_config.context_len.is_none() { + return Err("no resolvable context length (model_config.context_len)".into()); + } + if self.model_config.vocab_size.is_none() { + return Err("no resolvable vocab size (model_config.vocab_size)".into()); + } + Ok(()) + } + + /// Bind address `host:port`. `host` is expected to be an IP — the result is + /// parsed as a `SocketAddr`. + pub fn bind(&self) -> String { + format!("{}:{}", self.host, self.port) + } + + /// Whether the HTTP access log is emitted, mirroring the Python server: + /// uvicorn runs at `log_level_http or log_level` and prints access lines + /// only at info/debug. `--log-level-http warning` turns them off. + pub fn http_access_log_enabled(&self) -> bool { + let level = self + .log_level_http + .as_deref() + .filter(|s| !s.is_empty()) + .unwrap_or(&self.log_level); + matches!( + level.to_ascii_lowercase().as_str(), + "trace" | "debug" | "info" + ) + } + + /// Pinned API threads for the embedded HTTP api-server. Python `server_args` + /// has no such field — this is derived: enough to cover the widest pool. + pub fn api_worker_num(&self) -> usize { + 4.max(self.tokenizer_worker_num) + .max(self.detokenizer_worker_num) + } +} diff --git a/rust/sglang-server/src/runtime/runnable.rs b/rust/sglang-server/src/runtime/runnable.rs new file mode 100644 index 000000000000..66b7535d76d1 --- /dev/null +++ b/rust/sglang-server/src/runtime/runnable.rs @@ -0,0 +1,10 @@ +//! The [`Runnable`] stage trait — the one contract every pipeline stage +//! (CPU-bound worker or TM router) implements to be spawned by the runtime. + +/// A pipeline stage that owns its channel handles + config and runs a blocking +/// loop until its inbox closes. Lets the runtime spawn stages uniformly via +/// `threads::spawn_stage` / `threads::spawn_pool` instead of free `run_*` functions with +/// positional handles. Implemented by every CPU-bound worker and TM router. +pub trait Runnable: Send + 'static { + fn run(self); +} diff --git a/rust/sglang-server/src/tokenizer_manager.rs b/rust/sglang-server/src/tokenizer_manager.rs new file mode 100644 index 000000000000..fae3be1b46f4 --- /dev/null +++ b/rust/sglang-server/src/tokenizer_manager.rs @@ -0,0 +1,89 @@ +//! TokenizerManager — owns the request lifecycle across two isolated threads: +//! +//! * [`ingress`] — drives the ingress FSM (Received → Validating → +//! Normalizing → {Tokenizing | PreSendValidating}) and pushes tokenized requests to the +//! scheduler ring. +//! * [`egress`] — drains the scheduler-output ring and routes each chunk to +//! the owning detokenizer shard. +//! +//! The two run on separate pinned threads with no shared state, connected to +//! the rest of the pipeline only through `flume` channels: [`TmEvent`] into the +//! ingress loop, [`Senders`] fanning out to the pools. + +mod ingress; + +use crate::ids::Rid; +use crate::message::{DetokMsg, Request}; + +/// Blocking receive that also wakes on shutdown: returns `None` when `rx` closes +/// *or* the `shutdown` sender is dropped. +pub fn recv(rx: &flume::Receiver, shutdown: &flume::Receiver<()>) -> Option { + flume::Selector::new() + .recv(rx, |r| r.ok()) + .recv(shutdown, |_| None) + .wait() +} + +/// Events into the TokenizerManager ingress loop. API server + tokenizer pool +/// share this one inbox, keeping the loop a single consumer (no `select`). +pub enum TmEvent { + /// A freshly received request from the API server. + Ingress(Request), + /// A request back from the tokenizer pool: `PreSendValidating` (ids filled) on success, + /// or `Failed` on a tokenize error. `drive` handles both. + Tokenized(Request), +} + +/// Producer-side handles, cloned into every stage that needs to emit. +/// Who asked for an abort. Both variants do the same work in +/// [`Ingress::on_abort`](crate::tokenizer_manager::ingress::Ingress) — deregister +/// the detok entry, tell the scheduler to stop — and the source is kept for +/// diagnostics. +/// +/// There is no in-flight rid registry to keep consistent, and so no release +/// ordering to get wrong: [`Rid::from_client`] makes every client-supplied rid +/// internally unique, so a resubmit of the "same" rid is a different `Rid` and +/// cannot be tangled up with an abort still in flight for the original. +#[derive(Clone, Debug)] +pub enum AbortSource { + /// From an `AbortGuard` drop. Owns the release. + Guard(Rid), + /// From a detokenizer terminal path. Aborts the scheduler work. + Detok(Rid), +} + +impl AbortSource { + pub fn rid(&self) -> &Rid { + match self { + Self::Guard(rid) | Self::Detok(rid) => rid, + } + } +} + +#[derive(Clone)] +pub struct Senders { + /// → TokenizerManager ingress loop. + pub tm: flume::Sender, + /// → the same loop, but UNBOUNDED and abort-only. + /// + /// Aborts cannot share the bounded inbox. `try_send` there drops them exactly + /// when they matter most — under overload — leaving the scheduler generating + /// for a dead connection; and the caller then faces a false choice between + /// releasing the rid (a live entry can be overwritten by a resubmit) and + /// holding it (a permanent leak). An unbounded lane removes the dilemma: an + /// abort is a small `String` and is always accepted, so releases can be + /// unconditional again. It cannot grow without bound in practice — one entry + /// per in-flight request, each already bounded by the inbox that admitted it. + pub abort: flume::Sender, + /// → Tokenizer pool (CPU-bound, pinned threads). + pub tok: flume::Sender, + /// → Detokenizer shards, indexed by `Rid::shard(detok.len())`. + pub detok: Vec>, +} + +impl Senders { + #[inline] + pub fn detok_for(&self, rid: &Rid) -> &flume::Sender { + &self.detok[rid.shard(self.detok.len())] + } +} diff --git a/rust/sglang-server/src/tokenizer_manager/ingress.rs b/rust/sglang-server/src/tokenizer_manager/ingress.rs new file mode 100644 index 000000000000..57d425a490bd --- /dev/null +++ b/rust/sglang-server/src/tokenizer_manager/ingress.rs @@ -0,0 +1,1210 @@ +//! TokenizerManager — ingress side. +//! +//! [`Ingress`] is a single-consumer stage draining one inbox fed by both the API +//! server (fresh requests) and the Tokenizer pool (returned requests). It owns +//! the request while driving the ingress FSM and hands it off by *moving* it to +//! the next stage; nothing here is shared, so no locks. +//! +//! Edges driven here (from the design table): +//! Received → Validating +//! Validating → Normalizing (generate: sampling-param normalize/verify) +//! Validating → PreSendValidating (control: no tokenize, no sampling params) +//! Normalizing → {Encoding | Tokenizing | PreSendValidating} (by ValidationOutcome) +//! Tokenizing → PreSendValidating (on TokenizeDone, when the request returns) +//! PreSendValidating → Queued (checks needing the tokenized length) +//! Queued → ring (handed to the scheduler) +//! +//! The egress edges (Streaming/Finalizing/Completed) are driven on the egress +//! side (see `egress` + `detokenizer`). + +use bytes::Bytes; + +use crate::error::Error; +use crate::fsm::{Event, RequestState, ValidationOutcome}; + +use crate::message::{ + AbortReq, ControlRequest, DetokMsg, EgressItem, GenerateRequest, IngressMsg, Request, + RequestKind, +}; +use crate::ring::IngressProducer; +use crate::runtime::{Runnable, ServerArgs}; +use crate::tokenizer_manager::{AbortSource, Senders, TmEvent}; + +/// Ingress FSM dispatcher stage. Owns its inbox + downstream handles, so the +/// runtime spawns it as a [`Runnable`] rather than calling a free `run_*` fn +/// with positional arguments. +pub struct Ingress { + rx: flume::Receiver, + /// Unbounded abort lane (see [`Senders::abort`]). Selected against `rx` so an + /// abort is handled promptly even while the bounded inbox is saturated. + abort_rx: flume::Receiver, + senders: Senders, + ingress: IngressProducer, + limits: Limits, + shutdown: flume::Receiver<()>, +} + +/// Longest client-supplied rid accepted. It keys the detok table and travels on +/// every chunk, so its length is a recurring cost; Python mints 32-byte uuid hex. +const MAX_RID_LEN: usize = 128; + +/// What ingress admits, resolved once at boot from the scheduler's `server_args`. +/// A struct rather than more positional `new` arguments — these grew from two to +/// six, and every one of them is a `u64`/`bool` that would be trivial to swap at +/// a call site. +/// +/// NOT `Default`-able on purpose. `vocab_size` and `context_len` are mandatory, +/// and their zero value is the most restrictive setting there is — a derived +/// `Default` would silently build limits that reject every request rather than +/// failing loudly. Tests construct these explicitly (see `test_limits`). +#[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. Mandatory — + /// [`ServerArgs::validate_mandatory`](crate::runtime::ServerArgs) rejects a + /// boot without it, so ingress can check unconditionally. + pub vocab_size: u64, + /// `model_config.context_len`, the ceiling for input + `max_new_tokens`. + /// Mandatory, as above. + 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 TryFrom<&ServerArgs> for Limits { + type Error = Error; + + fn try_from(sa: &ServerArgs) -> Result { + Ok(Self { + skip_tokenizer_init: sa.skip_tokenizer_init, + vocab_size: sa + .model_config + .vocab_size + .ok_or_else(|| Error::Validation("vocab_size missing".into()))?, + context_len: sa + .model_config + .context_len + .ok_or_else(|| Error::Validation("context_len missing".into()))?, + num_reserved_tokens: sa.num_reserved_tokens, + allow_auto_truncate: sa.allow_auto_truncate, + enable_return_hidden_states: sa.enable_return_hidden_states, + }) + } +} + +impl Ingress { + pub fn new( + rx: flume::Receiver, + abort_rx: flume::Receiver, + senders: Senders, + ingress: IngressProducer, + limits: Limits, + shutdown: flume::Receiver<()>, + ) -> Self { + Self { + rx, + abort_rx, + senders, + ingress, + limits, + shutdown, + } + } +} + +/// Which lane produced the next item. +enum Lane { + Abort(AbortSource), + Event(TmEvent), +} + +impl Runnable for Ingress { + fn run(self) { + loop { + // Select, not a drain-then-block: an abort arriving while the inbox is + // idle must still be handled at once. + let next = flume::Selector::new() + .recv(&self.abort_rx, |r| r.ok().map(Lane::Abort)) + .recv(&self.rx, |r| r.ok().map(Lane::Event)) + .recv(&self.shutdown, |_| None) + .wait(); + match next { + Some(Lane::Abort(rid)) => self.on_abort(rid), + // A fresh request and one returning from the tokenizer pool. + Some(Lane::Event(TmEvent::Ingress(req) | TmEvent::Tokenized(req))) => { + self.drive(req) + } + None => { + // Shutdown, or the inbox closed. Drain whatever is still queued + // on the abort lane first: those requests are in flight on the + // scheduler, and the selector may report the closed inbox before + // it ever looks at a pending abort. + while let Ok(source) = self.abort_rx.try_recv() { + self.on_abort(source); + } + return; + } + } + } + } +} + +impl Ingress { + /// Reject a request: → `Failed`, notify the client, deregister (unconditional + /// — a no-op when nothing was registered). + /// `registered` says whether this request ever reached `register_detok`. It + /// must: `Deregister`'s handler is a bare `table.remove(&rid)`, so a + /// request rejected BEFORE registering would evict whatever entry currently + /// holds that key — a concurrent request's sink — leaving that client with no + /// terminal frame and a hung connection. Python cannot hit this because it + /// validates before `rid_to_state[obj.rid] = state`. + fn fail(&self, req: &mut Request, err: Error, registered: bool) { + // Log only server faults (500); 4xx/499/503 are expected and would spam. + if err.http_status() == 500 { + tracing::error!(rid = %req.rid, error = %err, "ingress rejected request"); + } + let _ = req.state.apply(Event::Error(err.clone())); + let _ = req.sink.try_send(EgressItem::Error(err)); // client may be gone + if registered { + let _ = self.senders.detok_for(&req.rid).send(DetokMsg::Deregister { + rid: req.rid.clone(), + }); + } + } + + /// Drive a request through its ingress states until it terminates (failed or + /// pushed to the ring) or is handed to the tokenizer pool (re-entering as a + /// `Tokenized` event). Each arm acts and advances the FSM; the loop + /// re-dispatches. The arms are the design table's states, `Failed` the single + /// reject path. + fn drive(&self, mut req: Request) { + // Flipped once `register_detok` succeeds; `fail` must not deregister before + // that (see `fail`). A pool return re-enters `drive` already registered. + let mut registered = !matches!(req.state, RequestState::Received); + loop { + match req.state.clone() { + // Validate, then register the sink before the request leaves Rust. + // Failures move to `Failed` and fall through to the reject arm. + RequestState::Received => { + if let Err(e) = validate(&mut req, &self.limits) { + let _ = req.state.apply(Event::Error(e)); // → Failed + continue; + } + if !self.register_detok(&req) { + let _ = req + .state + .apply(Event::Error(Error::Internal("detok shard gone".into()))); + continue; + } + registered = true; + // `validate` advanced Received → Validating; keep driving. + } + // Control skips normalization (no sampling params) straight to the + // pre-send checks; generate goes to Normalizing. + RequestState::Validating => match &req.kind { + RequestKind::Control(_) => { + let _ = req + .state + .apply(Event::Validated(ValidationOutcome::AlreadyTokenized)); + } + RequestKind::Generate(_) => { + let _ = req.state.apply(Event::NeedsNormalize); + } + }, + // Normalize + verify sampling params (off the scheduler loop), then + // pick the branch; a bad param becomes `Failed`. + RequestState::Normalizing => { + let outcome = { + let RequestKind::Generate(g) = &mut req.kind else { + // Unreachable (control never reaches here); reject so a + // bug can't leak/hang a registered request. + self.fail( + &mut req, + Error::Internal("non-generate request in Normalizing".into()), + registered, + ); + return; + }; + match g + .sampling_params + .normalize(self.limits.skip_tokenizer_init, self.limits.vocab_size) + { + Err(e) => Err(e), + // Client ids skip the pool; text goes to the tokenizer. + Ok(()) if g.already_tokenized() => { + Ok(ValidationOutcome::AlreadyTokenized) + } + Ok(()) => Ok(ValidationOutcome::NeedsTokenize), + } + }; + match outcome { + Err(e) => { + let _ = req.state.apply(Event::Error(e)); // → Failed + } + Ok(o) => { + // AlreadyTokenized → Queued, NeedsTokenize → Tokenizing. + let _ = req.state.apply(Event::Validated(o)); + } + } + } + // Hand off to the tokenizer pool; it returns the request as a + // `Tokenized` event (PreSendValidating, or Failed on error). + // Doesn't loop. + RequestState::Tokenizing => { + if let Err(err) = self.senders.tok.send(req) { + // Pool gone (workers exited); flume hands the request back. + let mut req = err.into_inner(); + // Past `Received`, so registration happened. + self.fail( + &mut req, + Error::Internal("tokenizer pool gone".into()), + true, + ); + } + return; + } + // The checks that need the final `input_ids`: every branch + // converges here (client ids arrive directly, text arrives from + // the tokenizer pool), so they run once per request regardless of + // how it was tokenized. `validate` runs too early — at `Received` + // a text request has no ids yet. + RequestState::PreSendValidating => { + if let RequestKind::Generate(g) = &mut req.kind + && let Err(e) = check_total_tokens(g, &self.limits) + { + let _ = req.state.apply(Event::Error(e)); // → Failed + continue; + } + let _ = req.state.apply(Event::PreSendValidated); // → Queued + } + // Push the wire message (control frame or generate payload) to the ring. + RequestState::Queued => { + // `matches!` reads the discriminant without holding a borrow, + // so `req` can be moved into the push below. + if matches!(req.kind, RequestKind::Generate(_)) { + self.push_to_ring(req); + } else { + self.push_control_to_ring(req); + } + return; + } + // The single reject path for every post-register failure. + RequestState::Failed(e) => { + self.fail(&mut req, e, registered); + return; + } + // Unreachable (egress states never reach here). Reject via `fail`/ + // return (not apply + continue, which would spin on a terminal state). + other => { + self.fail( + &mut req, + Error::Internal(format!("unexpected ingress state: {other:?}")), + registered, + ); + return; + } + } + } + } + + /// Register the egress sink with the owning detok shard (by id) so the response + /// has a home. Carries the per-request detok flags — `return_text_in_logprobs` + /// (decode logprob text on this shard) and `no_stop_trim` (keep the matched + /// stop in the output) — so the shard needs no back-reference to the request. + /// Returns `false` if the shard is gone. + fn register_detok(&self, req: &Request) -> bool { + let (decode_logprob_text, no_stop_trim) = match &req.kind { + RequestKind::Generate(g) => ( + g.return_text_in_logprobs.unwrap_or(false), + g.sampling_params.no_stop_trim, + ), + RequestKind::Control(_) => (false, false), + }; + self.senders + .detok_for(&req.rid) + .send(DetokMsg::Register { + rid: req.rid.clone(), + sink: req.sink.clone(), + decode_logprob_text, + no_stop_trim, + }) + .is_ok() + } + + /// Push a bare control request (`[tag, rid, nil]`) onto the ingress ring. The + /// scheduler dispatches it (e.g. `GetInternalStateReq`) and replies via the + /// egress ring as a single `Result`. + fn push_control_to_ring(&self, mut req: Request) { + let encode = match &req.kind { + RequestKind::Control(control) => control.encode(), + _ => Err(Error::Internal( + "non-control request reached push_control_to_ring".into(), + )), + }; + let header = match encode { + Ok(b) => b, + Err(e) => { + self.fail(&mut req, e, true); // on the push path: registered + return; + } + }; + // Control requests carry no tensor cell — empty `ids`. + if !self.ingress.try_push(IngressMsg { + header, + ids: Bytes::new(), + }) { + self.fail(&mut req, Error::QueueFull, true); // registered + } + } + + /// Client disconnected (or a detok terminal): deregister the sink, then push an + /// `AbortReq(rid)` so the scheduler stops generating for it. + /// + /// A failed push is logged, not retried: the scheduler keeps generating and the + /// chunks arrive for a rid no longer in the detok table, where they are dropped. + /// That wastes GPU work until the request finishes on its own, but it cannot be + /// misdelivered — the rid is unique to this request for the process's lifetime + /// ([`Rid::from_client`]), so no later request can ever answer to it. + fn on_abort(&self, source: AbortSource) { + let rid = source.rid().clone(); + let _ = self + .senders + .detok_for(&rid) + .send(DetokMsg::Deregister { rid: rid.clone() }); + + // The ring is BOUNDED and drops pushes under exactly the load this matters + // for, so report the miss rather than assuming the scheduler was told. + match ControlRequest::AbortReq(AbortReq::new(rid.as_str().to_string(), false)).encode() { + Ok(header) => { + if !self.ingress.try_push(IngressMsg { + header, + ids: Bytes::new(), + }) { + tracing::error!( + rid = %rid, + "abort dropped: ingress ring full; the scheduler keeps generating \ + for this request until it finishes on its own" + ); + } + } + Err(e) => tracing::error!(rid = %rid, error = %e, "abort encode failed"), + } + } + + /// Serialize the tokenized request to its `TokenizedGenerateReqInput` wire and + /// push it onto the ingress ring for the scheduler. On backpressure, fail it. + fn push_to_ring(&self, mut req: Request) { + // Only generate requests reach here (control uses `push_control_to_ring`). + // Validate + serialize while borrowing `g` immutably; the resulting `Bytes` + // own their data, so the borrow ends before any `fail(&mut req)`. + let serialized = match &req.kind { + RequestKind::Generate(g) if g.already_tokenized() => g + .encode_header() + .map(|header| (header, g.encode_data_buf())), + RequestKind::Generate(_) => Err(Error::Tokenize("empty input_ids".into())), + _ => Err(Error::Internal( + "non-generate request reached push_to_ring".into(), + )), + }; + let (header, ids) = match serialized { + Ok(v) => v, + Err(e) => { + self.fail(&mut req, e, true); // on the push path: registered + return; + } + }; + + if !self.ingress.try_push(IngressMsg { header, ids }) { + self.fail(&mut req, Error::QueueFull, true); // registered + } + // On success the scheduler owns the request (egress arrives by rid); we + // drop our `Request` here — the detok shard holds the sink. + } +} + +/// `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. +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})" + ))); + } + } + } + } + + // 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. +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(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::fsm::RequestState; + use crate::message::{EgressSink, GenerateRequest, SamplingParams}; + use crate::ring::{IngressConsumer, ingress_ring}; + use tokio::sync::mpsc; + + /// An `Ingress` plus its detok-shard receiver, ring consumer (keep alive — + /// dropping it closes the ring → false QueueFull), and tm inbox sender. + fn make_ingress() -> ( + Ingress, + flume::Receiver, + IngressConsumer, + flume::Sender, + ) { + make_ingress_with(test_limits()) + } + + fn make_ingress_with_abort( + abort_rx: flume::Receiver, + ) -> ( + Ingress, + flume::Receiver, + IngressConsumer, + flume::Sender, + ) { + make_ingress_inner(test_limits(), abort_rx) + } + + fn make_ingress_with( + limits: Limits, + ) -> ( + Ingress, + flume::Receiver, + IngressConsumer, + flume::Sender, + ) { + let (abort_tx, abort_rx) = flume::unbounded::(); + std::mem::forget(abort_tx); // keep the lane open; tests end by dropping tm_tx + make_ingress_inner(limits, abort_rx) + } + + fn make_ingress_inner( + limits: Limits, + abort_rx: flume::Receiver, + ) -> ( + Ingress, + flume::Receiver, + IngressConsumer, + flume::Sender, + ) { + let (tok_tx, _tok_rx) = flume::unbounded(); + let (detok_tx, detok_rx) = flume::unbounded(); + let senders = Senders { + tm: flume::unbounded().0, + abort: flume::unbounded().0, + tok: tok_tx, + detok: vec![detok_tx], + }; + let (ingress_producer, consumer) = ingress_ring(16); + let (tm_tx, tm_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 ingress = Ingress::new(tm_rx, abort_rx, senders, ingress_producer, limits, sd_rx); + (ingress, detok_rx, consumer, tm_tx) + } + + /// 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::(); + let (ingress_producer, consumer) = ingress_ring(16); + let (sd_tx, sd_rx) = flume::unbounded::<()>(); + std::mem::forget(sd_tx); + let ingress = Ingress::new( + flume::unbounded().1, + flume::unbounded().1, + Senders { + tm: flume::unbounded().0, + abort: flume::unbounded().0, + tok: flume::unbounded().0, + detok: vec![detok_tx], + }, + ingress_producer, + test_limits(), + sd_rx, + ); + + ingress.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: EgressSink::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 ingress 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 (ingress, detok_rx, consumer, _tm_tx) = make_ingress_with(Limits { + context_len: 4, + ..test_limits() + }); + ingress.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 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::(); + let senders = Senders { + tm: flume::unbounded().0, + abort: abort_tx, + tok: tok_tx, + detok: vec![detok_tx], + }; + let (producer, _consumer) = ingress_ring(1); + let (_tm_tx, tm_rx) = flume::unbounded(); + let (sd_tx, sd_rx) = flume::unbounded::<()>(); + std::mem::forget(sd_tx); + let ingress = Ingress::new(tm_rx, abort_rx, senders, producer, test_limits(), sd_rx); + + ingress.on_abort(AbortSource::Guard("pushed".into())); + ingress.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 (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + 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]); + } + ingress.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 (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + ingress.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 (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + // top_p = 2.0 is outside (0, 1], so `SamplingParams::normalize` rejects it. + let bad = SamplingParams { + top_p: 2.0, + ..Default::default() + }; + ingress.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 + /// ingress with a 400 — passed through, it reaches the embedding lookup + /// and kills the scheduler process (`make_ingress` bounds vocab at 1000). + #[test] + fn out_of_vocab_input_ids_rejected() { + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + 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]); + } + ingress.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 (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let mut req = generate_req(22, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut req.kind { + g.input_ids = Some(vec![-1]); + } + ingress.drive(req); + match detok_rx.try_recv() { + Err(_) | Ok(DetokMsg::Deregister { .. }) => {} + Ok(_) => panic!("negative token id must not be admitted"), + } + + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + let mut req = generate_req(23, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut req.kind { + g.token_ids_logprob = Some(vec![999_999]); + } + ingress.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 (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + // Empty map → all sampling defaults, passes normalization. + ingress.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_ingress() { + let (ingress, detok_rx, _consumer, tm_tx) = make_ingress(); + // 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); + ingress.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::(); + let (ingress, detok_rx, _consumer, tm_tx) = make_ingress_with_abort(abort_rx); + abort_tx.send(AbortSource::Guard("rid-13".into())).unwrap(); + drop(abort_tx); + drop(tm_tx); + ingress.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 (ingress, detok_rx, _consumer, tm_tx) = make_ingress(); + 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); + ingress.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() { + // `make_ingress` drops the tok receiver, so `tok.send` fails. + let (ingress, detok_rx, _consumer, _tm_tx) = make_ingress(); + // No ids → NeedsTokenize → Tokenizing branch. + let mut req = generate_req(21, SamplingParams::default()); + if let RequestKind::Generate(g) = &mut req.kind { + g.input_ids = None; + } + ingress.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"); + } +} diff --git a/rust/sglang-server/src/utils/regex.rs b/rust/sglang-server/src/utils/regex.rs index 38d730fecf25..65e57abc3bf3 100644 --- a/rust/sglang-server/src/utils/regex.rs +++ b/rust/sglang-server/src/utils/regex.rs @@ -5,6 +5,9 @@ //! and afford. See [`validate`] for the two rejection classes and why the //! invariant is one-directional. +use std::collections::HashMap; +use std::sync::{LazyLock, Mutex}; + use crate::error::Error; /// `MAX_LEN` from Python's `get_max_seq_length`: the bound for an *unbounded* stop @@ -209,6 +212,49 @@ fn check_escape(b: &[u8], i: usize) -> Result<(), String> { Ok(()) } +/// Entries kept in [`ADMISSION_CACHE`], mirroring CPython's `re._MAXCACHE`. +const ADMISSION_CACHE_CAP: usize = 512; + +/// Memo of admitted patterns → their bound. +/// +/// Admission is a pure function of the pattern text, and an expensive one: ~87% of +/// it is HIR translation, which expands `\w`/`\W` into large Unicode class unions. +/// A 256-byte `\W`-heavy pattern (exactly [`MAX_STOP_REGEX_LEN`]) measures 574 µs, +/// and a request may carry [`MAX_STOP_REGEX_COUNT`] of them — 18 ms of admission on +/// the single ingress thread, re-derived from scratch on every request. It +/// multiplies through a batch, because one `sampling_params` object broadcasts to +/// every item: a 13.6 KB body measured **1.01 s**, during which that thread serves +/// no other request, no abort and no health probe. +/// +/// Only successes are memoized. A rejected pattern fails inside [`validate`], which +/// is the cheap 8% — the expensive translate runs only after it passes — so the +/// hazard is entirely on the admitted side, and this keeps the entry a plain +/// `usize` rather than something that has to reconstruct an `Error` faithfully. +/// +/// Cleared wholesale when full rather than evicted one at a time: that is what +/// CPython's `re` does, and it keeps the hot path one lookup with no LRU +/// bookkeeping. The lock is held across a hash lookup and nothing else, and is +/// taken almost exclusively by the one ingress thread. +static ADMISSION_CACHE: LazyLock, usize>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +fn cached_bound(pattern: &str) -> Option { + ADMISSION_CACHE + .lock() + .ok() + .and_then(|c| c.get(pattern).copied()) +} + +fn cache_bound(pattern: &str, max_len: usize) { + let Ok(mut c) = ADMISSION_CACHE.lock() else { + return; + }; + if c.len() >= ADMISSION_CACHE_CAP { + c.clear(); + } + c.insert(pattern.into(), max_len); +} + /// A `stop_regex` that has been admitted, together with the bound derived while /// admitting it. /// @@ -239,6 +285,11 @@ impl<'a> RegexPattern<'a> { /// `Err` for anything CPython's `re` cannot compile, or cannot match cheaply /// enough to run on every decode step — see [`validate`]. fn build(pattern: &'a str) -> Result { + // Same pattern text ⇒ same verdict and same bound, so a repeat costs a hash + // lookup instead of a parse + translate. See [`ADMISSION_CACHE`]. + if let Some(max_len) = cached_bound(pattern) { + return Ok(Self { pattern, max_len }); + } let ast = validate(pattern)?; // Translate the AST `validate` already produced instead of re-parsing. The full // `regex_syntax::Parser` parses AND translates, so calling it here would parse a @@ -255,6 +306,7 @@ impl<'a> RegexPattern<'a> { )) })?; let max_len = hir_max_len(&hir); + cache_bound(pattern, max_len); Ok(Self { pattern, max_len }) } @@ -578,14 +630,53 @@ fn ast_len(ast: ®ex_syntax::ast::Ast) -> (u64, Option) { mod tests { use super::*; - /// Bound-only view of [`RegexPattern::new`], so the corpus rows read as + /// Bound-only view of [`RegexPattern`], so the corpus rows read as /// `pattern -> bound` without naming the type at every call. fn stop_regex_bound(pattern: &str) -> Result { RegexPattern::try_from(pattern).map(|r| r.max_len()) } - /// The newtype's contract: an admitted pattern keeps its own text alongside - /// its own bound, so the two cannot be paired up wrongly. + /// The admission memo must be indistinguishable from admitting afresh. + /// + /// It short-circuits the validator, so a wrong entry would admit a pattern + /// nobody checked or hand back another pattern's bound — and the bound sizes + /// the scheduler's match window, which is the under-estimate class of bug this + /// module exists to prevent. Three properties, one per way that could break: + /// a repeat agrees with a cold run, a rejection is never memoized, and the + /// wholesale clear at [`ADMISSION_CACHE_CAP`] loses nothing but the entries. + #[test] + fn admission_memo_agrees_with_admitting_afresh() { + // Distinct from any other test's patterns: the cache is process-wide, so a + // shared pattern would make this pass for the wrong reason. + let admitted = r"memo\d{3}[a-f]+"; + let cold = RegexPattern::try_from(admitted).expect("valid").max_len(); + let warm = RegexPattern::try_from(admitted).expect("valid").max_len(); + assert_eq!( + cold, warm, + "a memoized bound must equal a freshly derived one" + ); + + // Rejections are re-validated every time, so the memo can never turn one + // into an admission. + let rejected = r"memo(?:.|.)*Z"; + assert!(RegexPattern::try_from(rejected).is_err()); + assert!( + RegexPattern::try_from(rejected).is_err(), + "a rejected pattern must stay rejected on the second try" + ); + + // Overflow the cache, then re-check: clearing must not corrupt or stale a + // subsequent lookup. + for i in 0..=ADMISSION_CACHE_CAP { + let _ = RegexPattern::try_from(format!("memofill{i}").as_str()); + } + assert_eq!( + RegexPattern::try_from(admitted).expect("valid").max_len(), + cold, + "the bound must survive a cache clear" + ); + } + #[test] fn admitted_pattern_carries_its_own_text_and_bound() { let p = RegexPattern::try_from(r"\d{6}").expect("valid");