[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>

rust/Cargo.lock

@@ -5941,7 +5941,7 @@ dependencies = [
 
 [[package]]
 name = "vllm-proto"
-version = "0.3.0"
+version = "0.4.0"
 dependencies = [
  "prost",
  "prost-types",

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"

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

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,
             })),
         }
     }

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,
     }
 }
 

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),

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;

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],

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;

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<Vec<String>>,
     #[serde(default)]
-    pub new_sampling_mask: Option<OpaqueValue>,
+    pub new_sampling_mask: Option<MaybeWireSamplingMask>,
     /// 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(())
     }
 }

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<Vec<u32>>,
+}
+
+/// 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<WireNdArray>,
+    #[serde(default)]
+    cu_num_generated_tokens: Option<Vec<usize>>,
+}
+
+/// Sampling-mask field while it transitions from Python wire data to rows.
+#[derive(Debug, Clone, PartialEq, EnumAsInner)]
+pub enum MaybeWireSamplingMask {
+    Wire(Box<WireSamplingMask>),
+    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<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        WireSamplingMask::deserialize(deserializer).map(|value| Self::Wire(Box::new(value)))
+    }
+}
+
+impl Serialize for MaybeWireSamplingMask {
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    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<Self> {
+        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<Self, String> {
+        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::<Vec<_>>();
+        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<SamplingMask> {
+        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::<Vec<_>>();
+        let offsets = [0_i64, 2, 4].into_iter().flat_map(i64::to_le_bytes).collect::<Vec<_>>();
+
+        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]]);
+    }
+}

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.

rust/src/engine-core-client/src/tests/python_compat.py

@@ -234,14 +234,10 @@ sampling_mask_wire = [
     [
         "<i4",
         [5],
-        msgpack.ExtType(3, np.array([2, 12, 16, 17, 18], dtype=np.int32).tobytes()),
-    ],
-    [
-        "<i8",
-        [2],
-        msgpack.ExtType(3, np.array([0, 5], dtype=np.int64).tobytes()),
+        msgspec.msgpack.Ext(3, np.array([2, 12, 16, 17, 18], dtype=np.int32).tobytes()),
     ],
     None,
+    None,
 ]
 outputs_with_sampling_mask = EngineCoreOutputs(
     outputs=[

rust/src/llm/src/error.rs

@@ -12,4 +12,12 @@ pub enum Error {
     EmptyPromptTokenIds { request_id: String },
     #[error("engine-core error")]
     EngineCoreClient(#[from] vllm_engine_core_client::Error),
+    #[error(
+        "sampling mask for request `{request_id}` has {row_count} rows for {token_count} generated tokens"
+    )]
+    SamplingMaskTokenCountMismatch {
+        request_id: String,
+        token_count: usize,
+        row_count: usize,
+    },
 }

rust/src/llm/src/output.rs

@@ -12,6 +12,7 @@ use futures::{Stream, StreamExt as _, pin_mut};
 use serde::{Deserialize, Serialize};
 use vllm_engine_core_client::protocol::logprobs::Logprobs;
 use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason};
+use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
 use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
 
 use crate::error::Result;
@@ -44,6 +45,8 @@ pub struct CollectedGenerateOutput {
     /// Connector-specific encoder cache transfer parameters for disaggregated
     /// serving.
     pub ec_transfer_params: Option<serde_json::Value>,
+    /// Sampling support sets aligned one-to-one with generated token positions.
+    pub sampling_mask: Option<SamplingMask>,
 }
 
 /// 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<serde_json::Value>,
+    /// Sampling support sets aligned one-to-one with `token_ids`.
+    pub sampling_mask: Option<SamplingMask>,
 }
 
 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<T: Stream<Item = Result<GenerateOutput>> + 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<T: Stream<Item = Result<GenerateOutput>> + 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<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                         },
                         kv_transfer_params: None,
                         ec_transfer_params: None,
+                        sampling_mask,
                     });
                 }
 
@@ -387,6 +409,15 @@ impl<T: Stream<Item = Result<GenerateOutput>> + 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);
                 }
             }

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)]

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

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

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(

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,
         };
 

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,
                 }),
             )),
         ]);

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<serde_json::Value>,
+    /// Sampling support sets aligned with all generated token positions.
+    pub sampling_mask: Option<SamplingMask>,
 }
 
 /// 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<DecodedLogprobs> = None;
+    let mut sampling_mask: Option<SamplingMask> = 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;

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<serde_json::Value>,
+    /// Sampling support sets aligned with generated token positions.
+    pub sampling_mask: Option<SamplingMask>,
 }
 
 #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")]
@@ -90,6 +93,7 @@ impl<T: TextOutputStream> T {
                                 usage: vllm_llm::TokenUsage::default(),
                                 kv_transfer_params: None,
                                 ec_transfer_params: None,
+                                sampling_mask: None,
                             })
                         };
 
@@ -99,6 +103,7 @@ impl<T: TextOutputStream> 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,
                 })),
             }),
         ]);