From 3842970a82f0a9e52cdb563e4c4fd8a3b3d780fc Mon Sep 17 00:00:00 2001 From: ChenWenjing <54166744+Shirley125@users.noreply.github.com> Date: Wed, 3 Jun 2026 16:07:41 +0800 Subject: [PATCH] [bugfix] fix streaming input issue after rebase 0.22.0 (#4085) Signed-off-by: CHEN <116010019@link.cuhk.edu.cn> --- tests/core/sched/test_omni_ar_scheduler_streaming.py | 4 ++-- vllm_omni/core/sched/omni_ar_scheduler.py | 5 ++--- .../transfer_adapter/chunk_transfer_adapter.py | 3 ++- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index f6539a40a..41fb414cb 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -61,7 +61,7 @@ def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() sched._update_request_as_session(session, _make_update([10, 20])) - assert session.discard_latest_async_tokens is True + assert session.async_tokens_to_discard == 1 assert session.num_output_placeholders == 0 assert session.spec_token_ids == [] # The async placeholder makes token 9 unconfirmed, so only 7 and 8 are @@ -83,7 +83,7 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder() sched._update_request_as_session(session, _make_update([10, 20])) - assert session.discard_latest_async_tokens is False + assert getattr(session, "async_tokens_to_discard", 0) == 0 assert session.num_output_placeholders == 0 assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20] assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 9, 10, 20] diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 1fac5518e..a6d5edd8f 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -434,7 +434,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Downstream async-chunk stages receive real payloads from the # connector. This update only resumes polling for the next segment. self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id) - request.discard_latest_async_tokens = True + request.async_tokens_to_discard = 1 request.num_output_placeholders = 0 request.spec_token_ids = [] request._output_token_ids.clear() @@ -632,13 +632,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): """ req_id = session.request_id self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids) - session.discard_latest_async_tokens = False outstanding_async_tokens = getattr(session, "num_output_placeholders", 0) if outstanding_async_tokens > 0: # Async scheduling may already have sampled the previous # segment's next token. Drop that late token instead of # appending it to the new streaming segment. - session.discard_latest_async_tokens = True + session.async_tokens_to_discard = 1 session.num_computed_tokens -= session.num_output_placeholders session.num_output_placeholders = 0 session.spec_token_ids = [] 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 04fb81bbb..121fb79ec 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 @@ -542,7 +542,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # Requests that have loaded chunk from last round # of schedule, but have not scheduled continue - if request.request_id in self.finished_requests: + if self.is_done_receiving_chunks(request.request_id): + request.additional_information = None continue # Requests that waiting for chunk self.load_async(request)