[Feat] Support step-boundary abort in diffusion

[asukaqaq-s]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

based on：RFC 874 and

PR
#1368
#1625

merging their work and adding further modifications to support stepwise execution/scheduling and request abort.

Summary

Support aborting both running and waiting diffusion requests, with immediate cancellation and zero post-abort performance overhead.

Pipeline / Runner

• Add step-level execution support (--step-execution) so the abort queue is checked between each denoise step, not only after the entire request completes

Worker / Executor

• Add execute_step() / execute_stepwise() path that forwards per-step RunnerOutput (with step_index and finished fields) back to the engine
• execute_request() wraps full-request DiffusionOutput into RunnerOutput for unified handling at the engine layer
  Introduce DiffusionRequestAbortedError for structured abort signaling across layers

Scheduler / Engine / Entrypoints

• StepScheduler checks abort status between steps; engine-level guard short-circuits update_from_output for already-aborted requests to prevent infinite re-queue loop
• AsyncOmniDiffusion uses Future-based request tracking: queued requests are cancelled instantly via future.cancel(), running requests are forwarded to engine.abort()
• _request_id_to_sched_req_id mapping supports aborting batched requests by any constituent request ID

Benchmark (Qwen-Image, 50 steps, single GPU)

Configuration	Resolution	Baseline	Post-abort	Total	Overhead
Original vllm-omni	512x512	10.2s	27.7s	42.9s	+17.5s (+172%)
Original vllm-omni	1024x1024	21.2s	116.0s	142.3s	+94.8s (+447%)
Abort waiting (entrypoints rewrite)	512x512	6.8s	5.8s	17.6s	-1.0s (~0%)
Abort waiting (entrypoints rewrite)	1024x1024	20.3s	34.8s	60.2s	+14.5s (+71%)
Abort waiting + running (step-exec)	512x512	6.5s	5.6s	17.1s	-0.9s (~0%)
Abort waiting + running (step-exec)	1024x1024	20.2s	20.9s	46.1s	+0.7s (+3%)

Conclusions:

• Original 1024 post-abort takes 116s — severe blocking, abort is completely ineffective, the previous request must finish before the next one can start
• Entrypoints rewrite solves abort for queued requests (future.cancel()), 512 already shows zero overhead
• Step-execution solves abort for running requests (checks abort_queue between each denoise step), 1024 overhead drops from +71% to +3%
• No baseline performance regression (512 being faster is due to other optimizations on the dev-exp branch)

verify

Parallel Mode	Baseline (s)	Post-abort (s)	Abort Recovery	Status
--ring 2	14.14	14.36	5/5	Pass
--usp 2	12.87	13.10	5/5	Pass
--cfg-parallel-size 2	21.08	21.22	5/5	Pass

Key findings:
All three parallel strategies (Ring, Ulysses SP, CFG) work correctly with --step-execution.
Post-abort latency is consistent with baseline — no performance degradation after repeated aborts.
CFG parallel requires negative_prompt in the request to activate; without it the second GPU idles.

E2E test

server

baseline：
vllm serve Qwen/Qwen-Image --omni --port 8091
outrs：
vllm serve Qwen/Qwen-Image --omni --port 8091 --step-execution

client

curl -s http://localhost:8091/v1/chat/completions   -H "Content-Type: application/json"   -d '{
    "messages": [
      {"role": "user", "content": "A flower"}
    ],
    "extra_body": {
      "height": 512,
      "width": 512,
      "num_inference_steps": 50,
      "true_cfg_scale": 4.0,
      "seed": 42
    }
  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png

Test plan

• 18 unit tests pass (test_diffusion_stepmodel.py): scheduler, engine, async abort
• End-to-end curl abort test: 5x SIGINT interrupts followed by successful generation
• Verified future.cancel() for queued requests (never reach engine)
• Verified engine.abort() forwarding for running requests
• Verified post-abort recovery with zero performance overhead

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: dfcffff64f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Pass request-mode result in scheduler's expected type

In the default non-step path (step_execution=False), self.execute_fn is executor.execute_request, which returns a RunnerOutput, but this value is passed straight into RequestScheduler.update_from_output. RequestScheduler.update_from_output expects a DiffusionOutput and dereferences output.error, so the first normal request can fail with AttributeError before any result is returned.

Useful? React with 👍 / 👎.

[asukaqaq-s]

fixed

[chatgpt-codex-connector]

Call abort_request with the scheduler's actual signature

_process_aborts_queue builds a list of request IDs, then _abort_requests forwards that list to scheduler.abort_request with an extra status argument. Current scheduler implementations accept only a single req_id parameter, so any abort processing can raise TypeError instead of marking requests aborted.

Useful? React with 👍 / 👎.

[asukaqaq-s]

fixed

[david6666666]

@wtomin @SamitHuang @david6666666 PTAL, thx

[asukaqaq-s]

All conflicts resolved. Summary:

• data.py: Kept both — step_execution field (ours) + is_moe property (main)
• diffusion_engine.py: Kept both — exec_total_time timing (main) + abort check (ours)
• scheduler.py: Deleted (refactored into sched/); ported unpack_diffusion_output_shm from main into multiproc_executor.py at both dequeue paths

Note: The inline diffusion path introduced in #1715 bypasses AsyncOmniDiffusion,
so engine.abort() is never called on request cancellation. This breaks step-level
abort for --step-execution mode.

TO Fix: call self._inline_engine.engine.abort(request_id) in _generate_inline's
CancelledError handler. Will push a fix in the next commit.

[asukaqaq-s]

• Fixed inline diffusion request cancellation in AsyncOmni.

After merging origin/main, single-stage diffusion started using the new inline path, which called synchronous OmniDiffusion.generate() through run_in_executor(). That kept the stage-worker RPC optimization, but broke abort behavior: cancelling the asyncio task only cancelled the outer future, not the underlying running thread, and queued inline requests were not cleanly abortable.

The fix keeps inline execution fully in-process, but switches the inline engine to AsyncOmniDiffusion,which already has proper request-state tracking, queued-future cancellation, and forwarding of running-request aborts into DiffusionEngine. I also updated inline cancellation handling to explicitly abort on CancelledError, and changed inline AsyncOmni.abort() to call the async wrapper’s abort() instead of bypassing it and talking directly to the raw engine.

Result: inline diffusion now preserves the RPC-saving fast path while restoring both step-level interruption for running requests and clean cancellation for waiting requests.

[asukaqaq-s]

Bug fixed: CFG parallel previously deadlocked in step-execution mode because non-rank-0 workers received None from predict_noise_maybe_with_cfg() and skipped the collective broadcast in scheduler_step_maybe_with_cfg(). Fixed by always calling step_scheduler() regardless of noise_pred value.

[hsliuustc0106]

@ApsarasX PTAL

[Gaohan123]

@wtomin @ZJY0516 @SamitHuang PTAL

[yJader]

Merge conflicts with main have been resolved

[hsliuustc0106]

do we have a validation for this arg if it's accidentally triggered in non-diffusion models? or do we expect this will help in ar+dit model as well?

[asukaqaq-s]

1. If this argument is triggered in non-diffusion engine scenarios, it should not have any effect.
2. For AR+DiT models, I have also discussed this in other community issues. If they adopt the diffusion engine, step-execution will be effective, and I expect it will help as well.

[wtomin]

Thanks for your great contributions. I have a few comments:

1. Currently only qwen-image supports step-level execution, but it is an important feature that should be included in docs/user_guide/diffusion_features.md, and please give the feature compatibility information in this doc, too.
2. In tests/e2e/online_serving/test_qwen_image_expansion.py test script, please add --step-execution test case.

[wtomin]

LGTM