wehub-resource-sync
164 lines
6.7 KiB
Markdown
164 lines
6.7 KiB
Markdown
# 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
|
|
- per-request progress and completion remain independent
|
|
|
|
Here, "continuous batching" means the step-wise path enabled by
|
|
`step_execution=True`. Request-mode `DiffusionRequestBatch` is static
|
|
request-level batching for one full pipeline `forward()` call; it does not
|
|
admit or remove requests between denoise steps.
|
|
|
|
## 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 three 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
|
|
key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
|
|
targeting different adapters or scales run in separate batches and the
|
|
worker can activate exactly one adapter per step.
|
|
|
|
The scheduler batching unit is one logical `OmniDiffusionRequest`. In the
|
|
step-wise path, runtime tensor batching is represented as `StepInputBatch`. For
|
|
request-mode prompt semantics, see
|
|
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
|
|
|
|
## 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.
|
|
|
|
When `step_execution=True`, the engine routes work through the step-wise
|
|
executor path. The continuous batching behavior is defined by scheduler-side
|
|
compatibility gating and runner-side `StepInputBatch` packing.
|
|
|
|
## Current Limitations
|
|
|
|
- Experimental feature; use `max_num_seqs=1` for the older conservative path.
|
|
- Only native pipelines that already support `step_execution=True`.
|
|
- Only homogeneous batches keyed by `SamplingParamsKey` are supported.
|
|
- `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)
|