@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4a9d6f9572

[Bugfix] Preserve sampling masks in DELTA and TITO streaming outputs (#55935)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: fh3 <a@b> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: fh3 <a@b> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

aoshen02 committed Oct 1, 2026 at 14:43 UTC 4a9d6f957203902ea6bff0b275c26876e5a47e46

8 files changed +135 -6

docs/training/sampling_mask.md

+2 -2

@@ -53,6 +53,8 @@ The mask is also available via the `/inference/v1/generate` HTTP endpoint:

53 }

54 ```

56 +With `stream=true`, each chunk carries the masks of the tokens it emits.

57 +

58 ## Requirements

60 | Requirement | Reason |

@@ -145,5 +147,3 @@ Each position then carries the sampled token's logprob and the top-128

147 - **Engine-level flag:** `--return-sampling-mask` globally disables the

148 FlashInfer fused sampler. All requests pay the cost of the PyTorch sampling

149 path, even if they don't need the mask.

148 -- **No streaming support:** The mask is returned only in the final response,

149 - not in intermediate streaming chunks.

rust/src/server/src/routes/inference/generate.rs

+19 -1

@@ -231,6 +231,7 @@ async fn generate_chunk_stream(

231 logprobs,

232 finish_reason: finish_reason.map(|reason| reason.as_str().to_string()),

233 token_ids,

234 + sampling_mask: output.sampling_mask.map(|mask| mask.rows),

235 }],

236 usage: include_continuous_usage

237 .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)),

@@ -331,6 +332,7 @@ fn collect_generate(

332 logprobs,

333 finish_reason: Some(finish_reason),

334 token_ids: collected.token_ids,

335 + sampling_mask: collected.sampling_mask.map(|mask| mask.rows),

336 }],

337 prompt_logprobs,

338 prompt_token_id_logprobs: collected

@@ -516,6 +518,7 @@ mod tests {

518 use futures::{TryStreamExt as _, stream};

519 use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange};

520 use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;

521 + use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;

522 use vllm_llm::GeneratePromptInfo;

523

524 use super::*;

@@ -657,6 +660,18 @@ mod tests {

660 .expect("collect chunks")

661 }

662

663 + #[tokio::test]

664 + async fn generate_chunk_stream_returns_sampling_mask() {

665 + let mut output = stream_output(Some(&[11]), vec![33], Some(FinishReason::Length));

666 + output.sampling_mask = Some(SamplingMask {

667 + rows: vec![vec![33, 44]],

668 + });

669 +

670 + let chunks = collect_chunks(vec![output], false, None).await;

671 +

672 + assert_eq!(chunks[0].choices[0].sampling_mask, Some(vec![vec![33, 44]]));

673 + }

674 +

675 #[tokio::test]

676 async fn generate_chunk_stream_omits_prompt_metadata_by_default() {

677 let chunks = collect_chunks(

@@ -792,7 +807,9 @@ mod tests {

807 kv_transfer_params: None,

808 ec_transfer_params: None,

809 prompt_token_ids: vec![10, 20],

795 - sampling_mask: None,

810 + sampling_mask: Some(SamplingMask {

811 + rows: vec![vec![30, 40]],

812 + }),

813 spec_decode_metrics: None,

814 };

815

@@ -813,6 +830,7 @@ mod tests {

830

831 assert!(response.prompt_token_ids.is_none());

832 assert!(response.mm_placeholders.is_none());

833 + assert_eq!(response.choices[0].sampling_mask, Some(vec![vec![30, 40]]));

834 }

835

836 #[test]

rust/src/server/src/routes/inference/generate/types.rs

@@ -64,6 +64,7 @@ pub(super) struct GenerateResponseChoice {

64 pub logprobs: Option<ChatLogProbs>,

65 pub finish_reason: Option<String>,

66 pub token_ids: Vec<u32>,

67 + pub sampling_mask: Option<Vec<Vec<u32>>>,

68 }

70 /// Mirrors the Python vLLM `GenerateResponseStreamChoice` class.

@@ -74,6 +75,7 @@ pub(super) struct GenerateResponseStreamChoice {

75 pub logprobs: Option<ChatLogProbs>,

76 pub finish_reason: Option<String>,

77 pub token_ids: Vec<u32>,

78 + pub sampling_mask: Option<Vec<Vec<u32>>>,

79 }

81 /// Mirrors the Python vLLM `GenerateStreamResponse` class.

tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py

+42 -1

@@ -34,7 +34,7 @@ from vllm.multimodal.inputs import (

34 MultiModalKwargsItem,

35 PlaceholderRange,

36 )

37 -from vllm.outputs import CompletionOutput, RequestOutput

37 +from vllm.outputs import CompletionOutput, RequestOutput, SamplingMask

38 from vllm.renderers import renderer_from_config

39 from vllm.renderers.online_renderer import OnlineRenderer

40 from vllm.sampling_params import SamplingParams

@@ -1154,3 +1154,44 @@ async def test_stream_prompt_tokens_details_zero_cached():

1154 # Zero cached tokens must be present, not omitted

1155 assert usage_chunk["usage"]["prompt_tokens_details"] is not None

1156 assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 0

1157 +

1158 +

1159 +@pytest.mark.asyncio

1160 +@pytest.mark.parametrize("with_mask", [False, True])

1161 +async def test_stream_sampling_mask_matches_each_token_chunk(with_mask):

1162 + engine = _mock_engine()

1163 +

1164 + async def generate(*args, **kwargs):

1165 + for position, tokens in enumerate(([10], [20, 30])):

1166 + result = _make_request_output(

1167 + "req-mask",

1168 + token_ids=list(tokens),

1169 + finish_reason="length" if position else None,

1170 + finished=bool(position),

1171 + )

1172 + if with_mask:

1173 + result.outputs[0].sampling_mask = SamplingMask(

1174 + [[token, token + 1] for token in tokens]

1175 + )

1176 + yield result

1177 +

1178 + engine.generate = MagicMock(side_effect=generate)

1179 + serving = _build_serving_tokens(engine)

1180 + response = await serving.serve_tokens(

1181 + GenerateRequest(

1182 + model=MODEL_NAME, token_ids=[1, 2, 3], sampling_params={}, stream=True

1183 + )

1184 + )

1185 + chunks = _parse_sse_chunks([chunk async for chunk in response])

1186 + choices = [

1187 + choice

1188 + for chunk in chunks

1189 + if isinstance(chunk, dict)

1190 + for choice in chunk.get("choices", [])

1191 + ]

1192 + assert [choice["token_ids"] for choice in choices] == [[10], [20, 30]]

1193 + for choice in choices:

1194 + expected = (

1195 + [[token, token + 1] for token in choice["token_ids"]] if with_mask else None

1196 + )

1197 + assert choice["sampling_mask"] == expected

tests/v1/engine/test_output_processor.py

+56 -1

@@ -18,7 +18,7 @@ from tests.v1.engine.utils import (

18 from vllm import PoolingParams

19 from vllm.logprobs import FlatLogprobs, Logprob, PromptLogprobs, SampleLogprobs

20 from vllm.lora.request import LoRARequest

21 -from vllm.outputs import CompletionOutput, RequestOutput

21 +from vllm.outputs import CompletionOutput, RequestOutput, SamplingMask

22 from vllm.sampling_params import RequestOutputKind, SamplingParams

23 from vllm.tokenizers import TokenizerLike

24 from vllm.v1.engine import (

@@ -1575,3 +1575,58 @@ def test_abort_requests(runner: str, abort_by: str, dummy_test_vectors):

1575 output_processor.abort_requests([request.request_id], internal=True)

1576 else:

1577 output_processor.abort_requests([request.external_req_id], internal=False)

1578 +

1579 +

1580 +@pytest.mark.parametrize("output_kind", list(RequestOutputKind))

1581 +def test_sampling_masks_follow_output_kind(output_kind):

1582 + state = RequestState.__new__(RequestState)

1583 + state.detokenizer = MagicMock()

1584 + state.detokenizer.get_next_output_text.return_value = ""

1585 + state.logprobs_processor = MagicMock()

1586 + state.logprobs_processor.logprobs = None

1587 + state.logprobs_processor.cumulative_logprob = None

1588 + state.output_kind = output_kind

1589 + state.request_index = 0

1590 + state.sampling_mask_chunks = []

1591 + state.routed_experts_chunks = []

1592 + state.spec_decode_metrics = None

1593 +

1594 + supports = [[10, 11], [20], [30, 31]]

1595 + masks = []

1596 + for position, support in enumerate(supports):

1597 + state.sampling_mask_chunks.append(SamplingMaskLists(np.asarray(support)))

1598 + finish = FinishReason.LENGTH if position == len(supports) - 1 else None

1599 + mask = state._new_completion_output([support[0]], finish, None).sampling_mask

1600 + masks.append(mask.token_ids if mask is not None else None)

1601 +

1602 + if output_kind == RequestOutputKind.DELTA:

1603 + assert masks == [[support] for support in supports]

1604 + else:

1605 + assert masks == [None, None, supports]

1606 +

1607 +

1608 +def test_request_output_add_merges_delta_sampling_masks():

1609 + def delta(token_ids, sampling_mask):

1610 + completion = CompletionOutput(

1611 + index=0,

1612 + text="",

1613 + token_ids=token_ids,

1614 + cumulative_logprob=None,

1615 + logprobs=None,

1616 + sampling_mask=sampling_mask,

1617 + )

1618 + return RequestOutput(

1619 + request_id="request",

1620 + prompt=None,

1621 + prompt_token_ids=[1],

1622 + prompt_logprobs=None,

1623 + outputs=[completion],

1624 + finished=False,

1625 + )

1626 +

1627 + merged = delta([], None)

1628 + merged.add(delta([10], SamplingMask([[10, 11]])), aggregate=True)

1629 + merged.add(delta([20], SamplingMask([[20]])), aggregate=True)

1630 +

1631 + assert merged.outputs[0].token_ids == [10, 20]

1632 + assert merged.outputs[0].sampling_mask.token_ids == [[10, 11], [20]]

vllm/entrypoints/scale_out/token_in_token_out/serving.py

@@ -515,6 +515,10 @@ class ServingTokens(GenerateBaseServing):

515 else None

516 )

517

518 + sampling_mask = None

519 + if output.sampling_mask is not None:

520 + sampling_mask = output.sampling_mask.token_ids

521 +

522 chunk = GenerateStreamResponse(

523 request_id=request_id,

524 choices=[

@@ -524,6 +528,7 @@ class ServingTokens(GenerateBaseServing):

528 finish_reason=finish_reason,

529 token_ids=as_list(delta_token_ids),

530 routed_experts=routed_experts_b64,

531 + sampling_mask=sampling_mask,

532 )

533 ],

534 )

vllm/outputs.py

@@ -197,6 +197,12 @@ class RequestOutput:

197 if next_completion.logprobs:

198 assert completion.logprobs is not None

199 completion.logprobs.extend(next_completion.logprobs) # type: ignore[arg-type]

200 + if next_completion.sampling_mask is not None:

201 + if completion.sampling_mask is None:

202 + completion.sampling_mask = SamplingMask([])

203 + completion.sampling_mask.token_ids.extend(

204 + next_completion.sampling_mask.token_ids

205 + )

206 completion.cumulative_logprob = (

207 next_completion.cumulative_logprob

208 )

vllm/v1/engine/output_processor.py

+3 -1

@@ -439,7 +439,7 @@ class RequestState:

439 logprobs = logprobs[-num_new_tokens:] if num_new_tokens else logprobs[:0]

440

441 sampling_mask = None

442 - if finished and self.sampling_mask_chunks:

442 + if (delta or finished) and self.sampling_mask_chunks:

443 sampling_mask = SamplingMask(

444 [

445 position

@@ -447,6 +447,8 @@ class RequestState:

447 for position in chunk.to_nested_list()

448 ]

449 )

450 + if delta:

451 + self.sampling_mask_chunks.clear()

452

453 # Concatenate routed experts on finish

454 routed_experts = None