[Core][Frontend] Support request-level batching for diffusion pipelines

[yJader]

Purpose

Implements Phase 2 of RFC #3550. This completes the diffusion request contract cleanup by making OmniDiffusionRequest represent exactly one logical request with one prompt, while runtime batching is represented by scheduler output and runner-side batch structures.

In this PR:

Request contract

• OmniDiffusionRequest.prompts is replaced by a single prompt field; request identity remains the required single request_id. __post_init__ rejects list inputs defensively.
• Scheduler output sends NewRequestData(request_id, req), carrying the already-initialized OmniDiffusionRequest directly so the executor/worker forward it instead of rebuilding it (which would re-trigger __post_init__ and corrupt sentinel fields like guidance_scale_provided). Fields like prompt / sampling_params / kv_sender_info are read off req rather than duplicated on the payload.

Scheduler batching

• Scheduler batching is based on compatible independent requests via SamplingParamsKey, with conservative LoRA homogeneity constraints. num_outputs_per_prompt is added to the key so requests producing differently-shaped outputs do not get co-batched.

Request-level batch & execution path

• DiffusionRequestBatch is introduced as the request-level batch abstraction above the step/tensor-level InputBatch (StepInputBatch alias). It wraps list[OmniDiffusionRequest] and exposes prompts / sampling_params / request_id / kv_sender_info compatibility properties so pipeline forward bodies stay close to upstream. Worker-side DiffusionRequestState (single prompt) is retained for the stepwise path only.
• Pipeline forward is unified to forward(req: DiffusionRequestBatch) -> list[DiffusionOutput]. A class attribute supports_request_batch advertises whether the pipeline supports one batched forward.
• The executor exposes two request-mode entries — execute_request (one execute_model RPC per scheduled request, the preserved upstream path) and execute_batch (a single execute_model_batch RPC carrying the whole DiffusionSchedulerOutput). The engine resolves request-batch capability at initialization from the configured pipeline class, including custom pipeline classes, and binds batch-capable pipelines directly to execute_batch for all request-mode cycles, including single-request cycles.
• On the batch path the worker execute_model_batch builds the DiffusionRequestBatch and runs per-request setup (KV transfer, generator/seed), one cache refresh, and one LoRA activation per batch. Non-batch pipelines stay on the per-request execution path.
• BatchRunnerOutput routes results per request_id, including per-request success, error, and abort outputs.
• Diffusion IPC SHM packing now traverses both shallow RunnerOutput.result wrappers and nested BatchRunnerOutput.runner_outputs[*].result, so request-mode batch outputs with large tensors continue to use the SHM transfer path instead of falling back to pickle IPC.
• Request-mode batching has an opt-in admission wait window through request_batch_max_wait_ms. The engine can wait briefly before the first schedule() of a wave so bursty compatible requests can accumulate, while the default 0.0 keeps the existing no-wait behavior. The scheduler exposes waiting/running queue counters for this logic, and the serve CLI / default diffusion stage config forward the option into OmniDiffusionConfig.

API boundary (breaking change)

• The packed diffusion list-prompt submission path no longer constructs packed diffusion requests. AsyncOmni.generate(prompt=[...]) rejects early with a clear error when a diffusion stage is present; StagePool.submit_initial() / submit_update() keep a defensive list check.

Pipeline & test migration

• Diffusion model pipelines and focused tests are migrated to the DiffusionRequestBatch forward contract; non-batch models change only signature/return type, batch-capable models add supports_request_batch = True plus per-request output splitting.
• output_formatter.py also follows the single-request contract: it reads request.prompt, emits one OmniRequestOutput per request, and no longer keeps a request.prompts compatibility branch.

Compatibility/docs/release notes

• This is a user-facing breaking change for diffusion AsyncOmni.generate(prompt=[...]) packed-list behavior. Users should submit multiple independent requests to use automatic scheduler batching.
• vLLM-Omni docs now describe request-level diffusion batching, max_num_seqs, and the optional request_batch_max_wait_ms admission wait setting.

Test Plan

Core diffusion request-mode regression

Focused coverage for request contract, scheduler batching, runner behavior,
engine routing, multiprocess dispatch, and stage process integration:

python -m pytest \
  tests/diffusion/test_diffusion_request.py \
  tests/diffusion/test_diffusion_scheduler.py \
  tests/diffusion/test_diffusion_step_pipeline.py \
  tests/diffusion/test_diffusion_model_runner.py \
  tests/diffusion/test_diffusion_engine_cleanup.py \
  tests/diffusion/test_diffusion_engine_rpc_routing.py \
  tests/diffusion/test_multiproc_engine_concurrency.py \
  tests/diffusion/test_stage_diffusion_proc.py \
  -q

Touched backend and model coverage

Representative backend/model unit coverage for changed diffusion paths, excluding
hardware-dependent advanced model tests:

python -m pytest \
  tests/diffusion/diffusion_backend/test_diffusers_backend.py \
  tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py \
  tests/diffusion/models/dmd2/test_dmd2_scheduler.py \
  tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py \
  tests/diffusion/models/ovis_image/test_ovis_image.py \
  tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py \
  tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py \
  tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py \
  tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py \
  tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py \
  -m "not advanced_model" \
  -q

Output formatter regression

Focused coverage for the single-prompt output formatting contract and the removal of the old multi-prompt compatibility path:

python -m pytest \
  tests/diffusion/test_diffusion_output_formatter.py \
  -q

IPC coverage

Large-tensor SHM packing coverage for request-mode and step-pipeline outputs:

python -m pytest \
  tests/diffusion/test_diffusion_ipc.py \
  tests/diffusion/test_diffusion_step_pipeline.py::TestIPC::test_pack_unpack_runner_output_shm \
  -q

Request-batch capability routing

Focused coverage for the supports_request_batch route from engine dispatch to
worker-side RequestBatch execution:

python -m pytest \
  tests/diffusion/test_diffusion_engine.py \
  tests/diffusion/test_diffusion_model_runner.py \
  tests/diffusion/test_multiproc_engine_concurrency.py::TestRequestModeDispatch \
  -q

Request-batch admission and CLI forwarding

Focused coverage for request-batch capability detection, admission wait,
scheduler queue counters, and the request_batch_max_wait_ms CLI/config path:

python -m pytest \
  tests/diffusion/test_diffusion_engine.py::TestRequestBatchCapability \
  tests/diffusion/test_diffusion_engine.py::TestRequestBatchAdmission \
  tests/diffusion/test_diffusion_scheduler.py \
  tests/entrypoints/test_async_omni_diffusion_config.py \
  -q

Qwen-Image performance comparison

Compare Qwen-Image performance with the same setup: 512x512, 20 denoising
steps, FLASH_ATTN, single A100, 1 warmup run, and 30 measured runs. The
baseline uses merge-base packed prompt-list; the candidate and StepScheduler
checks use two concurrent single-prompt requests with max_num_seqs: 2.

Residual old-contract check

Run a residual old-contract grep before submission:

rg "req\\.prompts|request\\.prompts|OmniDiffusionRequest\\(.*prompts|add_batch_request_async" \
  vllm_omni tests

Expected remaining add_batch_request_async hits should be limited to the abstract interface or explicit rejection stubs. Remaining req.prompts / request.prompts hits should be in DiffusionRequestBatch compatibility properties, pre-processing helpers, comments, or rejection messages — not in single-request execution paths.

Test Result

• Core request-contract, scheduler, runner, engine-routing, stage, and multiproc coverage: 137 passed, 23 warnings in 66.91s
• Representative touched backend/model unit coverage: 177 passed, 2 deselected, 20 warnings in 3.37s
• IPC coverage: 10 passed, 16 warnings in 2.56s
• Request-batch capability routing coverage: 25 passed, 16 warnings in 3.87s
• Request-batch admission and CLI forwarding focused coverage: 72 passed, 20 warnings in 3.51s
• KV flow regression coverage after OpenPI request-field migration: 18 passed, 17 warnings in 0.84s
• Residual old-contract grep rerun. Remaining add_batch_request_async / req.prompts / request.prompts
  hits are in compatibility or rejection paths; focused prompts= check now only finds the
  StageDiffusionProc rejection test.

Static request-batch route checks:

• Qwen-Image declares supports_request_batch = True.
• Request mode keeps step_execution disabled, sets max_num_seqs: 2.
• DiffusionEngine selects execute_batch for request-batch-capable pipelines when step_execution is disabled.

Qwen-Image request batching performance comparison:

• Model: Qwen-Image
• Setup: 512x512, 20 steps, FLASH_ATTN, single A100, 1 warmup run, 30 measured runs
• Candidate uses request-mode scheduler batching with step_execution disabled. StepScheduler benchmark is intentionally not used for this PR result.

Path	Commit	ok/error	mean ms	p50 ms	p90 ms	amortized mean ms/image
Merge-base prompt-list	b8f68174	30/0	4657.41	4658.68	4676.52	2328.70
Current request scheduler batch	7771f3c2	30/0	4664.54	4660.43	4701.68	2332.27

Current request scheduler batch delta vs merge-base prompt-list: mean +0.15%, p50 +0.04%, p90 +0.54%, amortized mean +0.15%. No significant performance regression is observed with the new request-level batching path.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[SamitHuang]

how to control the batch size?

[yJader]

Use max_num_seqs, consistent with the step-wise batching path

[SamitHuang]

I think it's better to compare step-level batching and request-level batching on time performance

[SamitHuang]

This refactor introduces request-level batching but is missing some critical validation and architectural cleanup.

• vllm_omni/diffusion/request.py:35: If diffusion models later support multi-part single prompts (e.g., interleaved text and image dicts in a list), this check (isinstance(self.prompt, list)) will incorrectly reject them. Consider checking if the elements are valid single-prompt types rather than rejecting all lists.
• vllm_omni/diffusion/worker/diffusion_model_runner.py:310: Hardcoding tea_cache specific logic in the model runner breaks encapsulation. The cache backend interface should define whether it requires num_inference_steps, or gracefully handle None without runner-side workarounds.
• vllm_omni/entrypoints/async_omni.py:313: Since this is a user-facing API breaking change, ensure that this error message also points users to the updated documentation or an example script for the new independent request submission pattern.

[SamitHuang]

seems it does not support batching variable-length prompts in my test

[SamitHuang]

@knlnguyen1802 PTAL

[hsliuustc0106]

• qwen-image
• qwen-image-edit
• WAN 2.1/2.2
• SDXL
• LTX 2.3

any updates?

[yJader]

• qwen-image
• qwen-image-edit
• WAN 2.1/2.2
• SDXL
• LTX 2.3

any updates?

Currently request batch supports Qwen-Image, Flux, SD3.5, and LTX2.3. The implementation reuses/adapts the existing prompt-list batching path in pipelines, so we first enabled it for the models we care about most.

[hsliuustc0106]

I think this PR needs to polish for further abstraction, but let's wait for the next version