diff --git a/.github/workflows/e2e-gpu-job.yml b/.github/workflows/e2e-gpu-job.yml index b5abed138..b97a2995e 100644 --- a/.github/workflows/e2e-gpu-job.yml +++ b/.github/workflows/e2e-gpu-job.yml @@ -54,6 +54,11 @@ on: type: string default: "nixl" description: "KV transfer backend for vLLM PD workers: nixl or mooncake" + connection_mode: + required: false + type: string + default: "" + description: "Wire override for local backends: zmq runs the local cases over ZMQ" jobs: run: @@ -66,6 +71,7 @@ jobs: E2E_RUNTIME: ${{ inputs.engine }} E2E_GPU_TIER: ${{ inputs.gpu_tier }} E2E_VLLM_KV_BACKEND: ${{ inputs.vllm_kv_backend }} + E2E_CONNECTION_MODE: ${{ inputs.connection_mode }} ROUTER_LOCAL_MODEL_PATH: /models steps: - name: Checkout code @@ -166,6 +172,12 @@ jobs: if [ "${{ inputs.engine }}" = "vllm" ]; then SUFFIX="-${{ inputs.vllm_kv_backend }}" fi + # gRPC and ZMQ chat legs share engine/GPU/test-dir and differ only by + # wire mode; fold it in so their log artifacts don't collide. + MODE="${{ inputs.connection_mode }}" + if [ -n "$MODE" ]; then + SUFFIX="${SUFFIX}-${MODE}" + fi ARTIFACT="e2e-worker-logs-${{ inputs.engine }}-gpu${{ inputs.gpu_tier }}-${LABEL}${SUFFIX}" bash scripts/ci_dump_worker_logs.sh e2e-logs "$ARTIFACT" echo "artifact=${ARTIFACT}" >> "$GITHUB_OUTPUT" diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 9230953fc..fa8260a39 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -543,6 +543,45 @@ jobs: test_dirs: ${{ matrix.test_dirs || 'e2e_test/chat_completions' }} secrets: inherit + e2e-1gpu-chat-zmq: + name: e2e-1gpu-chat-zmq (${{ matrix.engine }}) + needs: [build-wheel, detect-changes] + if: >- + always() + && !cancelled() + && needs.build-wheel.result == 'success' + && (github.event_name != 'pull_request' + || (needs.detect-changes.result == 'success' + && (needs.detect-changes.outputs.common == 'true' + || needs.detect-changes.outputs.chat-completions == 'true'))) + # Same single-worker chat suite as e2e-1gpu-chat, driven over the ZMQ + # direct-backend wire. The collection hook deselects the gRPC-only + # families (PD, EPD, multi-worker) for a ZMQ lane, so this covers the + # local cases only. + strategy: + fail-fast: false + matrix: + include: + # The ZMQ lane restarts the engine per model group (66-90s of + # load+warmup each); the suite needs more headroom than the gRPC + # lane, whose pool amortizes worker bring-up. + - engine: vllm + timeout: 46 + test_timeout: 40 + - engine: tokenspeed + timeout: 50 + test_timeout: 40 + uses: ./.github/workflows/e2e-gpu-job.yml + with: + engine: ${{ matrix.engine }} + gpu_tier: "1" + runner: 1-gpu-h100 + timeout: ${{ matrix.timeout }} + test_timeout: ${{ matrix.test_timeout }} + test_dirs: e2e_test/chat_completions + connection_mode: zmq + secrets: inherit + e2e-1gpu-completions: name: e2e-1gpu-completions (${{ matrix.engine }}) needs: [build-wheel, detect-changes] @@ -1049,7 +1088,7 @@ jobs: path: benchmark_go_bindings/ finish: - needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e] + needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-chat-zmq, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e] if: always() runs-on: k8s-runner-cpu permissions: {} @@ -1064,6 +1103,7 @@ jobs: "${{ needs.unit-tests.result }}" == "failure" || \ "${{ needs.benchmarks.result }}" == "failure" || \ "${{ needs.e2e-1gpu-chat.result }}" == "failure" || \ + "${{ needs.e2e-1gpu-chat-zmq.result }}" == "failure" || \ "${{ needs.e2e-1gpu-completions.result }}" == "failure" || \ "${{ needs.e2e-1gpu-embeddings.result }}" == "failure" || \ "${{ needs.e2e-1gpu-gateway.result }}" == "failure" || \ diff --git a/crates/engine_zmq_client/src/codec/tensor.rs b/crates/engine_zmq_client/src/codec/tensor.rs index 4f3e959ea..dadba06f0 100644 --- a/crates/engine_zmq_client/src/codec/tensor.rs +++ b/crates/engine_zmq_client/src/codec/tensor.rs @@ -136,6 +136,42 @@ impl WireNdArray { Self::from_raw_bytes(dtype, shape, Bytes::from(data)) } + /// Build from little-endian `float32` bytes, casting each element to + /// `dtype`. Mirrors the model-dtype cast the engine's own frontend applies + /// to floating multimodal tensors before they reach the model. + pub fn from_f32_bytes_cast( + dtype: super::dtype::ModelDtype, + shape: Vec, + data: &[u8], + ) -> std::result::Result { + use super::dtype::ModelDtype; + if !data.len().is_multiple_of(4) { + return Err(format!( + "float32 buffer length {} is not a multiple of 4", + data.len() + )); + } + validate_element_count(&shape, data.len() / 4)?; + let floats = data + .as_chunks::<4>() + .0 + .iter() + .map(|c| f32::from_le_bytes(*c)); + Ok(match dtype { + ModelDtype::Float32 => Self::from_raw("float32", shape, data.to_vec()), + ModelDtype::Float16 => Self::from_raw_bytes( + "float16", + shape, + bytes_from_pod_vec(floats.map(f16::from_f32).collect::>()), + ), + ModelDtype::BFloat16 => Self::from_raw_bytes( + "bfloat16", + shape, + bytes_from_pod_vec(floats.map(bf16::from_f32).collect::>()), + ), + }) + } + /// Build from an owned immutable raw-view buffer. pub fn from_raw_bytes(dtype: impl Into, shape: Vec, data: Bytes) -> Self { Self { diff --git a/crates/engine_zmq_client/src/error.rs b/crates/engine_zmq_client/src/error.rs index ce8691983..77f3a6153 100644 --- a/crates/engine_zmq_client/src/error.rs +++ b/crates/engine_zmq_client/src/error.rs @@ -53,6 +53,8 @@ pub enum Error { context: &'static str, field: &'static str, }, + #[error("invalid structured outputs params: {message}")] + InvalidStructuredOutputsParams { message: String }, #[error("request `{request_id}` is already in flight")] DuplicateRequestId { request_id: String }, #[error("data parallel rank {rank} is out of range for {num_engines} engine(s)")] diff --git a/crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs b/crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs index 81ac5e8ea..7fe448588 100644 --- a/crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs +++ b/crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs @@ -54,13 +54,14 @@ pub struct SamplingParams { pub repetition_penalty: f64, /// Minimum number of tokens to generate before EOS / stop handling. pub min_new_tokens: u32, - /// Structured-output JSON schema. SMG rejects constraints upstream. + /// Structured-output JSON schema. Set from the request's `constraint`; + /// at most one of these four structured-output fields is populated. pub json_schema: Option, - /// Structured-output regex. SMG rejects constraints upstream. + /// Structured-output regex. Set from the request's `constraint`. pub regex: Option, - /// Structured-output EBNF grammar. SMG rejects constraints upstream. + /// Structured-output EBNF grammar. Set from the request's `constraint`. pub ebnf: Option, - /// Structured-output structural tag. SMG rejects constraints upstream. + /// Structured-output structural tag. Set from the request's `constraint`. pub structural_tag: Option, /// Ignore the EOS token and keep generating until another stop condition. pub ignore_eos: bool, diff --git a/crates/engine_zmq_client/src/protocol/vllm/mod.rs b/crates/engine_zmq_client/src/protocol/vllm/mod.rs index 44d96abdd..cec63b4db 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/mod.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/mod.rs @@ -7,20 +7,21 @@ //! shapes, field order, and `array_like` positional-tuple encoding are the wire //! contract with Python `EngineCoreProc` — do not reorder fields. //! -//! Text generation is typed fully. Multimodal features, structured outputs -//! (guided decoding), and pooling params are carried as [`crate::codec::OpaqueValue`] -//! for now — they serialize as `nil` on the text path and get strongly typed in -//! the multimodal phase. +//! Text generation, structured outputs (guided decoding), and multimodal +//! features are typed fully. Pooling params and prompt embeds are carried as +//! [`crate::codec::OpaqueValue`] — they serialize as `nil` on supported paths. // The startup handshake is engine-neutral (TokenSpeed speaks the same // protocol); re-exported here so existing `vllm::handshake` paths keep working. pub use crate::protocol::handshake; pub mod logprobs; pub mod lora; +pub mod multimodal; pub mod output; pub mod request; pub mod sampling; pub mod stats; +pub mod structured_outputs; use bytes::Bytes; diff --git a/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs b/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs new file mode 100644 index 000000000..50c65a757 --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs @@ -0,0 +1,324 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Ported from the Apache-2.0 reference `vllm-engine-core-client` +// (vllm-project/vllm): protocol/multimodal.rs. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use serde_tuple::{Deserialize_tuple, Serialize_tuple}; + +use crate::codec::tensor::WireTensor; + +/// Multimodal feature payload carried at `EngineCoreRequest.mm_features`. +/// +/// Python: `list[MultiModalFeatureSpec] | None` (`vllm/v1/engine/__init__.py`). +pub type MmFeatures = Vec; + +/// A single multimodal input with its processed data and metadata. A request +/// containing multiple multimodal items carries one `MmFeatureSpec` per item. +/// +/// Python: `MultiModalFeatureSpec` (`vllm/multimodal/inputs.py`), a dataclass — +/// encodes as a string-keyed msgpack map. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MmFeatureSpec { + /// Processed multimodal data for this item. `None` only when the engine's + /// receiver cache already holds `identifier` — an external frontend that + /// does not mirror that cache protocol must always send the full data. + pub data: Option, + + /// The input modality, e.g. `"image"`, `"audio"`, `"video"`. + pub modality: String, + + /// The hash for caching encoder outputs (with LoRA prefix if applicable). + pub identifier: String, + + /// The location of the `modality` tokens corresponding to this item in + /// the prompt. + pub mm_position: PlaceholderRange, + + /// The hash for caching processor outputs (without LoRA prefix). + #[serde(default)] + pub mm_hash: Option, +} + +/// Placeholder location information for one multimodal item. +/// +/// Python: `PlaceholderRange` (`vllm/multimodal/inputs.py`), a dataclass — +/// encodes as a string-keyed msgpack map. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct PlaceholderRange { + /// The start index of the placeholder in the prompt. + pub offset: usize, + + /// The length of the placeholder. + pub length: usize, + + /// A boolean mask of shape `(length,)` indicating which positions between + /// `offset` and `offset + length` receive embeddings. `None` means all. + #[serde(default)] + pub is_embed: Option, +} + +/// Processed keyword arguments for a single multimodal item, keyed by model +/// kwarg name (e.g. `pixel_values`). +/// +/// Python: `MultiModalKwargsItem` (`vllm/multimodal/inputs.py`) — encoded by +/// the serializer hooks as a string-keyed map. +pub type MmKwargsItem = BTreeMap; + +/// One processed keyword argument of a `MmKwargsItem`. +/// +/// Python: `MultiModalFieldElem` (`vllm/multimodal/inputs.py`). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MmFieldElem { + /// The keyword argument value passed to the model. `None` only when the + /// item is cached engine-side (see [`MmFeatureSpec::data`]). + pub data: Option, + + /// How this field's values combine with other items' for batching. + pub field: MmField, +} + +/// Processed multimodal keyword argument value (Python `NestedTensors`). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MmKwargValue { + Tensor(WireTensor), + Int(i64), + Float(f64), + List(Vec), +} + +/// How to interpret tensor data belonging to a keyword argument. +/// +/// Wire form is a 2-tuple `(factory_name, kwargs_map)` with factory names +/// `"batched"`, `"flat"`, `"shared"` — the serializer's +/// `MMF_CLASS_TO_FACTORY` encoding. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(try_from = "MmFieldWire", into = "MmFieldWire")] +pub enum MmField { + Batched(MmBatchedField), + Flat(MmFlatField), + Shared(MmSharedField), +} + +/// Python `MultiModalFieldConfig.batched`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmBatchedField { + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python `MultiModalFieldConfig.flat` / `flat_from_sizes`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmFlatField { + /// For each multimodal item, a slice (`dim=0`) or a tuple of slices + /// (`dim>0`) that extracts the data corresponding to it. + pub slices: Vec, + + /// The dimension to extract data from, default 0. + pub dim: i32, + + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python `MultiModalFieldConfig.shared`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmSharedField { + pub batch_size: usize, + + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python slice encoded as `(start, stop, step)`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)] +pub struct SliceSpec { + pub start: Option, + pub stop: Option, + pub step: Option, +} + +/// A single slice or a tuple of slices used by [`MmFlatField`]. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MmSlice { + Slice(SliceSpec), + Slices(Vec), +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)] +struct MmFieldWire { + name: String, + inner: MmFieldWireInner, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +enum MmFieldWireInner { + Batched(MmBatchedField), + Flat(MmFlatField), + Shared(MmSharedField), +} + +impl TryFrom for MmField { + type Error = String; + + fn try_from(value: MmFieldWire) -> Result { + match (value.name.as_str(), value.inner) { + ("batched", MmFieldWireInner::Batched(kwargs)) => Ok(Self::Batched(kwargs)), + ("flat", MmFieldWireInner::Flat(kwargs)) => Ok(Self::Flat(kwargs)), + ("shared", MmFieldWireInner::Shared(kwargs)) => Ok(Self::Shared(kwargs)), + (name, _) => Err(format!( + "mismatched or unknown multimodal field factory {name:?}" + )), + } + } +} + +impl From for MmFieldWire { + fn from(value: MmField) -> Self { + match value { + MmField::Batched(kwargs) => Self { + name: "batched".to_string(), + inner: MmFieldWireInner::Batched(kwargs), + }, + MmField::Flat(kwargs) => Self { + name: "flat".to_string(), + inner: MmFieldWireInner::Flat(kwargs), + }, + MmField::Shared(kwargs) => Self { + name: "shared".to_string(), + inner: MmFieldWireInner::Shared(kwargs), + }, + } + } +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use rmpv::Value; + + use super::*; + use crate::codec::encode_msgpack; + + fn encode_value(value: &T) -> Value { + let bytes = encode_msgpack(value).expect("encode value"); + rmpv::decode::read_value(&mut Cursor::new(bytes)).expect("decode value") + } + + #[test] + fn field_serializes_to_python_factory_tuple() { + let field = MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(1200), + step: None, + })], + dim: 0, + keep_on_cpu: false, + }); + + let value = encode_value(&field); + let Value::Array(items) = value else { + panic!("field should encode as a 2-tuple array"); + }; + assert_eq!(items.len(), 2); + assert_eq!(items[0].as_str(), Some("flat")); + let Value::Map(kwargs) = &items[1] else { + panic!("field kwargs should encode as a map"); + }; + for key in ["slices", "dim", "keep_on_cpu"] { + assert!( + kwargs.iter().any(|(k, _)| k.as_str() == Some(key)), + "missing kwarg {key}" + ); + } + } + + #[test] + fn field_round_trips_python_factory_tuple() { + for field in [ + MmField::Batched(MmBatchedField { keep_on_cpu: true }), + MmField::Shared(MmSharedField { + batch_size: 4, + keep_on_cpu: false, + }), + ] { + let encoded = encode_msgpack(&field).expect("encode field"); + let decoded: MmField = rmp_serde::from_slice(&encoded).expect("decode field"); + assert_eq!(decoded, field); + } + } + + #[test] + fn feature_spec_serializes_as_named_map_with_tensor_ext() { + let mut item = MmKwargsItem::new(); + item.insert( + "pixel_values".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Tensor( + WireTensor::from_f32(vec![2, 3], vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) + .expect("tensor built"), + )), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + ); + let spec = MmFeatureSpec { + data: Some(item), + modality: "image".to_string(), + identifier: "abc123".to_string(), + mm_position: PlaceholderRange { + offset: 5, + length: 6, + is_embed: None, + }, + mm_hash: Some("abc123".to_string()), + }; + + let value = encode_value(&spec); + let Value::Map(entries) = value else { + panic!("feature spec should encode as a map"); + }; + for key in ["data", "modality", "identifier", "mm_position", "mm_hash"] { + assert!( + entries.iter().any(|(k, _)| k.as_str() == Some(key)), + "missing key {key}" + ); + } + + // The tensor payload must reach the wire as the 3-tuple + // (dtype, shape, ext-3 raw view). + let data = entries + .iter() + .find(|(k, _)| k.as_str() == Some("data")) + .map(|(_, v)| v) + .expect("data present"); + let Value::Map(kwargs) = data else { + panic!("kwargs item should encode as a map"); + }; + let Value::Map(elem) = &kwargs[0].1 else { + panic!("field elem should encode as a map"); + }; + let tensor = elem + .iter() + .find(|(k, _)| k.as_str() == Some("data")) + .map(|(_, v)| v) + .expect("elem data present"); + let Value::Array(tuple) = tensor else { + panic!("tensor should encode as (dtype, shape, data)"); + }; + assert_eq!(tuple[0].as_str(), Some("float32")); + assert!(matches!(&tuple[2], Value::Ext(3, _))); + } +} diff --git a/crates/engine_zmq_client/src/protocol/vllm/request.rs b/crates/engine_zmq_client/src/protocol/vllm/request.rs index de7546f5e..8d3b0b0d1 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/request.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/request.rs @@ -2,9 +2,6 @@ // // Ported from the Apache-2.0 reference `vllm-engine-core-client` // (vllm-project/vllm): protocol/request.rs. -// -// `mm_features` is carried as OpaqueValue for now (typed in the multimodal -// phase); on the text path it is `nil` and carries no aux-frame tensors. use std::collections::{BTreeMap, HashMap}; @@ -15,7 +12,7 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple}; use crate::{ codec::OpaqueValue, - protocol::vllm::{lora, sampling::EngineCoreSamplingParams}, + protocol::vllm::{lora, multimodal::MmFeatures, sampling::EngineCoreSamplingParams}, Error, Result, }; @@ -72,8 +69,8 @@ pub struct ReasoningParserKwargs { pub struct EngineCoreRequest { pub request_id: String, pub prompt_token_ids: Option>, - /// Multimodal features (untyped for now; `nil` on the text path). - pub mm_features: Option, + /// Multimodal features, one per input item, sorted by placeholder offset. + pub mm_features: Option, pub sampling_params: Option, /// Pooling parameters, preserved in the schema but not yet strongly typed. pub pooling_params: Option, @@ -131,9 +128,10 @@ impl EngineCoreRequest { Ok(()) } - // NOTE: send-side aux-frame extraction (walking `mm_features` for large - // tensors) is added with the typed multimodal module. Text requests carry - // no tensors, so the transport send path appends no aux frames for now. + // NOTE: multimodal tensors are sent as inline ext-3 raw views in the + // request frame — valid at any size (the engine's aux-frame split is an + // encoder-side optimization only). Send-side aux extraction is a perf + // follow-up. } #[cfg(test)] diff --git a/crates/engine_zmq_client/src/protocol/vllm/sampling.rs b/crates/engine_zmq_client/src/protocol/vllm/sampling.rs index 13ed51a9a..ac89ce7cd 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/sampling.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/sampling.rs @@ -2,16 +2,13 @@ // // Ported from the Apache-2.0 reference `vllm-engine-core-client` // (vllm-project/vllm): protocol/sampling.rs. -// -// `structured_outputs` (guided decoding) is carried as OpaqueValue for now; it -// serializes as `nil` when absent and gets strongly typed in a later phase. use std::collections::{BTreeSet, HashMap}; use serde::{Deserialize, Serialize}; use serde_default::DefaultFromSerde; -use crate::codec::OpaqueValue; +use super::structured_outputs::StructuredOutputsParams; fn default_top_p() -> f32 { 1.0 @@ -115,9 +112,9 @@ pub struct EngineCoreSamplingParams { /// Tokenized bad words to avoid during generation. #[serde(rename = "_bad_words_token_ids")] pub bad_words_token_ids: Option>>, - /// Structured outputs (guided decoding). Carried untyped for now; `nil` on - /// the text path. - pub structured_outputs: Option, + /// Structured outputs (guided decoding). `None` (serialized `nil`) on the + /// unconstrained text path. + pub structured_outputs: Option, /// Specific token IDs for which log probabilities should be returned at each /// position, in addition to the sampled/scored token. pub logprob_token_ids: Option>, diff --git a/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs b/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs new file mode 100644 index 000000000..ffef05a19 --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs @@ -0,0 +1,336 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Ported from the Apache-2.0 reference `vllm-engine-core-client` +// (vllm-project/vllm): protocol/structured_outputs.rs. + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use serde_json::Value; + +use crate::error::{Error, Result}; + +/// Structured-output backend selected for EngineCore grammar compilation. +/// +/// Python stores this in `StructuredOutputsParams._backend` after request +/// validation. This client selects the backend per constraint: structural +/// tags require xgrammar (the triggered-tags format is not understood by +/// guidance's legacy structures/triggers parser); everything else lowers to +/// guidance. Peer-supplied `_backend` values are ignored. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum StructuredOutputBackend { + Xgrammar, + #[default] + Guidance, + Outlines, + LmFormatEnforcer, +} + +/// The single structured-output constraint selected for a request. +#[derive(Debug, Clone, PartialEq)] +pub enum StructuredOutputConstraint { + /// JSON schema (as a dict/object or JSON string) constraining the output. + Json(Value), + /// Regular expression the output must match. + Regex(String), + /// List of allowed output strings (the model must produce one of these). + Choice(Vec), + /// Context-free grammar (in EBNF-like notation) the output must conform to. + Grammar(String), + /// Output must be valid JSON (free-form, no schema). + JsonObject, + /// Structural tag configuration (JSON-encoded string). + StructuralTag(String), +} + +/// Additional structured-output options that do not select the constraint mode. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct StructuredOutputOptions { + /// Disable any additional whitespace in guided JSON output. + pub disable_any_whitespace: bool, + /// Disable `additionalProperties` in JSON schema output. + pub disable_additional_properties: bool, + /// Custom whitespace pattern for guided JSON output. + pub whitespace_pattern: Option, +} + +/// Parameters for configuring structured outputs (guided decoding). +/// +/// This is the semantic Rust representation: exactly one constraint mode is +/// always selected. The Python/msgpack product-shaped representation is kept in +/// the private wire type below and used only at serde boundaries. +#[derive(Debug, Clone, PartialEq)] +pub struct StructuredOutputsParams { + pub constraint: StructuredOutputConstraint, + pub options: StructuredOutputOptions, + /// Structured-output backend, mirroring Python's internal `_backend`. + /// + /// Peer-supplied values are ignored during deserialization. This matches the + /// Python request boundary, where `_backend` is set by validation rather + /// than accepted as a request-level backend selector. + pub backend: StructuredOutputBackend, +} + +impl StructuredOutputsParams { + pub fn json(json: Value) -> Self { + Self::from_constraint(StructuredOutputConstraint::Json(json)) + } + + pub fn regex(regex: impl Into) -> Self { + Self::from_constraint(StructuredOutputConstraint::Regex(regex.into())) + } + + pub fn choice(choice: Vec) -> Self { + Self::from_constraint(StructuredOutputConstraint::Choice(choice)) + } + + pub fn grammar(grammar: impl Into) -> Self { + Self::from_constraint(StructuredOutputConstraint::Grammar(grammar.into())) + } + + pub fn json_object() -> Self { + Self::from_constraint(StructuredOutputConstraint::JsonObject) + } + + pub fn structural_tag(structural_tag: impl Into) -> Self { + Self::from_constraint(StructuredOutputConstraint::StructuralTag( + structural_tag.into(), + )) + } + + fn from_constraint(constraint: StructuredOutputConstraint) -> Self { + // Structural tags use the triggered-tags format that only xgrammar + // compiles; guidance's parser expects the legacy structures/triggers + // shape and fails the request at grammar build. + let backend = match &constraint { + StructuredOutputConstraint::StructuralTag(_) => StructuredOutputBackend::Xgrammar, + _ => StructuredOutputBackend::default(), + }; + Self { + constraint, + options: StructuredOutputOptions::default(), + backend, + } + } +} + +/// `true` when a boolean is `false`; used to drop default-`false` flags from the +/// serialized map to match the sparse `omit_defaults` wire shape. Takes `&bool` +/// because serde's `skip_serializing_if` requires a by-reference predicate. +#[expect(clippy::trivially_copy_pass_by_ref)] +fn is_false(v: &bool) -> bool { + !*v +} + +/// Wire-compatible structured-output payload used by Python engine-core. +/// +/// Python models `StructuredOutputsParams` as a product-shaped dataclass with +/// several optional constraint fields, then validates that exactly one of those +/// fields is present. This client exposes [`StructuredOutputsParams`] as an +/// enum-backed domain type instead, while using this private wire type for +/// ser/de. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(default)] +struct WireStructuredOutputsParams { + json: Option, + regex: Option, + choice: Option>, + grammar: Option, + json_object: Option, + #[serde(skip_serializing_if = "is_false")] + disable_any_whitespace: bool, + #[serde(skip_serializing_if = "is_false")] + disable_additional_properties: bool, + whitespace_pattern: Option, + structural_tag: Option, + #[serde( + default, + rename = "_backend", + deserialize_with = "serde_with::rust::deserialize_ignore_any" + )] + backend: StructuredOutputBackend, +} + +impl TryFrom for StructuredOutputsParams { + type Error = Error; + + fn try_from(raw: WireStructuredOutputsParams) -> Result { + use StructuredOutputConstraint::*; + + let mut constraint = None; + + macro_rules! insert_constraint { + ($name:literal, $value:expr) => { + if let Some(value) = $value { + if let Some((existing, _)) = constraint { + return Err(Error::InvalidStructuredOutputsParams { + message: format!( + "multiple structured output constraints specified: {existing}, {}", + $name + ), + }); + } + constraint = Some(($name, value)); + } + }; + } + + insert_constraint!("json", raw.json.map(Json)); + insert_constraint!("regex", raw.regex.map(Regex)); + insert_constraint!("choice", raw.choice.map(Choice)); + insert_constraint!("grammar", raw.grammar.map(Grammar)); + match raw.json_object { + Some(true) => { + insert_constraint!("json_object", Some(JsonObject)); + } + Some(false) => { + return Err(Error::InvalidStructuredOutputsParams { + message: "structured_outputs.json_object must be true if set; omit structured_outputs to disable structured outputs".to_string(), + }); + } + None => {} + } + insert_constraint!("structural_tag", raw.structural_tag.map(StructuralTag)); + + Ok(Self { + constraint: constraint.map(|(_, c)| c).ok_or_else(|| { + Error::InvalidStructuredOutputsParams { + message: "missing structured output constraint".to_string(), + } + })?, + options: StructuredOutputOptions { + disable_any_whitespace: raw.disable_any_whitespace, + disable_additional_properties: raw.disable_additional_properties, + whitespace_pattern: raw.whitespace_pattern, + }, + backend: raw.backend, + }) + } +} + +impl From for WireStructuredOutputsParams { + fn from(params: StructuredOutputsParams) -> Self { + let mut raw = Self { + disable_any_whitespace: params.options.disable_any_whitespace, + disable_additional_properties: params.options.disable_additional_properties, + whitespace_pattern: params.options.whitespace_pattern, + backend: params.backend, + ..Self::default() + }; + + match params.constraint { + StructuredOutputConstraint::Json(json) => raw.json = Some(json), + StructuredOutputConstraint::Regex(regex) => raw.regex = Some(regex), + StructuredOutputConstraint::Choice(choice) => raw.choice = Some(choice), + StructuredOutputConstraint::Grammar(grammar) => raw.grammar = Some(grammar), + StructuredOutputConstraint::JsonObject => raw.json_object = Some(true), + StructuredOutputConstraint::StructuralTag(structural_tag) => { + raw.structural_tag = Some(structural_tag); + } + } + + raw + } +} + +impl Serialize for StructuredOutputsParams { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: Serializer, + { + WireStructuredOutputsParams::from(self.clone()).serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for StructuredOutputsParams { + fn deserialize(deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + WireStructuredOutputsParams::deserialize(deserializer)? + .try_into() + .map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn structured_outputs_backend_ignores_deserialized_value() { + let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({ + "json_object": true, + "_backend": "xgrammar", + })) + .unwrap(); + + assert_eq!(params.backend, StructuredOutputBackend::Guidance); + assert_eq!(params.constraint, StructuredOutputConstraint::JsonObject); + + let value = serde_json::to_value(params).unwrap(); + assert_eq!(value["_backend"], "guidance"); + } + + #[test] + fn structural_tag_selects_xgrammar_backend() { + // The triggered-tags format only compiles under xgrammar; guidance's + // legacy parser fails the request at grammar build. + let params = StructuredOutputsParams::structural_tag(r#"{"format":{}}"#); + assert_eq!(params.backend, StructuredOutputBackend::Xgrammar); + + let value = serde_json::to_value(params).unwrap(); + assert_eq!(value["_backend"], "xgrammar"); + assert_eq!(value["structural_tag"], r#"{"format":{}}"#); + } + + #[test] + fn structured_outputs_rejects_missing_constraint() { + let error = + serde_json::from_value::(serde_json::json!({})).unwrap_err(); + + assert!(error + .to_string() + .contains("missing structured output constraint")); + } + + #[test] + fn structured_outputs_rejects_multiple_constraints() { + let error = serde_json::from_value::(serde_json::json!({ + "json": {"type": "object"}, + "regex": ".*", + })) + .unwrap_err(); + + assert!(error + .to_string() + .contains("multiple structured output constraints specified: json, regex")); + } + + #[test] + fn structured_outputs_rejects_json_object_false() { + let error = serde_json::from_value::(serde_json::json!({ + "json_object": false, + })) + .unwrap_err(); + + assert!(error.to_string().contains("json_object must be true")); + } + + #[test] + fn structured_outputs_serializes_through_raw_shape() { + let params = StructuredOutputsParams { + constraint: StructuredOutputConstraint::StructuralTag( + r#"{"type":"structural_tag"}"#.to_string(), + ), + options: StructuredOutputOptions::default(), + backend: StructuredOutputBackend::Xgrammar, + }; + + let value = serde_json::to_value(params).unwrap(); + + assert_eq!(value["structural_tag"], r#"{"type":"structural_tag"}"#); + assert_eq!(value["_backend"], "xgrammar"); + assert!(value.get("json").is_none()); + } +} diff --git a/crates/grpc_client/src/tokenspeed_scheduler.rs b/crates/grpc_client/src/tokenspeed_scheduler.rs index bd428dd6b..9a543f947 100644 --- a/crates/grpc_client/src/tokenspeed_scheduler.rs +++ b/crates/grpc_client/src/tokenspeed_scheduler.rs @@ -178,12 +178,7 @@ impl TokenSpeedSchedulerClient { // ── Request builders ────────────────────────────────────────────── - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_generate_request_from_chat( - &self, request_id: String, body: &ChatCompletionRequest, processed_text: String, @@ -208,12 +203,7 @@ impl TokenSpeedSchedulerClient { }) } - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_plain_generate_request( - &self, request_id: String, body: &GenerateRequest, original_text: Option, @@ -242,12 +232,7 @@ impl TokenSpeedSchedulerClient { }) } - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_generate_request_from_responses( - &self, request_id: String, body: &ResponsesRequest, processed_text: String, @@ -267,12 +252,7 @@ impl TokenSpeedSchedulerClient { }) } - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_generate_request_from_messages( - &self, request_id: String, body: &CreateMessageRequest, processed_text: String, @@ -295,12 +275,7 @@ impl TokenSpeedSchedulerClient { }) } - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_generate_request_from_completion( - &self, request_id: String, body: &CompletionRequest, original_text: String, diff --git a/crates/tokenizer/src/mock.rs b/crates/tokenizer/src/mock.rs index f057aabe2..245e1968b 100644 --- a/crates/tokenizer/src/mock.rs +++ b/crates/tokenizer/src/mock.rs @@ -122,6 +122,11 @@ impl TokenizerTrait for MockTokenizer { self.reverse_vocab.get(&id).cloned() } + fn eos_token_ids(&self) -> &[u32] { + // `` in the mock vocab. + &[999] + } + fn as_any(&self) -> &dyn std::any::Any { self } diff --git a/crates/tokenizer/src/stop.rs b/crates/tokenizer/src/stop.rs index 820620cec..3b0ed3e6f 100644 --- a/crates/tokenizer/src/stop.rs +++ b/crates/tokenizer/src/stop.rs @@ -77,6 +77,9 @@ pub struct StopSequenceDecoder { jail_max_bytes: usize, /// Whether we've stopped stopped: bool, + /// The string stop sequence that triggered the stop, if any. Set only for + /// string-sequence matches; token-level stops leave this `None`. + matched_stop: Option, /// True when there are no string stop sequences (only token-level stops). /// In this mode the jail buffer is bypassed entirely for lower overhead. token_only: bool, @@ -139,6 +142,7 @@ impl StopSequenceDecoder { jail_buffer: String::new(), jail_max_bytes, stopped: false, + matched_stop: None, token_only, } } @@ -208,6 +212,7 @@ impl StopSequenceDecoder { let input = Input::new(&self.jail_buffer).span(search_start..self.jail_buffer.len()); if let Some(mat) = ac.find(input) { self.stopped = true; + self.matched_stop = Some(self.jail_buffer[mat.start()..mat.end()].to_string()); let is_visible = mat.pattern().as_usize() >= self.visible_boundary_idx; if is_visible { @@ -287,11 +292,18 @@ impl StopSequenceDecoder { self.stopped } + /// The string stop sequence that triggered the stop, if a string sequence + /// matched. `None` for token-level stops or when no stop has fired. + pub fn matched_stop(&self) -> Option<&str> { + self.matched_stop.as_deref() + } + /// Reset the decoder state pub fn reset(&mut self) { self.jail_buffer.clear(); self.sequence.clear(); self.stopped = false; + self.matched_stop = None; } } @@ -455,6 +467,40 @@ mod tests { )); } + #[test] + fn test_matched_stop_reports_matched_string() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_sequence("test"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // No match yet. + assert_eq!(decoder.matched_stop(), None); + + decoder.process_token(1).unwrap(); // "Hello" + decoder.process_token(2).unwrap(); // "world" + assert_eq!(decoder.matched_stop(), None); + + // "test" triggers the string stop; the matched string is captured. + decoder.process_token(3).unwrap(); // "test" + assert_eq!(decoder.matched_stop(), Some("test")); + + // Reset clears it. + decoder.reset(); + assert_eq!(decoder.matched_stop(), None); + } + + #[test] + fn test_matched_stop_none_for_token_stop() { + // A token-id stop is not a string match, so matched_stop stays None. + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_token(999); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + let result = decoder.process_token(999).unwrap(); + assert_eq!(result, SequenceDecoderOutput::Stopped); + assert_eq!(decoder.matched_stop(), None); + } + #[test] fn test_flush_after_partial() { let tokenizer = Arc::new(MockTokenizer::new()); diff --git a/e2e_test/fixtures/hooks.py b/e2e_test/fixtures/hooks.py index 3eb440864..8986252e1 100644 --- a/e2e_test/fixtures/hooks.py +++ b/e2e_test/fixtures/hooks.py @@ -14,10 +14,16 @@ import os import pytest -from infra import cleanup_pool, get_runtime +from infra import ConnectionMode, cleanup_pool, get_connection_mode_override, get_runtime from .markers import resolve_class_marker +# Local wires a plain (non-PD/EPD) test case can be authored with; these are the +# only backends a ZMQ lane reuses (mapped onto ZMQ by the ``setup_backend`` +# fixture). Anything else — ``pd_*``/``epd_*``, EPD topology tuples, cloud +# vendor names — is out of scope for ZMQ. +_ZMQ_LOCAL_WIRES = frozenset({"grpc", "http"}) + # --------------------------------------------------------------------------- # Marker registration # --------------------------------------------------------------------------- @@ -117,6 +123,66 @@ def _get_marker(item: pytest.Item, name: str): return resolve_class_marker(item, name) +def _setup_backend_param(item: pytest.Item): + """Return the ``setup_backend`` parametrize value for an item, or None.""" + callspec = getattr(item, "callspec", None) + if callspec is None: + return None + return getattr(callspec, "params", {}).get("setup_backend") + + +def _is_multi_worker(item: pytest.Item) -> bool: + """True when the item's ``workers`` marker asks for a PD/multi-worker topology.""" + marker = resolve_class_marker(item, "workers") + if marker is None: + return False + kwargs = marker.kwargs + if (kwargs.get("count") or 1) > 1: + return True + return bool(kwargs.get("prefill") or kwargs.get("decode")) + + +def _zmq_dedup_key(item: pytest.Item) -> tuple: + """Group key ignoring the ``setup_backend`` value. + + Lets us collapse a case parametrized on both ``grpc`` and ``http`` into a + single ZMQ run (they map onto the same wire) while keeping distinct + ``api_client`` (or other) parametrizations apart. + """ + callspec = getattr(item, "callspec", None) + params = getattr(callspec, "params", {}) or {} + others = tuple(sorted((k, repr(v)) for k, v in params.items() if k != "setup_backend")) + return (item.nodeid.split("[", 1)[0], others) + + +def _filter_zmq_items(items: list[pytest.Item]) -> tuple[list[pytest.Item], list[pytest.Item]]: + """Split items into (kept, deselected) for a ZMQ lane. + + Keeps single-worker local cases (``grpc``/``http`` map onto ZMQ) and drops + the gRPC-only families: PD (``pd_*``), EPD (``epd_*`` and topology tuples), + and multiple-worker topologies. A case authored for both ``grpc`` and + ``http`` is collapsed to one ZMQ run. Items without a ``setup_backend`` + parametrization are left untouched. + """ + kept: list[pytest.Item] = [] + deselected: list[pytest.Item] = [] + groups_with_grpc = {_zmq_dedup_key(it) for it in items if _setup_backend_param(it) == "grpc"} + for item in items: + param = _setup_backend_param(item) + if param is None: + kept.append(item) + continue + if param not in _ZMQ_LOCAL_WIRES or _is_multi_worker(item): + deselected.append(item) + continue + # Collapse the http twin when a grpc one covers the same ZMQ run. + if param == "http" and _zmq_dedup_key(item) in groups_with_grpc: + deselected.append(item) + continue + kept.append(item) + return kept, deselected + + def pytest_collection_modifyitems( config: pytest.Config, items: list[pytest.Item], @@ -124,7 +190,9 @@ def pytest_collection_modifyitems( """Filter + order collected tests. Filtering: env vars ``E2E_ENGINE``, ``E2E_VENDOR``, ``E2E_GPU_TIER`` - select the matching slice when set. + select the matching slice when set. When ``E2E_CONNECTION_MODE=zmq`` the + lane additionally drops the gRPC-only families (PD, EPD, multi-worker) and + collapses ``grpc``/``http`` twins onto a single ZMQ run. Ordering: items are sorted by ``(backend, model)`` so consecutive classes that share a backend cluster together. This is what lets @@ -157,6 +225,12 @@ def pytest_collection_modifyitems( selected.append(item) items[:] = selected + if get_connection_mode_override() == ConnectionMode.ZMQ: + kept, deselected = _filter_zmq_items(items) + if deselected: + config.hook.pytest_deselected(items=deselected) + items[:] = kept + items.sort(key=_pool_sort_key) diff --git a/e2e_test/fixtures/setup_backend.py b/e2e_test/fixtures/setup_backend.py index aedce60fa..8f63e3814 100644 --- a/e2e_test/fixtures/setup_backend.py +++ b/e2e_test/fixtures/setup_backend.py @@ -23,13 +23,16 @@ from infra import ( DEFAULT_MODEL, DEFAULT_ROUTER_TIMEOUT, + DEFAULT_STARTUP_TIMEOUT, ENV_MODEL, ENV_SKIP_BACKEND_SETUP, RUNTIME_LABELS, THIRD_PARTY_MODELS, ConnectionMode, Gateway, + Runtime, WorkerType, + get_connection_mode_override, get_runtime, launch_cloud_gateway, ) @@ -61,6 +64,24 @@ _worker_start_failures: dict[str, int] = {} # engine -> count _MAX_WORKER_START_FAILURES = 3 # fail fast after this many failures (matches --reruns 2) +# Engines that speak the direct-ZMQ backend wire in e2e. Pairing ZMQ with any +# other engine can't work, so we reject it up front instead of timing out on a +# worker that never becomes ready. +ZMQ_CAPABLE_ENGINES = frozenset({Runtime.VLLM.value, Runtime.TOKENSPEED.value}) + + +def _validate_connection_mode(connection_mode: ConnectionMode, engine: str) -> None: + """Reject connection-mode/engine pairings that cannot start. + + Raises ``ValueError`` when a lane selects ZMQ for an engine that does not + support the direct-ZMQ backend. + """ + if connection_mode == ConnectionMode.ZMQ and engine not in ZMQ_CAPABLE_ENGINES: + raise ValueError( + f"ConnectionMode.ZMQ is only supported for engines " + f"{sorted(ZMQ_CAPABLE_ENGINES)}, not {engine!r}" + ) + def _start_workers_tracked(**kwargs) -> list: """Start workers via the session pool and track failures for fail-fast. @@ -91,6 +112,25 @@ def _start_gateway(gateway: Gateway, gateway_config: dict, **mode_kwargs) -> Non ) +def _gateway_readiness_timeout( + connection_mode: ConnectionMode, model_id: str, base_timeout: float +) -> float: + """Effective gateway readiness timeout for the given connection mode. + + gRPC/HTTP workers are health-checked (model fully loaded) by the pool + before the gateway starts, so the gateway only has to connect — the short + router timeout suffices. ZMQ engines instead spawn and return immediately; + their model load happens *inside* the gateway's readiness gate, so that gate + must cover model load too. Use the model's ``startup_timeout`` (what the + worker gate would have applied), never shrinking an explicitly larger + gateway timeout. + """ + if connection_mode != ConnectionMode.ZMQ: + return base_timeout + startup_timeout = get_model_spec(model_id).get("startup_timeout", DEFAULT_STARTUP_TIMEOUT) + return max(base_timeout, startup_timeout) + + def _make_openai_client(gateway: Gateway) -> openai.OpenAI: return openai.OpenAI(base_url=f"{gateway.base_url}/v1", api_key="not-used") @@ -145,7 +185,13 @@ def setup_backend(request: pytest.FixtureRequest): is_pd = backend_name.startswith("pd_") protocol = backend_name.replace("epd_", "").replace("pd_", "") connection_mode = ConnectionMode(protocol) + # A lane can override the local wire (e.g. run grpc/http cases over ZMQ); + # PD/EPD keep their own wire since they are excluded from those lanes. + mode_override = get_connection_mode_override() + if mode_override is not None and not is_pd and not is_epd: + connection_mode = mode_override engine = get_runtime() + _validate_connection_mode(connection_mode, engine) model_path = get_model_spec(model_id)["model"] workers_config = get_marker_kwargs(request, "workers", defaults=_WORKER_DEFAULTS) log_dir = os.environ.get("E2E_LOG_DIR") or gateway_config.get("log_dir") @@ -232,18 +278,36 @@ def _setup_local( gpus=workers_config.get("gpus"), extra_engine_args=workers_config.get("extra_engine_args"), ) + # ZMQ engines dial this gateway's handshake sockets, so they cannot be + # reused by a later class's gateway — the pool starts them fresh and the + # caller owns their teardown (like the PD path). gRPC/HTTP workers stay + # in the pool and outlive the gateway. + is_zmq = connection_mode == ConnectionMode.ZMQ + # ZMQ engines load the model inside the gateway's readiness gate (the worker + # spawn returned immediately), so that gate must cover model load. + gateway_config = { + **gateway_config, + "timeout": _gateway_readiness_timeout(connection_mode, model_id, gateway_config["timeout"]), + } try: _start_gateway( gateway, gateway_config, worker_urls=[w.base_url for w in workers], model_path=model_path, + backend=engine if is_zmq else None, ) logger.info("%s backend ready at %s", backend_name, gateway.base_url) yield backend_name, model_path, _make_openai_client(gateway), gateway finally: - logger.info("Tearing down %s backend (workers stay in pool)", backend_name) + logger.info( + "Tearing down %s backend (%s)", + backend_name, + "stopping ZMQ workers" if is_zmq else "workers stay in pool", + ) gateway.shutdown() + if is_zmq: + stop_workers(workers) # --------------------------------------------------------------------------- @@ -463,20 +527,34 @@ def test_router_state(backend_router): backend_name = request.param model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL) connection_mode = ConnectionMode(backend_name) + mode_override = get_connection_mode_override() + if mode_override is not None: + connection_mode = mode_override + engine = get_runtime() + _validate_connection_mode(connection_mode, engine) model_path = get_model_spec(model_id)["model"] + is_zmq = connection_mode == ConnectionMode.ZMQ # Route through the pool so we evict any cached class-scope worker - # holding the GPUs we need. The pool retains ownership; we don't stop - # the workers ourselves. + # holding the GPUs we need. The pool retains ownership of gRPC/HTTP + # workers; ZMQ engines are bound to this gateway, so we stop them here. workers = get_pool().acquire( model_id=model_id, - engine=get_runtime(), + engine=engine, mode=connection_mode, count=1, ) gateway = Gateway() try: - gateway.start(worker_urls=[w.base_url for w in workers], model_path=model_path) + gateway.start( + worker_urls=[w.base_url for w in workers], + model_path=model_path, + backend=engine if is_zmq else None, + # ZMQ loads the model inside the gateway's readiness gate; cover it. + timeout=_gateway_readiness_timeout(connection_mode, model_id, DEFAULT_ROUTER_TIMEOUT), + ) yield gateway finally: gateway.shutdown() + if is_zmq: + stop_workers(workers) diff --git a/e2e_test/fixtures/test_connection_mode_validation.py b/e2e_test/fixtures/test_connection_mode_validation.py new file mode 100644 index 000000000..54ab0d782 --- /dev/null +++ b/e2e_test/fixtures/test_connection_mode_validation.py @@ -0,0 +1,34 @@ +"""Unit tests for connection-mode/engine validation (no GPU).""" + +from __future__ import annotations + +import pytest +from infra.constants import ConnectionMode, Runtime + +# setup_backend pulls in the cloud SDKs; skip if the env lacks them. +setup_backend = pytest.importorskip("fixtures.setup_backend") + + +@pytest.mark.parametrize("engine", [Runtime.VLLM.value, Runtime.TOKENSPEED.value]) +def test_zmq_allowed_for_capable_engines(engine): + # Does not raise. + setup_backend._validate_connection_mode(ConnectionMode.ZMQ, engine) + + +@pytest.mark.parametrize( + "engine", + [Runtime.SGLANG.value, Runtime.TRTLLM.value, Runtime.MLX.value], +) +def test_zmq_rejected_for_incapable_engines(engine): + with pytest.raises(ValueError, match="ConnectionMode.ZMQ is only supported"): + setup_backend._validate_connection_mode(ConnectionMode.ZMQ, engine) + + +@pytest.mark.parametrize("mode", [ConnectionMode.GRPC, ConnectionMode.HTTP]) +@pytest.mark.parametrize( + "engine", + [Runtime.SGLANG.value, Runtime.VLLM.value, Runtime.TOKENSPEED.value], +) +def test_non_zmq_modes_accept_any_engine(mode, engine): + # Non-ZMQ wires impose no engine restriction. + setup_backend._validate_connection_mode(mode, engine) diff --git a/e2e_test/fixtures/test_hooks_zmq_filter.py b/e2e_test/fixtures/test_hooks_zmq_filter.py new file mode 100644 index 000000000..b3237e67e --- /dev/null +++ b/e2e_test/fixtures/test_hooks_zmq_filter.py @@ -0,0 +1,127 @@ +"""Unit tests for the ZMQ-lane collection filter (no GPU).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from fixtures import hooks + + +class _FakeItem: + """Minimal stand-in for a pytest ``Item`` for the collection helpers. + + Exposes only what ``_filter_zmq_items`` / ``_zmq_dedup_key`` touch: + ``nodeid``, ``callspec.params``, ``cls`` (always None so marker resolution + falls back to ``get_closest_marker``), and a ``workers`` marker. + """ + + def __init__(self, nodeid, params=None, workers=None): + self.nodeid = nodeid + self.cls = None + self.callspec = SimpleNamespace(params=params) if params is not None else None + self._workers = workers + + def get_closest_marker(self, name): + if name == "workers": + return self._workers + return None + + +def _item(nodeid, setup_backend=None, extra_params=None, workers=None): + params = None + if setup_backend is not None or extra_params is not None: + params = {} + if setup_backend is not None: + params["setup_backend"] = setup_backend + if extra_params: + params.update(extra_params) + return _FakeItem(nodeid, params=params, workers=workers) + + +def _workers_marker(**kwargs): + return pytest.mark.workers(**kwargs).mark + + +# --------------------------------------------------------------------------- +# _zmq_dedup_key +# --------------------------------------------------------------------------- + + +def test_dedup_key_ignores_setup_backend_value(): + grpc = _item("t.py::test_x[grpc]", setup_backend="grpc") + http = _item("t.py::test_x[http]", setup_backend="http") + assert hooks._zmq_dedup_key(grpc) == hooks._zmq_dedup_key(http) + + +def test_dedup_key_keeps_other_params_apart(): + a = _item("t.py::test_x[grpc-a]", setup_backend="grpc", extra_params={"api_client": "a"}) + b = _item("t.py::test_x[grpc-b]", setup_backend="grpc", extra_params={"api_client": "b"}) + assert hooks._zmq_dedup_key(a) != hooks._zmq_dedup_key(b) + + +def test_dedup_key_splits_nodeid_at_bracket(): + key, _others = hooks._zmq_dedup_key(_item("t.py::test_x[grpc]", setup_backend="grpc")) + assert key == "t.py::test_x" + + +# --------------------------------------------------------------------------- +# _filter_zmq_items +# --------------------------------------------------------------------------- + + +def test_grpc_http_twins_collapse_to_grpc(): + grpc = _item("t.py::test_x[grpc]", setup_backend="grpc") + http = _item("t.py::test_x[http]", setup_backend="http") + kept, deselected = hooks._filter_zmq_items([grpc, http]) + assert kept == [grpc] + assert deselected == [http] + + +def test_http_only_case_is_retained(): + http = _item("t.py::test_x[http]", setup_backend="http") + kept, deselected = hooks._filter_zmq_items([http]) + assert kept == [http] + assert deselected == [] + + +def test_distinct_other_params_keep_both_wires(): + # A grpc/http pair that differs on another param is NOT a twin. + grpc = _item("t.py::test_x[grpc-a]", setup_backend="grpc", extra_params={"api_client": "a"}) + http = _item("t.py::test_x[http-b]", setup_backend="http", extra_params={"api_client": "b"}) + kept, deselected = hooks._filter_zmq_items([grpc, http]) + assert kept == [grpc, http] + assert deselected == [] + + +@pytest.mark.parametrize("param", ["pd_grpc", "epd_grpc", ("epd_grpc", (1, 1, 1)), ()]) +def test_non_local_wire_families_are_deselected(param): + item = _item("t.py::test_x[p]", setup_backend=param) + kept, deselected = hooks._filter_zmq_items([item]) + assert kept == [] + assert deselected == [item] + + +def test_multi_worker_case_is_deselected(): + item = _item("t.py::test_x[grpc]", setup_backend="grpc", workers=_workers_marker(count=2)) + kept, deselected = hooks._filter_zmq_items([item]) + assert kept == [] + assert deselected == [item] + + +def test_pd_worker_topology_is_deselected(): + item = _item( + "t.py::test_x[grpc]", + setup_backend="grpc", + workers=_workers_marker(prefill=1, decode=1), + ) + kept, deselected = hooks._filter_zmq_items([item]) + assert kept == [] + assert deselected == [item] + + +def test_items_without_setup_backend_are_untouched(): + item = _item("t.py::test_plain") # no callspec / params + kept, deselected = hooks._filter_zmq_items([item]) + assert kept == [item] + assert deselected == [] diff --git a/e2e_test/infra/__init__.py b/e2e_test/infra/__init__.py index 89044d746..bde475869 100644 --- a/e2e_test/infra/__init__.py +++ b/e2e_test/infra/__init__.py @@ -11,6 +11,7 @@ DEFAULT_RUNTIME, DEFAULT_STARTUP_TIMEOUT, ENV_BACKENDS, + ENV_CONNECTION_MODE, ENV_MODEL, ENV_MODELS, ENV_RUNTIME, @@ -32,6 +33,7 @@ ConnectionMode, Runtime, WorkerType, + get_connection_mode_override, get_runtime, is_mlx, is_sglang, @@ -102,6 +104,7 @@ "ENV_BACKENDS", "ENV_MODEL", "ENV_RUNTIME", + "ENV_CONNECTION_MODE", "ENV_STARTUP_TIMEOUT", "ENV_SKIP_MODEL_POOL", "ENV_SKIP_BACKEND_SETUP", @@ -109,6 +112,7 @@ "ENV_SHOW_WORKER_LOGS", # Runtime helpers "get_runtime", + "get_connection_mode_override", "is_vllm", "is_sglang", "is_trtllm", diff --git a/e2e_test/infra/constants.py b/e2e_test/infra/constants.py index ce4b523c6..caac79b26 100644 --- a/e2e_test/infra/constants.py +++ b/e2e_test/infra/constants.py @@ -9,6 +9,7 @@ class ConnectionMode(StrEnum): HTTP = "http" GRPC = "grpc" + ZMQ = "zmq" class WorkerType(StrEnum): @@ -35,7 +36,7 @@ class Runtime(StrEnum): # Convenience sets -LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC}) +LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC, ConnectionMode.ZMQ}) LOCAL_RUNTIMES = frozenset( {Runtime.SGLANG, Runtime.VLLM, Runtime.TRTLLM, Runtime.MLX, Runtime.TOKENSPEED} ) @@ -59,6 +60,9 @@ class Runtime(StrEnum): ENV_RUNTIME = ( "E2E_RUNTIME" # Runtime for gRPC tests — one of Runtime.{SGLANG,VLLM,TRTLLM,TOKENSPEED} ) +ENV_CONNECTION_MODE = ( + "E2E_CONNECTION_MODE" # Per-lane wire override — see get_connection_mode_override +) ENV_STARTUP_TIMEOUT = "E2E_STARTUP_TIMEOUT" ENV_SKIP_MODEL_POOL = "SKIP_MODEL_POOL" ENV_SKIP_BACKEND_SETUP = "SKIP_BACKEND_SETUP" @@ -125,6 +129,31 @@ def is_tokenspeed() -> bool: return get_runtime() == "tokenspeed" +def get_connection_mode_override() -> "ConnectionMode | None": + """Per-lane wire-protocol override for local backends. + + Set ``E2E_CONNECTION_MODE`` to run the existing local test cases over a + different wire (like ``E2E_RUNTIME`` picks the engine): a ``grpc``/``http`` + case then runs over that mode without a separate parametrize value. PD/EPD + backends keep their own wire. Returns ``None`` when the var is unset or + blank (the workflow always exports it and leaves it empty for non-override + lanes); a set-but-unrecognized value is a misconfiguration and raises. + """ + value = os.environ.get(ENV_CONNECTION_MODE) + if value is None: + return None + value = value.strip() + if not value: + return None + valid = [mode.value for mode in ConnectionMode] + try: + return ConnectionMode(value.lower()) + except ValueError: + raise ValueError( + f"{ENV_CONNECTION_MODE}={value!r} is not a valid connection mode; use one of {valid}" + ) from None + + ENV_VLLM_KV_BACKEND = "E2E_VLLM_KV_BACKEND" diff --git a/e2e_test/infra/gateway.py b/e2e_test/infra/gateway.py index 145e9d272..68a123772 100644 --- a/e2e_test/infra/gateway.py +++ b/e2e_test/infra/gateway.py @@ -122,6 +122,7 @@ def start( igw_mode: bool = False, cloud_backend: str | None = None, history_backend: str = "memory", + backend: str | None = None, policy: str = "round_robin", timeout: float = DEFAULT_ROUTER_TIMEOUT, show_output: bool | None = None, @@ -229,8 +230,13 @@ def start( self.model_path = model_path self.pd_mode = False self.igw_mode = False + mode_args = ["--model-path", model_path, "--worker-urls", *worker_urls] + # ZMQ workers share one wire across engine runtimes, so the router + # cannot probe the backend from the ipc:// URL — pin it explicitly. + if backend is not None: + mode_args += ["--backend", backend] self._launch( - mode_args=["--model-path", model_path, "--worker-urls", *worker_urls], + mode_args=mode_args, timeout=timeout, show_output=show_output, extra_args=extra_args, diff --git a/e2e_test/infra/test_connection_mode.py b/e2e_test/infra/test_connection_mode.py new file mode 100644 index 000000000..44262b6c3 --- /dev/null +++ b/e2e_test/infra/test_connection_mode.py @@ -0,0 +1,43 @@ +"""Unit tests for E2E connection-mode env parsing (no GPU).""" + +from __future__ import annotations + +import pytest +from infra.constants import ( + ENV_CONNECTION_MODE, + ConnectionMode, + get_connection_mode_override, +) + + +def test_unset_returns_none(monkeypatch): + monkeypatch.delenv(ENV_CONNECTION_MODE, raising=False) + assert get_connection_mode_override() is None + + +@pytest.mark.parametrize( + "value,expected", + [ + ("zmq", ConnectionMode.ZMQ), + ("ZMQ", ConnectionMode.ZMQ), + ("Grpc", ConnectionMode.GRPC), + (" http ", ConnectionMode.HTTP), + ], +) +def test_valid_values_are_case_insensitive(monkeypatch, value, expected): + monkeypatch.setenv(ENV_CONNECTION_MODE, value) + assert get_connection_mode_override() == expected + + +@pytest.mark.parametrize("value", ["", " "]) +def test_set_but_blank_returns_none(monkeypatch, value): + # The workflow always exports E2E_CONNECTION_MODE and leaves it empty for + # non-override lanes, so a blank value must mean "no override", not an error. + monkeypatch.setenv(ENV_CONNECTION_MODE, value) + assert get_connection_mode_override() is None + + +def test_invalid_value_raises(monkeypatch): + monkeypatch.setenv(ENV_CONNECTION_MODE, "bogus") + with pytest.raises(ValueError, match="not a valid connection mode"): + get_connection_mode_override() diff --git a/e2e_test/infra/test_zmq_cmd_builders.py b/e2e_test/infra/test_zmq_cmd_builders.py new file mode 100644 index 000000000..e12229e6a --- /dev/null +++ b/e2e_test/infra/test_zmq_cmd_builders.py @@ -0,0 +1,71 @@ +"""Unit tests for the ZMQ direct-backend command builders (no GPU).""" + +from __future__ import annotations + +import pytest +from infra.constants import ConnectionMode +from infra.worker import Worker + +_VLLM_MODEL = "meta-llama/Llama-3.2-1B-Instruct" +_TS_MODEL = "Qwen/Qwen3.5-9B" + + +@pytest.fixture +def serve(): + # The ZMQ builders delegate to smg.serve, so the wheel must be importable. + # Scoped to a fixture (not module import) so the gRPC test below still runs + # when the wheel is absent. + return pytest.importorskip("smg.serve") + + +def _worker(engine, port=50111): + return Worker( + model_id=_TS_MODEL if engine == "tokenspeed" else _VLLM_MODEL, + engine=engine, + port=port, + gpu_ids=[0], + mode=ConnectionMode.ZMQ, + ) + + +def test_zmq_base_url_matches_serve_helper(serve): + w = _worker("vllm", port=50123) + assert w.base_url == serve._zmq_ipc_url(50123) + assert w.base_url.startswith("ipc://") + + +def test_vllm_zmq_cmd_is_headless_with_derived_handshake_port(serve): + w = _worker("vllm") + cmd = w._build_vllm_zmq_cmd("/models/llama", 1, {"vllm_args": ["--max-model-len", "2048"]}) + assert "serve" in cmd + assert "--headless" in cmd + assert "/models/llama" in cmd + # The engine dials the same tcp port SMG derives from the ipc path. + expected_port = serve._zmq_handshake_port(serve._zmq_ipc_url(w.port)) + assert cmd[cmd.index("--data-parallel-rpc-port") + 1] == str(expected_port) + # Model-spec engine args ride through. + assert cmd[cmd.index("--max-model-len") + 1] == "2048" + + +def test_tokenspeed_zmq_cmd_is_headless_with_derived_handshake_port(serve): + w = _worker("tokenspeed") + cmd = w._build_tokenspeed_zmq_cmd( + "/models/qwen", 1, {"tokenspeed_args": ["--attention-backend", "fa3"]} + ) + assert "serve" in cmd + assert "--headless" in cmd + assert "/models/qwen" in cmd + expected_port = serve._zmq_handshake_port(serve._zmq_ipc_url(w.port)) + assert cmd[cmd.index("--data-parallel-rpc-port") + 1] == str(expected_port) + assert cmd[cmd.index("--attention-backend") + 1] == "fa3" + + +def test_grpc_worker_still_uses_grpc_url(): + w = Worker( + model_id=_VLLM_MODEL, + engine="vllm", + port=50111, + gpu_ids=[0], + mode=ConnectionMode.GRPC, + ) + assert w.base_url == "grpc://127.0.0.1:50111" diff --git a/e2e_test/infra/worker.py b/e2e_test/infra/worker.py index e1bdb5c7e..f7b2346aa 100644 --- a/e2e_test/infra/worker.py +++ b/e2e_test/infra/worker.py @@ -2,6 +2,7 @@ from __future__ import annotations +import argparse import json import logging import os @@ -54,6 +55,13 @@ class Worker: @property def base_url(self) -> str: """Base URL for this worker.""" + if self.mode == ConnectionMode.ZMQ: + # ipc:// worker URL the router binds; the engine dials the tcp + # handshake port SMG derives from it. Reuse serve's helper so the + # path format stays in lockstep with the launcher and the router. + from smg.serve import _zmq_ipc_url + + return _zmq_ipc_url(self.port) if self.mode == ConnectionMode.GRPC: return f"grpc://{DEFAULT_HOST}:{self.port}" return f"http://{DEFAULT_HOST}:{self.port}" @@ -105,6 +113,17 @@ def start( return # Wait for health check + if self.mode == ConnectionMode.ZMQ: + # SMG (the router) binds the ZMQ sockets and this engine dials in; + # there is no worker port to probe. The gateway's readiness gate + # (wait_for_workers_ready) covers the engine, so just proceed. + logger.info( + "Worker %s spawned at %s (PID %d) — ZMQ readiness gated by the gateway", + self.model_id, + self.base_url, + self.process.pid, + ) + return if self.mode == ConnectionMode.GRPC: self._wait_grpc_healthy(timeout) else: @@ -184,7 +203,9 @@ def _build_cmd(self) -> list[str]: if self.engine == "sglang": cmd = self._build_sglang_cmd(model_path, tp_size, features, spec) elif self.engine == "vllm": - if self.mode == ConnectionMode.GRPC: + if self.mode == ConnectionMode.ZMQ: + cmd = self._build_vllm_zmq_cmd(model_path, tp_size, spec) + elif self.mode == ConnectionMode.GRPC: cmd = self._build_vllm_grpc_cmd(model_path, tp_size, spec) else: cmd = self._build_vllm_http_cmd(model_path, tp_size, spec) @@ -193,12 +214,15 @@ def _build_cmd(self) -> list[str]: elif self.engine == "mlx": cmd = self._build_mlx_cmd(model_path, spec) elif self.engine == "tokenspeed": - if self.mode != ConnectionMode.GRPC: + if self.mode == ConnectionMode.ZMQ: + cmd = self._build_tokenspeed_zmq_cmd(model_path, tp_size, spec) + elif self.mode == ConnectionMode.GRPC: + cmd = self._build_tokenspeed_grpc_cmd(model_path, tp_size, spec) + else: raise ValueError( - "TokenSpeed e2e workers only support gRPC mode; " + "TokenSpeed e2e workers only support gRPC or ZMQ mode; " "HTTP mode would go through the existing OpenAI frontend." ) - cmd = self._build_tokenspeed_grpc_cmd(model_path, tp_size, spec) else: raise ValueError(f"Unsupported engine: {self.engine}") @@ -252,6 +276,21 @@ def _build_sglang_cmd( return cmd + def _build_vllm_zmq_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: + """Build the headless vLLM EngineCore command for the ZMQ direct backend. + + Delegates to the ``smg serve`` launcher so the engine flags and the + FNV-1a handshake port stay identical to the production launch path. + """ + from smg.serve import VllmWorkerLauncher + + args = argparse.Namespace( + connection_mode="zmq", model=model_path, tensor_parallel_size=tp_size + ) + return VllmWorkerLauncher().build_command( + args, list(spec.get("vllm_args", [])), DEFAULT_HOST, self.port + ) + def _build_vllm_grpc_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: """Build vLLM gRPC server command.""" return self._build_vllm_base_cmd("vllm.entrypoints.grpc_server", model_path, tp_size, spec) @@ -316,6 +355,21 @@ def _build_mlx_cmd(self, model_path: str, spec: dict) -> list[str]: cmd.extend(extra) return cmd + def _build_tokenspeed_zmq_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: + """Build the headless TokenSpeed command for the ZMQ direct backend. + + Delegates to the ``smg serve`` launcher so the engine flags and the + FNV-1a handshake port stay identical to the production launch path. + """ + from smg.serve import TokenspeedWorkerLauncher + + args = argparse.Namespace( + connection_mode="zmq", model=model_path, tensor_parallel_size=tp_size + ) + return TokenspeedWorkerLauncher().build_command( + args, list(spec.get("tokenspeed_args", [])), DEFAULT_HOST, self.port + ) + def _build_tokenspeed_grpc_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: """Build TokenSpeed gRPC server command. diff --git a/e2e_test/infra/worker_pool.py b/e2e_test/infra/worker_pool.py index 75f58c4a3..7070cd943 100644 --- a/e2e_test/infra/worker_pool.py +++ b/e2e_test/infra/worker_pool.py @@ -90,8 +90,10 @@ def acquire( # Non-REGULAR workers (PD prefill/decode) aren't cached, but we # still have to release any cached regular worker first — it - # holds the GPUs the caller is about to claim. - if worker_type != WorkerType.REGULAR: + # holds the GPUs the caller is about to claim. ZMQ workers are + # likewise uncached: the engine dials one gateway's handshake + # sockets and cannot be reused by the next class's gateway. + if worker_type != WorkerType.REGULAR or mode == ConnectionMode.ZMQ: if self._key is not None: logger.info( "WorkerPool: evicting %s to free GPUs for non-REGULAR %s/%s", diff --git a/model_gateway/src/routers/grpc/backend_client.rs b/model_gateway/src/routers/grpc/backend_client.rs index 908e0e989..226556ae9 100644 --- a/model_gateway/src/routers/grpc/backend_client.rs +++ b/model_gateway/src/routers/grpc/backend_client.rs @@ -13,16 +13,23 @@ use openai_protocol::{ messages::CreateMessageRequest, worker::WorkerLoadResponse, }; use smg_grpc_client::{ - common_proto, tokenizer_bundle::StreamBundle, SglangSchedulerClient, VllmEngineClient, + common_proto, tokenizer_bundle::StreamBundle, SglangSchedulerClient, TokenSpeedSchedulerClient, + VllmEngineClient, }; -use crate::routers::grpc::{ - client::{GenerateRequestBuildOptions, GrpcClient, HealthCheckResponse, ModelInfo, ServerInfo}, - proto_wrapper::{ - finish_vllm_request, ProtoEmbedComplete, ProtoEmbedRequest, ProtoGenerateRequest, - ProtoStream, +use crate::{ + routers::grpc::{ + client::{ + GenerateRequestBuildOptions, GrpcClient, HealthCheckResponse, ModelInfo, ServerInfo, + }, + proto_wrapper::{ + finish_tokenspeed_request, finish_vllm_request, ProtoEmbedComplete, ProtoEmbedRequest, + ProtoGenerateRequest, ProtoStream, + }, + zmq_client::ZmqEngineClient, + MultimodalData, }, - zmq_client::ZmqEngineClient, + worker::RuntimeType, }; /// A backend connection: gRPC (any engine) or direct ZMQ (vLLM EngineCore or @@ -35,19 +42,17 @@ pub enum BackendClient { impl BackendClient { /// Runtime type backing this client. - pub fn runtime_type(&self) -> crate::worker::RuntimeType { + pub fn runtime_type(&self) -> RuntimeType { match self { Self::Grpc(client) => client.runtime_type(), Self::Zmq(client) => client.runtime(), } } - /// True if this backend speaks the vLLM protocol (gRPC-vLLM or ZMQ). - pub fn is_vllm(&self) -> bool { - match self { - Self::Grpc(client) => client.is_vllm(), - Self::Zmq(_) => true, - } + /// True if this is a direct-ZMQ backend (the engine receives token ids only + /// and cannot match string stops itself). + pub fn is_zmq(&self) -> bool { + matches!(self, Self::Zmq(_)) } /// Local liveness. gRPC has no cheap local flag (it uses a health RPC), so @@ -88,14 +93,14 @@ impl BackendClient { pub async fn get_model_info(&self) -> Result { match self { Self::Grpc(client) => client.get_model_info().await, - Self::Zmq(client) => Ok(ModelInfo::Vllm(client.get_model_info())), + Self::Zmq(client) => Ok(client.get_model_info()), } } pub async fn get_server_info(&self) -> Result { match self { Self::Grpc(client) => client.get_server_info().await, - Self::Zmq(client) => Ok(ServerInfo::Vllm(client.get_server_info())), + Self::Zmq(client) => Ok(client.get_server_info()), } } @@ -168,14 +173,7 @@ impl BackendClient { ) -> Result { match self { Self::Grpc(client) => client.generate(req).await, - Self::Zmq(client) => match req { - ProtoGenerateRequest::Vllm(boxed_req) => { - Ok(ProtoStream::Zmq(client.generate(*boxed_req).await?)) - } - _ => Err(tonic::Status::internal( - "ZMQ backend expects a vLLM generate request", - )), - }, + Self::Zmq(client) => Ok(ProtoStream::Zmq(client.generate(req).await?)), } } @@ -191,6 +189,10 @@ impl BackendClient { } } + #[expect( + clippy::unreachable, + reason = "assembly stage guarantees matching MultimodalData variant for each backend" + )] pub fn build_chat_request( &self, request_id: String, @@ -203,22 +205,50 @@ impl BackendClient { Self::Grpc(client) => { client.build_chat_request(request_id, body, processed_text, token_ids, options) } - Self::Zmq(_) => { - reject_zmq_multimodal(&options)?; - finish_vllm_request(None, |mm| { - VllmEngineClient::build_generate_request_from_chat( - request_id, - body, - processed_text, - token_ids, - mm, - options.tool_constraints, - ) - }) - } + // A ZMQ backend speaks vLLM EngineCore or TokenSpeed directly; build + // the native request for its runtime, mirroring the gRPC per-engine + // dispatch in `GrpcClient::build_chat_request`. + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::TokenSpeed(data) => data.into_proto(true), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_tokenspeed_request(tokenspeed_mm, |mm| { + TokenSpeedSchedulerClient::build_generate_request_from_chat( + request_id, + body, + processed_text, + token_ids, + mm, + options.tool_constraints, + ) + }) + } + _ => { + let vllm_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::Vllm(data) => data.into_proto(), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_vllm_request(vllm_mm, |mm| { + VllmEngineClient::build_generate_request_from_chat( + request_id, + body, + processed_text, + token_ids, + mm, + options.tool_constraints, + ) + }) + } + }, } } + #[expect( + clippy::unreachable, + reason = "assembly stage guarantees matching MultimodalData variant for each backend" + )] pub fn build_messages_request( &self, request_id: String, @@ -231,19 +261,42 @@ impl BackendClient { Self::Grpc(client) => { client.build_messages_request(request_id, body, processed_text, token_ids, options) } - Self::Zmq(_) => { - reject_zmq_multimodal(&options)?; - finish_vllm_request(None, |mm| { - VllmEngineClient::build_generate_request_from_messages( - request_id, - body, - processed_text, - token_ids, - mm, - options.tool_constraints, - ) - }) - } + // Mirrors the gRPC per-engine dispatch: build the request natively for + // the ZMQ backend's runtime (vLLM EngineCore or TokenSpeed). + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::TokenSpeed(data) => data.into_proto(true), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_tokenspeed_request(tokenspeed_mm, |mm| { + TokenSpeedSchedulerClient::build_generate_request_from_messages( + request_id, + body, + processed_text, + token_ids, + mm, + options.tool_constraints, + ) + }) + } + _ => { + let vllm_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::Vllm(data) => data.into_proto(), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_vllm_request(vllm_mm, |mm| { + VllmEngineClient::build_generate_request_from_messages( + request_id, + body, + processed_text, + token_ids, + mm, + options.tool_constraints, + ) + }) + } + }, } } @@ -258,15 +311,26 @@ impl BackendClient { Self::Grpc(client) => { client.build_completion_request(request_id, body, original_text, token_ids) } - Self::Zmq(_) => { - let req = VllmEngineClient::build_generate_request_from_completion( - request_id, - body, - original_text, - token_ids, - )?; - Ok(ProtoGenerateRequest::Vllm(Box::new(req))) - } + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let req = TokenSpeedSchedulerClient::build_generate_request_from_completion( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::TokenSpeed(Box::new(req))) + } + _ => { + let req = VllmEngineClient::build_generate_request_from_completion( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::Vllm(Box::new(req))) + } + }, } } @@ -281,23 +345,26 @@ impl BackendClient { Self::Grpc(client) => { client.build_generate_request(request_id, body, original_text, token_ids) } - Self::Zmq(_) => { - let req = VllmEngineClient::build_plain_generate_request( - request_id, - body, - original_text, - token_ids, - )?; - Ok(ProtoGenerateRequest::Vllm(Box::new(req))) - } + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let req = TokenSpeedSchedulerClient::build_plain_generate_request( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::TokenSpeed(Box::new(req))) + } + _ => { + let req = VllmEngineClient::build_plain_generate_request( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::Vllm(Box::new(req))) + } + }, } } } - -/// ZMQ text path does not carry multimodal inputs yet. -fn reject_zmq_multimodal(options: &GenerateRequestBuildOptions) -> Result<(), String> { - if options.multimodal_inputs.is_some() { - return Err("ZMQ backend does not support multimodal inputs yet".to_string()); - } - Ok(()) -} diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index b420173ab..2bd1221b5 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -541,13 +541,13 @@ impl GrpcClient { )?; Ok(ProtoGenerateRequest::Mlx(Box::new(req))) } - Self::TokenSpeed(client) => { + Self::TokenSpeed(_) => { let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { MultimodalData::TokenSpeed(data) => data.into_proto(true), _ => unreachable!("caller guarantees matching variant"), }); finish_tokenspeed_request(tokenspeed_mm, |mm| { - client.build_generate_request_from_chat( + TokenSpeedSchedulerClient::build_generate_request_from_chat( request_id, body, processed_text, @@ -633,13 +633,13 @@ impl GrpcClient { )?; Ok(ProtoGenerateRequest::Mlx(Box::new(req))) } - Self::TokenSpeed(client) => { + Self::TokenSpeed(_) => { let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { MultimodalData::TokenSpeed(data) => data.into_proto(true), _ => unreachable!("caller guarantees matching variant"), }); finish_tokenspeed_request(tokenspeed_mm, |mm| { - client.build_generate_request_from_messages( + TokenSpeedSchedulerClient::build_generate_request_from_messages( request_id, body, processed_text, @@ -696,8 +696,8 @@ impl GrpcClient { )?; Ok(ProtoGenerateRequest::Mlx(Box::new(req))) } - Self::TokenSpeed(client) => { - let req = client.build_generate_request_from_completion( + Self::TokenSpeed(_) => { + let req = TokenSpeedSchedulerClient::build_generate_request_from_completion( request_id, body, original_text, @@ -752,8 +752,8 @@ impl GrpcClient { )?; Ok(ProtoGenerateRequest::Mlx(Box::new(req))) } - Self::TokenSpeed(client) => { - let req = client.build_plain_generate_request( + Self::TokenSpeed(_) => { + let req = TokenSpeedSchedulerClient::build_plain_generate_request( request_id, body, original_text, diff --git a/model_gateway/src/routers/grpc/common/stages/helpers.rs b/model_gateway/src/routers/grpc/common/stages/helpers.rs index 3cecfe4a6..d8f5331fd 100644 --- a/model_gateway/src/routers/grpc/common/stages/helpers.rs +++ b/model_gateway/src/routers/grpc/common/stages/helpers.rs @@ -2,6 +2,7 @@ use std::sync::Arc; +use llm_tokenizer::traits::Tokenizer; use rand::RngExt; use smg_grpc_client::{ mlx_proto, @@ -291,6 +292,118 @@ fn apply_tokenspeed_sampling_defaults( apply_opt!(repetition_penalty); } +/// Convert single-token stop strings into `stop_token_ids` entries so the engine +/// can halt generation early for the common case (e.g. `["."]`, `["\n"]`). +/// +/// The proto `stop_token_ids` field is a flat list of single token ids, so a +/// multi-token stop string cannot be represented there — pushing its sub-tokens +/// would stop far too eagerly (on any one of them). Multi-token, empty, and +/// unknown stops are therefore left to the router-side `StopSequenceDecoder`, +/// which detokenizes worker output and trims the stop text. Existing +/// `stop_token_ids` are preserved and deduped. +fn encode_single_token_stops( + stops: Vec, + stop_token_ids: &mut Vec, + tokenizer: Option<&Arc>, +) { + // Without a tokenizer we cannot encode (not expected on paths that resolve + // one to tokenize the prompt). Safe: the strings are already dropped by the + // caller, so the router-side decoder remains the source of truth. + let Some(tokenizer) = tokenizer else { + if !stops.is_empty() { + warn!( + "No tokenizer available to encode string stop sequences; \ + relying on router-side stop decoder only" + ); + } + return; + }; + + for stop in stops { + if stop.is_empty() { + continue; + } + // add_special_tokens=false: we want the literal token(s) for the stop + // string, not a BOS/EOS-wrapped encoding. + match tokenizer.encode(&stop, false) { + Ok(encoding) => match encoding.token_ids() { + [id] => { + if !stop_token_ids.contains(id) { + stop_token_ids.push(*id); + } + } + ids => debug!( + stop = %stop, + token_count = ids.len(), + "string stop is not single-token; handled by router-side stop decoder" + ), + }, + Err(e) => warn!( + stop = %stop, + error = %e, + "Failed to encode string stop sequence; relying on router-side stop decoder" + ), + } + } +} + +/// Router-authoritative string-`stop` resolution for backends whose engine +/// cannot match string stops itself. +/// +/// vLLM over gRPC detokenizes server-side (`detokenize=bool(stop)`), TRT-LLM +/// tokenizes stop words server-side, and MLX has no string-`stop` field — those +/// keep their strings untouched. Two paths cannot: +/// - SGLang gRPC workers run with `skip_tokenizer_init=True` and reject string +/// stops outright (a 400 for any request carrying `stop`); and +/// - every direct-ZMQ backend (vLLM EngineCore, TokenSpeed) receives token ids +/// only, so the engine never sees — and cannot match — a stop string. +/// +/// For both, the router owns the tokenizer and already matches string stops via +/// `StopSequenceDecoder` (it detokenizes worker output and trims), so the worker +/// never needs the raw strings. This drops the string `stop` list and forwards +/// any single-token stop as a `stop_token_ids` entry for early stopping; the +/// router-side decoder handles the rest. This is the single resolution point +/// shared by SGLang gRPC and every ZMQ backend. +pub(crate) fn resolve_string_stops( + request: &mut ProtoGenerateRequest, + tokenizer: Option<&Arc>, + is_zmq: bool, +) { + // SGLang always needs it; other protos only when talking to a ZMQ backend + // (which, on this path, always carries the vLLM proto). + match request { + ProtoGenerateRequest::Sglang(req) => { + if let Some(params) = req.sampling_params.as_mut() { + let stops = std::mem::take(&mut params.stop); + encode_single_token_stops(stops, &mut params.stop_token_ids, tokenizer); + } + } + ProtoGenerateRequest::Vllm(req) if is_zmq => { + if let Some(params) = req.sampling_params.as_mut() { + let stops = std::mem::take(&mut params.stop); + encode_single_token_stops(stops, &mut params.stop_token_ids, tokenizer); + // The engine only stops at EOS when the frontend supplies the + // ids, and the connect-time model-dir resolution has nothing + // to read when the worker's model id is a repo id rather than + // a local path. The tokenizer carries the merged EOS set, so + // fold it into the stop tokens as the always-available + // backstop — without it an uncapped request generates to the + // full context window. + if !params.ignore_eos { + if let Some(tokenizer) = tokenizer { + for &id in tokenizer.eos_token_ids() { + if !params.stop_token_ids.contains(&id) { + params.stop_token_ids.push(id); + } + } + } + } + } + } + _ => {} + } +} + /// Inject PD bootstrap metadata for SGLang if needed. /// /// SGLang uses DisaggregatedParams with bootstrap host/port/room. @@ -434,3 +547,162 @@ mod request_id_tests { assert!(id.starts_with("chatcmpl-")); } } + +#[cfg(test)] +mod stop_resolution_tests { + use std::sync::Arc; + + use llm_tokenizer::{mock::MockTokenizer, traits::Tokenizer}; + use smg_grpc_client::{sglang_proto, vllm_proto}; + + use super::{resolve_string_stops, ProtoGenerateRequest}; + + fn mock_tokenizer() -> Arc { + // MockTokenizer vocab: "." => 6, "Hello" => 1, "world" => 2. `encode` + // splits on whitespace, so "." => [6] (single) and "Hello world" => + // [1, 2] (multi); unknown words encode to []. + Arc::new(MockTokenizer::new()) + } + + fn sglang_request(stop: Vec<&str>, stop_token_ids: Vec) -> ProtoGenerateRequest { + ProtoGenerateRequest::Sglang(Box::new(sglang_proto::GenerateRequest { + sampling_params: Some(sglang_proto::SamplingParams { + stop: stop.into_iter().map(str::to_string).collect(), + stop_token_ids, + ..Default::default() + }), + ..Default::default() + })) + } + + fn vllm_request(stop: Vec<&str>, stop_token_ids: Vec) -> ProtoGenerateRequest { + ProtoGenerateRequest::Vllm(Box::new(vllm_proto::GenerateRequest { + sampling_params: Some(vllm_proto::SamplingParams { + stop: stop.into_iter().map(str::to_string).collect(), + stop_token_ids, + ..Default::default() + }), + ..Default::default() + })) + } + + fn sglang_params(req: &ProtoGenerateRequest) -> &sglang_proto::SamplingParams { + match req { + ProtoGenerateRequest::Sglang(r) => r.sampling_params.as_ref().unwrap(), + _ => panic!("expected SGLang request"), + } + } + + fn vllm_params(req: &ProtoGenerateRequest) -> &vllm_proto::SamplingParams { + match req { + ProtoGenerateRequest::Vllm(r) => r.sampling_params.as_ref().unwrap(), + _ => panic!("expected vLLM request"), + } + } + + #[test] + fn sglang_single_token_becomes_stop_token_id() { + let mut req = sglang_request(vec!["."], vec![]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty(), "string stop should be cleared"); + assert_eq!(params.stop_token_ids, vec![6]); + } + + #[test] + fn sglang_multi_token_relies_on_router_decoder() { + // "Hello world" => [1, 2]: can't be a flat stop_token_id, so it must not + // be forwarded (would over-eagerly stop on any subtoken). + let mut req = sglang_request(vec!["Hello world"], vec![]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty()); + assert!(params.stop_token_ids.is_empty()); + } + + #[test] + fn sglang_mixed_only_single_token_forwarded_and_dedups() { + let mut req = sglang_request(vec![".", "Hello world"], vec![6, 42]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty()); + assert_eq!( + params.stop_token_ids, + vec![6, 42], + "existing ids kept, no dup" + ); + } + + #[test] + fn sglang_without_tokenizer_still_clears_strings() { + let mut req = sglang_request(vec!["."], vec![]); + resolve_string_stops(&mut req, None, false); + + let params = sglang_params(&req); + assert!( + params.stop.is_empty(), + "strings dropped so worker won't 400" + ); + assert!(params.stop_token_ids.is_empty()); + } + + #[test] + fn vllm_resolved_only_over_zmq() { + // gRPC vLLM keeps its strings (the servicer detokenizes engine-side). + let mut grpc = vllm_request(vec!["."], vec![]); + resolve_string_stops(&mut grpc, Some(&mock_tokenizer()), false); + let params = vllm_params(&grpc); + assert_eq!( + params.stop, + vec![".".to_string()], + "gRPC vLLM stop preserved" + ); + assert!(params.stop_token_ids.is_empty()); + + // ZMQ vLLM (EngineCore sees token ids only) resolves like SGLang, and + // gains the tokenizer's EOS ids so generation always terminates. + let mut zmq = vllm_request(vec!["."], vec![]); + resolve_string_stops(&mut zmq, Some(&mock_tokenizer()), true); + let params = vllm_params(&zmq); + assert!(params.stop.is_empty(), "ZMQ vLLM stop cleared"); + assert_eq!(params.stop_token_ids, vec![6, 999]); + } + + #[test] + fn vllm_zmq_appends_tokenizer_eos_ids() { + let mut req = vllm_request(vec![], vec![7]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![7, 999]); + + // Already-present EOS ids are not duplicated. + let mut req = vllm_request(vec![], vec![999]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![999]); + } + + #[test] + fn vllm_zmq_ignore_eos_skips_injection() { + let mut req = ProtoGenerateRequest::Vllm(Box::new(vllm_proto::GenerateRequest { + sampling_params: Some(vllm_proto::SamplingParams { + stop_token_ids: vec![7], + ignore_eos: true, + ..Default::default() + }), + ..Default::default() + })); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![7]); + } + + #[test] + fn noop_when_no_string_stops() { + let mut req = sglang_request(vec![], vec![7]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert_eq!(params.stop_token_ids, vec![7], "unrelated ids untouched"); + } +} diff --git a/model_gateway/src/routers/grpc/harmony/mod.rs b/model_gateway/src/routers/grpc/harmony/mod.rs index 59896abfe..dd0c3641c 100644 --- a/model_gateway/src/routers/grpc/harmony/mod.rs +++ b/model_gateway/src/routers/grpc/harmony/mod.rs @@ -35,6 +35,7 @@ pub(crate) mod parser; pub(crate) mod processor; pub(crate) mod responses; pub(crate) mod stages; +pub(crate) mod stop; pub(crate) mod streaming; pub(crate) mod types; diff --git a/model_gateway/src/routers/grpc/harmony/processor.rs b/model_gateway/src/routers/grpc/harmony/processor.rs index 6a74a1bec..df975093a 100644 --- a/model_gateway/src/routers/grpc/harmony/processor.rs +++ b/model_gateway/src/routers/grpc/harmony/processor.rs @@ -16,6 +16,7 @@ use tracing::error; use super::{ builder::{convert_harmony_logprobs, try_harmony_encoding}, + stop::TextStopScanner, HarmonyParserAdapter, }; use crate::routers::{ @@ -39,11 +40,16 @@ impl HarmonyResponseProcessor { } /// Process a non-streaming Harmony chat response + /// + /// `router_stop_strings` is non-empty only when the router must enforce + /// string `stop` sequences itself (direct-ZMQ backends: the engine sees + /// token ids only). pub async fn process_non_streaming_chat_response( &self, execution_result: ExecutionResult, chat_request: Arc, dispatch: DispatchMetadata, + router_stop_strings: &[String], ) -> Result { let request_logprobs = chat_request.logprobs; @@ -103,16 +109,47 @@ impl HarmonyResponseProcessor { None }; + let mut analysis = parsed.analysis; + let mut final_text = parsed.final_text; + let mut tool_calls = parsed.commentary; + let mut finish_reason = parsed.finish_reason; + let mut matched_stop = matched_stop; + + // Router-enforced string stops: scan channel text in emission + // order (analysis before final), truncate after the first match + // (retaining the stop as suffix, like the token-forwarding gRPC + // path), and drop everything generated past it. + if !router_stop_strings.is_empty() { + let mut scanner = TextStopScanner::new(router_stop_strings.to_vec()); + if let Some(text) = analysis.take() { + let (out, stopped) = scanner.scan_complete(&text); + analysis = (!out.is_empty()).then_some(out); + if stopped { + final_text.clear(); + tool_calls = None; + } + } + if scanner.matched().is_none() && !final_text.is_empty() { + let (out, stopped) = scanner.scan_complete(&final_text); + final_text = out; + if stopped { + tool_calls = None; + } + } + if let Some(stop) = scanner.matched() { + finish_reason = "stop".to_string(); + matched_stop = Some(serde_json::Value::String(stop.to_string())); + } + } + // Build response message (assistant) let message = ChatCompletionMessage { role: "assistant".to_string(), - content: (!parsed.final_text.is_empty()).then_some(parsed.final_text), - tool_calls: parsed.commentary, - reasoning_content: parsed.analysis, + content: (!final_text.is_empty()).then_some(final_text), + tool_calls, + reasoning_content: analysis, }; - let finish_reason = parsed.finish_reason; - // Accumulate reasoning tokens across all responses total_reasoning_tokens += parsed.reasoning_token_count; diff --git a/model_gateway/src/routers/grpc/harmony/stages/request_building.rs b/model_gateway/src/routers/grpc/harmony/stages/request_building.rs index dda6e4426..568afb603 100644 --- a/model_gateway/src/routers/grpc/harmony/stages/request_building.rs +++ b/model_gateway/src/routers/grpc/harmony/stages/request_building.rs @@ -2,21 +2,24 @@ use async_trait::async_trait; use axum::response::Response; -use smg_grpc_client::{SglangGenerateRequestOptions, VllmEngineClient}; +use smg_grpc_client::{SglangGenerateRequestOptions, TokenSpeedSchedulerClient, VllmEngineClient}; use tracing::{debug, error}; -use crate::routers::{ - error, - grpc::{ - backend_client::BackendClient, - client::GrpcClient, - common::stages::{helpers, PipelineStage}, - context::{ - ClientSelection, ExecutionPlan, ExecutionPlanKind, PreparationOutput, RequestContext, - RequestType, +use crate::{ + routers::{ + error, + grpc::{ + backend_client::BackendClient, + client::GrpcClient, + common::stages::{helpers, PipelineStage}, + context::{ + ClientSelection, ExecutionPlan, ExecutionPlanKind, PreparationOutput, + RequestContext, RequestType, + }, + proto_wrapper::ProtoGenerateRequest, }, - proto_wrapper::ProtoGenerateRequest, }, + worker::RuntimeType, }; /// Harmony Request Building stage: Convert Harmony tokens to gRPC request @@ -299,12 +302,11 @@ impl PipelineStage for HarmonyRequestBuildingStage { }; ProtoGenerateRequest::Mlx(Box::new(req)) } - BackendClient::Grpc(GrpcClient::TokenSpeed(tokenspeed_client)) => { + BackendClient::Grpc(GrpcClient::TokenSpeed(_)) => { let req = match &ctx.input.request_type { RequestType::Chat(request) => { let body = modified_request.as_deref().unwrap_or_else(|| request.as_ref()); - tokenspeed_client - .build_generate_request_from_chat( + TokenSpeedSchedulerClient::build_generate_request_from_chat( request_id, body, placeholder_processed_text, @@ -317,8 +319,7 @@ impl PipelineStage for HarmonyRequestBuildingStage { error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) })? } - RequestType::Responses(request) => tokenspeed_client - .build_generate_request_from_responses( + RequestType::Responses(request) => TokenSpeedSchedulerClient::build_generate_request_from_responses( request_id, request.as_ref(), placeholder_processed_text, @@ -344,9 +345,58 @@ impl PipelineStage for HarmonyRequestBuildingStage { }; ProtoGenerateRequest::TokenSpeed(Box::new(req)) } + // A ZMQ worker speaks vLLM EngineCore or TokenSpeed directly; build the + // request natively for its runtime, mirroring the gRPC per-engine + // dispatch above. Both support the Harmony request types (Chat + + // Responses). + BackendClient::Zmq(zmq_client) if zmq_client.runtime() == RuntimeType::TokenSpeed => { + let req = match &ctx.input.request_type { + RequestType::Chat(request) => { + let body = modified_request + .as_deref() + .unwrap_or_else(|| request.as_ref()); + TokenSpeedSchedulerClient::build_generate_request_from_chat( + request_id, + body, + placeholder_processed_text, + token_ids, + None, // Harmony path: multimodal not yet wired + tool_constraints, + ) + .map_err(|e| { + error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TokenSpeed ZMQ generate request"); + error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) + })? + } + RequestType::Responses(request) => { + TokenSpeedSchedulerClient::build_generate_request_from_responses( + request_id, + request.as_ref(), + placeholder_processed_text, + token_ids, + tool_constraints, + ) + .map_err(|e| { + error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TokenSpeed ZMQ generate request from responses"); + error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) + })? + } + RequestType::Embedding(_) => { + return Err(error::bad_request( + "harmony_embedding_not_supported", + "Embedding requests are not supported with Harmony models".to_string(), + )); + } + _ => { + return Err(error::bad_request( + "unsupported_request_type", + "Unsupported request type for Harmony models".to_string(), + )); + } + }; + ProtoGenerateRequest::TokenSpeed(Box::new(req)) + } BackendClient::Zmq(_) => { - // A ZMQ worker is a vLLM engine, so it uses the same vLLM request - // builders and supports the same request types (Chat + Responses). let req = match &ctx.input.request_type { RequestType::Chat(request) => { let body = modified_request diff --git a/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs b/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs index f745ca589..488c16b4d 100644 --- a/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs +++ b/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs @@ -12,12 +12,38 @@ use crate::{ error, grpc::{ common::stages::PipelineStage, - context::{FinalResponse, RequestContext, RequestType}, + context::{ClientSelection, FinalResponse, RequestContext, RequestType}, }, }, worker::AttachedBody, }; +/// String `stop` sequences the ROUTER must enforce: only for direct-ZMQ +/// backends, where the engine receives token ids and never sees the strings. +/// Empty for gRPC backends (the engine matches stops itself). +fn router_stop_strings(ctx: &RequestContext) -> Vec { + let is_zmq = ctx + .state + .clients + .as_ref() + .is_some_and(|clients| match clients { + ClientSelection::Single { client } => client.is_zmq(), + ClientSelection::Disaggregated { decode, .. } => decode.is_zmq(), + }); + if !is_zmq { + return Vec::new(); + } + match &ctx.input.request_type { + RequestType::Chat(_) => ctx + .chat_request_arc() + .stop + .as_ref() + .map(|stop| stop.to_vec()) + .unwrap_or_default(), + _ => Vec::new(), + } +} + /// Harmony Response Processing stage: Parse and format Harmony responses /// /// Takes output tokens from execution and parses them using HarmonyParserAdapter @@ -80,6 +106,7 @@ impl PipelineStage for HarmonyResponseProcessingStage { execution_result, ctx.chat_request_arc(), dispatch, + router_stop_strings(ctx), ); // Attach load guards to response body for proper RAII lifecycle @@ -93,9 +120,15 @@ impl PipelineStage for HarmonyResponseProcessingStage { // For non-streaming, delegate to Harmony response processor to build ChatCompletionResponse let chat_request = ctx.chat_request_arc(); + let stops = router_stop_strings(ctx); let response = self .processor - .process_non_streaming_chat_response(execution_result, chat_request, dispatch) + .process_non_streaming_chat_response( + execution_result, + chat_request, + dispatch, + &stops, + ) .await?; ctx.state.response.final_response = Some(FinalResponse::Chat(response)); diff --git a/model_gateway/src/routers/grpc/harmony/stop.rs b/model_gateway/src/routers/grpc/harmony/stop.rs new file mode 100644 index 000000000..f6f49a72d --- /dev/null +++ b/model_gateway/src/routers/grpc/harmony/stop.rs @@ -0,0 +1,163 @@ +//! Incremental stop-string scanning for harmony channel text. +//! +//! The direct-ZMQ engine receives token ids only and cannot match string +//! `stop` sequences; the regular pipeline covers this with the router-side +//! `StopSequenceDecoder`, but harmony emits channel-parsed text rather than +//! raw tokens, so stops are matched here on the decoded text instead. The +//! matched stop stays in the output as its suffix, mirroring what the +//! token-forwarding gRPC path yields for harmony models. + +/// Result of feeding text through the scanner. +pub(crate) struct StopScan { + /// Text safe to emit: everything up to and including a match, or the + /// prefix that cannot be part of a future match. + pub emit: String, + /// A stop sequence completed inside this push. + pub stopped: bool, +} + +pub(crate) struct TextStopScanner { + stops: Vec, + /// Tail held back because it could still grow into a match. + pending: String, + matched: Option, +} + +impl TextStopScanner { + pub fn new(stops: Vec) -> Self { + Self { + stops, + pending: String::new(), + matched: None, + } + } + + pub fn matched(&self) -> Option<&str> { + self.matched.as_deref() + } + + /// Feed a text delta. Once a stop has matched, further pushes emit nothing. + pub fn push(&mut self, text: &str) -> StopScan { + if self.matched.is_some() { + return StopScan { + emit: String::new(), + stopped: true, + }; + } + self.pending.push_str(text); + + // Earliest match across all stops wins; ties prefer the longer stop. + let mut best: Option<(usize, &str)> = None; + for stop in &self.stops { + if let Some(at) = self.pending.find(stop.as_str()) { + let better = match best { + None => true, + Some((best_at, best_stop)) => { + at < best_at || (at == best_at && stop.len() > best_stop.len()) + } + }; + if better { + best = Some((at, stop)); + } + } + } + if let Some((at, stop)) = best { + let emit = self.pending[..at + stop.len()].to_string(); + self.matched = Some(stop.to_string()); + self.pending.clear(); + return StopScan { + emit, + stopped: true, + }; + } + + // Hold back the longest suffix that is a prefix of some stop; emit the + // rest (it can never become part of a match). + let max_hold = (self.stops.iter().map(String::len).max()) + .unwrap_or(1) + .saturating_sub(1); + let mut hold = 0; + for len in (1..=max_hold.min(self.pending.len())).rev() { + let Some(start) = self.pending.len().checked_sub(len) else { + continue; + }; + if !self.pending.is_char_boundary(start) { + continue; + } + let tail = &self.pending[start..]; + if self.stops.iter().any(|s| s.starts_with(tail)) { + hold = len; + break; + } + } + let emit = self.pending[..self.pending.len() - hold].to_string(); + self.pending.drain(..self.pending.len() - hold); + StopScan { + emit, + stopped: false, + } + } + + /// Emit whatever is still held back (end of stream, no match). + pub fn flush(&mut self) -> String { + std::mem::take(&mut self.pending) + } + + /// Scan a complete text: returns the (possibly truncated) text and whether + /// a stop matched inside it. + pub fn scan_complete(&mut self, text: &str) -> (String, bool) { + let scan = self.push(text); + let mut out = scan.emit; + if !scan.stopped { + out.push_str(&self.flush()); + } + (out, scan.stopped) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn complete_scan_truncates_after_first_match() { + let mut scanner = TextStopScanner::new(vec![",".to_string()]); + let (out, stopped) = scanner.scan_complete("1, 2, 3"); + assert_eq!(out, "1,"); + assert!(stopped); + assert_eq!(scanner.matched(), Some(",")); + } + + #[test] + fn streaming_holds_back_partial_matches() { + let mut scanner = TextStopScanner::new(vec!["STOP".to_string()]); + let scan = scanner.push("abc ST"); + assert_eq!(scan.emit, "abc "); + assert!(!scan.stopped); + let scan = scanner.push("OP tail"); + assert_eq!(scan.emit, "STOP"); + assert!(scan.stopped); + // Post-match pushes are swallowed. + let scan = scanner.push("more"); + assert_eq!(scan.emit, ""); + assert!(scan.stopped); + } + + #[test] + fn flush_releases_unmatched_tail() { + let mut scanner = TextStopScanner::new(vec!["END".to_string()]); + let scan = scanner.push("value EN"); + assert_eq!(scan.emit, "value "); + assert_eq!(scanner.flush(), "EN"); + assert_eq!(scanner.matched(), None); + } + + #[test] + fn earliest_match_wins_across_stops() { + let mut scanner = TextStopScanner::new(vec!["zz".to_string(), "b".to_string()]); + let (out, stopped) = scanner.scan_complete("abzz"); + assert_eq!(out, "ab"); + assert!(stopped); + assert_eq!(scanner.matched(), Some("b")); + } +} diff --git a/model_gateway/src/routers/grpc/harmony/streaming.rs b/model_gateway/src/routers/grpc/harmony/streaming.rs index cef6b643a..66d231452 100644 --- a/model_gateway/src/routers/grpc/harmony/streaming.rs +++ b/model_gateway/src/routers/grpc/harmony/streaming.rs @@ -27,6 +27,7 @@ use tracing::{debug, error}; use super::{ builder::{convert_harmony_logprobs, try_harmony_encoding}, processor::ResponsesIterationResult, + stop::TextStopScanner, types::HarmonyChannelDelta, HarmonyParserAdapter, }; @@ -100,11 +101,15 @@ impl HarmonyStreamingProcessor { clippy::disallowed_methods, reason = "streaming tasks are fire-and-forget by design; client disconnect terminates them" )] + /// `router_stop_strings` is non-empty only when the router must enforce + /// string `stop` sequences itself (direct-ZMQ backends: the engine sees + /// token ids only). pub fn process_streaming_chat_response( self: Arc, execution_result: context::ExecutionResult, chat_request: Arc, dispatch: context::DispatchMetadata, + router_stop_strings: Vec, ) -> Response { // Create SSE channel let (tx, rx) = mpsc::unbounded_channel::>(); @@ -113,8 +118,14 @@ impl HarmonyStreamingProcessor { match execution_result { context::ExecutionResult::Single { stream } => { tokio::spawn(async move { - let result = - Self::process_single_stream(stream, dispatch, chat_request, &tx).await; + let result = Self::process_single_stream( + stream, + dispatch, + chat_request, + &tx, + router_stop_strings, + ) + .await; if let Err(e) = result { error!("Harmony streaming error: {}", e); @@ -137,6 +148,7 @@ impl HarmonyStreamingProcessor { dispatch, chat_request, &tx, + router_stop_strings, ) .await; @@ -179,6 +191,7 @@ impl HarmonyStreamingProcessor { dispatch: context::DispatchMetadata, original_request: Arc, tx: &mpsc::UnboundedSender>, + router_stop_strings: Vec, ) -> Result<(), String> { let mut prompt_tokens = HashMap::new(); let mut cached_tokens = HashMap::new(); @@ -189,6 +202,7 @@ impl HarmonyStreamingProcessor { tx, &mut prompt_tokens, &mut cached_tokens, + &router_stop_strings, ) .await } @@ -200,6 +214,7 @@ impl HarmonyStreamingProcessor { dispatch: context::DispatchMetadata, original_request: Arc, tx: &mpsc::UnboundedSender>, + router_stop_strings: Vec, ) -> Result<(), String> { // Phase 1: Process prefill stream (collect metadata) let mut prompt_tokens: HashMap = HashMap::new(); @@ -222,6 +237,7 @@ impl HarmonyStreamingProcessor { tx, &mut prompt_tokens, &mut cached_tokens, + &router_stop_strings, ) .await?; @@ -244,6 +260,7 @@ impl HarmonyStreamingProcessor { tx: &mpsc::UnboundedSender>, prompt_tokens: &mut HashMap, cached_tokens: &mut HashMap, + router_stop_strings: &[String], ) -> Result<(), String> { // Timing for metrics let start_time = Instant::now(); @@ -253,6 +270,12 @@ impl HarmonyStreamingProcessor { let mut parsers: HashMap = HashMap::new(); let mut is_firsts: HashMap = HashMap::new(); let mut matched_stops: HashMap> = HashMap::new(); + // Router-enforced string stops (direct-ZMQ): per-index, per-channel + // scanners. Once an index stops, its further deltas are swallowed and + // the engine's own Complete is not re-emitted. + let mut analysis_scanners: HashMap = HashMap::new(); + let mut final_scanners: HashMap = HashMap::new(); + let mut router_stopped: std::collections::HashSet = std::collections::HashSet::new(); let mut completion_tokens = CompletionTokenTracker::new(); // Reusable SSE encoder shared across every chunk emitted for this stream. let mut encoder = SseEncoder::new(); @@ -303,21 +326,76 @@ impl HarmonyStreamingProcessor { .map_err(|e| format!("Parse error: {e}"))?; // Emit SSE event if there's a delta - if let Some(delta) = delta_result { + if let Some(mut delta) = delta_result { + if router_stopped.contains(&index) { + continue; + } + let mut stop_matched: Option = None; + if !router_stop_strings.is_empty() { + if let Some(text) = delta.analysis_delta.take() { + let scanner = analysis_scanners.entry(index).or_insert_with(|| { + TextStopScanner::new(router_stop_strings.to_vec()) + }); + let scan = scanner.push(&text); + delta.analysis_delta = (!scan.emit.is_empty()).then_some(scan.emit); + if scan.stopped { + stop_matched = scanner.matched().map(str::to_string); + delta.final_delta = None; + delta.commentary_delta = None; + } + } + if stop_matched.is_none() { + if let Some(text) = delta.final_delta.take() { + let scanner = + final_scanners.entry(index).or_insert_with(|| { + TextStopScanner::new(router_stop_strings.to_vec()) + }); + let scan = scanner.push(&text); + delta.final_delta = + (!scan.emit.is_empty()).then_some(scan.emit); + if scan.stopped { + stop_matched = scanner.matched().map(str::to_string); + delta.commentary_delta = None; + } + } + } + } + + let has_payload = delta.analysis_delta.is_some() + || delta.final_delta.is_some() + || delta.commentary_delta.is_some(); let is_first = is_firsts.get(&index).copied().unwrap_or(false); - Self::emit_chunk_delta( - &delta, - index, - is_first, - dispatch, - original_request, - tx, - &mut encoder, - chunk_logprobs, - )?; + if has_payload || is_first { + Self::emit_chunk_delta( + &delta, + index, + is_first, + dispatch, + original_request, + tx, + &mut encoder, + chunk_logprobs, + )?; + + if is_first { + is_firsts.insert(index, false); + } + } - if is_first { - is_firsts.insert(index, false); + // A router-side stop fired: emit the final chunk now + // and swallow the rest of this index's stream (the + // engine keeps generating until its own limits). + if let Some(stop) = stop_matched { + Self::emit_final_chunk( + index, + "stop", + Some(&serde_json::Value::String(stop)), + dispatch, + original_request, + tx, + &mut encoder, + )?; + router_stopped.insert(index); } } } @@ -341,6 +419,37 @@ impl HarmonyStreamingProcessor { let final_output = parser.finalize(complete_wrapper.finish_reason().to_string()); + // A router-side stop already closed this choice. + if router_stopped.contains(&index) { + continue; + } + + // Release scanner-held text that never became a match. + let flushed = HarmonyChannelDelta { + analysis_delta: analysis_scanners + .get_mut(&index) + .map(TextStopScanner::flush) + .filter(|s| !s.is_empty()), + commentary_delta: None, + final_delta: final_scanners + .get_mut(&index) + .map(TextStopScanner::flush) + .filter(|s| !s.is_empty()), + is_final: false, + }; + if flushed.analysis_delta.is_some() || flushed.final_delta.is_some() { + Self::emit_chunk_delta( + &flushed, + index, + false, + dispatch, + original_request, + tx, + &mut encoder, + None, + )?; + } + Self::emit_final_chunk( index, &final_output.finish_reason, diff --git a/model_gateway/src/routers/grpc/mod.rs b/model_gateway/src/routers/grpc/mod.rs index 6014c7b2f..6f374ba86 100644 --- a/model_gateway/src/routers/grpc/mod.rs +++ b/model_gateway/src/routers/grpc/mod.rs @@ -18,6 +18,7 @@ pub(crate) mod regular; pub(crate) mod router; // Used by routers/factory pub mod utils; // Used by routers/http and bindings/golang pub mod zmq_client; // ZMQ backend adapter behind the vLLM client surface +pub(crate) mod zmq_multimodal; // Proto mm inputs → EngineCore mm_features // Re-export for convenience pub use proto_wrapper::{MultimodalData, TensorBytes}; diff --git a/model_gateway/src/routers/grpc/multimodal/assemble.rs b/model_gateway/src/routers/grpc/multimodal/assemble.rs index ff839eec7..d08abcbb8 100644 --- a/model_gateway/src/routers/grpc/multimodal/assemble.rs +++ b/model_gateway/src/routers/grpc/multimodal/assemble.rs @@ -28,16 +28,19 @@ use super::{ transport::{mm_encoder_input_dtype, resolve_mm_shm_enabled, resolve_mm_shm_min_bytes}, MediaBatch, MultimodalIntermediate, PrecomputedMultimodalIntermediate, PromptBinding, }; -use crate::routers::grpc::{ - backend_client::BackendClient, - client::GrpcClient, - context::WorkerSelection, - proto_wrapper::{ - cleanup_tokenspeed_items_encoder_shm, SglangMultimodalData, TensorBytes, - TokenSpeedModality, TokenSpeedMultimodalData, TokenSpeedMultimodalItem, TokenSpeedTensor, - TrtllmMultimodalData, VllmMultimodalData, +use crate::{ + routers::grpc::{ + backend_client::BackendClient, + client::GrpcClient, + context::WorkerSelection, + proto_wrapper::{ + cleanup_tokenspeed_items_encoder_shm, SglangMultimodalData, TensorBytes, + TokenSpeedModality, TokenSpeedMultimodalData, TokenSpeedMultimodalItem, + TokenSpeedTensor, TrtllmMultimodalData, VllmMultimodalData, + }, + MultimodalData, }, - MultimodalData, + worker::RuntimeType, }; /// Assemble backend-specific multimodal data from the intermediate. @@ -108,7 +111,20 @@ async fn assemble_multimodal_data_impl( BackendClient::Grpc(GrpcClient::Mlx(_)) => { anyhow::bail!("MLX does not support multimodal inputs") } - BackendClient::Zmq(_) => anyhow::bail!("ZMQ backend does not support multimodal inputs"), + BackendClient::Zmq(client) => match client.runtime() { + RuntimeType::Vllm | RuntimeType::Unspecified => { + let batch = into_single_batch(intermediate, "vLLM")?; + let mut data = assemble_vllm(batch, workers)?; + // The ZMQ translate reads tensor bytes inline; this wire has no + // /dev/shm or RDMA pull on the engine side. + data.shm_enabled = false; + data.rdma_enabled = false; + Ok(MultimodalData::Vllm(data)) + } + runtime => anyhow::bail!( + "multimodal inputs are not supported over the {runtime} ZMQ backend yet" + ), + }, } } diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index 6d1d8d9e0..2275f3eb9 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -78,14 +78,19 @@ impl ResponseProcessor { // Accumulate text with early breaks let mut final_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => final_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { final_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -177,8 +182,14 @@ impl ResponseProcessor { } } - // Step 3: Use finish reason directly from proto (already OpenAI-compatible string) - let finish_reason_str = complete.finish_reason(); + // Step 3: Determine finish reason. A local stop-decoder match takes + // precedence over the engine's reason (which is "length" when stop + // strings are enforced gateway-side rather than by the backend). + let finish_reason_str = if stopped { + "stop" + } else { + complete.finish_reason() + }; // Override finish reason if we have tool calls let final_finish_reason_str = if tool_calls.is_some() { @@ -187,7 +198,12 @@ impl ResponseProcessor { finish_reason_str }; - let matched_stop = complete.matched_stop_json(); + // When the local decoder matched a stop string, surface it (the engine + // reports no stop_reason over the ZMQ path); otherwise use the engine's. + let matched_stop = stop_decoder + .matched_stop() + .map(|s| serde_json::Value::String(s.to_string())) + .or_else(|| complete.matched_stop_json()); // Step 4: Convert output logprobs if present let logprobs = complete.output_logprobs().map(|ref proto_logprobs| { @@ -579,14 +595,19 @@ impl ResponseProcessor { })?; let mut final_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => final_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { final_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -726,10 +747,20 @@ impl ResponseProcessor { } } - // Step 4: Determine stop_reason and stop_sequence (derived from same conditions) - let finish_reason_str = complete.finish_reason(); - let matched_stop = complete.matched_stop_json(); - let stop_sequence = matched_stop.and_then(|v| v.as_str().map(String::from)); + // Step 4: Determine stop_reason and stop_sequence (derived from same conditions). + // A local stop-decoder match takes precedence over the engine's reason + // (the backend has no stop-string detection over ZMQ), surfacing the + // matched sequence for a StopSequence result. + let finish_reason_str = if stopped { + "stop" + } else { + complete.finish_reason() + }; + let stop_sequence = stop_decoder.matched_stop().map(String::from).or_else(|| { + complete + .matched_stop_json() + .and_then(|v| v.as_str().map(String::from)) + }); let stop_reason = if tool_calls.is_some() || finish_reason_str == "tool_calls" { Some(messages::StopReason::ToolUse) @@ -832,14 +863,19 @@ impl ResponseProcessor { }; let mut decoded_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => decoded_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { decoded_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -851,7 +887,12 @@ impl ResponseProcessor { prompt_tokens = prompt_tokens.max(complete.prompt_tokens()); total_completion += complete.completion_tokens(); - let finish_reason = { + // A local stop-decoder match takes precedence over the engine's + // reason (which is "length" when stop strings are enforced + // gateway-side rather than by the backend). + let finish_reason = if stopped { + Some("stop".to_string()) + } else { let reason = complete.finish_reason(); if reason.is_empty() { None @@ -868,7 +909,13 @@ impl ResponseProcessor { } }; - let matched_stop = complete.matched_stop_json(); + // When the local decoder matched a stop string, surface it (the + // engine reports no stop_reason over the ZMQ path); otherwise use + // the engine's. + let matched_stop = stop_decoder + .matched_stop() + .map(|s| serde_json::Value::String(s.to_string())) + .or_else(|| complete.matched_stop_json()); let suffix_len = completion_req.suffix.as_ref().map_or(0, |s| s.len()); let echo_len = if completion_req.echo { diff --git a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs index 4967044e9..3dd71e948 100644 --- a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs @@ -147,6 +147,13 @@ impl PipelineStage for ChatRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs index 6477ccecc..8bd9b63da 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs @@ -141,6 +141,12 @@ impl PipelineStage for CompletionRequestBuildingStage { let disaggregated = matches!(clients, ClientSelection::Disaggregated { .. }); let request_type = &ctx.input.request_type; let workers = ctx.state.workers.as_ref(); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + let tokenizer = ctx.tokenizer_arc(); let plan = match items.as_slice() { [] => { @@ -149,9 +155,8 @@ impl PipelineStage for CompletionRequestBuildingStage { "No prompts prepared", )) } - [item] => ExecutionPlan::generate( - self.plan_kind, - self.build_proto_request( + [item] => { + let mut proto_request = self.build_proto_request( builder_client, helpers::resolve_request_id( request_type, @@ -163,8 +168,10 @@ impl PipelineStage for CompletionRequestBuildingStage { &completion_request, request_type, workers, - )?, - ), + )?; + helpers::resolve_string_stops(&mut proto_request, tokenizer.as_ref(), is_zmq); + ExecutionPlan::generate(self.plan_kind, proto_request) + } batch_items => { // The shared id (client rid or middleware request id) stays // clean for the response; per-sub engine ids get a uniqueness @@ -186,14 +193,16 @@ impl PipelineStage for CompletionRequestBuildingStage { } else { format!("{shared_request_id}-p{i}") }; - requests.push(self.build_proto_request( + let mut proto_request = self.build_proto_request( builder_client, sub_request_id, item, &completion_request, request_type, workers, - )?); + )?; + helpers::resolve_string_stops(&mut proto_request, tokenizer.as_ref(), is_zmq); + requests.push(proto_request); } ExecutionPlan::Batch { kind: self.plan_kind, diff --git a/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs index 3be932e43..2f7634632 100644 --- a/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs @@ -86,6 +86,13 @@ impl PipelineStage for GenerateRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs index 23d91b89a..aa73ff82a 100644 --- a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs @@ -148,6 +148,13 @@ impl PipelineStage for MessageRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index 505a8dc3b..b9d888f0c 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -271,6 +271,9 @@ impl StreamingProcessor { let mut stream_buffers: HashMap = HashMap::new(); let mut finish_reasons: HashMap = HashMap::new(); let mut matched_stops: HashMap> = HashMap::new(); + // Indices whose local stop decoder fired: their finish reason is pinned + // to "stop" and later engine output for the index is ignored. + let mut stopped_indices: HashSet = HashSet::new(); let mut prompt_tokens: HashMap = HashMap::new(); let mut completion_tokens = CompletionTokenTracker::new(); let mut cached_tokens: HashMap = HashMap::new(); @@ -384,6 +387,12 @@ impl StreamingProcessor { let index = chunk.index(); + // Once the local stop decoder has fired for an index, ignore + // any further engine output the backend emits for it. + if stopped_indices.contains(&index) { + continue; + } + completion_tokens.record_chunk(&chunk); // Get or create stop decoder for this index @@ -406,9 +415,26 @@ impl StreamingProcessor { }); // Process tokens through stop decoder - let (chunk_text, _should_stop) = + let (chunk_text, should_stop) = Self::process_chunk_tokens(stop_decoder, chunk.token_ids()); + if should_stop { + // Stop-decoder match takes precedence: pin "stop" even if + // the backend's eventual Complete carries "length" (the + // local stop sequence fired first). Any pre-stop text in + // `chunk_text` is still emitted below before the finish + // reason is flushed in Phase 4. + finish_reasons + .entry(index) + .or_insert_with(|| "stop".to_string()); + matched_stops.entry(index).or_insert_with(|| { + stop_decoder + .matched_stop() + .map(|s| Value::String(s.to_string())) + }); + stopped_indices.insert(index); + } + if chunk_text.is_empty() { continue; } @@ -568,9 +594,13 @@ impl StreamingProcessor { cached_tokens.insert(index, complete.cached_tokens()); reasoning_tokens.insert(index, complete.reasoning_tokens()); - finish_reasons.insert(index, complete.finish_reason().to_string()); - matched_stops.insert(index, complete.matched_stop_json()); + // A local stop-decoder match already pinned "stop" for this + // index; don't let the engine's finish reason overwrite it. + if !stopped_indices.contains(&index) { + finish_reasons.insert(index, complete.finish_reason().to_string()); + matched_stops.insert(index, complete.matched_stop_json()); + } // Don't break - continue reading all Complete messages for n>1 } @@ -1740,6 +1770,9 @@ impl StreamingProcessor { let mut prompt_tokens: u32 = 0; let mut finish_reason_str = String::new(); let mut matched_stop: Option = None; + // Set once the local stop decoder fires: pins "stop" and ignores later + // engine output (the backend has no stop-string detection over ZMQ). + let mut stopped = false; // Check parser availability once upfront. Run parser when the user explicitly // enabled thinking, or when the selected parser needs structural special tokens. @@ -1867,11 +1900,29 @@ impl StreamingProcessor { first_token_time = Some(Instant::now()); } + // Once the local stop decoder has fired, ignore further + // engine output for this (single-choice) request. + if stopped { + continue; + } + completion_tokens.record_chunk(&chunk); - let (chunk_text, _should_stop) = + let (chunk_text, should_stop) = Self::process_chunk_tokens(&mut stop_decoder, chunk.token_ids()); + if should_stop { + // Stop-decoder match takes precedence over the engine's + // eventual finish reason (the local stop sequence fired + // first). Pre-stop text in `chunk_text` is still emitted + // below; Phase 4 derives StopSequence from `matched_stop`. + stopped = true; + finish_reason_str = "stop".to_string(); + matched_stop = stop_decoder + .matched_stop() + .map(|s| Value::String(s.to_string())); + } + if chunk_text.is_empty() { continue; } @@ -2150,8 +2201,12 @@ impl StreamingProcessor { prompt_tokens = complete.prompt_tokens(); completion_tokens.record_complete(&complete); - finish_reason_str = complete.finish_reason().to_string(); - matched_stop = complete.matched_stop_json(); + // A local stop-decoder match already pinned "stop"; don't let + // the engine's finish reason overwrite it. + if !stopped { + finish_reason_str = complete.finish_reason().to_string(); + matched_stop = complete.matched_stop_json(); + } } ProtoResponseVariant::None => continue, } diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index e3fc0b692..269c36d00 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -13,12 +13,14 @@ // the request-execution stage is reused unchanged. use std::{ - collections::HashMap, + collections::{BTreeSet, HashMap}, + path::Path, sync::Arc, time::{Duration, SystemTime, UNIX_EPOCH}, }; use engine_zmq_client::{ + codec::dtype::ModelDtype, connect_handshake, connector::{EngineCoreClient, EngineCoreStream, TokenSpeedClient, TokenSpeedStream}, protocol::{ @@ -30,6 +32,7 @@ use engine_zmq_client::{ output::{EngineCoreFinishReason, EngineCoreOutput, StopReason}, request::EngineCoreRequest, sampling::EngineCoreSamplingParams, + structured_outputs::StructuredOutputsParams, }, EngineLoad, }, @@ -37,9 +40,16 @@ use engine_zmq_client::{ }; use futures::{stream::SelectAll, Stream, StreamExt}; use openai_protocol::worker::{SchedulerLoadSnapshot, WorkerLoadResponse}; -use smg_grpc_client::vllm_proto as vllm; +use smg_grpc_client::{tokenspeed_proto, vllm_proto as vllm}; -use crate::worker::RuntimeType; +use crate::{ + routers::grpc::{ + client::{ModelInfo, ServerInfo}, + proto_wrapper::ProtoGenerateRequest, + zmq_multimodal, + }, + worker::RuntimeType, +}; /// Loopback host for the same-host ZMQ transport (TCP handshake and local /// binds). Shared with the worker-side socket derivation. @@ -56,6 +66,58 @@ enum ZmqBackend { TokenSpeed(Arc), } +/// The model's EOS stop set, resolved from its local directory. EngineCore +/// has no tokenizer or model config — stopping at EOS is the frontend's job +/// (the ids ride each request), and without them generation only ends at +/// `max_tokens`. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct EosTokenIds { + /// Primary EOS id, carried as the request's `_eos_token_id`. + primary: Option, + /// Extra EOS ids (multi-EOS models), merged into `stop_token_ids`. + extra: Vec, +} + +impl EosTokenIds { + pub fn new(primary: Option, extra: Vec) -> Self { + Self { primary, extra } + } + + /// Resolve from `config.json` + `generation_config.json` in a local model + /// directory: primary = the model config's first id, extras = every other + /// listed id. Missing files or fields degrade to fewer ids. + pub fn from_model_dir(dir: &Path) -> Self { + let model_ids = eos_ids_from_file(&dir.join("config.json")); + let gen_ids = eos_ids_from_file(&dir.join("generation_config.json")); + let primary = (model_ids.first().or_else(|| gen_ids.first())).copied(); + let mut extra = Vec::new(); + for id in model_ids.into_iter().chain(gen_ids) { + if Some(id) != primary && !extra.contains(&id) { + extra.push(id); + } + } + Self { primary, extra } + } +} + +/// Read a config file's `eos_token_id`, which is a single id or a list. +fn eos_ids_from_file(path: &Path) -> Vec { + std::fs::read_to_string(path) + .ok() + .and_then(|text| serde_json::from_str::(&text).ok()) + .map(|config| eos_ids_from_value(config.get("eos_token_id"))) + .unwrap_or_default() +} + +fn eos_ids_from_value(value: Option<&serde_json::Value>) -> Vec { + let as_id = |v: &serde_json::Value| v.as_u64().and_then(|id| u32::try_from(id).ok()); + match value { + Some(serde_json::Value::Array(ids)) => ids.iter().filter_map(as_id).collect(), + Some(id) => as_id(id).into_iter().collect(), + None => Vec::new(), + } +} + /// Direct ZMQ connection to a same-host engine (vLLM EngineCore or TokenSpeed), /// presented behind the vLLM gRPC client surface. #[derive(Clone)] @@ -64,6 +126,9 @@ pub struct ZmqEngineClient { /// Model id advertised for metadata (the engine does not report it on the /// wire; it is configured at worker registration). model_id: String, + /// EOS ids attached to every vLLM request (the engine can't stop at EOS + /// without them). + eos: EosTokenIds, } impl ZmqEngineClient { @@ -74,12 +139,17 @@ impl ZmqEngineClient { /// engines connect to (chosen by SMG). `engine_count` is the number of DP /// ranks to await. `runtime` selects the wire protocol spoken over the shared /// transport (vLLM EngineCore vs TokenSpeed). + #[expect( + clippy::too_many_arguments, + reason = "transport constructor: endpoints, engine count, and runtime are all irreducible connection inputs" + )] pub async fn connect( handshake_address: &str, input_address: &str, output_address: &str, engine_count: usize, model_id: String, + eos: EosTokenIds, runtime: RuntimeType, timeout: Duration, ) -> Result> { @@ -126,7 +196,11 @@ impl ZmqEngineClient { // runtimes were rejected before the handshake. _ => ZmqBackend::Vllm(Arc::new(EngineCoreClient::new(transport))), }; - Ok(Self { backend, model_id }) + Ok(Self { + backend, + model_id, + eos, + }) } /// The engine runtime behind this connection (the wire protocol chosen at @@ -148,7 +222,10 @@ impl ZmqEngineClient { } /// Submit a generate request and return a stream of vLLM-proto responses. - /// The request is translated into the backend's wire protocol. + /// The request is the engine's native proto (vLLM for a vLLM backend, + /// TokenSpeed for a TokenSpeed backend — the [`BackendClient`] builders emit + /// the matching variant per runtime); it is translated into the backend's + /// wire protocol here. /// /// Over gRPC the engine-side frontend (e.g. vLLM's AsyncLLM) fans `n` out /// itself and multiplexes the choices onto one stream. The raw ZMQ wire has @@ -156,27 +233,56 @@ impl ZmqEngineClient { /// single-sample engine requests (see [`fan_out_requests`]); their outputs /// are merged back into one stream with each sub tagged via the proto /// `index` field, exactly like the gRPC contract. + /// + /// [`BackendClient`]: crate::routers::grpc::backend_client::BackendClient pub async fn generate( &self, - req: vllm::GenerateRequest, + req: ProtoGenerateRequest, ) -> Result { - let subs = fan_out_requests(req); // Sub-streams submitted before a mid-loop failure are dropped with the // error, which auto-aborts their engine-side requests. match &self.backend { ZmqBackend::Vllm(client) => { + let ProtoGenerateRequest::Vllm(req) = req else { + return Err(tonic::Status::internal( + "vLLM ZMQ backend expects a vLLM generate request", + )); + }; + // EngineCore needs a concrete `max_tokens`; vLLM's OpenAI frontend + // (which the ZMQ path bypasses) defaults an unset value to + // `max_model_len - prompt_len`. The context length comes from the + // engine's ready handshake, so a connected engine is required. + let (max_model_len, model_dtype) = client + .engines() + .first() + .map(|e| (e.ready_response.max_model_len, e.ready_response.dtype)) + .ok_or_else(|| tonic::Status::unavailable("no connected ZMQ engine"))?; let mut streams = SelectAll::new(); - for (index, sub) in subs.into_iter().enumerate() { - let request = - translate_request(sub).map_err(tonic::Status::invalid_argument)?; + for (index, sub) in fan_out_requests(*req).into_iter().enumerate() { + // The engine returns the sampled logprob plus up to + // `logprobs` ranked candidates per position; carry the + // requested count so the stream can shape `top_logprobs`. + let top_logprobs = sub + .sampling_params + .as_ref() + .and_then(|sp| sp.logprobs) + .filter(|&n| n > 0) + .map_or(0, |n| n as usize); + let request = translate_request(sub, max_model_len, model_dtype, &self.eos) + .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; - streams.push(VllmGenerateStream::new(stream, index as u32)); + streams.push(VllmGenerateStream::new(stream, index as u32, top_logprobs)); } Ok(ZmqGenerateStream::Vllm(streams)) } ZmqBackend::TokenSpeed(client) => { + let ProtoGenerateRequest::TokenSpeed(req) = req else { + return Err(tonic::Status::internal( + "TokenSpeed ZMQ backend expects a TokenSpeed generate request", + )); + }; let mut streams = SelectAll::new(); - for (index, sub) in subs.into_iter().enumerate() { + for (index, sub) in fan_out_tokenspeed_requests(*req).into_iter().enumerate() { let request = translate_request_tokenspeed(sub) .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; @@ -246,38 +352,55 @@ impl ZmqEngineClient { /// Model info derived from the handshake `EngineCoreReadyResponse` plus the /// configured model id (the engine does not report tokenizer/vocab metadata, - /// so those come from worker config). - pub fn get_model_info(&self) -> vllm::GetModelInfoResponse { + /// so those come from worker config). Returned as the runtime's native + /// metadata variant so the label mapping matches the gRPC path. + pub fn get_model_info(&self) -> ModelInfo { let max_context_length = self .engines() .first() .map(|e| e.ready_response.max_model_len) .unwrap_or(0); - vllm::GetModelInfoResponse { - model_path: self.model_id.clone(), - served_model_name: self.model_id.clone(), - tokenizer_path: self.model_id.clone(), - is_generation: true, - max_context_length: u32::try_from(max_context_length).unwrap_or(u32::MAX), - ..Default::default() + match &self.backend { + ZmqBackend::Vllm(_) => ModelInfo::Vllm(vllm::GetModelInfoResponse { + model_path: self.model_id.clone(), + served_model_name: self.model_id.clone(), + tokenizer_path: self.model_id.clone(), + is_generation: true, + max_context_length: u32::try_from(max_context_length).unwrap_or(u32::MAX), + ..Default::default() + }), + ZmqBackend::TokenSpeed(_) => { + ModelInfo::TokenSpeed(Box::new(tokenspeed_proto::GetModelInfoResponse { + model_path: self.model_id.clone(), + served_model_name: self.model_id.clone(), + tokenizer_path: self.model_id.clone(), + max_context_length: i32::try_from(max_context_length).unwrap_or(i32::MAX), + ..Default::default() + })) + } } } - /// Server info derived from the handshake response. - pub fn get_server_info(&self) -> vllm::GetServerInfoResponse { + /// Server info derived from the handshake response, as the runtime's native + /// metadata variant. + pub fn get_server_info(&self) -> ServerInfo { let data_parallel_size = self .engines() .first() .map(|e| e.ready_response.data_parallel_size) .unwrap_or(1); - let server_type = match &self.backend { - ZmqBackend::Vllm(_) => "vllm", - ZmqBackend::TokenSpeed(_) => "tokenspeed", - }; - vllm::GetServerInfoResponse { - data_parallel_size: i32::try_from(data_parallel_size).unwrap_or(i32::MAX), - server_type: server_type.to_string(), - ..Default::default() + match &self.backend { + ZmqBackend::Vllm(_) => ServerInfo::Vllm(vllm::GetServerInfoResponse { + data_parallel_size: i32::try_from(data_parallel_size).unwrap_or(i32::MAX), + server_type: "vllm".to_string(), + ..Default::default() + }), + // TokenSpeed's server-info proto carries no data-parallel size or + // server-type field; the ZMQ handshake supplies no `server_args` + // either, so only the fields it does expose are surfaced. + ZmqBackend::TokenSpeed(_) => { + ServerInfo::TokenSpeed(Box::::default()) + } } } } @@ -327,6 +450,9 @@ struct StreamState { /// `Complete`, so accumulate here and drain into `Complete`. output_logprobs_val: Vec, output_logprobs_idx: Vec, + /// Cumulative per-position ranked candidates (`top_logprobs`), accumulated + /// alongside the sampled logprobs and drained into the terminal `Complete`. + output_top_logprobs: Vec, } impl StreamState { @@ -336,7 +462,7 @@ impl StreamState { (!self.output_logprobs_val.is_empty()).then(|| vllm::OutputLogProbs { token_logprobs: std::mem::take(&mut self.output_logprobs_val), token_ids: std::mem::take(&mut self.output_logprobs_idx), - ..Default::default() + top_logprobs: std::mem::take(&mut self.output_top_logprobs), }) } } @@ -350,6 +476,10 @@ pub struct VllmGenerateStream { /// Choice index stamped on every chunk/complete (0 for n=1; the fan-out /// position for n>1) — the proto field the pipeline demuxes choices by. index: u32, + /// Number of ranked candidates the client requested per position; `0` when + /// only the sampled logprob (or nothing) was asked for, in which case no + /// `top_logprobs` are emitted. + top_logprobs: usize, /// Terminal `Complete` held back when the finish tick also carried new /// tokens: streaming frontends decode text/logprobs from chunks only, so /// the tick's delta goes out as a `Chunk` first. @@ -357,11 +487,12 @@ pub struct VllmGenerateStream { } impl VllmGenerateStream { - fn new(inner: EngineCoreStream, index: u32) -> Self { + fn new(inner: EngineCoreStream, index: u32, top_logprobs: usize) -> Self { Self { inner, state: StreamState::default(), index, + top_logprobs, pending: None, } } @@ -370,6 +501,7 @@ impl VllmGenerateStream { &mut self, output: EngineCoreOutput, ) -> Result { + let top_k = self.top_logprobs; let state = &mut self.state; if let Some(stats) = &output.prefill_stats { state.prompt_tokens = stats.num_prompt_tokens; @@ -379,11 +511,13 @@ impl VllmGenerateStream { state.completion_tokens += token_ids.len() as u32; state.output_ids.extend(token_ids.iter().copied()); - // Sampled-token logprobs (entry 0 per position), if requested. Chunks - // carry this tick's increment; the terminal `Complete` carries the - // cumulative set, so accumulate into `state` and drain it on finish. + // Sampled-token logprobs (entry 0 per position) plus the requested + // ranked candidates (`top_logprobs`). Chunks carry this tick's + // increment; the terminal `Complete` carries the cumulative set, so + // accumulate into `state` and drain it on finish. let mut tick_logprobs_val = Vec::new(); let mut tick_logprobs_idx = Vec::new(); + let mut tick_top_logprobs = Vec::new(); if let Some(logprobs) = &output.new_logprobs { let decoded = logprobs.as_direct().ok_or_else(|| { // The protocol layer resolves wire logprobs during decode, so @@ -391,21 +525,41 @@ impl VllmGenerateStream { tonic::Status::internal("unresolved wire logprobs in engine output") })?; for position in &decoded.positions { - if let Some(sampled) = position.entries.first() { - tick_logprobs_val.push(sampled.logprob); - tick_logprobs_idx.push(sampled.token_id); + let Some(sampled) = position.entries.first() else { + continue; + }; + tick_logprobs_val.push(sampled.logprob); + tick_logprobs_idx.push(sampled.token_id); + // The entries arrive sampled-first then rank-ordered; take the + // requested count so one ranked list lands per sampled token. + if top_k > 0 { + let mut top = vllm::TopLogProbs::default(); + for entry in position.entries.iter().take(top_k) { + top.values.push(entry.logprob); + top.token_ids.push(entry.token_id); + } + tick_top_logprobs.push(top); } } } let chunk_logprobs = (!tick_logprobs_val.is_empty()).then(|| vllm::OutputLogProbs { token_logprobs: tick_logprobs_val.clone(), token_ids: tick_logprobs_idx.clone(), - ..Default::default() + top_logprobs: tick_top_logprobs.clone(), }); state.output_logprobs_val.extend(tick_logprobs_val); state.output_logprobs_idx.extend(tick_logprobs_idx); + state.output_top_logprobs.extend(tick_top_logprobs); let response = match output.finish_reason { + // An engine-side request failure (e.g. grammar compilation) must + // surface as an error, not as a normal completion with empty + // output — that would produce a 200 with no content. + Some(EngineCoreFinishReason::Error) => { + return Err(tonic::Status::internal( + "engine finished the request with an error (see engine logs)", + )); + } Some(reason) => { let complete = vllm::GenerateResponse { response: Some(vllm::generate_response::Response::Complete( @@ -639,81 +793,63 @@ fn fan_out_requests(req: vllm::GenerateRequest) -> Vec { .collect() } -/// Translate a vLLM-proto generate request into a TokenSpeed +/// Split an `n > 1` TokenSpeed proto request into `n` single-sample +/// sub-requests, the TokenSpeed analogue of [`fan_out_requests`] (the wire has +/// no per-sample demux, so `generate` fans out here). An `n <= 1` request passes +/// through untouched. The TokenSpeed proto carries no seed, so the samples +/// differ by the engine's per-rid seeding alone — the suffixed request ids are +/// unique, so each rid seeds independently. +fn fan_out_tokenspeed_requests( + req: tokenspeed_proto::GenerateRequest, +) -> Vec { + let n = req.sampling_params.as_ref().map_or(1, |sp| sp.n.max(1)); + if n <= 1 { + return vec![req]; + } + (0..n) + .map(|i| { + let mut sub = req.clone(); + sub.request_id = format!("{}-{i}", req.request_id); + if let Some(sp) = sub.sampling_params.as_mut() { + sp.n = 1; + } + sub + }) + .collect() +} + +/// Translate a TokenSpeed proto `GenerateRequest` into the wire /// `TokenizedGenerateReqInput`. ZMQ mode requires pre-tokenized input (SMG /// tokenizes upstream). fn translate_request_tokenspeed( - req: vllm::GenerateRequest, + req: tokenspeed_proto::GenerateRequest, ) -> Result { - let input_ids = match req.input { - Some(vllm::generate_request::Input::Tokenized(tokenized)) => tokenized.input_ids, - Some(vllm::generate_request::Input::Text(_)) => { - return Err("ZMQ mode requires pre-tokenized input (TokenizedInput)".to_string()); - } + // The TokenSpeed ZMQ wire has no multimodal slot yet; reject loudly rather + // than silently dropping pixels (assembly also refuses upstream). + if req.mm_inputs.is_some() { + return Err( + "multimodal inputs are not supported over the TokenSpeed ZMQ backend".to_string(), + ); + } + let input_ids = match req.tokenized { + Some(tokenized) => tokenized.input_ids, None => { return Err("ZMQ mode requires pre-tokenized input; no input provided".to_string()); } }; - let stream = req.stream; - // Single-engine TokenSpeed: a pinned DP rank other than 0 cannot be honored. - if req.data_parallel_rank.is_some_and(|rank| rank != 0) { - return Err(format!( - "invalid data_parallel_rank {:?}: the TokenSpeed ZMQ backend is single-engine", - req.data_parallel_rank - )); + // Over the ZMQ wire TokenSpeed returns only the single sampled-token logprob + // per token: no top-k candidates (`top_logprobs_num > 1`) and no prompt + // logprobs (`token_ids_logprob`). Reject both rather than silently return + // fewer than asked. A bare `logprobs: true` (count 0/1) is the plain + // sampled-token logprob and is wired end-to-end via `return_logprob`. + if req.top_logprobs_num > 1 { + return Err("top_logprobs are not supported over the TokenSpeed ZMQ backend".to_string()); } - // TokenSpeed returns only the single sampled-token logprob per token - // (`top_logprobs_num = 0` on its wire) and no prompt logprobs. The vLLM - // sampling `logprobs` field is a count: the chat frontend maps a bare - // `logprobs: true` to `1` and `top_logprobs = k` to `k`, so a count above 1 - // (or `-1` = "all") is a top-k request that cannot be honored — reject it - // rather than silently return fewer. Counts of 0 or 1 are the plain - // sampled-token logprob and are wired end-to-end. Note the flip side: at - // this proto boundary a chat `top_logprobs: 1` is indistinguishable from a - // bare `logprobs: true` (both arrive as count 1), so it is accepted and its - // `top_logprobs` list simply stays empty. - if let Some(sp) = req.sampling_params.as_ref() { - if sp.logprobs.is_some_and(|n| !(0..=1).contains(&n)) { - return Err( - "top_logprobs are not supported over the TokenSpeed ZMQ backend".to_string(), - ); - } - if sp.prompt_logprobs.is_some() { - return Err( - "prompt logprobs are not supported over the TokenSpeed ZMQ backend".to_string(), - ); - } - // The response_format / forced-tool-choice constraint oneof is not - // translated onto the TokenSpeed structured-output fields yet; dropping - // it would return unconstrained text. - if sp.constraint.is_some() { - return Err( - "structured output constraints are not supported over the ZMQ backend yet" - .to_string(), - ); - } - // The TokenSpeed wire has no per-sample demux; n>1 is fanned out into - // single-sample sub-requests by `generate` before translation. - // Stop strings are not forwarded: the direct ZMQ path sends token ids - // only (the engine-side transport normalizes without a tokenizer) and - // the gateway's stop decoder does not enforce them, so they would be - // ignored and the request would run to max_tokens. - if !sp.stop.is_empty() { - return Err( - "stop strings are not supported over the TokenSpeed ZMQ backend yet; \ - use stop_token_ids" - .to_string(), - ); - } - // logit_bias is not translated onto the TokenSpeed wire either. - if !sp.logit_bias.is_empty() { - return Err("logit_bias is not supported over the TokenSpeed ZMQ backend".to_string()); - } + if !req.token_ids_logprob.is_empty() { + return Err( + "prompt logprobs are not supported over the TokenSpeed ZMQ backend".to_string(), + ); } - let return_logprob = req - .sampling_params - .as_ref() - .is_some_and(|sp| sp.logprobs.is_some()); Ok(TokenizedGenerateReqInput { rid: req.request_id, input_ids, @@ -725,53 +861,76 @@ fn translate_request_tokenspeed( params.normalize(); params }), - return_logprob, - stream, + return_logprob: req.return_logprob, + stream: req.stream, // Every other field keeps its neutral default (the fields after // `stream` are not even emitted; the engine fills them from defaults). ..TokenizedGenerateReqInput::default() }) } -/// Map vLLM-proto sampling params onto TokenSpeed's native `SamplingParams`, -/// in the normalized form: the engine skips its decode-time re-derivation once +/// Map TokenSpeed proto sampling params onto the wire `SamplingParams`, in the +/// normalized form: the engine skips its decode-time re-derivation once /// `is_normalized` is set, so [`TokenSpeedSamplingParams::normalize`] resolves /// the derived fields (top_k sentinel, greedy collapse) before encoding. -fn translate_sampling_tokenspeed(sp: vllm::SamplingParams) -> TokenSpeedSamplingParams { +/// +/// String `stop` sequences are not forwarded — the token-only engine cannot +/// match them; the router-side stop decoder trims them from the text instead. +fn translate_sampling_tokenspeed( + sp: tokenspeed_proto::SamplingParams, +) -> TokenSpeedSamplingParams { let mut params = TokenSpeedSamplingParams { - max_new_tokens: sp.max_tokens, + max_new_tokens: sp.max_new_tokens, stop_token_ids: (!sp.stop_token_ids.is_empty()).then_some(sp.stop_token_ids), temperature: f64::from(sp.temperature.unwrap_or(1.0)), - top_p: f64::from(sp.top_p), - // vLLM proto uses `0` for "all tokens"; TokenSpeed's API form is `-1` - // (`normalize` resolves it to the engine's disabled sentinel). - top_k: if sp.top_k == 0 { - -1 - } else { - i32::try_from(sp.top_k).unwrap_or(-1) - }, - min_p: f64::from(sp.min_p), - frequency_penalty: f64::from(sp.frequency_penalty), - presence_penalty: f64::from(sp.presence_penalty), - repetition_penalty: f64::from(sp.repetition_penalty), - min_new_tokens: sp.min_tokens, + top_p: f64::from(sp.top_p.unwrap_or(1.0)), + // The proto keeps the API convention `-1` = "all tokens" (and unset); + // `normalize` resolves it to the engine's disabled sentinel. + top_k: sp.top_k.unwrap_or(-1), + min_p: f64::from(sp.min_p.unwrap_or(0.0)), + frequency_penalty: f64::from(sp.frequency_penalty.unwrap_or(0.0)), + presence_penalty: f64::from(sp.presence_penalty.unwrap_or(0.0)), + repetition_penalty: f64::from(sp.repetition_penalty.unwrap_or(1.0)), + min_new_tokens: sp.min_new_tokens, ignore_eos: sp.ignore_eos, - // A negative seed is a "no seed" sentinel; drop it rather than wrap - // (the engine then derives a per-rid seed). - seed: sp.seed.and_then(|seed| u64::try_from(seed).ok()), + skip_special_tokens: sp.skip_special_tokens, + spaces_between_special_tokens: sp.spaces_between_special_tokens, + no_stop_trim: sp.no_stop_trim, // Proto `0` means unspecified; TokenSpeed expects at least one sample. - // The engine stores n without acting on it — n>1 is fanned out before - // translation, so this is always 1 on the wire. + // n>1 is fanned out before translation, so this is always 1 on the wire. n: sp.n.max(1), ..TokenSpeedSamplingParams::default() }; + apply_tokenspeed_constraint(&mut params, sp.constraint); params.normalize(); params } +/// Map the proto structured-output `constraint` oneof onto the wire's dedicated +/// fields. The oneof is single-valued, so at most one field is set; the rest +/// stay `None`. +fn apply_tokenspeed_constraint( + params: &mut TokenSpeedSamplingParams, + constraint: Option, +) { + use tokenspeed_proto::sampling_params::Constraint; + match constraint { + Some(Constraint::JsonSchema(schema)) => params.json_schema = Some(schema), + Some(Constraint::Regex(regex)) => params.regex = Some(regex), + Some(Constraint::EbnfGrammar(grammar)) => params.ebnf = Some(grammar), + Some(Constraint::StructuralTag(tag)) => params.structural_tag = Some(tag), + None => {} + } +} + /// Translate a vLLM-proto generate request into an `EngineCoreRequest`. ZMQ mode /// requires pre-tokenized input (SMG tokenizes upstream). -fn translate_request(req: vllm::GenerateRequest) -> Result { +fn translate_request( + req: vllm::GenerateRequest, + max_model_len: u64, + model_dtype: ModelDtype, + eos: &EosTokenIds, +) -> Result { let prompt_token_ids = match req.input { Some(vllm::generate_request::Input::Tokenized(tokenized)) => Some(tokenized.input_ids), Some(vllm::generate_request::Input::Text(_)) => { @@ -781,20 +940,24 @@ fn translate_request(req: vllm::GenerateRequest) -> Result1 is fanned out into // single-sample sub-requests by `generate` before translation. // The ZMQ renderer path has no prompt-logprob merge, so the engine's @@ -803,17 +966,44 @@ fn translate_request(req: vllm::GenerateRequest) -> Result EngineCoreSamplingParams { +fn translate_sampling( + sp: vllm::SamplingParams, + default_max_tokens: u32, + eos: &EosTokenIds, +) -> EngineCoreSamplingParams { + // Stopping at EOS is the frontend's duty here: the primary id rides + // `_eos_token_id`, extra ids merge into `stop_token_ids`, and the union + // feeds `_all_stop_token_ids` (engine-side `min_tokens` masking, built + // regardless of `ignore_eos`). + let mut stop_token_ids = sp.stop_token_ids; + if !sp.ignore_eos { + for id in &eos.extra { + if !stop_token_ids.contains(id) { + stop_token_ids.push(*id); + } + } + } + let mut all_stop_token_ids: BTreeSet = stop_token_ids.iter().copied().collect(); + all_stop_token_ids.extend(eos.primary); + all_stop_token_ids.extend(eos.extra.iter().copied()); let logit_bias = if sp.logit_bias.is_empty() { None } else { @@ -840,18 +1030,43 @@ fn translate_sampling(sp: vllm::SamplingParams) -> EngineCoreSamplingParams { frequency_penalty: sp.frequency_penalty, presence_penalty: sp.presence_penalty, repetition_penalty: sp.repetition_penalty, - max_tokens: sp.max_tokens.unwrap_or(16), + max_tokens: sp.max_tokens.unwrap_or(default_max_tokens), min_tokens: sp.min_tokens, - stop_token_ids: sp.stop_token_ids, + stop_token_ids, + eos_token_id: (!sp.ignore_eos).then_some(eos.primary).flatten(), + all_stop_token_ids, seed: sp.seed.map(i64::from), logprobs: sp.logprobs, // prompt_logprobs is rejected in `translate_request` (no renderer // support on the ZMQ path), so it is never forwarded. logit_bias, + structured_outputs: sp.constraint.and_then(translate_constraint), ..EngineCoreSamplingParams::default() } } +/// Map the proto `constraint` oneof onto typed structured-output params. The +/// backend defaults to guidance engine-side; `json_object=false` selects no +/// constraint (the caller opted out), so it maps to `None`. +fn translate_constraint( + constraint: vllm::sampling_params::Constraint, +) -> Option { + use vllm::sampling_params::Constraint; + match constraint { + Constraint::JsonSchema(schema) => Some(StructuredOutputsParams::json( + // The engine accepts a JSON schema object or a schema string; parse + // to preserve object shape, falling back to the raw string. + serde_json::from_str(&schema).unwrap_or(serde_json::Value::String(schema)), + )), + Constraint::Regex(regex) => Some(StructuredOutputsParams::regex(regex)), + Constraint::Grammar(grammar) => Some(StructuredOutputsParams::grammar(grammar)), + Constraint::StructuralTag(tag) => Some(StructuredOutputsParams::structural_tag(tag)), + Constraint::JsonObject(true) => Some(StructuredOutputsParams::json_object()), + Constraint::JsonObject(false) => None, + Constraint::Choice(choice) => Some(StructuredOutputsParams::choice(choice.choices)), + } +} + fn map_matched_stop(reason: StopReason) -> vllm::generate_complete::MatchedStop { match reason { StopReason::TokenId(id) => vllm::generate_complete::MatchedStop::MatchedTokenId(id), @@ -922,7 +1137,7 @@ mod tests { logprob: Option, finish: Option, ) -> EngineCoreOutputs { - let finished = finish.map(|_| std::collections::BTreeSet::from([request_id.to_string()])); + let finished = finish.map(|_| BTreeSet::from([request_id.to_string()])); let new_logprobs = logprob.map(|lp| { MaybeWireLogprobs::Direct(Logprobs { positions: vec![PositionLogprobs { @@ -963,6 +1178,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1020,7 +1236,10 @@ mod tests { stream: true, ..Default::default() }; - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::Vllm(Box::new(req))) + .await + .expect("generate"); let first = stream.next().await.expect("chunk item").expect("chunk ok"); match first.response { @@ -1065,6 +1284,144 @@ mod tests { engine_task.await.unwrap(); } + /// With `logprobs=k`, each position's ranked candidates are shaped into + /// `top_logprobs`, taking the sampled entry plus the leading candidates up + /// to the requested count (matching the gRPC servicer's `islice` behaviour). + #[tokio::test] + async fn generate_shapes_top_logprobs_to_requested_count() { + let dir = tempfile::tempdir().unwrap(); + let ep = |name: &str| format!("ipc://{}", dir.path().join(name).display()); + let (handshake, input, output) = (ep("hs.sock"), ep("in.sock"), ep("out.sock")); + + let (client, engine) = tokio::join!( + ZmqEngineClient::connect( + &handshake, + &input, + &output, + 1, + "m".to_string(), + EosTokenIds::default(), + RuntimeType::Vllm, + Duration::from_secs(10) + ), + connect_to_frontend( + &handshake, + EngineId::from_engine_index(0), + default_ready_response() + ), + ); + let client = client.expect("adapter connect"); + let engine = engine.expect("mock engine"); + + // One position with the sampled token (actual vocab rank) first, then + // the engine's ranked candidates. The wire carries `k + 1` entries. + let position = PositionLogprobs { + entries: vec![ + TokenLogprob { + token_id: 10, + logprob: -0.5, + rank: 5, + }, + TokenLogprob { + token_id: 20, + logprob: -0.1, + rank: 1, + }, + TokenLogprob { + token_id: 30, + logprob: -0.3, + rank: 2, + }, + ], + }; + let outputs = EngineCoreOutputs::RequestBatch(RequestBatchOutputs { + engine_index: 0, + outputs: vec![EngineCoreOutput { + request_id: "r1".to_string(), + new_token_ids: vec![10], + new_logprobs: Some(MaybeWireLogprobs::Direct(Logprobs { + positions: vec![position], + })), + finish_reason: Some(EngineCoreFinishReason::Length), + ..Default::default() + }], + finished_requests: Some(BTreeSet::from(["r1".to_string()])), + ..Default::default() + }); + + #[expect( + clippy::disallowed_methods, + reason = "engine task ends after responding" + )] + let engine_task = tokio::spawn(async move { + let (mut input, mut output) = engine.split(); + let inbound = input.recv().await.unwrap(); + let request = match inbound { + EngineInbound::Add(request) => request, + other => panic!("expected Add, got {other:?}"), + }; + assert_eq!(request.sampling_params.as_ref().unwrap().logprobs, Some(2)); + output.send_outputs(&outputs).await.unwrap(); + }); + + let req = vllm::GenerateRequest { + request_id: "r1".to_string(), + input: Some(vllm::generate_request::Input::Tokenized( + vllm::TokenizedInput { + original_text: String::new(), + input_ids: vec![1, 2, 3], + }, + )), + sampling_params: Some(vllm::SamplingParams { + max_tokens: Some(1), + logprobs: Some(2), + ..Default::default() + }), + stream: true, + ..Default::default() + }; + let mut stream = client + .generate(ProtoGenerateRequest::Vllm(Box::new(req))) + .await + .expect("generate"); + + // The requested count is 2, so `top_logprobs` keeps the sampled entry + // plus the leading candidate (the third entry is dropped). + let expected_top = vec![vllm::TopLogProbs { + values: vec![-0.5, -0.1], + token_ids: vec![10, 20], + }]; + + // The finish tick carried a token, so the delta streams as a chunk. + let chunk = stream.next().await.expect("chunk item").expect("chunk ok"); + match chunk.response { + Some(vllm::generate_response::Response::Chunk(chunk)) => { + let logprobs = chunk.output_logprobs.expect("chunk logprobs"); + assert_eq!(logprobs.token_logprobs, vec![-0.5]); + assert_eq!(logprobs.token_ids, vec![10]); + assert_eq!(logprobs.top_logprobs, expected_top); + } + other => panic!("expected chunk, got {other:?}"), + } + let complete = stream + .next() + .await + .expect("complete item") + .expect("complete ok"); + match complete.response { + Some(vllm::generate_response::Response::Complete(complete)) => { + let logprobs = complete.output_logprobs.expect("complete logprobs"); + assert_eq!(logprobs.token_logprobs, vec![-0.5]); + assert_eq!(logprobs.token_ids, vec![10]); + assert_eq!(logprobs.top_logprobs, expected_top); + } + other => panic!("expected complete, got {other:?}"), + } + assert!(stream.next().await.is_none()); + + engine_task.await.unwrap(); + } + /// End-to-end over ipc:// for a TokenSpeed backend: the adapter frames a /// tagged `TokenizedGenerateReqInput`, and maps `BatchTokenIDOutSlim` /// batches back to vLLM-proto responses. The mock engine speaks the shared @@ -1091,6 +1448,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::TokenSpeed, Duration::from_secs(10) ), @@ -1153,24 +1511,25 @@ mod tests { .unwrap(); }); - let req = vllm::GenerateRequest { + let req = tokenspeed_proto::GenerateRequest { request_id: "r1".to_string(), - input: Some(vllm::generate_request::Input::Tokenized( - vllm::TokenizedInput { - original_text: String::new(), - input_ids: vec![1, 2, 3], - }, - )), - sampling_params: Some(vllm::SamplingParams { - max_tokens: Some(2), - // Plain sampled-token logprob (count 1); must be wired through. - logprobs: Some(1), + tokenized: Some(tokenspeed_proto::TokenizedInput { + input_ids: vec![1, 2, 3], + original_text: String::new(), + }), + sampling_params: Some(tokenspeed_proto::SamplingParams { + max_new_tokens: Some(2), ..Default::default() }), + // Plain sampled-token logprob; must be wired through. + return_logprob: true, stream: true, ..Default::default() }; - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::TokenSpeed(Box::new(req))) + .await + .expect("generate"); let first = stream.next().await.expect("chunk item").expect("chunk ok"); match first.response { @@ -1230,35 +1589,33 @@ mod tests { } #[test] - fn tokenspeed_sampling_maps_top_k_sentinel_and_seed() { + fn tokenspeed_sampling_maps_top_k_sentinel_and_floors_n() { use engine_zmq_client::protocol::tokenspeed::sampling::TOP_K_DISABLED; - // Proto top_k=0 ("all tokens") normalizes to the engine's disabled - // sentinel; negative seed dropped; n floored to 1. - let mapped = translate_sampling_tokenspeed(vllm::SamplingParams { - top_k: 0, + // Unset top_k rides the API convention `-1` ("all tokens") and normalizes + // to the engine's disabled sentinel; n=0 floors to 1; max_new_tokens + // forwards. + let mapped = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + top_k: None, n: 0, - seed: Some(-1), - max_tokens: Some(8), + max_new_tokens: Some(8), ..Default::default() }); assert_eq!(mapped.top_k, TOP_K_DISABLED); assert_eq!(mapped.n, 1); - assert_eq!(mapped.seed, None); assert_eq!(mapped.max_new_tokens, Some(8)); // The wire form is always normalized (the engine skips re-derivation). assert!(mapped.is_normalized); - let mapped = translate_sampling_tokenspeed(vllm::SamplingParams { - top_k: 40, - seed: Some(7), + // An explicit top_k passes through unchanged. + let mapped = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + top_k: Some(40), ..Default::default() }); assert_eq!(mapped.top_k, 40); - assert_eq!(mapped.seed, Some(7)); // A near-zero temperature collapses to greedy on the wire. - let mapped = translate_sampling_tokenspeed(vllm::SamplingParams { + let mapped = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { temperature: Some(0.0), ..Default::default() }); @@ -1266,7 +1623,7 @@ mod tests { assert_eq!(mapped.top_k, 1); // Empty stop_token_ids ride as None (the normalized encoding). - let mapped = translate_sampling_tokenspeed(vllm::SamplingParams::default()); + let mapped = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams::default()); assert_eq!(mapped.stop_token_ids, None); } @@ -1285,110 +1642,289 @@ mod tests { } } - #[test] - fn tokenspeed_plain_logprobs_set_return_logprob() { - // Counts 0 and 1 are the plain sampled-token case; both accepted. - for count in [0, 1] { - let req = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - logprobs: Some(count), - ..Default::default() - })) - .expect("plain logprobs accepted"); - assert!(req.return_logprob); + fn ts_tokenized_req( + sampling: tokenspeed_proto::SamplingParams, + ) -> tokenspeed_proto::GenerateRequest { + tokenspeed_proto::GenerateRequest { + request_id: "r1".to_string(), + tokenized: Some(tokenspeed_proto::TokenizedInput { + input_ids: vec![1, 2, 3], + original_text: String::new(), + }), + sampling_params: Some(sampling), + stream: true, + ..Default::default() } + } - // No logprobs -> the flag stays false. - let req = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams::default())) - .expect("no logprobs accepted"); - assert!(!req.return_logprob); + #[test] + fn tokenspeed_return_logprob_flag_passes_through() { + // The request-level `return_logprob` drives the plain sampled-token + // logprob (count 0/1 in `top_logprobs_num` is the same case). + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.return_logprob = true; + let wire = translate_request_tokenspeed(req).expect("return_logprob accepted"); + assert!(wire.return_logprob); + + // Unset -> the flag stays false. + let wire = translate_request_tokenspeed(ts_tokenized_req( + tokenspeed_proto::SamplingParams::default(), + )) + .expect("no logprobs accepted"); + assert!(!wire.return_logprob); } #[test] - fn tokenspeed_rejects_top_k_and_prompt_logprobs() { - // Top-k (count > 1) cannot be honored. - assert!( - translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - logprobs: Some(5), - ..Default::default() - })) - .is_err() - ); - // "all" (count -1) cannot be honored. - assert!( - translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - logprobs: Some(-1), - ..Default::default() - })) - .is_err() - ); - // Prompt logprobs cannot be produced. - assert!( - translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - prompt_logprobs: Some(1), - ..Default::default() - })) - .is_err() - ); + fn tokenspeed_rejects_top_logprobs_and_prompt_logprobs() { + // Top-k logprobs (count > 1) cannot be honored over the wire. + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.top_logprobs_num = 5; + assert!(translate_request_tokenspeed(req).is_err()); + + // Prompt (input) logprobs cannot be produced. + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.token_ids_logprob = vec![1, 2]; + assert!(translate_request_tokenspeed(req).is_err()); + + // A bare count of 0/1 is the plain sampled-token case: accepted. + for count in [0, 1] { + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.top_logprobs_num = count; + assert!(translate_request_tokenspeed(req).is_ok()); + } } #[test] - fn tokenspeed_rejects_unsupported_sampling_features() { - // Structured-output constraints have no TokenSpeed wire slot. - let err = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - constraint: Some(vllm::sampling_params::Constraint::JsonObject(true)), + fn tokenspeed_maps_structured_output_constraints() { + // The `constraint` oneof maps 1:1 onto the wire's dedicated fields; the + // oneof is single-valued, so the other three stay unset. + use tokenspeed_proto::sampling_params::Constraint; + + let json = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + constraint: Some(Constraint::JsonSchema("{\"type\":\"object\"}".into())), ..Default::default() - })) - .expect_err("constraint rejected"); - assert!(err.contains("structured output"), "{err}"); + }); + assert_eq!(json.json_schema.as_deref(), Some("{\"type\":\"object\"}")); + assert_eq!(json.regex, None); + assert_eq!(json.ebnf, None); + assert_eq!(json.structural_tag, None); - // Stop strings have no wire slot and are not enforced by the gateway. - let err = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - stop: vec!["".to_string()], + let regex = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + constraint: Some(Constraint::Regex("[0-9]+".into())), ..Default::default() - })) - .expect_err("stop strings rejected"); - assert!(err.contains("stop_token_ids"), "{err}"); + }); + assert_eq!(regex.regex.as_deref(), Some("[0-9]+")); + assert_eq!(regex.json_schema, None); - // logit_bias has no wire slot. - let err = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - logit_bias: HashMap::from([(7, 1.0)]), + let ebnf = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + constraint: Some(Constraint::EbnfGrammar("root ::= \"a\"".into())), + ..Default::default() + }); + assert_eq!(ebnf.ebnf.as_deref(), Some("root ::= \"a\"")); + + let tag = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + constraint: Some(Constraint::StructuralTag("".into())), + ..Default::default() + }); + assert_eq!(tag.structural_tag.as_deref(), Some("")); + + // No constraint leaves all four structured-output fields unset. + let none = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams::default()); + assert_eq!(none.json_schema, None); + assert_eq!(none.regex, None); + assert_eq!(none.ebnf, None); + assert_eq!(none.structural_tag, None); + } + + #[test] + fn tokenspeed_forwards_stop_token_ids_and_drops_stop_strings() { + // String stops are resolved upstream; any that reach here are dropped + // (the token-only engine cannot match them) while stop token ids ride + // through and the router-side decoder trims residual text. + let req = translate_request_tokenspeed(ts_tokenized_req(tokenspeed_proto::SamplingParams { + stop: vec!["".to_string()], + stop_token_ids: vec![13], ..Default::default() })) - .expect_err("logit_bias rejected"); - assert!(err.contains("logit_bias"), "{err}"); + .expect("residual stop strings must not be rejected"); + assert_eq!(req.sampling_params.stop_token_ids, Some(vec![13])); + assert_eq!(req.sampling_params.stop, None); } #[test] - fn tokenspeed_rejects_nonzero_dp_rank() { - // Single-engine backend: only rank 0 (or none) is valid. - let mut req = tokenized_req(vllm::SamplingParams::default()); - req.data_parallel_rank = Some(1); + fn tokenspeed_rejects_multimodal_inputs() { + // The TokenSpeed ZMQ wire has no multimodal slot yet; reject rather than + // silently drop pixels. + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.mm_inputs = Some(tokenspeed_proto::MultimodalInputs::default()); assert!(translate_request_tokenspeed(req).is_err()); - - let mut req = tokenized_req(vllm::SamplingParams::default()); - req.data_parallel_rank = Some(0); - assert!(translate_request_tokenspeed(req).is_ok()); } #[test] fn vllm_rejects_unsupported_sampling_features() { - // Structured-output constraints are not translated onto the wire. - let err = translate_request(tokenized_req(vllm::SamplingParams { - constraint: Some(vllm::sampling_params::Constraint::JsonObject(true)), - ..Default::default() - })) - .expect_err("constraint rejected"); - assert!(err.contains("structured output"), "{err}"); - // Prompt logprobs have no renderer merge on the ZMQ path. - let err = translate_request(tokenized_req(vllm::SamplingParams { - prompt_logprobs: Some(1), - ..Default::default() - })) + let err = translate_request( + tokenized_req(vllm::SamplingParams { + prompt_logprobs: Some(1), + ..Default::default() + }), + 4096, + ModelDtype::BFloat16, + &EosTokenIds::default(), + ) .expect_err("prompt logprobs rejected"); assert!(err.contains("prompt logprobs"), "{err}"); } + #[test] + fn vllm_defaults_unset_max_tokens_to_remaining_context() { + let max_tokens = |sampling, max_model_len| { + translate_request( + tokenized_req(sampling), + max_model_len, + ModelDtype::BFloat16, + &EosTokenIds::default(), + ) + .expect("request translated") + .sampling_params + .expect("sampling params present") + .max_tokens + }; + + // Unset max_tokens defaults to `max_model_len - prompt_len` (prompt is + // 3 tokens), mirroring vLLM's bypassed OpenAI frontend. + assert_eq!(max_tokens(vllm::SamplingParams::default(), 100), 97); + // An explicit value is always honored. + assert_eq!( + max_tokens( + vllm::SamplingParams { + max_tokens: Some(8), + ..Default::default() + }, + 100, + ), + 8, + ); + } + + #[test] + fn vllm_attaches_eos_stop_ids() { + let eos = EosTokenIds::new(Some(5), vec![7]); + let sampling = |sp| { + translate_request(tokenized_req(sp), 4096, ModelDtype::BFloat16, &eos) + .expect("request translated") + .sampling_params + .expect("sampling params present") + }; + + // Primary rides `_eos_token_id`, extras merge into `stop_token_ids` + // without duplicating, and the union lands in `_all_stop_token_ids`. + let sp = sampling(vllm::SamplingParams { + stop_token_ids: vec![7, 9], + ..Default::default() + }); + assert_eq!(sp.eos_token_id, Some(5)); + assert_eq!(sp.stop_token_ids, vec![7, 9]); + assert_eq!(sp.all_stop_token_ids, BTreeSet::from([5, 7, 9])); + + // ignore_eos drops the EOS stops from the wire but keeps the + // bookkeeping set (mirrors the reference frontend). + let sp = sampling(vllm::SamplingParams { + stop_token_ids: vec![9], + ignore_eos: true, + ..Default::default() + }); + assert_eq!(sp.eos_token_id, None); + assert_eq!(sp.stop_token_ids, vec![9]); + assert_eq!(sp.all_stop_token_ids, BTreeSet::from([5, 7, 9])); + } + + #[test] + fn eos_token_ids_resolve_from_model_dir() { + let dir = tempfile::tempdir().expect("tempdir"); + std::fs::write(dir.path().join("config.json"), r#"{"eos_token_id": 5}"#).unwrap(); + std::fs::write( + dir.path().join("generation_config.json"), + r#"{"eos_token_id": [5, 7, 9]}"#, + ) + .unwrap(); + assert_eq!( + EosTokenIds::from_model_dir(dir.path()), + EosTokenIds::new(Some(5), vec![7, 9]), + ); + + // Missing files degrade to no ids, not an error. + let empty = tempfile::tempdir().expect("tempdir"); + assert_eq!( + EosTokenIds::from_model_dir(empty.path()), + EosTokenIds::default(), + ); + } + + #[test] + fn vllm_translates_structured_output_constraints() { + use engine_zmq_client::protocol::vllm::structured_outputs::{ + StructuredOutputBackend, StructuredOutputConstraint, + }; + + let translate = |constraint| { + translate_request( + tokenized_req(vllm::SamplingParams { + constraint: Some(constraint), + ..Default::default() + }), + 4096, + ModelDtype::BFloat16, + &EosTokenIds::default(), + ) + .expect("constraint translated") + .sampling_params + .expect("sampling params present") + .structured_outputs + }; + + // Each constraint mode maps onto its typed counterpart, always lowering + // to the guidance backend engine-side. + let json_object = translate(vllm::sampling_params::Constraint::JsonObject(true)) + .expect("json_object translated"); + assert_eq!( + json_object.constraint, + StructuredOutputConstraint::JsonObject + ); + assert_eq!(json_object.backend, StructuredOutputBackend::Guidance); + + let regex = translate(vllm::sampling_params::Constraint::Regex("a.*".to_string())) + .expect("regex translated"); + assert_eq!( + regex.constraint, + StructuredOutputConstraint::Regex("a.*".to_string()) + ); + + let choice = translate(vllm::sampling_params::Constraint::Choice( + vllm::ChoiceConstraint { + choices: vec!["yes".to_string(), "no".to_string()], + }, + )) + .expect("choice translated"); + assert_eq!( + choice.constraint, + StructuredOutputConstraint::Choice(vec!["yes".to_string(), "no".to_string()]) + ); + + // A JSON schema string is parsed to preserve object shape. + let json = translate(vllm::sampling_params::Constraint::JsonSchema( + r#"{"type":"object"}"#.to_string(), + )) + .expect("json schema translated"); + assert_eq!( + json.constraint, + StructuredOutputConstraint::Json(serde_json::json!({"type": "object"})) + ); + + // json_object=false means the caller opted out: no constraint. + assert!(translate(vllm::sampling_params::Constraint::JsonObject(false)).is_none()); + } + /// n=3 fans out into 3 single-sample wire requests with unique sub-rids. /// An explicit seed derives per-sub seeds (`seed + i`) so the samples /// differ deterministically; no seed stays `None` per sub (the engine @@ -1453,6 +1989,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1504,7 +2041,10 @@ mod tests { ..Default::default() }); req.request_id = "r1".to_string(); - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::Vllm(Box::new(req))) + .await + .expect("generate"); let mut completes = Vec::new(); while let Some(item) = stream.next().await { @@ -1550,6 +2090,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::TokenSpeed, Duration::from_secs(10) ), @@ -1578,12 +2119,15 @@ mod tests { let request: TokenizedGenerateReqInput = decode_msgpack(frames[1].as_ref()).unwrap(); assert_eq!(request.sampling_params.n, 1); - rids.push((request.rid.clone(), request.sampling_params.seed)); + // TokenSpeed has no seed on the wire; the engine derives one from + // the (unique) rid so all TP/DP ranks agree. + assert_eq!(request.sampling_params.seed, None); + rids.push(request.rid.clone()); } assert_eq!( rids, - vec![("r1-0".to_string(), Some(5)), ("r1-1".to_string(), Some(6))], - "sub-rids must be unique and seeds derived per sub" + vec!["r1-0".to_string(), "r1-1".to_string()], + "sub-rids must be unique per sub" ); // Both subs finish in one wire batch (the batch demux fans them // back out to their sub-streams). @@ -1603,13 +2147,15 @@ mod tests { .unwrap(); }); - let mut req = tokenized_req(vllm::SamplingParams { + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams { n: 2, - seed: Some(5), ..Default::default() }); req.request_id = "r1".to_string(); - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::TokenSpeed(Box::new(req))) + .await + .expect("generate"); let mut completes = Vec::new(); while let Some(item) = stream.next().await { @@ -1648,6 +2194,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1662,10 +2209,12 @@ mod tests { let (mut engine_input, _engine_output) = engine.split(); let stream = client - .generate(tokenized_req(vllm::SamplingParams { - n: 2, - ..Default::default() - })) + .generate(ProtoGenerateRequest::Vllm(Box::new(tokenized_req( + vllm::SamplingParams { + n: 2, + ..Default::default() + }, + )))) .await .expect("generate"); @@ -1680,7 +2229,7 @@ mod tests { drop(stream); // unfinished -> every sub auto-aborts - let mut aborted = std::collections::BTreeSet::new(); + let mut aborted = BTreeSet::new(); while aborted.len() < 2 { match engine_input.recv().await.unwrap() { EngineInbound::Abort(rids) => aborted.extend(rids), @@ -1689,7 +2238,7 @@ mod tests { } assert_eq!( aborted, - std::collections::BTreeSet::from(["r1-0".to_string(), "r1-1".to_string()]) + BTreeSet::from(["r1-0".to_string(), "r1-1".to_string()]) ); } diff --git a/model_gateway/src/routers/grpc/zmq_multimodal.rs b/model_gateway/src/routers/grpc/zmq_multimodal.rs new file mode 100644 index 000000000..f20bd3c08 --- /dev/null +++ b/model_gateway/src/routers/grpc/zmq_multimodal.rs @@ -0,0 +1,518 @@ +//! Proto multimodal inputs → EngineCore `mm_features` for the direct-ZMQ path. +//! +//! The gRPC servicer converts the batched proto tensors into per-item engine +//! structures Python-side (`_build_preprocessed_mm_inputs` + the engine's +//! `from_hf_inputs` split). The ZMQ path bypasses that process, so the same +//! split happens here: batched keys index row `i`, flat keys slice by the +//! cumulative sizes tensor, everything else is shared (replicated per item). +//! Floating tensors are cast to the model dtype — the engine applies no cast +//! on this path. + +use std::collections::{BTreeMap, HashMap, HashSet}; + +use bytes::Bytes; +use engine_zmq_client::{ + codec::{ + dtype::ModelDtype, + tensor::{WireArrayData, WireTensor}, + }, + protocol::vllm::multimodal::{ + MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargValue, + MmKwargsItem, MmSharedField, MmSlice, PlaceholderRange, SliceSpec, + }, +}; +use smg_grpc_client::{common_proto as common, vllm_proto as vllm}; + +/// A decoded (and dtype-cast) proto tensor ready for per-item slicing. +struct Decoded { + dtype: String, + shape: Vec, + bytes: Bytes, +} + +impl Decoded { + fn elem_size(&self) -> Result { + match self.dtype.as_str() { + "bool" => Ok(1), + "float16" | "bfloat16" => Ok(2), + "float32" | "uint32" | "int32" => Ok(4), + "int64" | "float64" => Ok(8), + other => Err(format!("unsupported multimodal tensor dtype {other:?}")), + } + } + + /// Bytes per index step along dim 0. + fn row_nbytes(&self) -> Result { + let inner: usize = self.shape.iter().skip(1).product(); + Ok(inner * self.elem_size()?) + } + + /// Zero-copy view of rows `[start, stop)` along dim 0. + fn slice_rows(&self, start: usize, stop: usize) -> Result { + let row = self.row_nbytes()?; + let (lo, hi) = (start * row, stop * row); + if hi > self.bytes.len() || start > stop { + return Err(format!( + "row slice {start}..{stop} out of bounds for tensor of {} bytes", + self.bytes.len() + )); + } + let mut shape = self.shape.clone(); + shape[0] = stop - start; + Ok(WireTensor::from_raw_bytes( + self.dtype.clone(), + shape, + self.bytes.slice(lo..hi), + )) + } + + fn whole(&self) -> WireTensor { + WireTensor::from_raw_bytes(self.dtype.clone(), self.shape.clone(), self.bytes.clone()) + } + + /// Flattened values as widened i64 (sizes tensors are int64 or uint32). + fn flat_i64(&self) -> Result, String> { + match self.dtype.as_str() { + "int64" => Ok((self.bytes.as_chunks::<8>().0.iter()) + .map(|c| i64::from_le_bytes(*c)) + .collect()), + "uint32" => Ok((self.bytes.as_chunks::<4>().0.iter()) + .map(|c| i64::from(u32::from_le_bytes(*c))) + .collect()), + other => Err(format!("flat sizes tensor has unsupported dtype {other:?}")), + } + } +} + +fn decode_tensor( + name: &str, + tensor: vllm::TensorData, + model_dtype: ModelDtype, +) -> Result { + let shape: Vec = tensor.shape.iter().map(|&d| d as usize).collect(); + let data = match tensor.payload { + Some(vllm::tensor_data::Payload::Inline(data)) => data, + Some(_) => { + return Err(format!( + "multimodal tensor {name:?} uses a non-inline payload; the ZMQ wire carries \ + tensors inline" + )); + } + None => return Err(format!("multimodal tensor {name:?} has no payload")), + }; + // Floating tensors arrive as float32 and are cast to the model dtype, + // mirroring the cast the engine's own frontend applies. + if tensor.dtype == "float32" { + let cast = WireTensor::from_f32_bytes_cast(model_dtype, shape.clone(), &data)?; + let WireArrayData::RawView(bytes) = cast.data else { + return Err(format!("cast tensor {name:?} lost its raw view")); + }; + return Ok(Decoded { + dtype: cast.dtype, + shape, + bytes, + }); + } + Ok(Decoded { + dtype: tensor.dtype, + shape, + bytes: Bytes::from(data), + }) +} + +/// Rename generic keys for video inputs, mirroring the servicer's `mm_key`. +fn mm_key(key: &str, is_video: bool) -> String { + if is_video && key == "pixel_values" { + "pixel_values_videos".to_string() + } else { + key.to_string() + } +} + +/// Build per-item `mm_features` from batched proto multimodal inputs. +pub(crate) fn build_mm_features( + mm: vllm::MultimodalInputs, + prompt_token_ids: &[u32], + model_dtype: ModelDtype, +) -> Result { + let num_items = mm.mm_placeholders.len(); + if num_items == 0 { + return Ok(Vec::new()); + } + if mm.mm_hashes.len() != num_items { + return Err(format!( + "multimodal hash count {} does not match placeholder count {num_items}", + mm.mm_hashes.len() + )); + } + let is_video = mm.modality == common::Modality::Video as i32; + let modality = if is_video { "video" } else { "image" }; + + // Decode every tensor once, applying the video key rename. + let mut tensors: BTreeMap = BTreeMap::new(); + if let Some(pixel_values) = mm.pixel_values { + tensors.insert( + mm_key("pixel_values", is_video), + decode_tensor("pixel_values", pixel_values, model_dtype)?, + ); + } + for (key, tensor) in mm.model_specific_tensors { + let decoded = decode_tensor(&key, tensor, model_dtype)?; + tensors.insert(mm_key(&key, is_video), decoded); + } + + let batched: HashSet = mm + .batched_keys + .iter() + .map(|k| mm_key(k, is_video)) + .collect(); + let flat: HashMap = mm + .flat_keys + .iter() + .map(|(k, v)| (mm_key(k, is_video), mm_key(v, is_video))) + .collect(); + let keep_on_cpu: HashSet = mm + .keep_on_cpu_keys + .iter() + .map(|k| mm_key(k, is_video)) + .collect(); + + // Split every kwarg into per-item elems. + let mut items: Vec = vec![MmKwargsItem::new(); num_items]; + for (key, decoded) in &tensors { + let on_cpu = keep_on_cpu.contains(key); + if batched.contains(key) { + if decoded.shape.first() != Some(&num_items) { + return Err(format!( + "batched tensor {key:?} has leading dim {:?}, expected {num_items} items", + decoded.shape.first() + )); + } + for (i, item) in items.iter_mut().enumerate() { + let mut tensor = decoded.slice_rows(i, i + 1)?; + tensor.shape.remove(0); + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(tensor)), + field: MmField::Batched(MmBatchedField { + keep_on_cpu: on_cpu, + }), + }, + ); + } + } else if let Some(sizes_key) = flat.get(key) { + let sizes = tensors + .get(sizes_key) + .ok_or_else(|| format!("flat sizes tensor {sizes_key:?} missing for {key:?}"))? + .flat_i64()?; + if sizes.len() != num_items { + return Err(format!( + "flat sizes tensor {sizes_key:?} has {} entries, expected {num_items}", + sizes.len() + )); + } + // Cumulative row offsets, and the full per-item slice list every + // elem carries (the engine's flat field serializes all slices). + let mut bounds = Vec::with_capacity(num_items + 1); + let mut total = 0usize; + bounds.push(total); + for size in &sizes { + let size = usize::try_from(*size) + .map_err(|_| format!("negative size in flat sizes tensor {sizes_key:?}"))?; + total += size; + bounds.push(total); + } + if decoded.shape.first() != Some(&total) { + return Err(format!( + "flat tensor {key:?} has leading dim {:?}, expected {total} total rows", + decoded.shape.first(), + )); + } + let slices: Vec = bounds + .windows(2) + .map(|w| { + MmSlice::Slice(SliceSpec { + start: Some(w[0] as isize), + stop: Some(w[1] as isize), + step: None, + }) + }) + .collect(); + for (i, item) in items.iter_mut().enumerate() { + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor( + decoded.slice_rows(bounds[i], bounds[i + 1])?, + )), + field: MmField::Flat(MmFlatField { + slices: slices.clone(), + dim: 0, + keep_on_cpu: on_cpu, + }), + }, + ); + } + } else { + // Shared: the full tensor replicated per item (the servicer's + // fallback for keys in neither batched nor flat sets). + for item in &mut items { + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(decoded.whole())), + field: MmField::Shared(MmSharedField { + batch_size: num_items, + keep_on_cpu: false, + }), + }, + ); + } + } + } + + // One feature per placeholder, in prompt-offset order. + let mut features: MmFeatures = Vec::with_capacity(num_items); + for ((placeholder, item), hash) in mm + .mm_placeholders + .iter() + .zip(items) + .zip(mm.mm_hashes.iter()) + { + let offset = placeholder.offset as usize; + let length = placeholder.length as usize; + features.push(MmFeatureSpec { + data: Some(item), + modality: modality.to_string(), + identifier: hash.clone(), + mm_position: PlaceholderRange { + offset, + length, + is_embed: is_embed_mask(prompt_token_ids, offset, length, mm.im_token_id)?, + }, + mm_hash: Some(hash.clone()), + }); + } + features.sort_by_key(|f| f.mm_position.offset); + Ok(features) +} + +/// Boolean embed mask over a placeholder range: `true` where the prompt token +/// is the image token, excluding structural tokens (vision start/end markers) +/// from the embedding scatter. `None` when every position is an embed slot. +fn is_embed_mask( + prompt_token_ids: &[u32], + offset: usize, + length: usize, + im_token_id: Option, +) -> Result, String> { + let Some(im_token_id) = im_token_id else { + return Ok(None); + }; + let end = offset + .checked_add(length) + .filter(|&end| end <= prompt_token_ids.len()) + .ok_or_else(|| { + format!( + "placeholder range {offset}+{length} exceeds prompt of {} tokens", + prompt_token_ids.len() + ) + })?; + let mask: Vec = prompt_token_ids[offset..end] + .iter() + .map(|&id| id == im_token_id) + .collect(); + if mask.iter().all(|&m| m) { + return Ok(None); + } + Ok(Some(WireTensor::from_bool(vec![length], mask)?)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn inline_tensor(shape: Vec, dtype: &str, data: Vec) -> vllm::TensorData { + vllm::TensorData { + shape, + dtype: dtype.to_string(), + payload: Some(vllm::tensor_data::Payload::Inline(data)), + } + } + + fn f32_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + fn i64_bytes(values: &[i64]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + fn placeholders(ranges: &[(u32, u32)]) -> Vec { + ranges + .iter() + .map(|&(offset, length)| vllm::PlaceholderRange { offset, length }) + .collect() + } + + fn base_inputs() -> vllm::MultimodalInputs { + vllm::MultimodalInputs { + pixel_values: Some(inline_tensor( + vec![2, 4], + "float32", + f32_bytes(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]), + )), + model_specific_tensors: Default::default(), + im_token_id: None, + mm_placeholders: placeholders(&[(1, 3), (6, 3)]), + mm_hashes: vec!["h0".to_string(), "h1".to_string()], + batched_keys: vec!["pixel_values".to_string()], + flat_keys: Default::default(), + keep_on_cpu_keys: vec![], + modality: common::Modality::Image as i32, + } + } + + fn tensor_of(elem: &MmFieldElem) -> &WireTensor { + match elem.data.as_ref().expect("data present") { + MmKwargValue::Tensor(tensor) => tensor, + other => panic!("expected tensor, got {other:?}"), + } + } + + #[test] + fn batched_keys_split_per_row_and_cast_to_model_dtype() { + let features = build_mm_features(base_inputs(), &[], ModelDtype::BFloat16).expect("built"); + assert_eq!(features.len(), 2); + + for (i, feature) in features.iter().enumerate() { + assert_eq!(feature.modality, "image"); + assert_eq!(feature.identifier, format!("h{i}")); + assert_eq!(feature.mm_hash.as_deref(), Some(format!("h{i}").as_str())); + let item = feature.data.as_ref().expect("item present"); + let tensor = tensor_of(&item["pixel_values"]); + // Row i of the [2, 4] float32 batch, cast to bfloat16. + assert_eq!(tensor.dtype, "bfloat16"); + assert_eq!(tensor.shape, vec![4]); + assert!(matches!( + item["pixel_values"].field, + MmField::Batched(MmBatchedField { keep_on_cpu: false }) + )); + } + assert_eq!(features[0].mm_position.offset, 1); + assert_eq!(features[1].mm_position.offset, 6); + } + + #[test] + fn flat_keys_slice_by_cumulative_sizes() { + let mut mm = base_inputs(); + mm.pixel_values = Some(inline_tensor(vec![5, 2], "float32", f32_bytes(&[0.0; 10]))); + mm.batched_keys = vec!["patches_per_image".to_string()]; + mm.flat_keys = [("pixel_values".to_string(), "patches_per_image".to_string())].into(); + mm.model_specific_tensors = [( + "patches_per_image".to_string(), + inline_tensor(vec![2], "int64", i64_bytes(&[2, 3])), + )] + .into(); + + let features = build_mm_features(mm, &[], ModelDtype::Float32).expect("built"); + let item0 = features[0].data.as_ref().expect("item 0"); + let item1 = features[1].data.as_ref().expect("item 1"); + assert_eq!(tensor_of(&item0["pixel_values"]).shape, vec![2, 2]); + assert_eq!(tensor_of(&item1["pixel_values"]).shape, vec![3, 2]); + + // Every elem carries the full per-item slice list. + let expected_slices = vec![ + MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(2), + step: None, + }), + MmSlice::Slice(SliceSpec { + start: Some(2), + stop: Some(5), + step: None, + }), + ]; + for item in [item0, item1] { + let MmField::Flat(flat) = &item["pixel_values"].field else { + panic!("expected flat field"); + }; + assert_eq!(flat.slices, expected_slices); + assert_eq!(flat.dim, 0); + } + } + + #[test] + fn unlisted_keys_are_shared_and_replicated() { + let mut mm = base_inputs(); + mm.model_specific_tensors = [( + "video_second_per_grid".to_string(), + inline_tensor(vec![2], "int64", i64_bytes(&[1, 1])), + )] + .into(); + + let features = build_mm_features(mm, &[], ModelDtype::BFloat16).expect("built"); + for feature in &features { + let item = feature.data.as_ref().expect("item present"); + let elem = &item["video_second_per_grid"]; + assert_eq!(tensor_of(elem).shape, vec![2]); + assert!(matches!( + elem.field, + MmField::Shared(MmSharedField { + batch_size: 2, + keep_on_cpu: false, + }) + )); + } + } + + #[test] + fn is_embed_masks_structural_tokens() { + let mut mm = base_inputs(); + mm.im_token_id = Some(7); + // Placeholder 0 covers tokens [7, 7, 5] (mixed); placeholder 1 covers + // [7, 7, 7] (all image tokens). + let prompt = [9, 7, 7, 5, 9, 9, 7, 7, 7]; + + let features = build_mm_features(mm, &prompt, ModelDtype::BFloat16).expect("built"); + let mask = features[0] + .mm_position + .is_embed + .as_ref() + .expect("mixed range keeps a mask"); + assert_eq!(mask.dtype, "bool"); + assert_eq!(mask.shape, vec![3]); + assert!(features[1].mm_position.is_embed.is_none()); + } + + #[test] + fn video_renames_pixel_values() { + let mut mm = base_inputs(); + mm.modality = common::Modality::Video as i32; + + let features = build_mm_features(mm, &[], ModelDtype::BFloat16).expect("built"); + let item = features[0].data.as_ref().expect("item present"); + assert!(item.contains_key("pixel_values_videos")); + assert!(!item.contains_key("pixel_values")); + assert_eq!(features[0].modality, "video"); + } + + #[test] + fn rejects_hash_mismatch_and_non_inline_payloads() { + let mut mm = base_inputs(); + mm.mm_hashes.pop(); + let err = build_mm_features(mm, &[], ModelDtype::BFloat16).expect_err("hash mismatch"); + assert!(err.contains("hash count"), "{err}"); + + let mut mm = base_inputs(); + mm.pixel_values = Some(vllm::TensorData { + shape: vec![2, 4], + dtype: "float32".to_string(), + payload: Some(vllm::tensor_data::Payload::Shm(Default::default())), + }); + let err = build_mm_features(mm, &[], ModelDtype::BFloat16).expect_err("shm rejected"); + assert!(err.contains("inline"), "{err}"); + } +} diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index a117fc545..d34b9a4e3 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -35,7 +35,7 @@ use crate::{ grpc::{ backend_client::BackendClient, client::GrpcClient, - zmq_client::{ZmqEngineClient, ZMQ_LOOPBACK_HOST}, + zmq_client::{EosTokenIds, ZmqEngineClient, ZMQ_LOOPBACK_HOST}, }, }, }; @@ -200,6 +200,19 @@ async fn connect_zmq_backend( let (handshake, input, output) = zmq_socket_addresses(&base_url, handshake_override.as_deref())?; ensure_ipc_socket_dir(&base_url).await?; + // The engine can't stop at EOS on its own (it has no tokenizer or model + // config); resolve the EOS ids from the local model dir so every request + // carries them. + let model_dir = std::path::Path::new(&model_id); + let eos = if model_dir.is_dir() { + EosTokenIds::from_model_dir(model_dir) + } else { + tracing::warn!( + "ZMQ worker model id '{model_id}' is not a local model directory; EOS ids \ + unavailable — generation stops only at max_tokens or explicit stops" + ); + EosTokenIds::default() + }; tracing::info!("Binding ZMQ client for worker {base_url} (handshake={handshake})"); match ZmqEngineClient::connect( &handshake, @@ -207,6 +220,7 @@ async fn connect_zmq_backend( &output, 1, model_id, + eos, runtime, ZMQ_CONNECT_TIMEOUT, ) @@ -2826,6 +2840,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), diff --git a/scripts/ci_install_tokenspeed.sh b/scripts/ci_install_tokenspeed.sh index 1b1d4f5aa..352f94d5e 100755 --- a/scripts/ci_install_tokenspeed.sh +++ b/scripts/ci_install_tokenspeed.sh @@ -22,7 +22,7 @@ fi # engine-watch workflow files an issue when this drifts) rather than # floating against ``main`` — upstream has renamed APIs before and the # gRPC servicer broke until we caught up. -TOKENSPEED_REF="${TOKENSPEED_REF:-0f68676069141857a605b8805c13131d9f53e901}" +TOKENSPEED_REF="${TOKENSPEED_REF:-57b49061a1aa258a7c27d8e25f32beb91a4ed896}" TOKENSPEED_REPO="${TOKENSPEED_REPO:-https://github.com/lightseekorg/tokenspeed.git}" TOKENSPEED_DIR="${TOKENSPEED_DIR:-/tmp/tokenspeed-src}"