From e50a0665485178f428e8e37ad7a8cff2fc413484 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Sun, 3 May 2026 20:02:22 +0800 Subject: [PATCH] [BugFix]: Fix async scheduer transfer exceed KV cache (#3318) Signed-off-by: princepride Co-authored-by: Hongsheng Liu --- vllm_omni/core/sched/omni_ar_scheduler.py | 40 ++++++++++++++++--- .../model_executor/models/bagel/bagel.py | 8 +++- 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 677d8f671..ccdcf1643 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStat from vllm.logger import init_logger from vllm.v1.core.sched.async_scheduler import AsyncScheduler as VLLMScheduler from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.scheduler import Scheduler as SyncScheduler from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs from vllm.v1.metrics.perf import PerfStats @@ -80,6 +81,31 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Snapshot prompt length for each streaming input update self._new_prompt_len_snapshot: dict[str, int] = {} + def _get_confirmed_num_computed_tokens(self, request: Request) -> int: + """num_computed_tokens minus async placeholders (KV actually on GPU).""" + return request.num_computed_tokens - request.num_output_placeholders + + def _update_request_with_output(self, request: Request, new_token_ids: list[int]) -> tuple[list[int], bool]: + """Append output tokens, then cache blocks up to the confirmed count + so KV transfer never sees blocks whose data has not been computed yet. + """ + if request.discard_latest_async_tokens: + request.discard_latest_async_tokens = False + return [], False + + status_before_update = request.status + + new_token_ids, stopped = SyncScheduler._update_request_with_output(self, request, new_token_ids) + + request.num_output_placeholders -= len(new_token_ids) + assert request.num_output_placeholders >= 0 + + if status_before_update == RequestStatus.RUNNING: + confirmed = self._get_confirmed_num_computed_tokens(request) + self.kv_cache_manager.cache_blocks(request, confirmed) + + return new_token_ids, stopped + def _get_kv_transfer_criteria(self) -> dict | None: # Note: vllm_config is available in Scheduler after super().__init__ if not hasattr(self, "vllm_config"): @@ -146,10 +172,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return True return False + # seq_len for KV transfer must exclude async placeholders. + confirmed_computed = self._get_confirmed_num_computed_tokens(request) + if criteria_type == "prefill_finished": - if request.num_computed_tokens >= request.num_prompt_tokens: + if confirmed_computed >= request.num_prompt_tokens: self.transfer_triggered_requests.add(request.request_id) - self._mark_request_for_kv_transfer(request.request_id, request.num_computed_tokens) + self._mark_request_for_kv_transfer(request.request_id, confirmed_computed) actually_queued = request.request_id in self.requests_needing_kv_transfer if stop_decode_on_trigger and actually_queued: @@ -169,9 +198,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): try: idx = new_token_ids.index(target_token_id) tokens_to_exclude = len(new_token_ids) - (idx + 1) - snapshot_len = request.num_computed_tokens - tokens_to_exclude + snapshot_len = confirmed_computed - tokens_to_exclude except ValueError: - snapshot_len = request.num_computed_tokens + snapshot_len = confirmed_computed self._mark_request_for_kv_transfer(request.request_id, snapshot_len) actually_queued = request.request_id in self.requests_needing_kv_transfer @@ -622,7 +651,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): ) else: self.waiting_for_transfer_free.add(request_id) - self._mark_request_for_kv_transfer(request_id, request.num_computed_tokens) + confirmed_computed = self._get_confirmed_num_computed_tokens(request) + self._mark_request_for_kv_transfer(request_id, confirmed_computed) # Return KV transfer metadata so it propagates to RequestOutput if request_id in self.requests_needing_kv_transfer: transfer_data = self.requests_needing_kv_transfer[request_id] diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index bb1c59e58..a805291dd 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -567,7 +567,12 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): return None if num_computed_tokens is not None and "image_shape" in meta: prefill_rope = meta["ropes"][0] if meta.get("ropes") else 0 - if num_computed_tokens > prefill_rope: + prefill_position_count = meta.get("prefill_position_count") + if prefill_position_count is not None: + num_decoded = num_computed_tokens - prefill_position_count + if num_decoded > 0: + meta["ropes"] = [prefill_rope + num_decoded] + elif num_computed_tokens > prefill_rope: meta["ropes"] = [num_computed_tokens] return meta @@ -849,6 +854,7 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): { "ropes": [rope], "image_shape": [img_H, img_W], + "prefill_position_count": int(end - start), } ) img2img_idx += 1