From eb11446b7f2e30ca582f8aff3afe12e9a2e66f6c Mon Sep 17 00:00:00 2001 From: Inesh Reddy Chappidi Date: Mon, 31 Aug 2026 11:27:13 -0400 Subject: [PATCH] [Bugfix][Core] Do not advance the segment watermark past an unmaintained request counter (#6834) Signed-off-by: Inesh Reddy Chappidi Co-authored-by: Gao Han --- .../test_chunk_transfer_adapter.py | 31 +++++++++++++++++++ .../chunk_transfer_adapter.py | 10 +++++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index def66971f..0e3194190 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -2959,3 +2959,34 @@ def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter) ) assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long)) + + +def test_save_async_boundary_holds_generation_without_request_counter(build_adapter): + """A boundary must not advance past a generation nothing can catch up to. + + When the producing stage maintains ``_omni_segment_generation`` on the + request, a boundary may start the next segment's watermark at + ``generation + 1`` because the following frame will carry it. On paths + where the request has no such counter, that speculative advance makes every + subsequent frame look stale and starves the stream permanently + (vllm-project/vllm-omni#6816). + """ + adapter, _ = build_adapter(stage_id=1) + + request = _req("req-nogen", RequestStatus.WAITING, external_req_id="ext-nogen") + request.resumable = True + request._omni_segment_generation = None + + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True) + assert adapter._segment_generation["ext-nogen"] == 0 + + follow_up = _req("req-nogen-next", RequestStatus.WAITING, external_req_id="ext-nogen") + follow_up.resumable = True + follow_up.num_computed_tokens = 3 + follow_up._omni_segment_generation = None + queued_before = len(adapter._pending_save_reqs) + + adapter.save_async(multimodal_output=None, request=follow_up, is_segment_finished=False) + + assert len(adapter._pending_save_reqs) == queued_before + 1 + assert adapter._pending_save_reqs[-1]["request"] is follow_up diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index c19159e58..b9b03c9c5 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -401,7 +401,15 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if is_segment_finished: # The queued FIFO item now owns the old segment. Start the next # segment's deduplication watermark before the worker sends it. - self._segment_generation[external_req_id] = generation + 1 + # Only advance past the current generation when the producer + # maintains a request-side counter that can catch up. On paths + # where the request carries no `_omni_segment_generation`, a + # speculative +1 makes every subsequent frame look stale and + # starves the stream permanently (see #6816). + if getattr(request, "_omni_segment_generation", None) is not None: + self._segment_generation[external_req_id] = generation + 1 + else: + self._segment_generation[external_req_id] = generation self.requests_num_chunks_sent.pop(external_req_id, None) if reject_reason is not None: logger.error("Cannot enqueue %s: %s", external_req_id, reject_reason)