@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 95d723700b

[Rust Frontend] Return sampling masks over gRPC (#56777)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Biswa Panda <biswa.panda@gmail.com>

Biswa Panda committed Sep 18, 2026 at 01:26 UTC 95d723700b1df07b20263c696c0342b92c188bfe
23 files changed +570 -27
rust/Cargo.lock
+1 -1
@@ -5941,7 +5941,7 @@ dependencies = [
5941
5942 [[package]]
5943 name = "vllm-proto"
5944 -version = "0.3.0"
5944 +version = "0.4.0"
5945 dependencies = [
5946 "prost",
5947 "prost-types",
rust/proto/Cargo.toml
+1 -1
@@ -1,6 +1,6 @@
1 [package]
2 name = "vllm-proto"
3 -version = "0.3.0"
3 +version = "0.4.0"
4 edition.workspace = true
5 license.workspace = true
6 description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend"
rust/proto/inference.proto
+2
@@ -160,6 +160,8 @@ message SequenceOutput {
160
161 // Only present in final output for this sequence
162 optional FinishInfo finish_info = 8;
163 + // Only present in final output when sampling-mask capture is enabled.
164 + repeated TokenIds sampling_mask = 10;
165 }
166
167 // Prompt info, returned in the first response
rust/src/chat/src/output/default/unified.rs
+1
@@ -438,6 +438,7 @@ mod tests {
438 finish_reason: crate::FinishReason::Stop(None),
439 kv_transfer_params: None,
440 ec_transfer_params: None,
441 + sampling_mask: None,
442 })),
443 }
444 }
rust/src/chat/src/output/harmony/tests.rs
+1
@@ -58,6 +58,7 @@ fn finished() -> Finished {
58 finish_reason: FinishReason::stop_eos(),
59 kv_transfer_params: None,
60 ec_transfer_params: None,
61 + sampling_mask: None,
62 }
63 }
64
rust/src/chat/tests/roundtrip.rs
+2
@@ -806,6 +806,7 @@ fn decoded_completion_stream(
806 finish_reason: FinishReason::stop_eos(),
807 kv_transfer_params: None,
808 ec_transfer_params: None,
809 + sampling_mask: None,
810 })),
811 }
812 });
@@ -817,6 +818,7 @@ fn decoded_completion_stream(
818 finish_reason: FinishReason::stop_eos(),
819 kv_transfer_params: None,
820 ec_transfer_params: None,
821 + sampling_mask: None,
822 });
823 events.push(DecodedTextEvent::TextDelta {
824 decoded: vllm_text::DecodedText::unattributed(chunk.delta),
rust/src/engine-core-client/src/protocol/logprobs.rs
+1 -1
@@ -1,7 +1,7 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -mod array;
4 +pub(crate) mod array;
5 #[cfg(test)]
6 mod tests;
7 mod wire;
rust/src/engine-core-client/src/protocol/logprobs/array.rs
+1 -1
@@ -50,7 +50,7 @@ pub(super) fn decode_array2_u32(
50 })
51 }
52
53 -pub(super) fn decode_array1_u32(
53 +pub(crate) fn decode_array1_u32(
54 value: WireNdArray,
55 field: &str,
56 frames: &[Bytes],
rust/src/engine-core-client/src/protocol/mod.rs
+1
@@ -30,6 +30,7 @@ pub mod multimodal;
30 pub mod output;
31 pub mod request;
32 pub mod sampling;
33 +pub mod sampling_mask;
34 mod serde_utils;
35 pub mod stats;
36 pub mod structured_outputs;
rust/src/engine-core-client/src/protocol/output.rs
+5 -1
@@ -15,6 +15,7 @@ use super::serde_utils::AllowTrailingFields;
15 use super::utility::UtilityOutput;
16 use crate::error::{Error, Result, ext_value_decode};
17 use crate::protocol::logprobs::MaybeWireLogprobs;
18 +use crate::protocol::sampling_mask::MaybeWireSamplingMask;
19 use crate::protocol::stats::{PrefillStats, SchedulerStats};
20 use crate::protocol::{OpaqueValue, decode_msgpack};
21
@@ -127,7 +128,7 @@ pub struct EngineCoreOutput {
128 #[serde(default)]
129 pub mm_cache_miss_hashes: Option<Vec<String>>,
130 #[serde(default)]
130 - pub new_sampling_mask: Option<OpaqueValue>,
131 + pub new_sampling_mask: Option<MaybeWireSamplingMask>,
132 /// Per-request speculative-decoding acceptance metrics, set on the final
133 /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here;
134 /// the Rust frontend does not yet surface it in responses.
@@ -150,6 +151,9 @@ impl EngineCoreOutput {
151 self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
152 .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
153 .transpose()?;
154 + self.new_sampling_mask = (self.new_sampling_mask.take())
155 + .map(|value| value.resolve(frames, "new_sampling_mask"))
156 + .transpose()?;
157 Ok(())
158 }
159 }
rust/src/engine-core-client/src/protocol/sampling_mask.rs new
+237
@@ -0,0 +1,237 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::ops::{Deref, DerefMut};
5 +
6 +use bytes::Bytes;
7 +use enum_as_inner::EnumAsInner;
8 +use serde::{Deserialize, Deserializer, Serialize};
9 +use serde_tuple::{Deserialize_tuple, Serialize_tuple};
10 +
11 +use crate::error::{Error, Result, bail_ext_value_decode};
12 +use crate::protocol::logprobs::array::decode_array1_u32;
13 +use crate::protocol::tensor::WireNdArray;
14 +
15 +/// Decoded sampling-mask support sets, one row per generated token position.
16 +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
17 +pub struct SamplingMask {
18 + pub rows: Vec<Vec<u32>>,
19 +}
20 +
21 +/// Python `SamplingMaskLists` tuple before ndarray raw views are resolved.
22 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
23 +pub struct WireSamplingMask {
24 + token_ids: WireNdArray,
25 + #[serde(default)]
26 + offsets: Option<WireNdArray>,
27 + #[serde(default)]
28 + cu_num_generated_tokens: Option<Vec<usize>>,
29 +}
30 +
31 +/// Sampling-mask field while it transitions from Python wire data to rows.
32 +#[derive(Debug, Clone, PartialEq, EnumAsInner)]
33 +pub enum MaybeWireSamplingMask {
34 + Wire(Box<WireSamplingMask>),
35 + Direct(SamplingMask),
36 +}
37 +
38 +impl Deref for MaybeWireSamplingMask {
39 + type Target = SamplingMask;
40 +
41 + fn deref(&self) -> &Self::Target {
42 + match self {
43 + Self::Wire(_) => panic!("sampling mask is still in wire format"),
44 + Self::Direct(value) => value,
45 + }
46 + }
47 +}
48 +
49 +impl DerefMut for MaybeWireSamplingMask {
50 + fn deref_mut(&mut self) -> &mut Self::Target {
51 + match self {
52 + Self::Wire(_) => panic!("sampling mask is still in wire format"),
53 + Self::Direct(value) => value,
54 + }
55 + }
56 +}
57 +
58 +impl<'de> Deserialize<'de> for MaybeWireSamplingMask {
59 + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
60 + where
61 + D: Deserializer<'de>,
62 + {
63 + WireSamplingMask::deserialize(deserializer).map(|value| Self::Wire(Box::new(value)))
64 + }
65 +}
66 +
67 +impl Serialize for MaybeWireSamplingMask {
68 + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
69 + where
70 + S: serde::Serializer,
71 + {
72 + match self {
73 + Self::Wire(value) => value.serialize(serializer),
74 + Self::Direct(value) => WireSamplingMask::from_direct(value)
75 + .map_err(serde::ser::Error::custom)?
76 + .serialize(serializer),
77 + }
78 + }
79 +}
80 +
81 +impl MaybeWireSamplingMask {
82 + pub(super) fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result<Self> {
83 + match self {
84 + Self::Direct(value) => Ok(Self::Direct(value)),
85 + Self::Wire(value) => value.resolve(frames, field_prefix).map(Self::Direct),
86 + }
87 + }
88 +}
89 +
90 +impl WireSamplingMask {
91 + fn from_direct(value: &SamplingMask) -> std::result::Result<Self, String> {
92 + let [row] = value.rows.as_slice() else {
93 + return Err(format!(
94 + "per-request engine-core sampling mask must contain exactly one row, got {}",
95 + value.rows.len()
96 + ));
97 + };
98 + let token_ids = row.iter().map(|&token_id| i64::from(token_id)).collect::<Vec<_>>();
99 + Ok(Self {
100 + token_ids: WireNdArray::from_i64(vec![token_ids.len()], token_ids)?,
101 + offsets: None,
102 + cu_num_generated_tokens: None,
103 + })
104 + }
105 +
106 + fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result<SamplingMask> {
107 + if let Some(indices) = self.cu_num_generated_tokens {
108 + bail_ext_value_decode!(
109 + "{field_prefix}.cu_num_generated_tokens: expected None for per-request engine-core sampling-mask payload, got {indices:?}"
110 + );
111 + }
112 +
113 + let token_ids =
114 + decode_array1_u32(self.token_ids, &format!("{field_prefix}.token_ids"), frames)?;
115 + let Some(offsets) = self.offsets else {
116 + return Ok(SamplingMask {
117 + rows: vec![token_ids],
118 + });
119 + };
120 + let offsets = decode_array1_u32(offsets, &format!("{field_prefix}.offsets"), frames)?;
121 + if offsets.first().copied() != Some(0) {
122 + bail_ext_value_decode!("{field_prefix}.offsets: first offset must be zero");
123 + }
124 + if offsets.windows(2).any(|pair| pair[0] > pair[1]) {
125 + bail_ext_value_decode!("{field_prefix}.offsets: offsets must be nondecreasing");
126 + }
127 + if offsets.last().copied().map(|value| value as usize) != Some(token_ids.len()) {
128 + bail_ext_value_decode!(
129 + "{field_prefix}.offsets: final offset {:?} does not match token_ids length {}",
130 + offsets.last(),
131 + token_ids.len()
132 + );
133 + }
134 +
135 + let rows = offsets
136 + .windows(2)
137 + .map(|pair| token_ids[pair[0] as usize..pair[1] as usize].to_vec())
138 + .collect();
139 + Ok(SamplingMask { rows })
140 + }
141 +}
142 +
143 +#[cfg(test)]
144 +mod tests {
145 + use super::*;
146 + use crate::protocol::dtype::{NumpyDtype, TensorDtype};
147 + use crate::protocol::output::{
148 + EngineCoreOutput, RequestBatchOutputs, decode_engine_core_outputs,
149 + };
150 + use crate::protocol::tensor::WireArrayData;
151 + use crate::protocol::{decode_msgpack, encode_msgpack};
152 +
153 + #[test]
154 + fn decodes_single_position_without_offsets() {
155 + let wire = WireSamplingMask {
156 + token_ids: WireNdArray::from_i64(vec![3], vec![2, 12, 16]).unwrap(),
157 + offsets: None,
158 + cu_num_generated_tokens: None,
159 + };
160 + let encoded = encode_msgpack(&wire).unwrap();
161 + let decoded: MaybeWireSamplingMask = decode_msgpack(&encoded).unwrap();
162 + let decoded = decoded.resolve(&[], "new_sampling_mask").unwrap();
163 +
164 + assert_eq!(decoded.into_direct().unwrap().rows, vec![vec![2, 12, 16]]);
165 + }
166 +
167 + #[test]
168 + fn decodes_csr_rows_and_rejects_bad_terminal_offset() {
169 + let wire = WireSamplingMask {
170 + token_ids: WireNdArray::from_i64(vec![4], vec![2, 12, 16, 18]).unwrap(),
171 + offsets: Some(WireNdArray::from_i64(vec![3], vec![0, 2, 4]).unwrap()),
172 + cu_num_generated_tokens: None,
173 + };
174 + let decoded = MaybeWireSamplingMask::Wire(Box::new(wire))
175 + .resolve(&[], "new_sampling_mask")
176 + .unwrap();
177 + assert_eq!(
178 + decoded.into_direct().unwrap().rows,
179 + vec![vec![2, 12], vec![16, 18]]
180 + );
181 +
182 + let malformed = WireSamplingMask {
183 + token_ids: WireNdArray::from_i64(vec![2], vec![2, 12]).unwrap(),
184 + offsets: Some(WireNdArray::from_i64(vec![2], vec![0, 1]).unwrap()),
185 + cu_num_generated_tokens: None,
186 + };
187 + let error = MaybeWireSamplingMask::Wire(Box::new(malformed))
188 + .resolve(&[], "new_sampling_mask")
189 + .unwrap_err();
190 + assert!(error.to_string().contains("final offset"));
191 + }
192 +
193 + #[test]
194 + fn decodes_sampling_mask_from_multipart_aux_frames() {
195 + let output = EngineCoreOutput {
196 + request_id: "req-multipart-mask".to_string(),
197 + new_token_ids: vec![16, 18],
198 + new_sampling_mask: Some(MaybeWireSamplingMask::Wire(Box::new(WireSamplingMask {
199 + token_ids: WireNdArray {
200 + dtype: NumpyDtype::little(TensorDtype::I32),
201 + shape: vec![4],
202 + data: WireArrayData::AuxIndex(1),
203 + },
204 + offsets: Some(WireNdArray {
205 + dtype: NumpyDtype::little(TensorDtype::I64),
206 + shape: vec![3],
207 + data: WireArrayData::AuxIndex(2),
208 + }),
209 + cu_num_generated_tokens: None,
210 + }))),
211 + ..Default::default()
212 + };
213 + let primary = encode_msgpack(&crate::protocol::output::EngineCoreOutputs::from(
214 + RequestBatchOutputs {
215 + outputs: vec![output],
216 + ..Default::default()
217 + },
218 + ))
219 + .unwrap();
220 + let token_ids =
221 + [2_i32, 12, 16, 18].into_iter().flat_map(i32::to_le_bytes).collect::<Vec<_>>();
222 + let offsets = [0_i64, 2, 4].into_iter().flat_map(i64::to_le_bytes).collect::<Vec<_>>();
223 +
224 + let decoded = decode_engine_core_outputs(&[
225 + Bytes::from(primary),
226 + Bytes::from(token_ids),
227 + Bytes::from(offsets),
228 + ])
229 + .unwrap();
230 + let mask = decoded.as_request_batch().unwrap().outputs[0]
231 + .new_sampling_mask
232 + .as_ref()
233 + .unwrap();
234 +
235 + assert_eq!(mask.rows, vec![vec![2, 12], vec![16, 18]]);
236 + }
237 +}
rust/src/engine-core-client/src/tests/client.rs
+6 -6
@@ -2712,15 +2712,15 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2712 decode_value(&rmp_serde::to_vec_named(&expected_multimodal_request.mm_features).unwrap());
2713 assert_eq!(python_mm_features, rust_mm_features);
2714
2715 - let decoded_sampling_mask_outputs: EngineCoreOutputs =
2716 - rmp_serde::from_slice(&sampling_mask_outputs_bytes).unwrap();
2715 + let decoded_sampling_mask_outputs =
2716 + decode_engine_core_outputs(&[bytes::Bytes::from(sampling_mask_outputs_bytes)]).unwrap();
2717 let sampling_mask_output =
2718 &decoded_sampling_mask_outputs.as_request_batch().unwrap().outputs[0];
2719 assert!(sampling_mask_output.mm_cache_miss_hashes.is_none());
2720 - assert!(matches!(
2721 - sampling_mask_output.new_sampling_mask.as_ref(),
2722 - Some(Value::Array(fields)) if fields.len() == 3
2723 - ));
2720 + assert_eq!(
2721 + sampling_mask_output.new_sampling_mask.as_ref().unwrap().rows,
2722 + vec![vec![2, 12, 16, 17, 18]]
2723 + );
2724
2725 let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap();
2726 // Match msgspec's base-schema result for the same extended Python message.
rust/src/engine-core-client/src/tests/python_compat.py
+2 -6
@@ -234,14 +234,10 @@ sampling_mask_wire = [
234 [
235 "<i4",
236 [5],
237 - msgpack.ExtType(3, np.array([2, 12, 16, 17, 18], dtype=np.int32).tobytes()),
238 - ],
239 - [
240 - "<i8",
241 - [2],
242 - msgpack.ExtType(3, np.array([0, 5], dtype=np.int64).tobytes()),
237 + msgspec.msgpack.Ext(3, np.array([2, 12, 16, 17, 18], dtype=np.int32).tobytes()),
238 ],
239 None,
240 + None,
241 ]
242 outputs_with_sampling_mask = EngineCoreOutputs(
243 outputs=[
rust/src/llm/src/error.rs
+8
@@ -12,4 +12,12 @@ pub enum Error {
12 EmptyPromptTokenIds { request_id: String },
13 #[error("engine-core error")]
14 EngineCoreClient(#[from] vllm_engine_core_client::Error),
15 + #[error(
16 + "sampling mask for request `{request_id}` has {row_count} rows for {token_count} generated tokens"
17 + )]
18 + SamplingMaskTokenCountMismatch {
19 + request_id: String,
20 + token_count: usize,
21 + row_count: usize,
22 + },
23 }
rust/src/llm/src/output.rs
+31
@@ -12,6 +12,7 @@ use futures::{Stream, StreamExt as _, pin_mut};
12 use serde::{Deserialize, Serialize};
13 use vllm_engine_core_client::protocol::logprobs::Logprobs;
14 use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason};
15 +use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
16 use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
17
18 use crate::error::Result;
@@ -44,6 +45,8 @@ pub struct CollectedGenerateOutput {
45 /// Connector-specific encoder cache transfer parameters for disaggregated
46 /// serving.
47 pub ec_transfer_params: Option<serde_json::Value>,
48 + /// Sampling support sets aligned one-to-one with generated token positions.
49 + pub sampling_mask: Option<SamplingMask>,
50 }
51
52 /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for
@@ -155,6 +158,8 @@ pub struct GenerateOutput {
158 /// Connector-specific encoder cache transfer parameters for disaggregated
159 /// serving.
160 pub ec_transfer_params: Option<serde_json::Value>,
161 + /// Sampling support sets aligned one-to-one with `token_ids`.
162 + pub sampling_mask: Option<SamplingMask>,
163 }
164
165 impl GenerateOutput {
@@ -202,6 +207,7 @@ impl GenerateOutput {
207 cached_token_count: 0,
208 kv_transfer_params: None,
209 ec_transfer_params: None,
210 + sampling_mask: None,
211 }
212 }
213 }
@@ -271,6 +277,16 @@ impl Stream for GenerateOutputStream {
277 }
278
279 let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap());
280 + let sampling_mask = raw.new_sampling_mask.map(|value| value.into_direct().unwrap());
281 + if let Some(mask) = sampling_mask.as_ref()
282 + && mask.rows.len() != raw.new_token_ids.len()
283 + {
284 + return Poll::Ready(Some(Err(crate::Error::SamplingMaskTokenCountMismatch {
285 + request_id: raw.request_id,
286 + token_count: raw.new_token_ids.len(),
287 + row_count: mask.rows.len(),
288 + })));
289 + }
290 let cached_token_count = raw
291 .prefill_stats
292 .as_ref()
@@ -291,6 +307,7 @@ impl Stream for GenerateOutputStream {
307 cached_token_count,
308 kv_transfer_params: raw.kv_transfer_params,
309 ec_transfer_params: raw.ec_transfer_params,
310 + sampling_mask,
311 };
312
313 Poll::Ready(Some(Ok(output)))
@@ -341,6 +358,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
358
359 while let Some(output) = stream.next().await.transpose()? {
360 cached_token_count = cached_token_count.max(output.cached_token_count);
361 + let sampling_mask = output.sampling_mask;
362 if let Some(info) = output.prompt_info {
363 if prompt_token_ids.is_none() {
364 prompt_token_ids = Some(info.prompt_token_ids.to_vec());
@@ -359,6 +377,9 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
377 existing.logprobs = Some(step_logprobs);
378 }
379 }
380 + if let Some(mut mask) = sampling_mask {
381 + existing.sampling_mask.get_or_insert_default().rows.append(&mut mask.rows);
382 + }
383 } else {
384 collected = Some(CollectedGenerateOutput {
385 request_id: output.request_id,
@@ -374,6 +395,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
395 },
396 kv_transfer_params: None,
397 ec_transfer_params: None,
398 + sampling_mask,
399 });
400 }
401
@@ -387,6 +409,15 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
409 };
410 collected.kv_transfer_params = output.kv_transfer_params;
411 collected.ec_transfer_params = output.ec_transfer_params;
412 + if let Some(mask) = collected.sampling_mask.as_ref()
413 + && mask.rows.len() != collected.token_ids.len()
414 + {
415 + return Err(crate::Error::SamplingMaskTokenCountMismatch {
416 + request_id: collected.request_id,
417 + token_count: collected.token_ids.len(),
418 + row_count: mask.rows.len(),
419 + });
420 + }
421 return Ok(collected);
422 }
423 }
rust/src/llm/tests/generate.rs
+59 -10
@@ -2,7 +2,7 @@
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use std::collections::BTreeSet;
5 -use std::sync::Once;
5 +use std::sync::{Arc, Once};
6 use std::time::Duration;
7
8 use futures::StreamExt as _;
@@ -18,6 +18,7 @@ use vllm_engine_core_client::protocol::output::{
18 };
19 use vllm_engine_core_client::protocol::request::EngineCoreRequest;
20 use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
21 +use vllm_engine_core_client::protocol::sampling_mask::{MaybeWireSamplingMask, SamplingMask};
22 use vllm_engine_core_client::protocol::stats::PrefillStats;
23 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
24 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
@@ -326,6 +327,11 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
327 num_local_cached_tokens: 1,
328 ..Default::default()
329 }),
330 + new_sampling_mask: Some(MaybeWireSamplingMask::Direct(
331 + SamplingMask {
332 + rows: vec![vec![1, 33, 99]],
333 + },
334 + )),
335 ..request_output_with_logprobs(
336 &request.request_id,
337 vec![33],
@@ -334,15 +340,22 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
340 Some(prompt_logprobs()),
341 )
342 },
337 - request_output_with_logprobs_and_kv(
338 - &request.request_id,
339 - vec![44],
340 - Some(EngineCoreFinishReason::Stop),
341 - Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)),
342 - None,
343 - Some(serde_json::json!({"connector": "x"})),
344 - None,
345 - ),
343 + EngineCoreOutput {
344 + new_sampling_mask: Some(MaybeWireSamplingMask::Direct(
345 + SamplingMask {
346 + rows: vec![vec![2, 44, 88]],
347 + },
348 + )),
349 + ..request_output_with_logprobs_and_kv(
350 + &request.request_id,
351 + vec![44],
352 + Some(EngineCoreFinishReason::Stop),
353 + Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)),
354 + None,
355 + Some(serde_json::json!({"connector": "x"})),
356 + None,
357 + )
358 + },
359 ],
360 ..Default::default()
361 }
@@ -375,6 +388,42 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
388 collected.kv_transfer_params,
389 Some(serde_json::json!({"connector": "x"}))
390 );
391 + assert_eq!(
392 + collected.sampling_mask,
393 + Some(SamplingMask {
394 + rows: vec![vec![1, 33, 99], vec![2, 44, 88]],
395 + })
396 + );
397 +}
398 +
399 +#[tokio::test]
400 +async fn collect_output_rejects_partial_sampling_mask() {
401 + let output = vllm_llm::GenerateOutput {
402 + request_id: "req-partial-mask".to_string(),
403 + prompt_info: Some(GeneratePromptInfo {
404 + prompt_token_ids: Arc::from([11_u32, 22]),
405 + prompt_logprobs: None,
406 + }),
407 + token_ids: vec![33, 44],
408 + logprobs: None,
409 + finish_reason: Some(FinishReason::Length),
410 + cached_token_count: 0,
411 + kv_transfer_params: None,
412 + ec_transfer_params: None,
413 + sampling_mask: Some(SamplingMask {
414 + rows: vec![vec![1, 33, 99]],
415 + }),
416 + };
417 +
418 + let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err();
419 + assert!(matches!(
420 + error,
421 + Error::SamplingMaskTokenCountMismatch {
422 + token_count: 2,
423 + row_count: 1,
424 + ..
425 + }
426 + ));
427 }
428
429 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
rust/src/server/Cargo.toml
+1
@@ -19,6 +19,7 @@ itertools.workspace = true
19 libc.workspace = true
20 llm-multimodal.workspace = true
21 openssl.workspace = true
22 +prost.workspace = true
23 prost-types.workspace = true
24 rmp-serde.workspace = true
25 rmpv.workspace = true
rust/src/server/src/grpc/convert.rs
+44
@@ -316,6 +316,11 @@ pub fn to_sequence_output(
316 _ => (vec![], vec![], vec![]),
317 };
318
319 + let sampling_mask = finished
320 + .and_then(|finished| finished.sampling_mask.as_ref())
321 + .map(|mask| mask.rows.iter().map(|row| pb::TokenIds { ids: row.clone() }).collect())
322 + .unwrap_or_default();
323 +
324 Ok(pb::SequenceOutput {
325 index: 0, // TODO: multi-sequence (n > 1) not supported
326 text: if opts.output_text {
@@ -333,6 +338,7 @@ pub fn to_sequence_output(
338 ranks: rank_values,
339 candidate_tokens: candidates,
340 finish_info,
341 + sampling_mask,
342 })
343 }
344
@@ -524,6 +530,7 @@ impl ResponseOpts {
530 mod tests {
531 use prost::Message as _;
532 use vllm_engine_core_client::protocol::output::StopReason;
533 + use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
534 use vllm_text::{
535 FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params,
536 };
@@ -720,6 +727,7 @@ mod tests {
727 finish_reason: reason,
728 kv_transfer_params: None,
729 ec_transfer_params: None,
730 + sampling_mask: None,
731 }
732 }
733
@@ -804,4 +812,40 @@ mod tests {
812 assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32);
813 assert_eq!(finish.stop_reason, Some(PbStopReason::EosTokenId(30)));
814 }
815 +
816 + #[test]
817 + fn sequence_output_only_carries_sampling_mask_on_terminal_output() {
818 + let mut fin = finished(FinishReason::Length);
819 + fin.sampling_mask = Some(SamplingMask {
820 + rows: vec![vec![1, 10], vec![2, 20]],
821 + });
822 +
823 + let terminal =
824 + to_sequence_output("", &[10, 20], None, Some(&fin), &ResponseOpts::default())
825 + .expect("terminal output");
826 + let intermediate = to_sequence_output("", &[9], None, None, &ResponseOpts::default())
827 + .expect("intermediate output");
828 +
829 + assert_eq!(
830 + terminal.sampling_mask,
831 + vec![
832 + pb::TokenIds { ids: vec![1, 10] },
833 + pb::TokenIds { ids: vec![2, 20] }
834 + ]
835 + );
836 + assert!(intermediate.sampling_mask.is_empty());
837 + }
838 +
839 + #[test]
840 + fn sampling_mask_uses_additive_proto_tag_ten() {
841 + let response = pb::SequenceOutput {
842 + sampling_mask: vec![pb::TokenIds { ids: vec![7, 8] }],
843 + ..Default::default()
844 + };
845 +
846 + assert_eq!(
847 + response.encode_to_vec(),
848 + vec![0x52, 0x04, 0x0a, 0x02, 0x07, 0x08]
849 + );
850 + }
851 }
rust/src/server/src/grpc/inference.rs
+1
@@ -257,6 +257,7 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
257 finish_reason: collected.finish_reason,
258 kv_transfer_params: collected.kv_transfer_params,
259 ec_transfer_params: collected.ec_transfer_params,
260 + sampling_mask: collected.sampling_mask,
261 };
262
263 let outputs = convert::to_sequence_output(
rust/src/server/src/routes/inference/generate.rs
+6
@@ -481,6 +481,7 @@ mod tests {
481 cached_token_count: 0,
482 kv_transfer_params: None,
483 ec_transfer_params: None,
484 + sampling_mask: None,
485 }),
486 Ok(GenerateOutput {
487 request_id: String::new(),
@@ -494,6 +495,7 @@ mod tests {
495 cached_token_count: 2,
496 kv_transfer_params: None,
497 ec_transfer_params: None,
498 + sampling_mask: None,
499 }),
500 ]);
501
@@ -575,6 +577,7 @@ mod tests {
577 cached_token_count: 0,
578 kv_transfer_params: None,
579 ec_transfer_params: None,
580 + sampling_mask: None,
581 }
582 }
583
@@ -658,6 +661,7 @@ mod tests {
661 kv_transfer_params: None,
662 ec_transfer_params: None,
663 prompt_token_ids: vec![10, 20],
664 + sampling_mask: None,
665 };
666
667 let response = collect_generate(
@@ -744,6 +748,7 @@ mod tests {
748 kv_transfer_params: None,
749 ec_transfer_params: None,
750 prompt_token_ids: vec![10, 20],
751 + sampling_mask: None,
752 };
753
754 let response = collect_generate(
@@ -776,6 +781,7 @@ mod tests {
781 },
782 kv_transfer_params: None,
783 ec_transfer_params: None,
784 + sampling_mask: None,
785 prompt_token_ids,
786 };
787
rust/src/server/src/routes/openai/completions.rs
+6
@@ -694,6 +694,7 @@ mod tests {
694 ))),
695 kv_transfer_params: None,
696 ec_transfer_params: None,
697 + sampling_mask: None,
698 }),
699 )),
700 ]);
@@ -806,6 +807,7 @@ mod tests {
807 finish_reason: FinishReason::Length,
808 kv_transfer_params: None,
809 ec_transfer_params: None,
810 + sampling_mask: None,
811 }),
812 )),
813 ]);
@@ -858,6 +860,7 @@ mod tests {
860 finish_reason: FinishReason::Length,
861 kv_transfer_params: None,
862 ec_transfer_params: None,
863 + sampling_mask: None,
864 }),
865 )),
866 ]);
@@ -913,6 +916,7 @@ mod tests {
916 finish_reason: FinishReason::Length,
917 kv_transfer_params: None,
918 ec_transfer_params: None,
919 + sampling_mask: None,
920 }),
921 )),
922 ]);
@@ -985,6 +989,7 @@ mod tests {
989 finish_reason: FinishReason::Length,
990 kv_transfer_params: None,
991 ec_transfer_params: None,
992 + sampling_mask: None,
993 }),
994 )),
995 ]);
@@ -1059,6 +1064,7 @@ mod tests {
1064 finish_reason: FinishReason::Length,
1065 kv_transfer_params: None,
1066 ec_transfer_params: None,
1067 + sampling_mask: None,
1068 }),
1069 )),
1070 ]);
rust/src/text/src/output/decoded.rs
+146
@@ -9,6 +9,7 @@ use serde::{Deserialize, Serialize};
9 use tracing::{Level, debug, trace};
10 use vllm_engine_core_client::AbortCause;
11 use vllm_engine_core_client::protocol::output::StopReason;
12 +use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
13 use vllm_llm::{FinishReason, GenerateOutput, TokenUsage};
14 use vllm_tokenizer::{DecodedText, DynTokenizer, IncrementalDecoder};
15
@@ -50,6 +51,8 @@ pub struct Finished {
51 /// Connector-specific encoder cache transfer parameters for disaggregated
52 /// serving.
53 pub ec_transfer_params: Option<serde_json::Value>,
54 + /// Sampling support sets aligned with all generated token positions.
55 + pub sampling_mask: Option<SamplingMask>,
56 }
57
58 /// Sample metadata emitted by one engine output update.
@@ -120,6 +123,7 @@ pub async fn decoded_text_event_stream(
123 let mut sampled_token_ids = Vec::new();
124 let mut output_token_count: usize = 0;
125 let mut sampled_logprobs: Option<DecodedLogprobs> = None;
126 + let mut sampling_mask: Option<SamplingMask> = None;
127
128 while let Some(next) = raw_stream.next().await {
129 let output = next?;
@@ -213,6 +217,7 @@ pub async fn decoded_text_event_stream(
217
218 let mut new_token_ids = output.token_ids;
219 let mut new_logprobs = output.logprobs;
220 + let mut new_sampling_mask = output.sampling_mask;
221
222 // Trim tokens and logprobs if we matched stop string.
223 if let Some(num_tokens) = truncate_tokens_to {
@@ -220,9 +225,15 @@ pub async fn decoded_text_event_stream(
225 if let Some(logprobs) = &mut new_logprobs {
226 logprobs.positions.truncate(num_tokens);
227 }
228 + if let Some(mask) = &mut new_sampling_mask {
229 + mask.rows.truncate(num_tokens);
230 + }
231 }
232
233 output_token_count += new_token_ids.len();
234 + if let Some(mut mask) = new_sampling_mask {
235 + sampling_mask.get_or_insert_default().rows.append(&mut mask.rows);
236 + }
237
238 let decoded_logprobs = new_logprobs
239 .as_ref()
@@ -246,6 +257,16 @@ pub async fn decoded_text_event_stream(
257 }
258
259 if let Some(reason) = finish_reason {
260 + if let Some(mask) = sampling_mask.as_ref()
261 + && mask.rows.len() != output_token_count
262 + {
263 + return Err(vllm_llm::Error::SamplingMaskTokenCountMismatch {
264 + request_id,
265 + token_count: output_token_count,
266 + row_count: mask.rows.len(),
267 + }
268 + .into());
269 + }
270 // Flush any remaining buffered text.
271 let (last_chunk, full_decoded) = decoder.flush(truncate_output_to)?;
272 let text_len = full_decoded.text.len();
@@ -293,6 +314,7 @@ pub async fn decoded_text_event_stream(
314 finish_reason: reason,
315 kv_transfer_params,
316 ec_transfer_params,
317 + sampling_mask,
318 })),
319 })
320 .await;
@@ -354,6 +376,7 @@ mod tests {
376
377 use futures::{Stream, StreamExt as _, stream};
378 use vllm_engine_core_client::AbortCause;
379 + use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob};
380 use vllm_llm::GenerateOutput;
381 use vllm_tokenizer::{TokenAnchor, TokenAttribution, test_utils::TestTokenizer};
382
@@ -384,6 +407,73 @@ mod tests {
407 s.bytes().map(u32::from).collect()
408 }
409
410 + #[tokio::test]
411 + async fn sampling_masks_survive_text_decode_and_collection() {
412 + let prompt: Arc<[u32]> = Arc::from([]);
413 + let mut first = GenerateOutput::for_test(Some(prompt), vec![b'a' as u32], None);
414 + first.sampling_mask = Some(SamplingMask {
415 + rows: vec![vec![1, b'a' as u32]],
416 + });
417 + let mut terminal =
418 + GenerateOutput::for_test(None, vec![b'b' as u32], Some(FinishReason::Length));
419 + terminal.sampling_mask = Some(SamplingMask {
420 + rows: vec![vec![2, b'b' as u32]],
421 + });
422 + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
423 +
424 + let collected = decoded_text_event_stream(
425 + "test-mask".into(),
426 + tokenizer,
427 + stream::iter([Ok(first), Ok(terminal)]),
428 + TextDecodeOptions::default(),
429 + false,
430 + )
431 + .collect_output()
432 + .await
433 + .unwrap();
434 +
435 + assert_eq!(
436 + collected.sampling_mask,
437 + Some(SamplingMask {
438 + rows: vec![vec![1, b'a' as u32], vec![2, b'b' as u32]],
439 + })
440 + );
441 + }
442 +
443 + #[tokio::test]
444 + async fn sampling_mask_mismatch_reuses_llm_error() {
445 + let prompt: Arc<[u32]> = Arc::from([]);
446 + let mut terminal = GenerateOutput::for_test(
447 + Some(prompt),
448 + vec![b'a' as u32, b'b' as u32],
449 + Some(FinishReason::Length),
450 + );
451 + terminal.sampling_mask = Some(SamplingMask {
452 + rows: vec![vec![1, b'a' as u32]],
453 + });
454 + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
455 +
456 + let error = decoded_text_event_stream(
457 + "test-mask-mismatch".into(),
458 + tokenizer,
459 + stream::iter([Ok(terminal)]),
460 + TextDecodeOptions::default(),
461 + false,
462 + )
463 + .collect_output()
464 + .await
465 + .unwrap_err();
466 +
467 + assert!(matches!(
468 + error,
469 + crate::Error::Llm(vllm_llm::Error::SamplingMaskTokenCountMismatch {
470 + token_count: 2,
471 + row_count: 1,
472 + ..
473 + })
474 + ));
475 + }
476 +
477 fn opts(stop: &[&str], min_tokens: u32) -> TextDecodeOptions {
478 TextDecodeOptions {
479 stop_strings: Some(stop.iter().map(|s| s.to_string()).collect()),
@@ -472,6 +562,7 @@ mod tests {
562 finish_reason: FinishReason::Length,
563 kv_transfer_params: None,
564 ec_transfer_params: None,
565 + sampling_mask: None,
566 }),
567 ),
568 ],
@@ -503,6 +594,7 @@ mod tests {
594 ))),
595 kv_transfer_params: None,
596 ec_transfer_params: None,
597 + sampling_mask: None,
598 }),
599 ),
600 ],
@@ -576,6 +668,60 @@ mod tests {
668 assert!(output.finish_reason.is_stop());
669 }
670
671 + #[tokio::test]
672 + async fn stream_stop_string_mid_chunk_keeps_token_logprob_mask_alignment() {
673 + let token_ids = ascii_tokens("abcxyz");
674 + let mut raw = GenerateOutput::for_test(
675 + Some(Arc::from([])),
676 + token_ids.clone(),
677 + Some(FinishReason::Length),
678 + );
679 + raw.logprobs = Some(Logprobs {
680 + positions: token_ids
681 + .iter()
682 + .map(|&token_id| PositionLogprobs {
683 + entries: vec![TokenLogprob {
684 + token_id,
685 + logprob: -0.25,
686 + rank: 1,
687 + }],
688 + })
689 + .collect(),
690 + });
691 + raw.sampling_mask = Some(SamplingMask {
692 + rows: token_ids.iter().map(|&token_id| vec![token_id, token_id + 100]).collect(),
693 + });
694 + let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
695 +
696 + let output = decoded_text_event_stream(
697 + "test-mid-chunk-stop".into(),
698 + tokenizer,
699 + stream::iter([Ok(raw)]),
700 + opts(&["cx"], 0),
701 + false,
702 + )
703 + .collect_output()
704 + .await
705 + .unwrap();
706 +
707 + assert_eq!(output.text, "ab");
708 + assert!(output.finish_reason.is_stop());
709 + assert_eq!(output.token_ids, ascii_tokens("abcx"));
710 + assert_eq!(output.logprobs.as_ref().unwrap().positions.len(), 4);
711 + assert_eq!(
712 + output.sampling_mask,
713 + Some(SamplingMask {
714 + rows: vec![
715 + vec![b'a' as u32, b'a' as u32 + 100],
716 + vec![b'b' as u32, b'b' as u32 + 100],
717 + vec![b'c' as u32, b'c' as u32 + 100],
718 + vec![b'x' as u32, b'x' as u32 + 100],
719 + ],
720 + })
721 + );
722 + assert_eq!(output.usage.output_token_count, 4);
723 + }
724 +
725 #[tokio::test]
726 async fn stream_stop_string_at_end() {
727 let output = run_to_completion(ascii_tokens("abcxyz"), opts(&["xyz"], 0)).await;
rust/src/text/src/output/mod.rs
+7
@@ -17,6 +17,7 @@ mod logprobs;
17 use std::sync::Arc;
18
19 use futures::{StreamExt as _, pin_mut};
20 +use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
21
22 use crate::{Error, FinishReason, Result, TextOutputStream};
23
@@ -35,6 +36,8 @@ pub struct CollectedTextOutput {
36 /// Connector-specific encoder cache transfer parameters for disaggregated
37 /// serving.
38 pub ec_transfer_params: Option<serde_json::Value>,
39 + /// Sampling support sets aligned with generated token positions.
40 + pub sampling_mask: Option<SamplingMask>,
41 }
42
43 #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")]
@@ -90,6 +93,7 @@ impl<T: TextOutputStream> T {
93 usage: vllm_llm::TokenUsage::default(),
94 kv_transfer_params: None,
95 ec_transfer_params: None,
96 + sampling_mask: None,
97 })
98 };
99
@@ -99,6 +103,7 @@ impl<T: TextOutputStream> T {
103 collected.usage = finished.usage;
104 collected.kv_transfer_params = finished.kv_transfer_params;
105 collected.ec_transfer_params = finished.ec_transfer_params;
106 + collected.sampling_mask = finished.sampling_mask;
107 return Ok(collected);
108 }
109 }
@@ -173,6 +178,7 @@ mod tests {
178 finish_reason: FinishReason::stop_eos(),
179 kv_transfer_params: None,
180 ec_transfer_params: None,
181 + sampling_mask: None,
182 })),
183 }),
184 ]);
@@ -295,6 +301,7 @@ mod tests {
301 finish_reason: FinishReason::stop_eos(),
302 kv_transfer_params: None,
303 ec_transfer_params: None,
304 + sampling_mask: None,
305 })),
306 }),
307 ]);