diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 601e11c054..dee3f30789 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5941,7 +5941,7 @@ dependencies = [ [[package]] name = "vllm-proto" -version = "0.3.0" +version = "0.4.0" dependencies = [ "prost", "prost-types", diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml index cac268bc74..492cb5a06d 100644 --- a/rust/proto/Cargo.toml +++ b/rust/proto/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vllm-proto" -version = "0.3.0" +version = "0.4.0" edition.workspace = true license.workspace = true description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend" diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 46313f915a..4eaf6f92f4 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -160,6 +160,8 @@ message SequenceOutput { // Only present in final output for this sequence optional FinishInfo finish_info = 8; + // Only present in final output when sampling-mask capture is enabled. + repeated TokenIds sampling_mask = 10; } // Prompt info, returned in the first response diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index 964f0bf35a..023e3c3ffa 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -438,6 +438,7 @@ mod tests { finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, })), } } diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index 8dcf84dd98..035aef3cd4 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -58,6 +58,7 @@ fn finished() -> Finished { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index db3559c44f..9cba9f0aa2 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -806,6 +806,7 @@ fn decoded_completion_stream( finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, })), } }); @@ -817,6 +818,7 @@ fn decoded_completion_stream( finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }); events.push(DecodedTextEvent::TextDelta { decoded: vllm_text::DecodedText::unattributed(chunk.delta), diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index 646c106a04..72d852ba2c 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -mod array; +pub(crate) mod array; #[cfg(test)] mod tests; mod wire; diff --git a/rust/src/engine-core-client/src/protocol/logprobs/array.rs b/rust/src/engine-core-client/src/protocol/logprobs/array.rs index 4adf1ae975..919ef80f46 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/array.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/array.rs @@ -50,7 +50,7 @@ pub(super) fn decode_array2_u32( }) } -pub(super) fn decode_array1_u32( +pub(crate) fn decode_array1_u32( value: WireNdArray, field: &str, frames: &[Bytes], diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs index 8f4ee06127..742fed59bd 100644 --- a/rust/src/engine-core-client/src/protocol/mod.rs +++ b/rust/src/engine-core-client/src/protocol/mod.rs @@ -30,6 +30,7 @@ pub mod multimodal; pub mod output; pub mod request; pub mod sampling; +pub mod sampling_mask; mod serde_utils; pub mod stats; pub mod structured_outputs; diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index c8a8705435..a45cbe3731 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -15,6 +15,7 @@ use super::serde_utils::AllowTrailingFields; use super::utility::UtilityOutput; use crate::error::{Error, Result, ext_value_decode}; use crate::protocol::logprobs::MaybeWireLogprobs; +use crate::protocol::sampling_mask::MaybeWireSamplingMask; use crate::protocol::stats::{PrefillStats, SchedulerStats}; use crate::protocol::{OpaqueValue, decode_msgpack}; @@ -127,7 +128,7 @@ pub struct EngineCoreOutput { #[serde(default)] pub mm_cache_miss_hashes: Option>, #[serde(default)] - pub new_sampling_mask: Option, + pub new_sampling_mask: Option, /// Per-request speculative-decoding acceptance metrics, set on the final /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here; /// the Rust frontend does not yet surface it in responses. @@ -150,6 +151,9 @@ impl EngineCoreOutput { self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take()) .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors")) .transpose()?; + self.new_sampling_mask = (self.new_sampling_mask.take()) + .map(|value| value.resolve(frames, "new_sampling_mask")) + .transpose()?; Ok(()) } } diff --git a/rust/src/engine-core-client/src/protocol/sampling_mask.rs b/rust/src/engine-core-client/src/protocol/sampling_mask.rs new file mode 100644 index 0000000000..7f5c4592a7 --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/sampling_mask.rs @@ -0,0 +1,237 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::ops::{Deref, DerefMut}; + +use bytes::Bytes; +use enum_as_inner::EnumAsInner; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_tuple::{Deserialize_tuple, Serialize_tuple}; + +use crate::error::{Error, Result, bail_ext_value_decode}; +use crate::protocol::logprobs::array::decode_array1_u32; +use crate::protocol::tensor::WireNdArray; + +/// Decoded sampling-mask support sets, one row per generated token position. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct SamplingMask { + pub rows: Vec>, +} + +/// Python `SamplingMaskLists` tuple before ndarray raw views are resolved. +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)] +pub struct WireSamplingMask { + token_ids: WireNdArray, + #[serde(default)] + offsets: Option, + #[serde(default)] + cu_num_generated_tokens: Option>, +} + +/// Sampling-mask field while it transitions from Python wire data to rows. +#[derive(Debug, Clone, PartialEq, EnumAsInner)] +pub enum MaybeWireSamplingMask { + Wire(Box), + Direct(SamplingMask), +} + +impl Deref for MaybeWireSamplingMask { + type Target = SamplingMask; + + fn deref(&self) -> &Self::Target { + match self { + Self::Wire(_) => panic!("sampling mask is still in wire format"), + Self::Direct(value) => value, + } + } +} + +impl DerefMut for MaybeWireSamplingMask { + fn deref_mut(&mut self) -> &mut Self::Target { + match self { + Self::Wire(_) => panic!("sampling mask is still in wire format"), + Self::Direct(value) => value, + } + } +} + +impl<'de> Deserialize<'de> for MaybeWireSamplingMask { + fn deserialize(deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + WireSamplingMask::deserialize(deserializer).map(|value| Self::Wire(Box::new(value))) + } +} + +impl Serialize for MaybeWireSamplingMask { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + match self { + Self::Wire(value) => value.serialize(serializer), + Self::Direct(value) => WireSamplingMask::from_direct(value) + .map_err(serde::ser::Error::custom)? + .serialize(serializer), + } + } +} + +impl MaybeWireSamplingMask { + pub(super) fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result { + match self { + Self::Direct(value) => Ok(Self::Direct(value)), + Self::Wire(value) => value.resolve(frames, field_prefix).map(Self::Direct), + } + } +} + +impl WireSamplingMask { + fn from_direct(value: &SamplingMask) -> std::result::Result { + let [row] = value.rows.as_slice() else { + return Err(format!( + "per-request engine-core sampling mask must contain exactly one row, got {}", + value.rows.len() + )); + }; + let token_ids = row.iter().map(|&token_id| i64::from(token_id)).collect::>(); + Ok(Self { + token_ids: WireNdArray::from_i64(vec![token_ids.len()], token_ids)?, + offsets: None, + cu_num_generated_tokens: None, + }) + } + + fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result { + if let Some(indices) = self.cu_num_generated_tokens { + bail_ext_value_decode!( + "{field_prefix}.cu_num_generated_tokens: expected None for per-request engine-core sampling-mask payload, got {indices:?}" + ); + } + + let token_ids = + decode_array1_u32(self.token_ids, &format!("{field_prefix}.token_ids"), frames)?; + let Some(offsets) = self.offsets else { + return Ok(SamplingMask { + rows: vec![token_ids], + }); + }; + let offsets = decode_array1_u32(offsets, &format!("{field_prefix}.offsets"), frames)?; + if offsets.first().copied() != Some(0) { + bail_ext_value_decode!("{field_prefix}.offsets: first offset must be zero"); + } + if offsets.windows(2).any(|pair| pair[0] > pair[1]) { + bail_ext_value_decode!("{field_prefix}.offsets: offsets must be nondecreasing"); + } + if offsets.last().copied().map(|value| value as usize) != Some(token_ids.len()) { + bail_ext_value_decode!( + "{field_prefix}.offsets: final offset {:?} does not match token_ids length {}", + offsets.last(), + token_ids.len() + ); + } + + let rows = offsets + .windows(2) + .map(|pair| token_ids[pair[0] as usize..pair[1] as usize].to_vec()) + .collect(); + Ok(SamplingMask { rows }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::dtype::{NumpyDtype, TensorDtype}; + use crate::protocol::output::{ + EngineCoreOutput, RequestBatchOutputs, decode_engine_core_outputs, + }; + use crate::protocol::tensor::WireArrayData; + use crate::protocol::{decode_msgpack, encode_msgpack}; + + #[test] + fn decodes_single_position_without_offsets() { + let wire = WireSamplingMask { + token_ids: WireNdArray::from_i64(vec![3], vec![2, 12, 16]).unwrap(), + offsets: None, + cu_num_generated_tokens: None, + }; + let encoded = encode_msgpack(&wire).unwrap(); + let decoded: MaybeWireSamplingMask = decode_msgpack(&encoded).unwrap(); + let decoded = decoded.resolve(&[], "new_sampling_mask").unwrap(); + + assert_eq!(decoded.into_direct().unwrap().rows, vec![vec![2, 12, 16]]); + } + + #[test] + fn decodes_csr_rows_and_rejects_bad_terminal_offset() { + let wire = WireSamplingMask { + token_ids: WireNdArray::from_i64(vec![4], vec![2, 12, 16, 18]).unwrap(), + offsets: Some(WireNdArray::from_i64(vec![3], vec![0, 2, 4]).unwrap()), + cu_num_generated_tokens: None, + }; + let decoded = MaybeWireSamplingMask::Wire(Box::new(wire)) + .resolve(&[], "new_sampling_mask") + .unwrap(); + assert_eq!( + decoded.into_direct().unwrap().rows, + vec![vec![2, 12], vec![16, 18]] + ); + + let malformed = WireSamplingMask { + token_ids: WireNdArray::from_i64(vec![2], vec![2, 12]).unwrap(), + offsets: Some(WireNdArray::from_i64(vec![2], vec![0, 1]).unwrap()), + cu_num_generated_tokens: None, + }; + let error = MaybeWireSamplingMask::Wire(Box::new(malformed)) + .resolve(&[], "new_sampling_mask") + .unwrap_err(); + assert!(error.to_string().contains("final offset")); + } + + #[test] + fn decodes_sampling_mask_from_multipart_aux_frames() { + let output = EngineCoreOutput { + request_id: "req-multipart-mask".to_string(), + new_token_ids: vec![16, 18], + new_sampling_mask: Some(MaybeWireSamplingMask::Wire(Box::new(WireSamplingMask { + token_ids: WireNdArray { + dtype: NumpyDtype::little(TensorDtype::I32), + shape: vec![4], + data: WireArrayData::AuxIndex(1), + }, + offsets: Some(WireNdArray { + dtype: NumpyDtype::little(TensorDtype::I64), + shape: vec![3], + data: WireArrayData::AuxIndex(2), + }), + cu_num_generated_tokens: None, + }))), + ..Default::default() + }; + let primary = encode_msgpack(&crate::protocol::output::EngineCoreOutputs::from( + RequestBatchOutputs { + outputs: vec![output], + ..Default::default() + }, + )) + .unwrap(); + let token_ids = + [2_i32, 12, 16, 18].into_iter().flat_map(i32::to_le_bytes).collect::>(); + let offsets = [0_i64, 2, 4].into_iter().flat_map(i64::to_le_bytes).collect::>(); + + let decoded = decode_engine_core_outputs(&[ + Bytes::from(primary), + Bytes::from(token_ids), + Bytes::from(offsets), + ]) + .unwrap(); + let mask = decoded.as_request_batch().unwrap().outputs[0] + .new_sampling_mask + .as_ref() + .unwrap(); + + assert_eq!(mask.rows, vec![vec![2, 12], vec![16, 18]]); + } +} diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 739c34d6b7..eb8e713fb4 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2712,15 +2712,15 @@ fn python_msgpack_fixtures_match_rust_encoding() { decode_value(&rmp_serde::to_vec_named(&expected_multimodal_request.mm_features).unwrap()); assert_eq!(python_mm_features, rust_mm_features); - let decoded_sampling_mask_outputs: EngineCoreOutputs = - rmp_serde::from_slice(&sampling_mask_outputs_bytes).unwrap(); + let decoded_sampling_mask_outputs = + decode_engine_core_outputs(&[bytes::Bytes::from(sampling_mask_outputs_bytes)]).unwrap(); let sampling_mask_output = &decoded_sampling_mask_outputs.as_request_batch().unwrap().outputs[0]; assert!(sampling_mask_output.mm_cache_miss_hashes.is_none()); - assert!(matches!( - sampling_mask_output.new_sampling_mask.as_ref(), - Some(Value::Array(fields)) if fields.len() == 3 - )); + assert_eq!( + sampling_mask_output.new_sampling_mask.as_ref().unwrap().rows, + vec![vec![2, 12, 16, 17, 18]] + ); let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap(); // Match msgspec's base-schema result for the same extended Python message. diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 5013f9e2a5..a19dc9fdaa 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -234,14 +234,10 @@ sampling_mask_wire = [ [ ", + /// Sampling support sets aligned one-to-one with generated token positions. + pub sampling_mask: Option, } /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for @@ -155,6 +158,8 @@ pub struct GenerateOutput { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Sampling support sets aligned one-to-one with `token_ids`. + pub sampling_mask: Option, } impl GenerateOutput { @@ -202,6 +207,7 @@ impl GenerateOutput { cached_token_count: 0, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, } } } @@ -271,6 +277,16 @@ impl Stream for GenerateOutputStream { } let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap()); + let sampling_mask = raw.new_sampling_mask.map(|value| value.into_direct().unwrap()); + if let Some(mask) = sampling_mask.as_ref() + && mask.rows.len() != raw.new_token_ids.len() + { + return Poll::Ready(Some(Err(crate::Error::SamplingMaskTokenCountMismatch { + request_id: raw.request_id, + token_count: raw.new_token_ids.len(), + row_count: mask.rows.len(), + }))); + } let cached_token_count = raw .prefill_stats .as_ref() @@ -291,6 +307,7 @@ impl Stream for GenerateOutputStream { cached_token_count, kv_transfer_params: raw.kv_transfer_params, ec_transfer_params: raw.ec_transfer_params, + sampling_mask, }; Poll::Ready(Some(Ok(output))) @@ -341,6 +358,7 @@ impl> + Send> T { while let Some(output) = stream.next().await.transpose()? { cached_token_count = cached_token_count.max(output.cached_token_count); + let sampling_mask = output.sampling_mask; if let Some(info) = output.prompt_info { if prompt_token_ids.is_none() { prompt_token_ids = Some(info.prompt_token_ids.to_vec()); @@ -359,6 +377,9 @@ impl> + Send> T { existing.logprobs = Some(step_logprobs); } } + if let Some(mut mask) = sampling_mask { + existing.sampling_mask.get_or_insert_default().rows.append(&mut mask.rows); + } } else { collected = Some(CollectedGenerateOutput { request_id: output.request_id, @@ -374,6 +395,7 @@ impl> + Send> T { }, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask, }); } @@ -387,6 +409,15 @@ impl> + Send> T { }; collected.kv_transfer_params = output.kv_transfer_params; collected.ec_transfer_params = output.ec_transfer_params; + if let Some(mask) = collected.sampling_mask.as_ref() + && mask.rows.len() != collected.token_ids.len() + { + return Err(crate::Error::SamplingMaskTokenCountMismatch { + request_id: collected.request_id, + token_count: collected.token_ids.len(), + row_count: mask.rows.len(), + }); + } return Ok(collected); } } diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 2d09aff889..09b5199767 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeSet; -use std::sync::Once; +use std::sync::{Arc, Once}; use std::time::Duration; use futures::StreamExt as _; @@ -18,6 +18,7 @@ use vllm_engine_core_client::protocol::output::{ }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::sampling_mask::{MaybeWireSamplingMask, SamplingMask}; use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; @@ -326,6 +327,11 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { num_local_cached_tokens: 1, ..Default::default() }), + new_sampling_mask: Some(MaybeWireSamplingMask::Direct( + SamplingMask { + rows: vec![vec![1, 33, 99]], + }, + )), ..request_output_with_logprobs( &request.request_id, vec![33], @@ -334,15 +340,22 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { Some(prompt_logprobs()), ) }, - request_output_with_logprobs_and_kv( - &request.request_id, - vec![44], - Some(EngineCoreFinishReason::Stop), - Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)), - None, - Some(serde_json::json!({"connector": "x"})), - None, - ), + EngineCoreOutput { + new_sampling_mask: Some(MaybeWireSamplingMask::Direct( + SamplingMask { + rows: vec![vec![2, 44, 88]], + }, + )), + ..request_output_with_logprobs_and_kv( + &request.request_id, + vec![44], + Some(EngineCoreFinishReason::Stop), + Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)), + None, + Some(serde_json::json!({"connector": "x"})), + None, + ) + }, ], ..Default::default() } @@ -375,6 +388,42 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { collected.kv_transfer_params, Some(serde_json::json!({"connector": "x"})) ); + assert_eq!( + collected.sampling_mask, + Some(SamplingMask { + rows: vec![vec![1, 33, 99], vec![2, 44, 88]], + }) + ); +} + +#[tokio::test] +async fn collect_output_rejects_partial_sampling_mask() { + let output = vllm_llm::GenerateOutput { + request_id: "req-partial-mask".to_string(), + prompt_info: Some(GeneratePromptInfo { + prompt_token_ids: Arc::from([11_u32, 22]), + prompt_logprobs: None, + }), + token_ids: vec![33, 44], + logprobs: None, + finish_reason: Some(FinishReason::Length), + cached_token_count: 0, + kv_transfer_params: None, + ec_transfer_params: None, + sampling_mask: Some(SamplingMask { + rows: vec![vec![1, 33, 99]], + }), + }; + + let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); + assert!(matches!( + error, + Error::SamplingMaskTokenCountMismatch { + token_count: 2, + row_count: 1, + .. + } + )); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index caccbdefc5..8990ba2f3b 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -19,6 +19,7 @@ itertools.workspace = true libc.workspace = true llm-multimodal.workspace = true openssl.workspace = true +prost.workspace = true prost-types.workspace = true rmp-serde.workspace = true rmpv.workspace = true diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 5f6ce5d42b..0f28e147ef 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -316,6 +316,11 @@ pub fn to_sequence_output( _ => (vec![], vec![], vec![]), }; + let sampling_mask = finished + .and_then(|finished| finished.sampling_mask.as_ref()) + .map(|mask| mask.rows.iter().map(|row| pb::TokenIds { ids: row.clone() }).collect()) + .unwrap_or_default(); + Ok(pb::SequenceOutput { index: 0, // TODO: multi-sequence (n > 1) not supported text: if opts.output_text { @@ -333,6 +338,7 @@ pub fn to_sequence_output( ranks: rank_values, candidate_tokens: candidates, finish_info, + sampling_mask, }) } @@ -524,6 +530,7 @@ impl ResponseOpts { mod tests { use prost::Message as _; use vllm_engine_core_client::protocol::output::StopReason; + use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_text::{ FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params, }; @@ -720,6 +727,7 @@ mod tests { finish_reason: reason, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, } } @@ -804,4 +812,40 @@ mod tests { assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32); assert_eq!(finish.stop_reason, Some(PbStopReason::EosTokenId(30))); } + + #[test] + fn sequence_output_only_carries_sampling_mask_on_terminal_output() { + let mut fin = finished(FinishReason::Length); + fin.sampling_mask = Some(SamplingMask { + rows: vec![vec![1, 10], vec![2, 20]], + }); + + let terminal = + to_sequence_output("", &[10, 20], None, Some(&fin), &ResponseOpts::default()) + .expect("terminal output"); + let intermediate = to_sequence_output("", &[9], None, None, &ResponseOpts::default()) + .expect("intermediate output"); + + assert_eq!( + terminal.sampling_mask, + vec![ + pb::TokenIds { ids: vec![1, 10] }, + pb::TokenIds { ids: vec![2, 20] } + ] + ); + assert!(intermediate.sampling_mask.is_empty()); + } + + #[test] + fn sampling_mask_uses_additive_proto_tag_ten() { + let response = pb::SequenceOutput { + sampling_mask: vec![pb::TokenIds { ids: vec![7, 8] }], + ..Default::default() + }; + + assert_eq!( + response.encode_to_vec(), + vec![0x52, 0x04, 0x0a, 0x02, 0x07, 0x08] + ); + } } diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 72943b6af0..d56415c747 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -257,6 +257,7 @@ impl pb::inference_server::Inference for InferenceServiceImpl { finish_reason: collected.finish_reason, kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, + sampling_mask: collected.sampling_mask, }; let outputs = convert::to_sequence_output( diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index dfb7e6cbb1..f82059dffd 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -481,6 +481,7 @@ mod tests { cached_token_count: 0, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), Ok(GenerateOutput { request_id: String::new(), @@ -494,6 +495,7 @@ mod tests { cached_token_count: 2, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), ]); @@ -575,6 +577,7 @@ mod tests { cached_token_count: 0, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, } } @@ -658,6 +661,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, prompt_token_ids: vec![10, 20], + sampling_mask: None, }; let response = collect_generate( @@ -744,6 +748,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, prompt_token_ids: vec![10, 20], + sampling_mask: None, }; let response = collect_generate( @@ -776,6 +781,7 @@ mod tests { }, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, prompt_token_ids, }; diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index a122afbcd9..06874a43bd 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -694,6 +694,7 @@ mod tests { ))), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -806,6 +807,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -858,6 +860,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -913,6 +916,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -985,6 +989,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); @@ -1059,6 +1064,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), )), ]); diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index a438c8e267..d3dc422656 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -9,6 +9,7 @@ use serde::{Deserialize, Serialize}; use tracing::{Level, debug, trace}; use vllm_engine_core_client::AbortCause; use vllm_engine_core_client::protocol::output::StopReason; +use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_llm::{FinishReason, GenerateOutput, TokenUsage}; use vllm_tokenizer::{DecodedText, DynTokenizer, IncrementalDecoder}; @@ -50,6 +51,8 @@ pub struct Finished { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Sampling support sets aligned with all generated token positions. + pub sampling_mask: Option, } /// Sample metadata emitted by one engine output update. @@ -120,6 +123,7 @@ pub async fn decoded_text_event_stream( let mut sampled_token_ids = Vec::new(); let mut output_token_count: usize = 0; let mut sampled_logprobs: Option = None; + let mut sampling_mask: Option = None; while let Some(next) = raw_stream.next().await { let output = next?; @@ -213,6 +217,7 @@ pub async fn decoded_text_event_stream( let mut new_token_ids = output.token_ids; let mut new_logprobs = output.logprobs; + let mut new_sampling_mask = output.sampling_mask; // Trim tokens and logprobs if we matched stop string. if let Some(num_tokens) = truncate_tokens_to { @@ -220,9 +225,15 @@ pub async fn decoded_text_event_stream( if let Some(logprobs) = &mut new_logprobs { logprobs.positions.truncate(num_tokens); } + if let Some(mask) = &mut new_sampling_mask { + mask.rows.truncate(num_tokens); + } } output_token_count += new_token_ids.len(); + if let Some(mut mask) = new_sampling_mask { + sampling_mask.get_or_insert_default().rows.append(&mut mask.rows); + } let decoded_logprobs = new_logprobs .as_ref() @@ -246,6 +257,16 @@ pub async fn decoded_text_event_stream( } if let Some(reason) = finish_reason { + if let Some(mask) = sampling_mask.as_ref() + && mask.rows.len() != output_token_count + { + return Err(vllm_llm::Error::SamplingMaskTokenCountMismatch { + request_id, + token_count: output_token_count, + row_count: mask.rows.len(), + } + .into()); + } // Flush any remaining buffered text. let (last_chunk, full_decoded) = decoder.flush(truncate_output_to)?; let text_len = full_decoded.text.len(); @@ -293,6 +314,7 @@ pub async fn decoded_text_event_stream( finish_reason: reason, kv_transfer_params, ec_transfer_params, + sampling_mask, })), }) .await; @@ -354,6 +376,7 @@ mod tests { use futures::{Stream, StreamExt as _, stream}; use vllm_engine_core_client::AbortCause; + use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob}; use vllm_llm::GenerateOutput; use vllm_tokenizer::{TokenAnchor, TokenAttribution, test_utils::TestTokenizer}; @@ -384,6 +407,73 @@ mod tests { s.bytes().map(u32::from).collect() } + #[tokio::test] + async fn sampling_masks_survive_text_decode_and_collection() { + let prompt: Arc<[u32]> = Arc::from([]); + let mut first = GenerateOutput::for_test(Some(prompt), vec![b'a' as u32], None); + first.sampling_mask = Some(SamplingMask { + rows: vec![vec![1, b'a' as u32]], + }); + let mut terminal = + GenerateOutput::for_test(None, vec![b'b' as u32], Some(FinishReason::Length)); + terminal.sampling_mask = Some(SamplingMask { + rows: vec![vec![2, b'b' as u32]], + }); + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new()); + + let collected = decoded_text_event_stream( + "test-mask".into(), + tokenizer, + stream::iter([Ok(first), Ok(terminal)]), + TextDecodeOptions::default(), + false, + ) + .collect_output() + .await + .unwrap(); + + assert_eq!( + collected.sampling_mask, + Some(SamplingMask { + rows: vec![vec![1, b'a' as u32], vec![2, b'b' as u32]], + }) + ); + } + + #[tokio::test] + async fn sampling_mask_mismatch_reuses_llm_error() { + let prompt: Arc<[u32]> = Arc::from([]); + let mut terminal = GenerateOutput::for_test( + Some(prompt), + vec![b'a' as u32, b'b' as u32], + Some(FinishReason::Length), + ); + terminal.sampling_mask = Some(SamplingMask { + rows: vec![vec![1, b'a' as u32]], + }); + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new()); + + let error = decoded_text_event_stream( + "test-mask-mismatch".into(), + tokenizer, + stream::iter([Ok(terminal)]), + TextDecodeOptions::default(), + false, + ) + .collect_output() + .await + .unwrap_err(); + + assert!(matches!( + error, + crate::Error::Llm(vllm_llm::Error::SamplingMaskTokenCountMismatch { + token_count: 2, + row_count: 1, + .. + }) + )); + } + fn opts(stop: &[&str], min_tokens: u32) -> TextDecodeOptions { TextDecodeOptions { stop_strings: Some(stop.iter().map(|s| s.to_string()).collect()), @@ -472,6 +562,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), ), ], @@ -503,6 +594,7 @@ mod tests { ))), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }), ), ], @@ -576,6 +668,60 @@ mod tests { assert!(output.finish_reason.is_stop()); } + #[tokio::test] + async fn stream_stop_string_mid_chunk_keeps_token_logprob_mask_alignment() { + let token_ids = ascii_tokens("abcxyz"); + let mut raw = GenerateOutput::for_test( + Some(Arc::from([])), + token_ids.clone(), + Some(FinishReason::Length), + ); + raw.logprobs = Some(Logprobs { + positions: token_ids + .iter() + .map(|&token_id| PositionLogprobs { + entries: vec![TokenLogprob { + token_id, + logprob: -0.25, + rank: 1, + }], + }) + .collect(), + }); + raw.sampling_mask = Some(SamplingMask { + rows: token_ids.iter().map(|&token_id| vec![token_id, token_id + 100]).collect(), + }); + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new()); + + let output = decoded_text_event_stream( + "test-mid-chunk-stop".into(), + tokenizer, + stream::iter([Ok(raw)]), + opts(&["cx"], 0), + false, + ) + .collect_output() + .await + .unwrap(); + + assert_eq!(output.text, "ab"); + assert!(output.finish_reason.is_stop()); + assert_eq!(output.token_ids, ascii_tokens("abcx")); + assert_eq!(output.logprobs.as_ref().unwrap().positions.len(), 4); + assert_eq!( + output.sampling_mask, + Some(SamplingMask { + rows: vec![ + vec![b'a' as u32, b'a' as u32 + 100], + vec![b'b' as u32, b'b' as u32 + 100], + vec![b'c' as u32, b'c' as u32 + 100], + vec![b'x' as u32, b'x' as u32 + 100], + ], + }) + ); + assert_eq!(output.usage.output_token_count, 4); + } + #[tokio::test] async fn stream_stop_string_at_end() { let output = run_to_completion(ascii_tokens("abcxyz"), opts(&["xyz"], 0)).await; diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index a415bb8900..1293ac6025 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -17,6 +17,7 @@ mod logprobs; use std::sync::Arc; use futures::{StreamExt as _, pin_mut}; +use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use crate::{Error, FinishReason, Result, TextOutputStream}; @@ -35,6 +36,8 @@ pub struct CollectedTextOutput { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Sampling support sets aligned with generated token positions. + pub sampling_mask: Option, } #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")] @@ -90,6 +93,7 @@ impl T { usage: vllm_llm::TokenUsage::default(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, }) }; @@ -99,6 +103,7 @@ impl T { collected.usage = finished.usage; collected.kv_transfer_params = finished.kv_transfer_params; collected.ec_transfer_params = finished.ec_transfer_params; + collected.sampling_mask = finished.sampling_mask; return Ok(collected); } } @@ -173,6 +178,7 @@ mod tests { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, })), }), ]); @@ -295,6 +301,7 @@ mod tests { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + sampling_mask: None, })), }), ]);