From b0a309e5e7908b3c2d70d60b1a74f26f9545ef8a Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Fri, 18 Sep 2026 14:11:22 +0800 Subject: [PATCH] Isolate Qwen-Image 2.1 prefill and decode scheduling (#7770) Signed-off-by: Tianyu Guo --- recipes/Qwen/Qwen-Image-2.1.md | 6 ++++++ .../qwen_image_21/pipeline_qwen_image_21.py | 6 ++++-- vllm_omni/diffusion/request.py | 2 ++ vllm_omni/diffusion/sched/step_scheduler.py | 16 +++++++++++++++- 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/recipes/Qwen/Qwen-Image-2.1.md b/recipes/Qwen/Qwen-Image-2.1.md index 33ab34743..056b21471 100644 --- a/recipes/Qwen/Qwen-Image-2.1.md +++ b/recipes/Qwen/Qwen-Image-2.1.md @@ -96,6 +96,12 @@ vllm serve Qwen/Qwen-Image-2.1 --omni \ --max-num-seqs 8 ``` +Step execution batches compatible requests at the same KV-cache phase. Once +running requests have completed their first denoising step, newly arriving +requests wait for that batch to finish before starting their own prefill. +`--max-num-seqs` controls batch capacity; mixed prefill/decode admission is not +supported yet. + ### Verification For a direct API smoke test: diff --git a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py index 75c191d0d..73d57a542 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py +++ b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py @@ -135,6 +135,7 @@ def get_qwen_image_21_pre_process_func( request: OmniDiffusionRequest, ): """Pre-process requests for QwenImage21Pipeline.""" + request.allow_mixed_step_phases = False prompt = request.prompt multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None @@ -1213,11 +1214,12 @@ class QwenImage21Pipeline( img_mask = self._append_target_slots(torch.cat(mask_rows, dim=0), latents.shape[1]) negative_img_mask = None if input_batch.do_true_cfg: + negative_seq_len = input_batch.negative_prompt_embeds.shape[1] neg_rows = [] for state in states: row = state.extra["negative_image_pad_mask"] - if row.shape[1] < max_seq_len: - row = F.pad(row, (0, max_seq_len - row.shape[1]), value=False) + if row.shape[1] < negative_seq_len: + row = F.pad(row, (0, negative_seq_len - row.shape[1]), value=False) neg_rows.append(row) negative_img_mask = self._append_target_slots(torch.cat(neg_rows, dim=0), latents.shape[1]) diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 9e9a4bfbb..f90f843a5 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -73,6 +73,8 @@ class OmniDiffusionRequest: # A model preprocessor may keep selected requests on the legacy full-forward # path even when the engine is globally configured for step execution. use_step_execution: bool = True + # Whether first and later denoising steps may share a step batch. + allow_mixed_step_phases: bool = True # KV-recv wall-clock (ms), set by the runner's _prepare_request_for_forward # and carried to DiffusionOutput for the vllm_omni:diffusion_kv_load_s metric. kv_recv_ms: float = 0.0 diff --git a/vllm_omni/diffusion/sched/step_scheduler.py b/vllm_omni/diffusion/sched/step_scheduler.py index a5a7a04ae..3b5f515ef 100644 --- a/vllm_omni/diffusion/sched/step_scheduler.py +++ b/vllm_omni/diffusion/sched/step_scheduler.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -13,6 +13,7 @@ from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler from vllm_omni.diffusion.sched.interface import ( DiffusionRequestStatus, DiffusionSchedulerOutput, + SchedulerRequestState, ) if TYPE_CHECKING: @@ -62,6 +63,19 @@ class StepScheduler(BaseScheduler): ) return request_id + def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool: + if not super()._can_schedule_waiting(state): + return False + + is_first_step = self._request_progress[state.request_id].current_step == 0 + for request_id in self._running: + running = self._request_states[request_id] + if state.req.allow_mixed_step_phases and running.req.allow_mixed_step_phases: + continue + if (self._request_progress[request_id].current_step == 0) != is_first_step: + return False + return True + def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]: scheduled_request_ids = sched_output.scheduled_request_ids if not scheduled_request_ids: