[Bugfix][Frontend] Preserve streaming chunk boundaries in RequestOutputCollector (#58941)

Signed-off-by: David Cheung <d7cheung@gmail.com>
Co-authored-by: Benedikt Falk <bfalk@nvidia.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>

tests/v1/engine/test_output_processor.py

@@ -42,6 +42,7 @@ from vllm.v1.engine.output_processor import (
     OutputProcessor,
     RequestOutputCollector,
     RequestState,
+    StreamingInputOutputCollector,
 )
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.loggers import PrometheusStatLogger
@@ -1642,6 +1643,53 @@ async def test_cumulative_output_collector_n():
     assert third[0].text == "c"
 
 
+def _make_boundary_output(
+    token_id: int, finish_reason: str | None, finished: bool = False
+) -> RequestOutput:
+    return RequestOutput(
+        request_id="my-request-id",
+        prompt=None,
+        prompt_token_ids=[1, 2, 3],
+        prompt_logprobs=None,
+        outputs=[
+            CompletionOutput(
+                index=0,
+                text="a",
+                token_ids=[token_id],
+                cumulative_logprob=None,
+                logprobs=None,
+                finish_reason=finish_reason,
+            )
+        ],
+        finished=finished,
+    )
+
+
+@pytest.mark.asyncio
+async def test_streaming_input_collector_preserves_chunk_boundary():
+    """A pending chunk-boundary marker (finish_reason on a non-finished
+    output) must not be erased by merging the next chunk's output into it."""
+    collector = StreamingInputOutputCollector(
+        RequestOutputKind.DELTA, request_id="my-request-id-int"
+    )
+
+    collector.put(_make_boundary_output(0, "stop"))
+    collector.put(_make_boundary_output(1, None))
+    collector.put(_make_boundary_output(2, None))
+
+    boundary = await collector.get()
+    assert boundary.outputs[0].finish_reason == "stop"
+    assert boundary.outputs[0].token_ids == [0]
+
+    following = await collector.get()
+    assert following.outputs[0].finish_reason is None
+    assert following.outputs[0].token_ids == [1, 2]
+
+    assert collector.output is None
+    assert not collector._backlog
+    assert not collector.ready.is_set()
+
+
 def _make_streaming_input_request(
     request_id: str, prompt_token_ids: list[int]
 ) -> EngineCoreRequest:

vllm/v1/engine/async_llm.py

@@ -53,7 +53,11 @@ from vllm.v1.engine.admission_control import SharedAdmissionStats
 from vllm.v1.engine.core_client import EngineCoreClient
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 from vllm.v1.engine.input_processor import InputProcessor
-from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector
+from vllm.v1.engine.output_processor import (
+    OutputProcessor,
+    RequestOutputCollector,
+    StreamingInputOutputCollector,
+)
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.executor import Executor
 from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult
@@ -620,7 +624,9 @@ class AsyncLLM(EngineClient):
         self.input_processor.assign_request_id(final_req)
         internal_req_id = final_req.request_id
 
-        queue = RequestOutputCollector(sampling_params.output_kind, internal_req_id)
+        queue = StreamingInputOutputCollector(
+            sampling_params.output_kind, internal_req_id
+        )
 
         async def handle_inputs():
             cancelled = False

vllm/v1/engine/output_processor.py

@@ -113,6 +113,53 @@ class RequestOutputCollector:
             self._input_stream_task = None
 
 
+def _is_chunk_boundary(
+    output: RequestOutput | PoolingRequestOutput | Exception | None,
+) -> bool:
+    """Whether the output marks the end of a streaming-input chunk: a
+    finish_reason on a request that has not finished."""
+    return (
+        isinstance(output, RequestOutput)
+        and not output.finished
+        and any(completion.finished() for completion in output.outputs)
+    )
+
+
+class StreamingInputOutputCollector(RequestOutputCollector):
+    """Collector for streaming-input requests. Doesn't merge across input
+    chunk boundaries (a finish_reason on a non-finished output)."""
+
+    def __init__(self, output_kind: RequestOutputKind, request_id: str):
+        super().__init__(output_kind, request_id)
+        self._backlog: deque[RequestOutput] = deque()
+
+    def put(self, output: RequestOutput | PoolingRequestOutput | Exception) -> None:
+        if isinstance(output, Exception):
+            self._backlog.clear()
+        elif self._backlog or _is_chunk_boundary(self.output):
+            assert isinstance(output, RequestOutput)
+            if self._backlog and not _is_chunk_boundary(self._backlog[-1]):
+                self._backlog[-1].add(output, aggregate=self.aggregate)
+            else:
+                self._backlog.append(output)
+            return
+        super().put(output)
+
+    async def get(self) -> RequestOutput | PoolingRequestOutput:
+        output = await super().get()
+        if self._backlog:
+            self.output = self._backlog.popleft()
+            self.ready.set()
+        return output
+
+    def get_nowait(self) -> RequestOutput | PoolingRequestOutput | None:
+        output = super().get_nowait()
+        if self._backlog:
+            self.output = self._backlog.popleft()
+            self.ready.set()
+        return output
+
+
 @dataclass
 class OutputProcessorOutput:
     request_outputs: list[RequestOutput | PoolingRequestOutput]