sglang rust server tm ingress

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"
+            );
+        }
+    }
+}

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::*;

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<ChunkEvent>),
+    /// 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 },
+}