Isolate Qwen-Image 2.1 prefill and decode scheduling (#7770)

Signed-off-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>

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:

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])
 

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

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: