[Bugfix] Preserve sampling masks in DELTA and TITO streaming outputs

tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py

@@ -603,3 +603,46 @@ async def mock_generate(*args, **kwargs):
     # Zero cached tokens must be present, not omitted
     assert usage_chunk["usage"]["prompt_tokens_details"] is not None
     assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("with_mask", [False, True])
+async def test_stream_sampling_mask_matches_each_token_chunk(with_mask):
+    from vllm.outputs import SamplingMask
+
+    engine = _mock_engine()
+
+    async def generate(*args, **kwargs):
+        for position, tokens in enumerate(([10], [20, 30])):
+            result = _make_request_output(
+                "req-mask",
+                token_ids=list(tokens),
+                finish_reason="length" if position else None,
+                finished=bool(position),
+            )
+            if with_mask:
+                result.outputs[0].sampling_mask = SamplingMask(
+                    [[token, token + 1] for token in tokens]
+                )
+            yield result
+
+    engine.generate = MagicMock(side_effect=generate)
+    serving = _build_serving_tokens(engine)
+    response = await serving.serve_tokens(
+        GenerateRequest(
+            model=MODEL_NAME, token_ids=[1, 2, 3], sampling_params={}, stream=True
+        )
+    )
+    chunks = _parse_sse_chunks([chunk async for chunk in response])
+    choices = [
+        choice
+        for chunk in chunks
+        if isinstance(chunk, dict)
+        for choice in chunk.get("choices", [])
+    ]
+    assert [choice["token_ids"] for choice in choices] == [[10], [20, 30]]
+    for choice in choices:
+        expected = (
+            [[token, token + 1] for token in choice["token_ids"]] if with_mask else None
+        )
+        assert choice["sampling_mask"] == expected

tests/v1/engine/test_output_processor.py

@@ -1448,3 +1448,76 @@ def test_abort_requests(runner: str, abort_by: str, dummy_test_vectors):
             output_processor.abort_requests([request.request_id], internal=True)
         else:
             output_processor.abort_requests([request.external_req_id], internal=False)
+
+
+@pytest.mark.parametrize("output_kind", list(RequestOutputKind))
+@pytest.mark.parametrize("coalesce", [False, True])
+@pytest.mark.parametrize("stream_interval", [1, 2])
+def test_sampling_masks_follow_output_token_boundaries(
+    output_kind, coalesce, stream_interval
+):
+    import numpy as np
+
+    from vllm.v1.outputs import SamplingMaskLists
+
+    state = RequestState.__new__(RequestState)
+    state.detokenizer = MagicMock()
+    state.detokenizer.get_next_output_text.return_value = ""
+    state.detokenizer.output_token_ids = []
+    state.detokenizer.num_output_tokens.side_effect = lambda: len(
+        state.detokenizer.output_token_ids
+    )
+    state.stream_interval = stream_interval
+    state.sent_tokens_offset = 0
+    state.external_req_id = "request"
+    state.parent_req = None
+    state.prompt = None
+    state.prompt_token_ids = [1]
+    state.lora_request = None
+    state.num_cached_tokens = 0
+    state.num_cache_creation_tokens = 0
+    state.stats = None
+    state.logprobs_processor = MagicMock()
+    state.logprobs_processor.logprobs = None
+    state.logprobs_processor.cumulative_logprob = None
+    state.output_kind = output_kind
+    state.request_index = 0
+    state.sampling_mask_chunks = []
+    state.routed_experts_chunks = []
+    state.spec_decode_metrics = None
+    collector = RequestOutputCollector(output_kind, "request")
+
+    expected = [[10, 11], [20], [30, 31]]
+    received = []
+    for position, support in enumerate(expected):
+        state.detokenizer.output_token_ids.append(support[0])
+        state.sampling_mask_chunks.append(SamplingMaskLists(np.asarray(support)))
+        finished = position == len(expected) - 1
+        result = state.make_request_output(
+            [support[0]], None, FinishReason.LENGTH if finished else None, None
+        )
+        if result is None:
+            continue
+        collector.put(result)
+        if not coalesce:
+            received.append(collector.get_nowait())
+    if coalesce:
+        received.append(collector.get_nowait())
+
+    if output_kind == RequestOutputKind.DELTA:
+        masks = [
+            support
+            for result in received
+            for support in result.outputs[0].sampling_mask.token_ids
+        ]
+        assert masks == expected
+        for result in received:
+            completion = result.outputs[0]
+            assert len(completion.sampling_mask.token_ids) == len(completion.token_ids)
+        assert state.sampling_mask_chunks == []
+        assert (
+            state._new_completion_output([], FinishReason.LENGTH, None).sampling_mask
+            is None
+        )
+    else:
+        assert received[-1].outputs[0].sampling_mask.token_ids == expected

vllm/entrypoints/scale_out/token_in_token_out/serving.py

@@ -450,6 +450,11 @@ async def serve_tokens_stream_generator(
                                 finish_reason=finish_reason,
                                 token_ids=as_list(delta_token_ids),
                                 routed_experts=routed_experts_b64,
+                                sampling_mask=(
+                                    output.sampling_mask.token_ids
+                                    if output.sampling_mask is not None
+                                    else None
+                                ),
                             )
                         ],
                     )

vllm/outputs.py

@@ -189,6 +189,13 @@ def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
                         if next_completion.logprobs:
                             assert completion.logprobs is not None
                             completion.logprobs.extend(next_completion.logprobs)  # type: ignore[arg-type]
+                        if next_completion.sampling_mask is not None:
+                            if completion.sampling_mask is None:
+                                completion.sampling_mask = next_completion.sampling_mask
+                            else:
+                                completion.sampling_mask.token_ids.extend(
+                                    next_completion.sampling_mask.token_ids
+                                )
                         completion.cumulative_logprob = (
                             next_completion.cumulative_logprob
                         )

vllm/v1/engine/output_processor.py

@@ -418,10 +418,12 @@ def _new_completion_output(
             logprobs = logprobs[-num_new_tokens:] if num_new_tokens else logprobs[:0]
 
         sampling_mask = None
-        if finished and self.sampling_mask_chunks:
+        if self.sampling_mask_chunks and (delta or finished):
             sampling_mask = SamplingMask(
                 [chunk.token_ids.tolist() for chunk in self.sampling_mask_chunks]
             )
+            if delta:
+                self.sampling_mask_chunks.clear()
 
         # Concatenate routed experts on finish
         routed_experts = None