[Feat] DiffusionEngine Support async batch infer (#2729)
Signed-off-by: Semmer <semmer@live.cn> Signed-off-by: jader <yjader@foxmail.com> Signed-off-by: asukaqaq-s <1311722138@qq.com> Co-authored-by: asukaqaq-s <1311722138@qq.com> Co-authored-by: jader <yjader@foxmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
This commit is contained in:
32 files changed
+1951
-190
No files matched your search
@@ -112,6 +112,7 @@ nav:
|
||||
- design/feature/async_chunk.md
|
||||
- design/feature/vae_parallel.md
|
||||
- design/feature/diffusion_step_execution.md
|
||||
- design/feature/diffusion_continuous_batching.md
|
||||
- Module Design:
|
||||
- design/module/ar_module.md
|
||||
- design/module/dit_module.md
|
||||
|
||||
@@ -737,6 +737,16 @@ step-scoped phases. The reference implementation is
|
||||
Do not enable `step_execution=True` until those four methods are implemented
|
||||
and validated against the request-level path.
|
||||
|
||||
If you want the pipeline to work with the experimental batched step-wise path
|
||||
(`max_num_seqs > 1`), also see:
|
||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
||||
|
||||
If you expose this in example scripts or recipes, keep it opt-in. Surface
|
||||
runtime features like `step_execution` as optional flags instead of silently
|
||||
turning them on. For Qwen-Image-style serving examples, document
|
||||
`--step-execution` as the feature gate and `--max-num-seqs N` as the
|
||||
companion batching knob.
|
||||
|
||||
### Cache Acceleration
|
||||
|
||||
#### TeaCache
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
# Continuous Batching for Step-Wise Diffusion
|
||||
|
||||
!!! warning "Experimental Feature"
|
||||
This feature is experimental. It currently applies only to native
|
||||
diffusion pipelines running with `step_execution=True`.
|
||||
|
||||
This document describes the batching extension built on top of
|
||||
[Diffusion Step Execution](diffusion_step_execution.md). The base
|
||||
step-execution contract is unchanged. The batching work is mainly in the
|
||||
scheduler and runner layers.
|
||||
|
||||
## Why It Helps
|
||||
|
||||
Step-wise execution breaks a long denoise loop into scheduler-visible units.
|
||||
That gives the runtime a place to admit other compatible requests between
|
||||
steps instead of waiting for an entire request to finish.
|
||||
|
||||
This matters most in low-MFU or bursty serving scenarios:
|
||||
|
||||
- one request's denoise step may not fully saturate the GPU
|
||||
- several compatible requests can share the same denoise forward pass
|
||||
- throughput and device utilization can improve without changing request-local
|
||||
scheduler state
|
||||
|
||||
This is **not** a guaranteed single-request latency win. The main benefit is
|
||||
usually higher utilization and better throughput when the workload contains
|
||||
multiple in-flight compatible requests.
|
||||
|
||||
## Overview
|
||||
|
||||
With continuous batching enabled:
|
||||
|
||||
- the scheduler may keep multiple compatible requests active at the same time
|
||||
- the runner packs request-local step state into one `InputBatch`
|
||||
- `denoise_step()` runs on that batch
|
||||
- `step_scheduler()` and `post_decode()` still run per request
|
||||
|
||||
The current implementation is conservative:
|
||||
|
||||
- only compatible requests are batched together
|
||||
- request-mode diffusion still runs with `max_num_seqs=1`
|
||||
- per-request progress and completion remain independent
|
||||
|
||||
## Enablement
|
||||
|
||||
Use `--step-execution` as the feature gate, then increase `--max-num-seqs`
|
||||
above `1` if you want batching:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
```
|
||||
|
||||
`--max-num-seqs 1` keeps the step-wise path without enabling batching.
|
||||
|
||||
For a reproducible replay flow using the bundled serving benchmark, see the
|
||||
Qwen-Image replay commands in
|
||||
[`benchmarks/diffusion/README.md`](gh-file:benchmarks/diffusion/README.md)
|
||||
and
|
||||
[`benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md`](gh-file:benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md).
|
||||
|
||||
## Scheduler
|
||||
|
||||
The scheduler derives its batch capacity from `max_num_seqs` through
|
||||
`max_num_running_reqs`.
|
||||
|
||||
Batch admission is gated by
|
||||
[`SamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
|
||||
which is built from shape-sensitive and CFG-sensitive sampling fields. This is
|
||||
the core correctness rule for batching: requests are only co-batched when they
|
||||
share the same denoise tensor contract.
|
||||
|
||||
There are two important details:
|
||||
|
||||
- `num_inference_steps` is not part of the key, so requests with different
|
||||
total step counts can still share a batch
|
||||
- requests also do not need to be at the same current denoise progress; active
|
||||
requests can continue batching even when their current step indices diverge
|
||||
- admission is still FIFO, so an incompatible request at the head of the
|
||||
waiting queue blocks later compatible requests
|
||||
|
||||
Today that compatibility rule is still shape-sensitive. `height`, `width`,
|
||||
`num_frames`, and CFG-related fields remain part of the key, so different
|
||||
resolutions or incompatible guidance settings do **not** co-batch yet.
|
||||
|
||||
The current batching unit is one `OmniDiffusionRequest`. Requests with
|
||||
multiple prompts do not participate in batching today.
|
||||
|
||||
## Runner
|
||||
|
||||
The runner keeps persistent per-request execution state in
|
||||
[`DiffusionRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
|
||||
while the scheduler owns a separate lightweight request state for queueing and
|
||||
lifecycle tracking.
|
||||
|
||||
For each step, the runner builds an
|
||||
[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the
|
||||
active request states:
|
||||
|
||||
- prompt embeddings and masks are normalized and padded
|
||||
- dynamic tensors such as `latents` and `timesteps` are gathered each step
|
||||
- buffers are reused when batch composition stays the same
|
||||
|
||||
The step-wise batched path is:
|
||||
|
||||
1. Run `prepare_encode()` for newly admitted requests.
|
||||
2. Build or refresh `InputBatch`.
|
||||
3. Run one batched `denoise_step(input_batch)`.
|
||||
4. Slice the batched `noise_pred` back per request.
|
||||
5. Run per-request `step_scheduler()`.
|
||||
6. Run `post_decode()` only for requests that finished denoising.
|
||||
7. Scatter updated latents back into persistent request state with
|
||||
[`scatter_latents()`](gh-file:vllm_omni/diffusion/worker/input_batch.py).
|
||||
|
||||
This keeps the shared work limited to the denoise forward pass while preserving
|
||||
request-local scheduler state and outputs.
|
||||
|
||||
## Engine
|
||||
|
||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) provides
|
||||
the background loop and async add-request path needed for multiple requests to
|
||||
accumulate in the scheduler.
|
||||
|
||||
This is supporting infrastructure, not the main design point. The batching
|
||||
behavior is defined by scheduler-side compatibility gating and runner-side
|
||||
batch packing.
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- Experimental feature; use `max_num_seqs=1` for the older conservative path.
|
||||
- Only native pipelines that already support `step_execution=True`.
|
||||
- Request-mode diffusion still clamps `max_num_seqs` back to `1`.
|
||||
- Only homogeneous batches keyed by `SamplingParamsKey` are supported.
|
||||
- Multi-prompt requests are not batched.
|
||||
- `cache_backend`, KV transfer, and other request-mode extras are not wired
|
||||
into the batched step-wise path yet.
|
||||
- Future work can relax the current same-shape restriction with richer
|
||||
heterogeneous batching policies such as bucketing or padded execution for
|
||||
different resolutions.
|
||||
|
||||
## Related Files
|
||||
|
||||
- Scheduler base:
|
||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
||||
- Scheduler interface:
|
||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Step scheduler:
|
||||
[`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
||||
- Runner:
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Input batch:
|
||||
[`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
||||
- Tests:
|
||||
[`tests/diffusion/test_diffusion_scheduler.py`](gh-file:tests/diffusion/test_diffusion_scheduler.py)
|
||||
@@ -1,4 +1,4 @@
|
||||
# Diffusion Step Execution
|
||||
# Adding Step Execution Support for Diffusion Pipelines
|
||||
|
||||
This guide documents vLLM-Omni's stepwise diffusion contract for model authors
|
||||
and contributors implementing `step_execution=True` support for a diffusion
|
||||
@@ -7,12 +7,20 @@ pipeline.
|
||||
For end-user enablement, supported models, and current limitations, see
|
||||
[Step Execution](../../user_guide/diffusion/step_execution.md).
|
||||
|
||||
This document describes the base step-execution contract only. For the
|
||||
experimental batching policy layered on top of the step-wise path, see
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
|
||||
## Current Support Scope
|
||||
|
||||
`step_execution` is **not** a generic diffusion toggle. It only works for
|
||||
pipelines that implement the segmented stateful contract in
|
||||
[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
|
||||
|
||||
This page is intentionally author-facing. Treat runtime enablement
|
||||
(`step_execution=True` in Python or `--step-execution` in serving) as an
|
||||
opt-in user knob layered on top of the implementation contract below.
|
||||
|
||||
Current in-tree support:
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
@@ -22,7 +30,9 @@ Current in-tree support:
|
||||
|
||||
Current engine/runtime limitations:
|
||||
|
||||
- `StepScheduler` only schedules `batch_size=1`.
|
||||
- Continuous batching with `max_num_seqs > 1` is experimental and documented in
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
Keep `max_num_seqs=1` if you want the older conservative behavior.
|
||||
- `cache_backend` is not supported in step mode.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
- Unsupported pipelines now fail early during model loading instead of failing on the first request.
|
||||
|
||||
@@ -10,6 +10,8 @@ This section contains design documents and architecture specifications for vLLM-
|
||||
|
||||
- [Disaggregated Inference](feature/disaggregated_inference.md)
|
||||
- [Ray-based Execution](feature/ray_based_execution.md)
|
||||
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
||||
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
||||
|
||||
## Module Design Documents
|
||||
|
||||
|
||||
@@ -201,7 +201,7 @@ class _BaseScheduler(SchedulerInterface):
|
||||
|
||||
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
|
||||
|
||||
- **Current constraint**: `max_num_running_reqs` remains `1` because the current engine path is still synchronous request-mode execution.
|
||||
- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
|
||||
|
||||
#### 2.4 Current `RequestScheduler` Policy
|
||||
|
||||
@@ -217,7 +217,7 @@ class RequestScheduler(_BaseScheduler):
|
||||
|
||||
- **FIFO request scheduling**: Waiting requests are promoted in queue order.
|
||||
|
||||
- **Single-request admission**: The current policy only admits one active request at a time.
|
||||
- **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
|
||||
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
|
||||
|
||||
@@ -237,7 +237,7 @@ while True:
|
||||
|
||||
- **No scheduler-owned IPC**: Scheduler no longer talks to workers directly.
|
||||
|
||||
- **Conservative concurrency**: The current request-mode implementation still allows only one active request at a time.
|
||||
- **Split concurrency model**: Request-mode diffusion remains single-active-request, while the step-wise path can keep multiple compatible requests running and advance them independently between denoise steps.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -8,6 +8,8 @@ implement the stepwise contract support it today.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -25,6 +27,18 @@ outputs = omni.generate(
|
||||
)
|
||||
```
|
||||
|
||||
### Serving
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
```
|
||||
|
||||
For serving, `--step-execution` enables the step-wise runtime. Continuous
|
||||
batching only becomes relevant when `--max-num-seqs > 1`.
|
||||
|
||||
## Supported Pipelines
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
@@ -32,9 +46,17 @@ outputs = omni.generate(
|
||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
||||
|
||||
!!! warning "Experimental continuous batching"
|
||||
When `--step-execution` is enabled and `max_num_seqs > 1` is configured,
|
||||
the step-wise path can batch
|
||||
compatible requests together. This is experimental. Requests with
|
||||
incompatible sampling parameters are intentionally kept in separate batches,
|
||||
and `max_num_seqs=1` remains the conservative default.
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- `step_execution` currently supports `batch_size=1` only.
|
||||
- Continuous batching under `step_execution` is experimental and only batches
|
||||
compatible requests.
|
||||
- `cache_backend` is not supported together with step execution.
|
||||
- Unsupported pipelines fail early during model loading.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
@@ -45,6 +67,11 @@ Use step execution only when you specifically need the pipeline to run through
|
||||
its stepwise request state machine. For normal diffusion inference, leave it
|
||||
disabled unless your workflow depends on this mode.
|
||||
|
||||
For Qwen-Image online serving, the usual progression is:
|
||||
|
||||
- start with `--step-execution --max-num-seqs 1` if you only need the step-wise path
|
||||
- increase `--max-num-seqs` after that if you want the experimental compatible-request batching behavior
|
||||
|
||||
If you are looking for general diffusion speedups, see
|
||||
[Diffusion Features Overview](../diffusion_features.md).
|
||||
|
||||
@@ -59,3 +86,7 @@ pipeline does not support step execution.
|
||||
If you want to add step execution support to a new diffusion pipeline, see the
|
||||
implementation guide:
|
||||
[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).
|
||||
|
||||
If you also want that pipeline to participate in the experimental batched
|
||||
step-wise path, see:
|
||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
||||
Reference in new issue
Block a user