anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feat] DiffusionEngine Support async batch infer (#2729)

Browse Source 
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:
Ting FUasukaqaq-sjaderHongsheng LiuGao Han
 authored and GitHub committed 2026-05-06 09:47:13 +12:00
1 parent 2b85474c23
commit e49fbd8a2d
32 files changed
+1951 -190
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

benchmarks/diffusion/README.md
+12
View File
@@ -130,9 +130,21 @@ Warmup flags:
- `--warmup-requests`: Number of warmup requests.
- `--warmup-num-inference-steps`: Steps used during warmup.
- `--warmup-concurrency`: Maximum concurrent warmup requests. Use this to warm
the same batch shape as the measured run instead of warming only batch=`1`.
- For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
Traffic / concurrency flags:
- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
### Batched warmup note
For batched serving runs, warm the same in-flight shape you plan to measure.
For example, a run with `--max-concurrency 8` should usually also use
`--warmup-requests 8 --warmup-concurrency 8`; otherwise the first measured
batch may still pay compile or CUDA-graph capture cost.
For a Qwen-Image continuous-batching replay example, see
[`performance_dashboard/qwen_image_serving_performance.md`](./performance_dashboard/qwen_image_serving_performance.md).
benchmarks/diffusion/diffusion_benchmark_serving.py
+59 -18
View File
@@ -751,6 +751,52 @@ async def iter_requests(
yield req
def _make_warmup_request(
requests_list: list[RequestFuncInput],
index: int,
args,
) -> RequestFuncInput:
warm_req = requests_list[index % len(requests_list)]
if args.warmup_num_inference_steps is not None:
warm_req = replace(
warm_req,
num_inference_steps=args.warmup_num_inference_steps,
)
if args.task == "t2v":
warm_req = replace(warm_req, num_frames=1)
return warm_req
async def _run_warmups(
requests_list: list[RequestFuncInput],
args,
session: aiohttp.ClientSession,
request_func,
) -> list[tuple[RequestFuncInput, RequestFuncOutput]]:
if not args.warmup_requests or not requests_list:
return []
warmup_requests = [_make_warmup_request(requests_list, i, args) for i in range(args.warmup_requests)]
warmup_concurrency = min(int(args.warmup_concurrency), len(warmup_requests))
warmup_semaphore = asyncio.Semaphore(warmup_concurrency)
print(
f"Running {len(warmup_requests)} warmup request(s) "
f"with num_inference_steps={args.warmup_num_inference_steps} "
f"and warmup_concurrency={warmup_concurrency}..."
)
async def limited_warmup_request_func(
req: RequestFuncInput,
) -> RequestFuncOutput:
async with warmup_semaphore:
return await request_func(req, session, None)
warmup_tasks = [asyncio.create_task(limited_warmup_request_func(req)) for req in warmup_requests]
warmup_outputs = await asyncio.gather(*warmup_tasks)
return list(zip(warmup_requests, warmup_outputs))
def calculate_metrics(
outputs: list[RequestFuncOutput],
total_duration: float,
@@ -900,23 +946,12 @@ async def benchmark(args):
pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)
async with aiohttp.ClientSession() as session:
warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]] = []
if args.warmup_requests and requests_list:
print(
f"Running {args.warmup_requests} warmup request(s) \
with num_inference_steps={args.warmup_num_inference_steps}..."
)
for i in range(args.warmup_requests):
warm_req = requests_list[i % len(requests_list)]
if args.warmup_num_inference_steps is not None:
warm_req = replace(
warm_req,
num_inference_steps=args.warmup_num_inference_steps,
)
if args.task == "t2v":
warm_req = replace(warm_req, num_frames=1)
warm_out = await limited_request_func(warm_req, session, None)
warmup_pairs.append((warm_req, warm_out))
warmup_pairs = await _run_warmups(
requests_list=requests_list,
args=args,
session=session,
request_func=request_func,
)
if args.slo:
# Prefer trace-provided per-request slo_ms. Only populate when missing.
@@ -1072,6 +1107,13 @@ if __name__ == "__main__":
default=1,
help="num_inference_steps used for warmup requests.",
)
parser.add_argument(
"--warmup-concurrency",
type=int,
default=1,
help="Maximum number of warmup requests to run concurrently. "
"Set this to match the real batch shape when warming up torch.compile or CUDA graphs.",
)
parser.add_argument("--width", type=int, default=None, help="Image/Video width.")
parser.add_argument("--height", type=int, default=None, help="Image/Video height.")
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames (for video).")
@@ -1135,5 +1177,4 @@ if __name__ == "__main__":
)
args = parser.parse_args()
asyncio.run(benchmark(args))
benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md
+14
View File
@@ -33,6 +33,12 @@ vllm serve Qwen/Qwen-Image --omni \
--port 8091
```
To replay step-wise continuous batching, compare
`vllm serve Qwen/Qwen-Image --omni --port 8089 --step-execution --max-num-seqs 1`
against the same command with `--max-num-seqs 8`. `--step-execution` is the
feature gate, and increasing `--max-num-seqs` above `1` lets the scheduler
keep more compatible requests active at once.
## 3.2 Key Parameters
| Parameter | Description |
@@ -70,6 +76,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
| `--task` | Task type (e.g., `t2i`) |
| `--num-prompts` | Total number of requests |
| `--max-concurrency` | Client-side concurrency |
| `--warmup-concurrency` | Warmup concurrency for batch-shape prewarming |
| `--random-request-config`| JSON string defining random request |
---
@@ -133,6 +140,12 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
]'
```
For a continuous-batching A/B replay, run the same benchmark command twice
against the two serving configs above and keep traffic plus warmup settings
identical across both runs. If the measured run uses `--max-concurrency 8`,
warm with `--warmup-requests 8 --warmup-concurrency 8` so the first measured
batch does not include compile or CUDA-graph capture overhead.
---
# 6. Performance Metrics
@@ -163,6 +176,7 @@ To ensure consistent and comparable benchmark results:
* Record parallel configuration
* Record benchmark parameters (resolution, concurrency, number of prompts)
* Ensure no background workload on GPUs during testing
* When comparing continuous batching, keep warmup settings the same across runs
---
docs/.nav.yml
+1
View File
@@ -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
docs/contributing/model/adding_diffusion_model.md
+10
View File
@@ -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
docs/design/feature/diffusion_continuous_batching.md
+156
View File
@@ -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)
docs/design/feature/diffusion_step_execution.md
+12 -2
View File
@@ -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.
docs/design/index.md
+2
View File
@@ -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
docs/design/module/dit_module.md
+3 -3
View File
@@ -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.
---
docs/user_guide/diffusion/step_execution.md
+32 -1
View File
@@ -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).
recipes/Qwen/Qwen-Image.md
+130
View File
@@ -0,0 +1,130 @@
# Qwen-Image for text-to-image serving with optional continuous batching on 1x A100 80GB
## Summary
- Vendor: Qwen
- Model: `Qwen/Qwen-Image`
- Task: Text-to-image generation
- Mode: Online serving with optional step-wise continuous batching
- Maintainer: Community
## When to use this recipe
Use this recipe when you want a known-good starting point for serving
`Qwen/Qwen-Image` on a single 80 GB A100, validate the normal online-serving
path, and optionally replay the benefit of step-wise continuous batching with
the benchmark assets already bundled in this repository.
## References
- Upstream or canonical docs:
[`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
- Related example under `examples/`:
[`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
- Related benchmark:
[`benchmarks/diffusion/diffusion_benchmark_serving.py`](../../benchmarks/diffusion/diffusion_benchmark_serving.py)
## Hardware Support
This recipe currently documents one CUDA GPU serving configuration. Extend it
with more hardware sections as community validation lands.
## GPU
### 1x A100 80GB
#### Environment
- OS: Linux
- Python: 3.10+
- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU
- vLLM version: Match the repository requirements for your checkout
- vLLM-Omni version or commit: Use the commit you are deploying from
#### Command
Start the baseline server:
```bash
vllm serve Qwen/Qwen-Image --omni --port 8091
```
To enable the step-wise runtime without batching:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--step-execution \
--max-num-seqs 1
```
To enable experimental compatible-request batching:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--step-execution \
--max-num-seqs 8
```
You can also use the example launcher and pass the extra flags through:
```bash
bash examples/online_serving/text_to_image/run_server.sh --step-execution --max-num-seqs 8
```
#### Verification
Run the existing client example after the server is ready:
```bash
python examples/online_serving/text_to_image/openai_chat_client.py \
--server http://localhost:8091 \
--prompt "A ceramic teapot on a wooden table" \
--output /tmp/qwen_image_recipe.png
```
For a direct API smoke test:
```bash
curl http://localhost:8091/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"model": "Qwen/Qwen-Image",
"prompt": "A ceramic teapot on a wooden table",
"size": "1024x1024",
"num_inference_steps": 20,
"seed": 42
}'
```
To replay the batching benefit with matched warmup:
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni \
--dataset vbench \
--task t2i \
--model Qwen/Qwen-Image \
--num-prompts 10 \
--height 1024 \
--width 1024 \
--num-inference-steps 50 \
--seed 42 \
--port 8089 \
--max-concurrency 8 \
--warmup-requests 8 \
--warmup-concurrency 8 \
--warmup-num-inference-steps 3 \
--disable-tqdm
```
Run that once against `--max-num-seqs 1`, then rerun it against `--max-num-seqs 8`
and compare the output JSON or terminal metrics.
#### Notes
- Memory usage: keep headroom for first-request compile and image decode overhead.
- Key flags: `--step-execution` enables the step-wise runtime; `--max-num-seqs` controls how many compatible requests may stay active together.
- Keep `--max-num-seqs 1` when you want the more conservative path, when traffic is mostly single-request, or when you are debugging correctness before measuring throughput.
- Current batching is still shape-sensitive: different step progress can co-batch, but different resolutions do not yet co-batch.
recipes/README.md
+2
View File
@@ -25,6 +25,8 @@ recipes/
## Available Recipes
- [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md): text-to-image serving recipe for
Qwen-Image on `1x A100 80GB`, including optional step-wise continuous batching replay
- [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md): online serving recipe for
multimodal chat on `1x A100 80GB`
- [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md): image-to-video serving
tests/diffusion/test_diffusion_engine.py
+118
View File
@@ -0,0 +1,118 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import queue
import threading
import time
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Any
import pytest
@dataclass
class DiffusionSchedulerOutput:
step_id: int
scheduled_new_reqs: list = field(default_factory=list)
scheduled_cached_reqs: Any = None
finished_req_ids: set = field(default_factory=set)
num_running_reqs: int = 0
num_waiting_reqs: int = 0
@property
def scheduled_req_ids(self):
ids = [req.request_id for req in self.scheduled_new_reqs]
if self.scheduled_cached_reqs:
ids.extend(self.scheduled_cached_reqs.sched_req_ids)
return ids
@property
def is_empty(self):
return len(self.scheduled_req_ids) == 0
class MockScheduler:
def __init__(self):
self._waiting_queue = []
self._step_id = 0
def add_request(self, request):
self._waiting_queue.append(request)
return request.request_id
def has_requests(self):
return len(self._waiting_queue) > 0
def schedule(self) -> DiffusionSchedulerOutput:
if not self._waiting_queue:
return DiffusionSchedulerOutput(step_id=self._step_id)
batch = []
while self._waiting_queue:
req = self._waiting_queue.pop(0)
batch.append(SimpleNamespace(request_id=req.request_id))
output = DiffusionSchedulerOutput(step_id=self._step_id, scheduled_new_reqs=batch)
self._step_id += 1
return output
def update_from_output(self, sched_output, runner_output):
# assume all new req finished
return [req.request_id for req in sched_output.scheduled_new_reqs]
@pytest.mark.asyncio
async def test_async_add_req_and_wait_for_response():
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
engine = object.__new__(DiffusionEngine)
engine.scheduler = MockScheduler()
engine._out_queue = {}
engine.abort_queue: queue.Queue[str] = queue.Queue()
engine._rpc_lock = threading.Lock()
engine._finalize_finished_request = lambda rid, out, err: out.result
def mock_execute_batch(sched_output):
req_ids = sched_output.scheduled_req_ids
time.sleep(1)
class MockRunnerOutput:
def __init__(self, ids):
self.req_id = ids
self.step_index = 0
self.finished = True
self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
def get_req_output(self, rid):
return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
return MockRunnerOutput(req_ids)
engine.execute_fn = mock_execute_batch
engine.stop_event = threading.Event()
engine.start_background_loop()
async def run_task(rid):
req = SimpleNamespace(request_id=rid)
start = time.time()
res = await engine.async_add_req_and_wait_for_response(req)
return rid, res, time.time() - start
task_ids = [f"req_{i}" for i in range(5)]
tasks = [run_task(rid) for rid in task_ids]
results = await asyncio.gather(*tasks)
engine.stop_event.set()
engine.worker_thread.join(timeout=5)
assert len(results) == 5
for rid, res, elapsed in results:
assert rid in res.result_data
eps = 0.5
latencies = [r[2] for r in results]
assert max(latencies) - min(latencies) < eps
tests/diffusion/test_diffusion_scheduler.py
+114 -6
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import queue
import threading
from types import SimpleNamespace
@@ -20,6 +21,7 @@ from vllm_omni.diffusion.sched import (
StepScheduler,
)
from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -34,7 +36,7 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
return SimpleNamespace(
return RunnerOutput(
req_id=req_id,
step_index=None,
finished=finished,
@@ -49,7 +51,7 @@ def _make_step_output(
finished: bool = False,
error: str | None = None,
):
return SimpleNamespace(
return RunnerOutput(
req_id=req_id,
step_index=step_index,
finished=finished,
@@ -91,6 +93,7 @@ class _StubScheduler(SchedulerInterface):
self._sched_req_id = request.request_ids[0]
self._state = None
self._scheduled = False
self.max_num_running_reqs = 1
def initialize(self, od_config) -> None:
self.initialized_with = od_config
@@ -218,6 +221,46 @@ class TestRequestScheduler:
assert third.num_running_reqs == 1
assert third.num_waiting_reqs == 0
def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
req_id_a = scheduler.add_request(_make_request("a"))
req_id_b = scheduler.add_request(_make_request("b"))
sched_output = scheduler.schedule()
assert _new_ids(sched_output) == [req_id_a, req_id_b]
assert sched_output.num_running_reqs == 2
assert sched_output.num_waiting_reqs == 0
def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=3))
req_id_a = scheduler.add_request(_make_request("a"))
req_id_b = scheduler.add_request(
OmniDiffusionRequest(
prompts=["prompt_b"],
sampling_params=OmniDiffusionSamplingParams(width=768),
request_ids=["b"],
)
)
scheduler.add_request(_make_request("c"))
first = scheduler.schedule()
assert _new_ids(first) == [req_id_a]
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 2
scheduler.update_from_output(first, _make_request_output(req_id_a))
second = scheduler.schedule()
assert _new_ids(second) == [req_id_b]
assert second.num_running_reqs == 1
assert second.num_waiting_reqs == 1
def test_abort_request_for_waiting_and_running(self) -> None:
req_id_a = self.scheduler.add_request(_make_request("a"))
req_id_b = self.scheduler.add_request(_make_request("b"))
@@ -348,18 +391,25 @@ class TestDiffusionEngine:
assert req_id in finished
assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
@pytest.mark.asyncio
async def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine._loop_started = False
engine._init_lock = asyncio.Lock()
engine.main_loop = asyncio.get_running_loop()
engine.stop_event = threading.Event()
engine.pre_process_func = None
engine.add_req_and_wait_for_response = mocker.Mock(
engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
)
with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
engine.step(_make_request("req-abort"))
await engine.step(_make_request("req-abort"))
def test_abort_queue_marks_request_finished_aborted(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(SimpleNamespace())
engine.abort_queue = queue.Queue()
@@ -519,7 +569,7 @@ class TestStepScheduler:
sched_output = self.scheduler.schedule()
finished = self.scheduler.update_from_output(
sched_output,
SimpleNamespace(
RunnerOutput(
req_id=req_id,
step_index=None,
finished=True,
@@ -576,6 +626,64 @@ class TestStepScheduler:
assert finished == {req_id}
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_ABORTED
def test_batches_compatible_step_requests(self) -> None:
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
req_a = scheduler.add_request(_make_step_request("a"))
req_b = scheduler.add_request(_make_step_request("b"))
sched_output = scheduler.schedule()
assert _new_ids(sched_output) == [req_a, req_b]
assert sched_output.num_running_reqs == 2
assert sched_output.num_waiting_reqs == 0
def test_step_batch_allows_different_num_inference_steps(self) -> None:
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
req_a = scheduler.add_request(_make_step_request("a", num_inference_steps=2))
req_b = scheduler.add_request(_make_step_request("b", num_inference_steps=4))
sched_output = scheduler.schedule()
assert _new_ids(sched_output) == [req_a, req_b]
assert sched_output.num_running_reqs == 2
assert sched_output.num_waiting_reqs == 0
def test_step_batch_rejects_different_sampling_key(self) -> None:
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=3))
req_a = scheduler.add_request(_make_step_request("a"))
req_b = scheduler.add_request(
_make_step_request(
"b",
sampling_params=OmniDiffusionSamplingParams(
height=768,
num_inference_steps=4,
),
)
)
scheduler.add_request(_make_step_request("c"))
sched_output = scheduler.schedule()
assert _new_ids(sched_output) == [req_a]
assert sched_output.num_running_reqs == 1
assert sched_output.num_waiting_reqs == 2
scheduler.update_from_output(
sched_output,
_make_step_output(req_a, step_index=4, finished=True),
)
second = scheduler.schedule()
assert _new_ids(second) == [req_b]
assert second.num_running_reqs == 1
assert second.num_waiting_reqs == 1
def test_preempt_request_preserves_step_index(self) -> None:
request = _make_step_request("preempt", num_inference_steps=3)
req_id = self.scheduler.add_request(request)
tests/diffusion/test_diffusion_step_pipeline.py
+20 -15
View File
@@ -77,12 +77,12 @@ class _StepPipeline:
self.prepare_calls += 1
state.timesteps = [torch.tensor(10), torch.tensor(5)]
state.latents = torch.tensor([0.0])
state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
return state
def denoise_step(self, state, **kwargs):
del state, kwargs
def denoise_step(self, input_batch, **kwargs):
self.denoise_calls += 1
return torch.tensor([1.0])
return torch.full_like(input_batch.prompt_embeds, fill_value=0.5)
def step_scheduler(self, state, noise_pred, **kwargs):
del noise_pred, kwargs
@@ -145,6 +145,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
state.step_index = 0
state.scheduler = self.scheduler
state.do_true_cfg = self.mode == "cfg"
state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
return state
def denoise_step(self, state, **kwargs):
@@ -291,6 +292,7 @@ def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
engine.scheduler = scheduler
engine.execute_fn = execute_fn
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
engine.abort_queue = queue.Queue()
return engine
@@ -327,10 +329,11 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
raise ValueError(f"Unsupported distributed test mode: {mode}")
runner = _make_distributed_runner(mode, device)
output = DiffusionModelRunner.execute_stepwise(
result = DiffusionModelRunner.execute_stepwise(
runner,
_make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
)
output = result.get_req_output("req-1")
assert output.finished is True
assert output.result is not None
@@ -354,14 +357,16 @@ class TestRunner:
req = _make_step_request()
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
first = result.get_req_output("req-1")
assert first.req_id == "req-1"
assert first.step_index == 1
assert first.finished is False
assert first.result is None
assert "req-1" in runner.state_cache
second = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
second = result.get_req_output("req-1")
assert second.req_id == "req-1"
assert second.step_index == 2
assert second.finished is True
@@ -393,8 +398,8 @@ class TestRunner:
num_waiting_reqs=0,
)
with pytest.raises(ValueError, match="batch_size=1"):
DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
assert len(result) == 2
def test_rejects_missing_cached_state(self):
runner = _make_runner()
@@ -408,8 +413,8 @@ class TestRunner:
req = _make_step_request()
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
output = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
output = result.get_req_output("req-1")
assert output.req_id == "req-1"
assert output.step_index == 0
assert output.finished is True
@@ -580,7 +585,7 @@ class TestEngine:
)
def test_step_engine_returns_error(self, execute_fn, expected_error, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(mocker.Mock())
scheduler.initialize(SimpleNamespace())
engine = _make_engine(scheduler, execute_fn=execute_fn)
output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2))
@@ -590,7 +595,7 @@ class TestEngine:
def test_step_execution_completes(self, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(mocker.Mock())
scheduler.initialize(SimpleNamespace())
engine = _make_engine(scheduler)
request = _make_engine_request("req-step", num_inference_steps=2)
@@ -616,7 +621,7 @@ class TestEngine:
def test_step_abort_stops_rescheduling_after_first_step(self, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(mocker.Mock())
scheduler.initialize(SimpleNamespace())
engine = _make_engine(scheduler)
request = _make_engine_request("req-stop", num_inference_steps=4)
@@ -641,7 +646,7 @@ class TestEngine:
def test_step_abort_after_reschedule_returns_aborted_output(self, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(mocker.Mock())
scheduler.initialize(SimpleNamespace())
engine = _make_engine(scheduler)
request = _make_engine_request("req-mid", num_inference_steps=4)
@@ -668,7 +673,7 @@ class TestEngine:
def test_finished_step_without_result_returns_error(self, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(mocker.Mock())
scheduler.initialize(SimpleNamespace())
engine = _make_engine(
scheduler,
execute_fn=lambda _: RunnerOutput(
tests/diffusion/test_multiproc_engine_concurrency.py
+7 -2
View File
@@ -19,6 +19,7 @@ from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
from vllm_omni.diffusion.sched import RequestScheduler
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
@@ -34,7 +35,11 @@ def _tagged_output(tag: str) -> DiffusionOutput:
def _mock_request(tag: str):
"""Return a lightweight request object identifiable by *tag*."""
return SimpleNamespace(request_ids=[tag])
return SimpleNamespace(
request_ids=[tag],
prompts=[f"prompt_{tag}"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
)
def _make_executor(num_gpus: int = 1):
@@ -406,7 +411,7 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
t.start()
# Wait until request execution is truly inside the lock and blocking.
add_req_blocked.wait(5)
assert add_req_blocked.wait(5)
# collective_rpc should time out at lock acquisition, not hang.
with pytest.raises(TimeoutError):
tests/diffusion/test_stage_diffusion_proc.py
+75 -2
View File
@@ -2,8 +2,9 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict
from dataclasses import asdict, dataclass
from types import SimpleNamespace
import pytest
@@ -18,7 +19,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
async def run_test():
captured = {}
def step(request):
async def step(request):
captured["request"] = request
return [
SimpleNamespace(
@@ -73,3 +74,75 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
assert result.images == ["img-1", "img-2"]
asyncio.run(run_test())
@dataclass
class MockOmniRequestOutput:
request_id: str = ""
status: str = "success"
BASE_HEIGHT = 512
BASE_WIDTH = 512
BASE_INFER_STEPS = 10
DELAY_BASE = 2
class MockDiffusionEngine:
async def step(self, request):
def simulate_step_delay(height, width, num_inference_steps) -> float:
return (height / BASE_HEIGHT) * (width / BASE_WIDTH) * (num_inference_steps / BASE_INFER_STEPS)
DELAY_BASE = 2
delay_scale = simulate_step_delay(
request.sampling_params.height, request.sampling_params.width, request.sampling_params.num_inference_steps
)
delay = DELAY_BASE + delay_scale * DELAY_BASE
await asyncio.sleep(delay)
return [MockOmniRequestOutput(request_id=request.request_id)]
@pytest.mark.asyncio
async def test_proc_process_request_with_batching_async_output():
stage_proc = object.__new__(StageDiffusionProc)
stage_proc._engine = MockDiffusionEngine()
test_requests = [
{
"request_id": "req_1",
"prompt": "prompt1",
"params": {"height": BASE_HEIGHT * 2, "width": BASE_WIDTH * 2, "num_inference_steps": BASE_INFER_STEPS * 1},
},
{
"request_id": "req_2",
"prompt": "prompt2",
"params": {"height": BASE_HEIGHT * 2, "width": BASE_WIDTH * 2, "num_inference_steps": BASE_INFER_STEPS * 2},
},
{
"request_id": "req_3",
"prompt": "prompt3",
"params": {"height": BASE_HEIGHT * 2, "width": BASE_WIDTH * 2, "num_inference_steps": BASE_INFER_STEPS * 3},
},
]
async def run_task(req_data):
start_time = time.time()
result = await stage_proc._process_request(
request_id=req_data["request_id"], prompt=req_data["prompt"], sampling_params_dict=req_data["params"]
)
end_time = time.time()
return result, end_time - start_time
coros = [run_task(req) for req in test_requests]
results = await asyncio.gather(*coros)
assert len(results) == len(test_requests)
base_time = DELAY_BASE
time_gap_std = DELAY_BASE * 2 * 2 * 1 # height/width/steps infer time scale
eps = 0.5
for i, (res, elapsed_time) in enumerate(results):
assert res.request_id == test_requests[i]["request_id"]
assert isinstance(res, MockOmniRequestOutput)
time_gap = elapsed_time - base_time
assert time_gap > time_gap_std - eps and time_gap < time_gap_std + eps
base_time = elapsed_time
vllm_omni/diffusion/diffusion_engine.py
+133 -9
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
import inspect
import queue
import threading
@@ -30,7 +31,7 @@ from vllm_omni.diffusion.registry import (
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
from vllm_omni.outputs import OmniRequestOutput
@@ -99,7 +100,17 @@ class DiffusionEngine:
StepScheduler() if self.step_execution else RequestScheduler()
)
self.scheduler.initialize(od_config)
if self.scheduler.max_num_running_reqs > 1 and not self.step_execution:
max_num_seqs = self.scheduler.max_num_running_reqs
self.scheduler.max_num_running_reqs = 1
logger.warning(f"Non-stepwise-execution does not support max-num-seqs={max_num_seqs}, set it to 1.")
self.main_loop: asyncio.AbstractEventLoop | None = None
self.stop_event: threading.Event | None = None
self._loop_started = False
self._init_lock = asyncio.Lock()
self._rpc_lock = threading.RLock()
self._cv = threading.Condition(self._rpc_lock)
self._out_queue: dict[str, asyncio.Future] = {}
self.abort_queue: queue.Queue[str] = queue.Queue()
self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request
@@ -110,7 +121,24 @@ class DiffusionEngine:
self.close()
raise e
def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
async def _check_and_start_background_loop(self):
if self._loop_started:
return
async with self._init_lock:
# double check, in case of lock queue issue
if self._loop_started:
return
self.main_loop = asyncio.get_running_loop()
self.stop_event = threading.Event()
self.worker_thread = threading.Thread(target=self._busy_loop)
self.worker_thread.start()
self._loop_started = True
async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
await self._check_and_start_background_loop()
diffusion_engine_start_time = time.perf_counter()
# Apply pre-processing if available
@@ -122,7 +150,7 @@ class DiffusionEngine:
logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
exec_start_time = time.perf_counter()
output = self.add_req_and_wait_for_response(request)
output = await self.async_add_req_and_wait_for_response(request)
exec_total_time = time.perf_counter() - exec_start_time
if output.aborted:
@@ -336,6 +364,62 @@ class DiffusionEngine:
return results
def _busy_loop(self):
while not self.stop_event.is_set():
self._process_aborts_queue()
with self._cv:
while not self.scheduler.has_requests() and not self.stop_event.is_set():
self._cv.wait(timeout=1.0)
if self.stop_event.is_set():
break
sched_output = self.scheduler.schedule()
if sched_output.is_empty:
self._handle_finished_requests(sched_output.finished_req_ids, None)
continue
try:
runner_output = self.execute_fn(sched_output)
except Exception as exc:
logger.error(
"Execution failed for diffusion requests %s", sched_output.scheduled_req_ids, exc_info=True
)
runner_output = BatchRunnerOutput.from_list(
[
RunnerOutput(
req_id=req_id,
step_index=None,
finished=True,
result=DiffusionOutput(error=str(exc)),
)
for req_id in sched_output.scheduled_req_ids
]
)
self._process_aborts_queue()
finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
self._handle_finished_requests(finished_req_ids, runner_output)
def _handle_finished_requests(
self,
finished_ids: set[str],
runner_output: RunnerOutput | None = None,
missing_result_error: str = "Diffusion execution finished without a final output",
):
for rid in finished_ids:
if rid in self._out_queue:
fut = self._out_queue.pop(rid)
if runner_output is not None:
_output = runner_output.get_req_output(rid)
else:
_output = None
out = self._finalize_finished_request(rid, _output, missing_result_error)
if not fut.done():
self.main_loop.call_soon_threadsafe(fut.set_result, out)
@staticmethod
def make_engine(
config: OmniDiffusionConfig,
@@ -351,6 +435,30 @@ class DiffusionEngine:
"""
return DiffusionEngine(config, scheduler=scheduler)
def add_request(self, request: OmniDiffusionRequest) -> str:
fut = self.main_loop.create_future()
with self._cv:
sched_req_id = self.scheduler.add_request(request)
self._out_queue[sched_req_id] = fut
self._cv.notify_all()
return sched_req_id
async def get_result(self, request_id: str) -> DiffusionOutput:
fut = self._out_queue.get(request_id)
if fut is None:
raise RuntimeError(f"Request {request_id} not found in output queue.")
try:
return await fut
except Exception as e:
logger.error(f"Wait for response failed: {e}")
raise
async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
sched_req_id = self.add_request(request)
return await self.get_result(sched_req_id)
def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
with self._rpc_lock:
target_sched_req_id = self.scheduler.add_request(request)
@@ -366,10 +474,8 @@ class DiffusionEngine:
raise RuntimeError("Diffusion scheduler has no runnable requests.")
continue
# NOTE: add_req_and_wait_for_response() is synchronous, and
# the scheduler currently enforces _max_batch_size = 1 (see
# vllm_omni/diffusion/sched/base_scheduler.py), so we directly
# take the single scheduled request here.
# NOTE: add_req_and_wait_for_response() is synchronous, will be only called
# within _dummy_run, only one request will be scheduled
sched_req_id = sched_output.scheduled_req_ids[0]
try:
runner_output = self.execute_fn(sched_output)
@@ -387,7 +493,12 @@ class DiffusionEngine:
self._process_aborts_queue()
finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
# sync func should receive one result
if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
raise ValueError("Sync func should receive one result at one time")
if target_sched_req_id in finished_req_ids:
runner_output = runner_output.get_req_output(target_sched_req_id)
return self._finalize_finished_request(
target_sched_req_id,
runner_output=runner_output,
@@ -510,13 +621,26 @@ class DiffusionEngine:
def close(self) -> None:
if hasattr(self, "scheduler"):
self.scheduler.close()
if hasattr(self, "stop_event") and self.stop_event is not None:
with self._cv:
self.stop_event.set()
self._cv.notify_all()
if hasattr(self, "worker_thread") and self.worker_thread.is_alive():
self.worker_thread.join(timeout=10)
if self.worker_thread.is_alive():
logger.warning("Worker thread did not terminate within 10s, resources may not be released.")
else:
self._loop_started = False
if hasattr(self, "executor"):
self.executor.shutdown()
def abort(self, request_id: str | Iterable[str]) -> None:
request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
for req_id in request_ids:
self.abort_queue.put(req_id)
with self._cv:
for req_id in request_ids:
self.abort_queue.put(req_id)
self._cv.notify_all()
def _process_aborts_queue(self) -> None:
if self.abort_queue.empty():
vllm_omni/diffusion/executor/abstract.py
+3 -3
View File
@@ -10,7 +10,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
if TYPE_CHECKING:
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
class DiffusionExecutor(ABC):
@@ -74,12 +74,12 @@ class DiffusionExecutor(ABC):
pass
@abstractmethod
def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
"""Execute request-mode work from a scheduler output."""
pass
@abstractmethod
def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
"""Execute step-mode work from a scheduler output."""
pass
vllm_omni/diffusion/executor/multiproc_executor.py
+5 -5
View File
@@ -22,7 +22,7 @@ from vllm_omni.diffusion.worker import WorkerProc
if TYPE_CHECKING:
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
logger = init_logger(__name__)
@@ -273,7 +273,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
logger.error(f"Generate call failed: {e}")
raise
def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
"""Adapt request-mode scheduler output to worker execute_model RPC."""
from vllm_omni.diffusion.worker.utils import RunnerOutput
@@ -301,9 +301,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
result=result,
)
def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
"""Forward step-mode scheduler output to worker execute_stepwise RPC."""
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, RunnerOutput
self._ensure_open()
result = self.collective_rpc(
@@ -313,7 +313,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
exec_all_ranks=True,
)
if isinstance(result, RunnerOutput):
if isinstance(result, BaseRunnerOutput):
return result
# TODO: Remove this fallback; DiffusionOutput cannot faithfully represent
# failed multi-request step batches.
vllm_omni/diffusion/inline_stage_diffusion_client.py
+2 -4
View File
@@ -108,8 +108,7 @@ class InlineStageDiffusionClient:
kv_sender_info=kv_sender_info,
)
loop = asyncio.get_running_loop()
results = await loop.run_in_executor(self._executor, self._engine.step, request)
results = await self._engine.step(request)
result = results[0]
if not result.request_id:
result.request_id = request_id
@@ -161,8 +160,7 @@ class InlineStageDiffusionClient:
kv_sender_info=kv_sender_info,
)
loop = asyncio.get_running_loop()
results = await loop.run_in_executor(self._executor, self._engine.step, request)
results = await self._engine.step(request)
all_images: list = []
merged_mm: dict[str, Any] = {}
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+20 -25
View File
@@ -44,6 +44,7 @@ from vllm_omni.diffusion.utils.size_utils import (
from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -864,54 +865,48 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
def denoise_step(
self,
state: "DiffusionRequestState",
input_batch: "InputBatch",
**kwargs: Any,
) -> torch.Tensor | None:
"""One denoise step: read from *state*, delegate to CFGParallelMixin.
"""One denoise step: read from *input_batch*, delegate to CFGParallelMixin.
Reuses ``predict_noise_maybe_with_cfg`` so that CFG-parallel,
sequential-CFG, and no-CFG paths are handled identically to
``diffuse()``.
"""
del kwargs
if self.interrupt:
return None
t = state.current_timestep
t = input_batch.timesteps
self._current_timestep = t
self.transformer.do_true_cfg = state.do_true_cfg
# Normalize timestep to [batch_size] tensor
if not torch.is_tensor(t):
t = torch.tensor(t, device=state.latents.device, dtype=state.latents.dtype)
self.transformer.do_true_cfg = input_batch.do_true_cfg
positive_kwargs, negative_kwargs, output_slice = self._build_denoise_kwargs(
latents=state.latents,
latents=input_batch.latents,
timestep=t,
guidance=state.guidance,
prompt_embeds=state.prompt_embeds,
prompt_embeds_mask=state.prompt_embeds_mask,
img_shapes=state.img_shapes,
txt_seq_lens=state.txt_seq_lens,
do_true_cfg=state.do_true_cfg,
negative_prompt_embeds=state.negative_prompt_embeds,
negative_prompt_embeds_mask=state.negative_prompt_embeds_mask,
negative_txt_seq_lens=state.negative_txt_seq_lens,
image_latents=state.sampling.image_latent,
guidance=input_batch.guidance,
prompt_embeds=input_batch.prompt_embeds,
prompt_embeds_mask=input_batch.prompt_embeds_mask,
img_shapes=input_batch.img_shapes,
txt_seq_lens=input_batch.txt_seq_lens,
do_true_cfg=input_batch.do_true_cfg,
negative_prompt_embeds=input_batch.negative_prompt_embeds,
negative_prompt_embeds_mask=input_batch.negative_prompt_embeds_mask,
negative_txt_seq_lens=input_batch.negative_txt_seq_lens,
image_latents=input_batch.image_latents,
extra_transformer_kwargs={
"attention_kwargs": self.attention_kwargs,
"return_dict": False,
},
)
true_cfg_scale = state.sampling.true_cfg_scale or 4.0
cfg_normalize = state.sampling.cfg_normalize
return self.predict_noise_maybe_with_cfg(
state.do_true_cfg,
true_cfg_scale,
input_batch.do_true_cfg,
input_batch.true_cfg_scale,
positive_kwargs,
negative_kwargs,
cfg_normalize,
input_batch.cfg_normalize,
output_slice,
)
vllm_omni/diffusion/sched/base_scheduler.py
+48 -8
View File
@@ -4,6 +4,7 @@
from __future__ import annotations
from collections import deque
from dataclasses import fields
from vllm.logger import init_logger
@@ -15,11 +16,23 @@ from vllm_omni.diffusion.sched.interface import (
DiffusionRequestStatus,
DiffusionSchedulerOutput,
NewRequestData,
SamplingParamsKey,
SchedulerInterface,
)
logger = init_logger(__name__)
_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey))
def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey | None:
"""Build a batch-compatibility key from the request's sampling params."""
if len(request.prompts) != 1:
return None
sampling = request.sampling_params
return SamplingParamsKey(**{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES})
class _BaseScheduler(SchedulerInterface):
"""Shared queue/state bookkeeping for diffusion schedulers."""
@@ -31,8 +44,9 @@ class _BaseScheduler(SchedulerInterface):
self._step_id: int = 0
self._waiting: deque[str] = deque()
self._running: list[str] = []
self._running_sampling_params_key: SamplingParamsKey | None = None
self._finished_req_ids: set[str] = set()
self._max_batch_size: int = 1
self.max_num_running_reqs: int = 1
def initialize(self, od_config: OmniDiffusionConfig) -> None:
self.od_config = od_config
@@ -41,11 +55,13 @@ class _BaseScheduler(SchedulerInterface):
self._step_id = 0
self._waiting.clear()
self._running.clear()
self._running_sampling_params_key = None
self._finished_req_ids.clear()
# The current DiffusionEngine execution mode does not support real
# request batching well, so we keep this fixed at 1 for now.
# TODO: Add support for multiple concurrent requests
self.max_num_running_reqs = 1
max_num_seqs = getattr(od_config, "max_num_seqs", 1)
try:
self.max_num_running_reqs = max(1, int(max_num_seqs))
except (TypeError, ValueError):
self.max_num_running_reqs = 1
self._reset_scheduler_state()
def add_request(self, request: OmniDiffusionRequest) -> str:
@@ -53,7 +69,7 @@ class _BaseScheduler(SchedulerInterface):
return self._add_request_with_sched_req_id(sched_req_id, request)
def _add_request_with_sched_req_id(self, sched_req_id: str, request: OmniDiffusionRequest) -> str:
state = DiffusionRequestState(sched_req_id=sched_req_id, req=request)
state = self._make_request_state(sched_req_id, request)
self._request_states[sched_req_id] = state
self._register_request_ids(request.request_ids, sched_req_id)
self._waiting.append(sched_req_id)
@@ -82,6 +98,8 @@ class _BaseScheduler(SchedulerInterface):
self._waiting.popleft()
was_new_request = state.status == DiffusionRequestStatus.WAITING
if not self._running:
self._running_sampling_params_key = state.sampling_params_key
state.status = DiffusionRequestStatus.RUNNING
self._running.append(sched_req_id)
if was_new_request:
@@ -124,6 +142,8 @@ class _BaseScheduler(SchedulerInterface):
return False
if sched_req_id in self._running:
self._running.remove(sched_req_id)
if not self._running:
self._running_sampling_params_key = None
self._waiting.appendleft(sched_req_id)
self._request_states[sched_req_id].status = DiffusionRequestStatus.PREEMPTED
return True
@@ -140,6 +160,7 @@ class _BaseScheduler(SchedulerInterface):
self._request_id_to_sched_req_id.clear()
self._waiting.clear()
self._running.clear()
self._running_sampling_params_key = None
self._finished_req_ids.clear()
self._reset_scheduler_state()
@@ -169,6 +190,8 @@ class _BaseScheduler(SchedulerInterface):
if running_to_remove:
self._running = [sched_req_id for sched_req_id in self._running if sched_req_id not in running_to_remove]
if not self._running:
self._running_sampling_params_key = None
if waiting_to_remove:
self._waiting = deque(
sched_req_id for sched_req_id in self._waiting if sched_req_id not in waiting_to_remove
@@ -208,9 +231,26 @@ class _BaseScheduler(SchedulerInterface):
def _pop_extra_request_state(self, sched_req_id: str) -> None:
"""Remove subclass-owned per-request state before popping request state."""
def _make_request_state(self, sched_req_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
return DiffusionRequestState(
sched_req_id=sched_req_id,
req=request,
sampling_params_key=get_sampling_params_key(request),
)
def _can_schedule_waiting(self, state: DiffusionRequestState) -> bool:
del state
return True
if not self._running:
return True
current_key = self._current_sampling_params_key()
return current_key is not None and current_key == state.sampling_params_key
def _current_sampling_params_key(self) -> SamplingParamsKey | None:
if self._running_sampling_params_key is not None or not self._running:
return self._running_sampling_params_key
state = self._request_states.get(self._running[0])
self._running_sampling_params_key = None if state is None else state.sampling_params_key
return self._running_sampling_params_key
def _register_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
for request_id in request_ids:
vllm_omni/diffusion/sched/interface.py
+30 -1
View File
@@ -38,6 +38,34 @@ class DiffusionRequestStatus(enum.IntEnum):
return status >= DiffusionRequestStatus.FINISHED_COMPLETED
@dataclass(frozen=True, eq=True)
class SamplingParamsKey:
"""Batch-compatibility key derived from ``OmniDiffusionSamplingParams``.
Only requests with the same key can be batched together.
Fields not included here are treated as request-local and do not
participate in the current homogeneous batching policy.
"""
# Spatial / temporal shape.
height: int | None = None
width: int | None = None
num_frames: int = 1
resolution: int | str | None = None
fps: int | None = None
frame_rate: float | None = None
boundary_ratio: float | None = None
# CFG / guidance.
do_classifier_free_guidance: bool = False
guidance_scale: float = 0.0
guidance_scale_provided: bool = False
guidance_scale_2: float | None = None
guidance_rescale: float = 0.0
true_cfg_scale: float | None = None
cfg_normalize: bool = False
@dataclass
class DiffusionRequestState:
"""Scheduler-owned state for one queued OmniDiffusionRequest."""
@@ -47,6 +75,7 @@ class DiffusionRequestState:
# TODO: Align this with OmniDiffusionRequest.request_ids once scheduler batching is supported.
sched_req_id: str
req: OmniDiffusionRequest
sampling_params_key: SamplingParamsKey | None = None
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
error: str | None = None
@@ -81,7 +110,7 @@ class CachedRequestData:
class DiffusionSchedulerOutput:
"""Output of a single scheduling cycle."""
step_id: int
step_id: int # global step index
scheduled_new_reqs: list[NewRequestData]
scheduled_cached_reqs: CachedRequestData
finished_req_ids: set[str]
vllm_omni/diffusion/sched/request_scheduler.py
+2 -1
View File
@@ -32,11 +32,12 @@ class RequestScheduler(_BaseScheduler):
terminal_statuses: dict[str, DiffusionRequestStatus] = {}
terminal_errors: dict[str, str | None] = {}
result = output.result
for sched_req_id in scheduled_req_ids:
state = self._request_states.get(sched_req_id)
if state is None or state.is_finished():
continue
req_output = output.get_req_output(sched_req_id)
result = req_output.result if req_output is not None else None
if result is None:
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
terminal_errors[sched_req_id] = "No output result"
vllm_omni/diffusion/sched/step_scheduler.py
+15 -5
View File
@@ -72,19 +72,29 @@ class StepScheduler(_BaseScheduler):
terminal_statuses: dict[str, DiffusionRequestStatus] = {}
terminal_errors: dict[str, str | None] = {}
output_error = output.result.error if output.result is not None else None
for sched_req_id in scheduled_req_ids:
state = self._request_states.get(sched_req_id)
progress = self._request_progress.get(sched_req_id)
if state is None or progress is None or state.is_finished():
continue
req_output = output.get_req_output(sched_req_id)
if req_output is None:
logger.warning(
"No RunnerOutput for request %s, treating as error",
sched_req_id,
)
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
terminal_errors[sched_req_id] = "No output for request"
continue
req_result = req_output.result
output_error = req_result.error if req_result is not None else None
if output_error is not None:
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
terminal_errors[sched_req_id] = output_error
continue
if output.step_index is None:
if req_output.step_index is None:
logger.warning(
"Received RunnerOutput with no step_index for request %s, treating as error",
sched_req_id,
@@ -94,9 +104,9 @@ class StepScheduler(_BaseScheduler):
continue
# We assume that the decoding stage is executed immediately after the denoising stage completes.
progress.current_step = output.step_index
state.req.sampling_params.step_index = output.step_index
if output.finished:
progress.current_step = req_output.step_index
state.req.sampling_params.step_index = req_output.step_index
if req_output.finished:
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
terminal_errors[sched_req_id] = None
else:
vllm_omni/diffusion/stage_diffusion_proc.py
+2 -4
View File
@@ -103,8 +103,7 @@ class StageDiffusionProc:
kv_sender_info=kv_sender_info,
)
loop = asyncio.get_running_loop()
results = await loop.run_in_executor(self._executor, self._engine.step, request)
results = await self._engine.step(request)
result = results[0]
if not result.request_id:
result.request_id = request_id
@@ -134,8 +133,7 @@ class StageDiffusionProc:
kv_sender_info=kv_sender_info,
)
loop = asyncio.get_running_loop()
results = await loop.run_in_executor(self._executor, self._engine.step, request)
results = await self._engine.step(request)
# Merge per-prompt results into a single combined output.
all_images: list = []
vllm_omni/diffusion/worker/diffusion_model_runner.py
+141 -71
View File
@@ -14,6 +14,7 @@ import copy
import time
from collections.abc import Iterable
from contextlib import nullcontext
from typing import Any
import torch
from torch.profiler import record_function
@@ -32,7 +33,8 @@ from vllm_omni.diffusion.offloader import get_offload_backend
from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput
from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, DiffusionRequestState, RunnerOutput
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
from vllm_omni.platforms import current_omni_platform
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
@@ -297,50 +299,107 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
"""Return whether current pipeline supports step execution."""
return self.pipeline is not None and supports_step_execution(self.pipeline)
def _update_states(self, scheduler_output: DiffusionSchedulerOutput) -> tuple[DiffusionRequestState, bool]:
def _update_states(
self, scheduler_output: DiffusionSchedulerOutput
) -> tuple[list[DiffusionRequestState], list[str]]:
"""Step-before update: cleanup finished requests and get/create one running state."""
for req_id in scheduler_output.finished_req_ids:
self.state_cache.pop(req_id, None)
if scheduler_output.num_scheduled_reqs != 1:
raise ValueError(
"Step mode currently supports batch_size=1, "
f"but got {scheduler_output.num_scheduled_reqs} scheduled requests."
)
resolved: list[DiffusionRequestState] = []
new_req_id: list[str] = []
try:
# process new requests
for sched_new_req in scheduler_output.scheduled_new_reqs:
# new_req_data = scheduler_output.scheduled_new_reqs[0]
req_id = sched_new_req.sched_req_id
req = sched_new_req.req
new_req_id.append(req_id)
if req_id in self.state_cache:
raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
request_ids = req.request_ids or [req_id]
if len(request_ids) != len(req.prompts):
raise ValueError(
f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
)
new_state = DiffusionRequestState(
req_id=req_id,
sampling=copy.deepcopy(req.sampling_params),
prompts=req.prompts,
)
self.state_cache[req_id] = new_state
resolved.append(new_state)
if scheduler_output.scheduled_new_reqs:
new_req_data = scheduler_output.scheduled_new_reqs[0]
req_id = new_req_data.sched_req_id
req = new_req_data.req
if req_id in self.state_cache:
raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
else:
req_id = scheduler_output.scheduled_cached_reqs.sched_req_ids[0]
state = self.state_cache.get(req_id)
if state is None:
raise ValueError(f"Missing cached state for request {req_id}.")
return state, False
# process cached requests
for req_id in scheduler_output.scheduled_cached_reqs.sched_req_ids:
state = self.state_cache.get(req_id)
if state is None:
raise ValueError(f"Missing cached state for request {req_id}.")
resolved.append(state)
except Exception:
for req_id in new_req_id:
self.state_cache.pop(req_id, None)
raise
request_ids = req.request_ids or [req_id]
if len(request_ids) != len(req.prompts):
raise ValueError(
f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
)
return resolved, new_req_id
state = DiffusionRequestState(
req_id=req_id,
sampling=copy.deepcopy(req.sampling_params),
prompts=req.prompts,
def _prepare_batch_inputs(self, states: list[DiffusionRequestState], new_request_ids: list[str]) -> InputBatch:
# process new reqs
for state in states:
if state.req_id in new_request_ids:
# set generator
if state.sampling.generator is None and state.sampling.seed is not None:
if state.sampling.generator_device is not None:
gen_device = state.sampling.generator_device
elif self.device.type == "cpu":
gen_device = "cpu"
else:
gen_device = self.device
state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
# encode
self.pipeline.prepare_encode(state)
input_batch = InputBatch.make_batch(
states,
cached_batch=getattr(self, "input_batch", None),
)
self.state_cache[req_id] = state
return state, True
self.input_batch = input_batch
return input_batch
def _update_states_after(self, state: DiffusionRequestState, finished: bool) -> None:
def _update_states_after(
self,
states: list[DiffusionRequestState],
input_batch: InputBatch,
interrupted: bool = False,
):
"""Step-after update: clear cached state for completed request."""
if finished:
self.state_cache.pop(state.req_id, None)
gathered_latents = torch.cat([state.latents for state in states], dim=0)
if (
input_batch.latents.size() == gathered_latents.size()
and input_batch.latents.dtype == gathered_latents.dtype
and input_batch.latents.device == gathered_latents.device
):
input_batch.latents.copy_(gathered_latents)
else:
input_batch.latents = gathered_latents.clone()
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
self.input_batch = input_batch
scatter_latents(states, input_batch)
for state in states:
if interrupted or state.denoise_completed:
self.state_cache.pop(state.req_id, None)
def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any:
model_state = getattr(self, "model_state", None)
if model_state is None:
return {}
prepare_attn = getattr(model_state, "prepare_attn", None)
if not callable(prepare_attn):
return {}
return prepare_attn(input_batch)
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput:
"""Execute one step for one scheduled request and return runner output."""
assert self.pipeline is not None, "Model not loaded. Call load_model() first."
if not self.supports_step_mode():
@@ -354,46 +413,57 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
use_hsdp = self.od_config.parallel_config.use_hsdp
grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
with grad_context:
state, is_new_request = self._update_states(scheduler_output)
states, new_request_ids = self._update_states(scheduler_output)
input_batch = self._prepare_batch_inputs(states, new_request_ids)
attn_metadata = self._prepare_attn_metadata(input_batch)
if is_new_request:
# TODO: support kv manager recv
# TODO: support cache backend
if state.sampling.generator is None and state.sampling.seed is not None:
if state.sampling.generator_device is not None:
gen_device = state.sampling.generator_device
elif self.device.type == "cpu":
gen_device = "cpu"
else:
gen_device = self.device
state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
with set_forward_context(
vllm_config=self.vllm_config,
omni_diffusion_config=self.od_config,
attn_metadata=attn_metadata,
):
noise_pred = self.pipeline.denoise_step(input_batch)
with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
# step0/new request: encode
if is_new_request:
self.pipeline.prepare_encode(state)
runner_output_list = []
pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
if noise_pred is None and pipeline_interrupted:
for state in states:
runner_output_list.append(
RunnerOutput(
req_id=state.req_id,
step_index=state.step_index,
finished=True,
result=DiffusionOutput(error="stepwise denoise interrupted"),
)
)
noise_pred = self.pipeline.denoise_step(state)
finished = False
# In CFG parallel mode, only rank 0 gets the actual noise_pred; non-rank-0 workers receive None.
# A true interrupt (all ranks return None) is detected by checking self.pipeline.interrupt.
if noise_pred is None and getattr(self.pipeline, "interrupt", False):
finished = True
result = DiffusionOutput(error="stepwise denoise interrupted")
else:
self.pipeline.step_scheduler(state, noise_pred)
finished = state.denoise_completed
if finished:
result = self.pipeline.post_decode(state)
else:
result = None
offset = 0
for req in states:
row_num = req.latents.shape[0]
self.pipeline.step_scheduler(
req, noise_pred[offset : offset + row_num] if noise_pred is not None else None
)
offset = offset + row_num
if req.denoise_completed:
result = self.pipeline.post_decode(req)
else:
result = None
runner_output_list.append(
RunnerOutput(
req_id=req.req_id,
step_index=req.step_index,
finished=req.denoise_completed,
result=result,
)
)
self._update_states_after(state, finished)
if noise_pred is not None and offset != noise_pred.shape[0]:
raise ValueError(
f"Stepwise noise_pred consumed {offset} rows, "
f"but batched noise_pred has {noise_pred.shape[0]} rows."
)
return RunnerOutput(
req_id=state.req_id,
step_index=state.step_index,
finished=finished,
result=result,
)
self._update_states_after(states, input_batch, pipeline_interrupted)
return BatchRunnerOutput.from_list(runner_output_list)
vllm_omni/diffusion/worker/diffusion_worker.py
+3 -3
View File
@@ -45,7 +45,7 @@ from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
from vllm_omni.lora.request import LoRARequest
from vllm_omni.platforms import current_omni_platform
from vllm_omni.profiler import OmniTorchProfilerWrapper, create_omni_profiler
@@ -262,7 +262,7 @@ class DiffusionWorker:
profiler.step()
return output
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
"""Execute one diffusion step by delegating to the model runner."""
assert self.model_runner is not None, "Model runner not initialized"
if self.lora_manager is not None:
@@ -834,7 +834,7 @@ class WorkerWrapperBase:
"""
return self.worker.execute_model(reqs, od_config)
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
"""Execute one diffusion step."""
return self.worker.execute_stepwise(scheduler_output)
vllm_omni/diffusion/worker/input_batch.py
+740
View File
@@ -0,0 +1,740 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Diffusion input-batch structures following the MRV2-style vLLM layout.
Request states remain the only persistent source of truth. Static tensors are
normalized/padded onto the request state once, while :class:`InputBatch`
assembles an ephemeral step-local view. Dynamic tensors are re-gathered every
step, and step outputs are scattered back into request states by
``scatter_latents()`` using ``idx_mapping``.
"""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
import numpy as np
import torch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
def _normalize_prompt_embeds(x: torch.Tensor) -> torch.Tensor:
if x.ndim == 2:
return x.unsqueeze(0)
if x.ndim == 3:
return x
raise ValueError(f"prompt_embeds must be 2D or 3D, got shape={tuple(x.shape)}")
def _normalize_mask(x: torch.Tensor) -> torch.Tensor:
if x.ndim == 1:
x = x.unsqueeze(0)
elif x.ndim != 2:
raise ValueError(f"prompt mask must be 1D or 2D, got shape={tuple(x.shape)}")
if x.dtype != torch.bool:
x = x != 0
return x
def _pad_prompt_embeds(x: torch.Tensor, target_seq_len: int) -> torch.Tensor:
x = _normalize_prompt_embeds(x)
bsz, seq_len, hidden = x.shape
if seq_len == target_seq_len:
return x
out = x.new_zeros((bsz, target_seq_len, hidden))
out[:, :seq_len] = x
return out
def _pad_mask(x: torch.Tensor, target_seq_len: int) -> torch.Tensor:
x = _normalize_mask(x)
bsz, seq_len = x.shape
if seq_len == target_seq_len:
return x
out = torch.zeros((bsz, target_seq_len), dtype=torch.bool, device=x.device)
out[:, :seq_len] = x
return out
def _select_states(
states: Sequence[DiffusionRequestState],
idx_mapping: torch.Tensor | None,
) -> tuple[list[DiffusionRequestState], torch.Tensor, np.ndarray]:
if not states:
raise ValueError("Cannot build InputBatch from empty states.")
if idx_mapping is None:
device = states[0].latents.device if states[0].latents is not None else None
idx_mapping = torch.arange(len(states), dtype=torch.int32, device=device)
else:
if idx_mapping.ndim != 1:
raise ValueError("idx_mapping must be a 1D tensor.")
idx_mapping = idx_mapping.to(dtype=torch.int32)
selected_states: list[DiffusionRequestState] = []
for batch_idx, state_idx in enumerate(idx_mapping.tolist()):
if state_idx < 0 or state_idx >= len(states):
raise ValueError(f"idx_mapping[{batch_idx}]={state_idx} is out of range for states.")
selected_states.append(states[state_idx])
return selected_states, idx_mapping, idx_mapping.detach().cpu().numpy()
def _prepare_req_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
return [state.req_id for state in states]
def _prepare_prompt_field_on_state(
state: DiffusionRequestState,
*,
embeds_attr: str,
mask_attr: str,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
embeds = getattr(state, embeds_attr)
if embeds is None:
return None, None
normalized_embeds = _normalize_prompt_embeds(embeds)
if normalized_embeds is not embeds:
setattr(state, embeds_attr, normalized_embeds)
embeds = normalized_embeds
mask = getattr(state, mask_attr)
if mask is None:
return embeds, None
normalized_mask = _normalize_mask(mask)
if normalized_mask is not mask:
setattr(state, mask_attr, normalized_mask)
return embeds, normalized_mask
def _prepare_reused_buffer(
current: torch.Tensor | None,
*,
shape: tuple[int, ...],
dtype: torch.dtype,
device: torch.device,
) -> torch.Tensor:
if current is not None and tuple(current.shape) == shape and current.dtype == dtype and current.device == device:
return current
return torch.empty(shape, dtype=dtype, device=device)
def _validate_gather_tensors(
values: Sequence[torch.Tensor],
*,
field_name: str,
) -> tuple[torch.dtype, torch.device, tuple[int, ...], int]:
if not values:
raise ValueError(f"Cannot gather empty tensor list for {field_name}.")
first = values[0]
dtype = first.dtype
device = first.device
suffix_shape = tuple(first.shape[1:])
total_rows = 0
for value in values:
if value.dtype != dtype:
raise ValueError(f"Mixed dtypes in {field_name} batch.")
if value.device != device:
raise ValueError(f"Mixed devices in {field_name} batch.")
if tuple(value.shape[1:]) != suffix_shape:
raise ValueError(
f"Mixed trailing shapes in {field_name} batch: expected {suffix_shape}, got {tuple(value.shape[1:])}."
)
total_rows += int(value.shape[0])
return dtype, device, suffix_shape, total_rows
def _gather_tensor_rows(
values: Sequence[torch.Tensor],
*,
field_name: str,
out: torch.Tensor | None = None,
) -> torch.Tensor:
dtype, device, suffix_shape, total_rows = _validate_gather_tensors(
values,
field_name=field_name,
)
gathered = _prepare_reused_buffer(
out,
shape=(total_rows, *suffix_shape),
dtype=dtype,
device=device,
)
row_offset = 0
for value in values:
next_row_offset = row_offset + int(value.shape[0])
gathered[row_offset:next_row_offset].copy_(value)
row_offset = next_row_offset
return gathered
def _get_seq_lens_from_mask(mask: torch.Tensor) -> list[int]:
return mask.sum(dim=1, dtype=torch.int32).tolist()
def _get_request_prompt_seq_lens(
state: DiffusionRequestState,
*,
embeds_attr: str,
mask_attr: str,
seq_lens_attr: str,
) -> list[int]:
embeds, mask = _prepare_prompt_field_on_state(
state,
embeds_attr=embeds_attr,
mask_attr=mask_attr,
)
if embeds is None:
raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
if mask is not None:
if mask.shape[0] != embeds.shape[0]:
raise ValueError(f"{mask_attr} batch dimension does not match {embeds_attr} for request {state.req_id}.")
return _get_seq_lens_from_mask(mask)
seq_lens = getattr(state, seq_lens_attr)
if seq_lens is not None:
return [int(value) for value in seq_lens]
return [int(embeds.shape[1])] * int(embeds.shape[0])
def _prepare_request_prompt_field(
state: DiffusionRequestState,
*,
embeds_attr: str,
mask_attr: str,
seq_lens_attr: str,
target_seq_len: int,
) -> tuple[torch.Tensor, torch.Tensor | None]:
embeds, mask = _prepare_prompt_field_on_state(
state,
embeds_attr=embeds_attr,
mask_attr=mask_attr,
)
if embeds is None:
raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
actual_seq_lens = _get_request_prompt_seq_lens(
state,
embeds_attr=embeds_attr,
mask_attr=mask_attr,
seq_lens_attr=seq_lens_attr,
)
max_actual_seq_len = max(actual_seq_lens) if actual_seq_lens else 0
if mask is None and max_actual_seq_len != target_seq_len:
raise ValueError(
f"Variable-length {embeds_attr} in batch but {mask_attr} is None. "
f"Provide masks or ensure {embeds_attr} have the same seq_len."
)
current_seq_len = int(embeds.shape[1])
if current_seq_len < target_seq_len:
embeds = _pad_prompt_embeds(embeds, target_seq_len)
setattr(state, embeds_attr, embeds)
if mask is not None:
mask = _pad_mask(mask, target_seq_len)
setattr(state, mask_attr, mask)
current_seq_len = target_seq_len
if current_seq_len > target_seq_len:
if max_actual_seq_len > target_seq_len:
raise ValueError(
f"{embeds_attr} for request {state.req_id} requires seq_len "
f"{max_actual_seq_len}, got target {target_seq_len}."
)
return embeds[:, :target_seq_len], None if mask is None else mask[:, :target_seq_len]
return embeds, mask
def _prepare_padded_prompt_fields(
states: Sequence[DiffusionRequestState],
*,
embeds_attr: str,
mask_attr: str,
seq_lens_attr: str,
embeds_out: torch.Tensor | None = None,
mask_out: torch.Tensor | None = None,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
prepared_fields = [
_prepare_prompt_field_on_state(
state,
embeds_attr=embeds_attr,
mask_attr=mask_attr,
)
for state in states
]
embeds_values = [embeds for embeds, _ in prepared_fields]
if not any(embeds is not None for embeds in embeds_values):
return None, None
if not all(embeds is not None for embeds in embeds_values):
raise ValueError(f"Mixed {embeds_attr} in batch.")
mask_values = [mask for _, mask in prepared_fields]
if any(mask is None for mask in mask_values):
if any(mask is not None for mask in mask_values):
raise ValueError(f"Mixed {mask_attr} in batch.")
target_seq_len = max(
max(
_get_request_prompt_seq_lens(
state,
embeds_attr=embeds_attr,
mask_attr=mask_attr,
seq_lens_attr=seq_lens_attr,
)
)
for state in states
)
request_embeds: list[torch.Tensor] = []
request_masks: list[torch.Tensor] = []
for state in states:
prepared_embeds, prepared_mask = _prepare_request_prompt_field(
state,
embeds_attr=embeds_attr,
mask_attr=mask_attr,
seq_lens_attr=seq_lens_attr,
target_seq_len=target_seq_len,
)
request_embeds.append(prepared_embeds)
if prepared_mask is not None:
request_masks.append(prepared_mask)
gathered_embeds = _gather_tensor_rows(
request_embeds,
field_name=embeds_attr,
out=embeds_out,
)
if not request_masks:
return gathered_embeds, None
gathered_masks = _gather_tensor_rows(
request_masks,
field_name=mask_attr,
out=mask_out,
)
return gathered_embeds, gathered_masks
def _prepare_latents(
states: Sequence[DiffusionRequestState],
*,
out: torch.Tensor | None = None,
) -> torch.Tensor:
latents_values = [state.latents for state in states]
if any(latents is None for latents in latents_values):
raise ValueError("All requests must have `latents` initialized.")
return _gather_tensor_rows(
[latents for latents in latents_values if latents is not None],
field_name="latents",
out=out,
)
def _require_state_latents(
state: DiffusionRequestState,
*,
for_field: str,
) -> torch.Tensor:
latents = state.latents
if latents is None:
raise ValueError(f"Request {state.req_id} has no latents while preparing {for_field}.")
return latents
def _expand_scalar_or_vector(
value: torch.Tensor,
*,
num_rows: int,
field_name: str,
) -> torch.Tensor:
if value.ndim == 0:
return value.reshape(1).expand(num_rows)
if value.ndim != 1:
raise ValueError(f"{field_name} must be scalar or 1D, got ndim={value.ndim}.")
if value.shape[0] == num_rows:
return value
if value.shape[0] == 1:
return value.expand(num_rows)
raise ValueError(
f"Per-request {field_name} must have either 1 element or {num_rows} elements; got {value.shape[0]}."
)
def _prepare_timesteps(
states: Sequence[DiffusionRequestState],
*,
out: torch.Tensor | None = None,
) -> torch.Tensor:
timestep_values: list[torch.Tensor] = []
for state in states:
timestep = state.current_timestep
if timestep is None:
raise ValueError("All requests must have a current timestep initialized.")
if not torch.is_tensor(timestep):
raise ValueError("InputBatch expects tensor timesteps; normalize them before batching.")
latents = _require_state_latents(state, for_field="timesteps")
timestep_values.append(
_expand_scalar_or_vector(
timestep,
num_rows=int(latents.shape[0]),
field_name="timestep tensor",
)
)
return _gather_tensor_rows(
timestep_values,
field_name="timesteps",
out=out,
)
def _prepare_cfg_scalars(
states: Sequence[DiffusionRequestState],
) -> tuple[bool, float, bool]:
def _cfg_scalars(state: DiffusionRequestState) -> tuple[bool, float, bool]:
true_cfg_scale = getattr(state.sampling, "true_cfg_scale", None) or 4.0
cfg_normalize = bool(getattr(state.sampling, "cfg_normalize", False))
return state.do_true_cfg, true_cfg_scale, cfg_normalize
scalars = _cfg_scalars(states[0])
for state in states[1:]:
if _cfg_scalars(state) != scalars:
raise ValueError("Mixed CFG settings in one diffusion batch are not supported.")
return scalars
def _prepare_guidance(
states: Sequence[DiffusionRequestState],
*,
out: torch.Tensor | None = None,
) -> torch.Tensor | None:
guidance_values = [state.guidance for state in states]
if all(guidance is None for guidance in guidance_values):
return None
if any(guidance is None for guidance in guidance_values):
raise ValueError("Mixed guidance in one diffusion batch are not supported.")
gathered_guidance: list[torch.Tensor] = []
for state in states:
latents = _require_state_latents(state, for_field="guidance")
guidance = state.guidance
assert guidance is not None
guidance_tensor = torch.as_tensor(
guidance,
device=latents.device,
dtype=latents.dtype,
)
gathered_guidance.append(
_expand_scalar_or_vector(
guidance_tensor,
num_rows=int(latents.shape[0]),
field_name="guidance tensor",
)
)
return _gather_tensor_rows(
gathered_guidance,
field_name="guidance",
out=out,
)
def _prepare_image_latents(
states: Sequence[DiffusionRequestState],
*,
out: torch.Tensor | None = None,
) -> torch.Tensor | None:
image_latents = [getattr(state.sampling, "image_latent", None) for state in states]
if all(image_latent is None for image_latent in image_latents):
return None
if any(image_latent is None for image_latent in image_latents):
raise ValueError("Mixed image_latent presence in one diffusion batch is not supported.")
return _gather_tensor_rows(
[image_latent for image_latent in image_latents if image_latent is not None],
field_name="image_latents",
out=out,
)
def _prepare_seq_lens(
states: Sequence[DiffusionRequestState],
attr_name: str,
) -> list[int] | None:
values = [getattr(state, attr_name) for state in states]
if all(value is None for value in values):
return None
if any(value is None for value in values):
raise ValueError(f"Mixed {attr_name} in batch.")
return [int(value[0]) for value in values if value is not None]
def _prepare_img_shapes(states: Sequence[DiffusionRequestState]) -> list | None:
values = [state.img_shapes for state in states]
if all(value is None for value in values):
return None
if any(value is None for value in values):
raise ValueError("Mixed img_shapes in batch.")
return [value[0] if value else [] for value in values if value is not None]
def _prepare_prompt_embeds(
states: Sequence[DiffusionRequestState],
*,
embeds_out: torch.Tensor | None = None,
mask_out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
prompt_embeds, prompt_embeds_mask = _prepare_padded_prompt_fields(
states,
embeds_attr="prompt_embeds",
mask_attr="prompt_embeds_mask",
seq_lens_attr="txt_seq_lens",
embeds_out=embeds_out,
mask_out=mask_out,
)
if prompt_embeds is None:
raise ValueError("All requests must have `prompt_embeds` initialized.")
return prompt_embeds, prompt_embeds_mask
def _prepare_negative_prompt_embeds(
states: Sequence[DiffusionRequestState],
*,
embeds_out: torch.Tensor | None = None,
mask_out: torch.Tensor | None = None,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
return _prepare_padded_prompt_fields(
states,
embeds_attr="negative_prompt_embeds",
mask_attr="negative_prompt_embeds_mask",
seq_lens_attr="negative_txt_seq_lens",
embeds_out=embeds_out,
mask_out=mask_out,
)
def _same_composition(
cached_batch: InputBatch | None,
req_ids: list[str],
idx_mapping_np: np.ndarray,
) -> bool:
if cached_batch is None:
return False
if cached_batch.req_ids != req_ids:
return False
return np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np)
def _scatter_batch_tensor_by_mapping(
states: Sequence[DiffusionRequestState],
idx_mapping_np: np.ndarray,
*,
attr_name: str,
value: torch.Tensor,
) -> None:
row_offset = 0
for batch_idx, state_idx in enumerate(idx_mapping_np.tolist()):
if state_idx < 0 or state_idx >= len(states):
raise ValueError(f"idx_mapping[{batch_idx}]={state_idx} is out of range for states.")
state = states[state_idx]
state_value = getattr(state, attr_name)
num_rows = 1 if state_value is None else int(state_value.shape[0])
next_row_offset = row_offset + num_rows
value_slice = value[row_offset:next_row_offset]
if state_value is None:
setattr(state, attr_name, value_slice.clone())
elif (
tuple(state_value.shape) != tuple(value_slice.shape)
or state_value.dtype != value_slice.dtype
or state_value.device != value_slice.device
):
setattr(state, attr_name, value_slice.clone())
else:
state_value.copy_(value_slice)
row_offset = next_row_offset
if row_offset != int(value.shape[0]):
raise ValueError(
f"Scatter for {attr_name} consumed {row_offset} rows, but batch has {int(value.shape[0])} rows."
)
@dataclass
class InputBatch:
"""Ephemeral step-level batch view.
Static request-local tensors are normalized and padded onto
``DiffusionRequestState`` itself, making the request state the persistent
source of truth. ``InputBatch`` only assembles a contiguous view for the
current step and refreshes dynamic fields in-place when composition is
unchanged.
"""
req_ids: list[str]
num_reqs: int
num_reqs_after_padding: int
idx_mapping: torch.Tensor
idx_mapping_np: np.ndarray
latents: torch.Tensor
timesteps: torch.Tensor
prompt_embeds: torch.Tensor
prompt_embeds_mask: torch.Tensor | None
negative_prompt_embeds: torch.Tensor | None
negative_prompt_embeds_mask: torch.Tensor | None
guidance: torch.Tensor | None = None
do_true_cfg: bool = False
true_cfg_scale: float = 4.0
cfg_normalize: bool = False
image_latents: torch.Tensor | None = None
img_shapes: list | None = None
txt_seq_lens: list[int] | None = None
negative_txt_seq_lens: list[int] | None = None
def __post_init__(self) -> None:
if len(self.req_ids) != int(self.idx_mapping.numel()):
raise ValueError("`req_ids` and `idx_mapping` must have the same length.")
if self.num_reqs != len(self.req_ids):
raise ValueError("`num_reqs` must match the number of request ids.")
if self.num_reqs_after_padding < self.num_reqs:
raise ValueError("`num_reqs_after_padding` must be >= `num_reqs`.")
def _refresh_dynamic_fields(
self,
selected_states: Sequence[DiffusionRequestState],
) -> None:
self.latents = _prepare_latents(selected_states, out=self.latents)
self.timesteps = _prepare_timesteps(selected_states, out=self.timesteps)
def _refresh_static_fields(
self,
states: Sequence[DiffusionRequestState],
) -> None:
self.do_true_cfg, self.true_cfg_scale, self.cfg_normalize = _prepare_cfg_scalars(states)
self.guidance = _prepare_guidance(states, out=self.guidance)
self.image_latents = _prepare_image_latents(states, out=self.image_latents)
self.prompt_embeds, self.prompt_embeds_mask = _prepare_prompt_embeds(
states,
embeds_out=self.prompt_embeds,
mask_out=self.prompt_embeds_mask,
)
(
self.negative_prompt_embeds,
self.negative_prompt_embeds_mask,
) = _prepare_negative_prompt_embeds(
states,
embeds_out=self.negative_prompt_embeds,
mask_out=self.negative_prompt_embeds_mask,
)
self.img_shapes = _prepare_img_shapes(states)
self.txt_seq_lens = _prepare_seq_lens(states, "txt_seq_lens")
self.negative_txt_seq_lens = _prepare_seq_lens(states, "negative_txt_seq_lens")
def _repack_dynamic_fields(
self,
selected_states: Sequence[DiffusionRequestState],
) -> None:
self._refresh_dynamic_fields(selected_states)
def _rebuild(
self,
selected_states: Sequence[DiffusionRequestState],
idx_mapping: torch.Tensor,
idx_mapping_np: np.ndarray,
req_ids: list[str],
) -> InputBatch:
self.req_ids = req_ids
self.num_reqs = len(req_ids)
self.num_reqs_after_padding = len(req_ids)
self.idx_mapping = idx_mapping
self.idx_mapping_np = idx_mapping_np
self.latents = _prepare_latents(selected_states, out=self.latents)
self.timesteps = _prepare_timesteps(selected_states, out=self.timesteps)
self._refresh_static_fields(selected_states)
self.__post_init__()
return self
@classmethod
def make_batch(
cls,
states: Sequence[DiffusionRequestState],
idx_mapping: torch.Tensor | None = None,
cached_batch: InputBatch | None = None,
) -> InputBatch:
"""Build a temporary step-local batch view from request states."""
selected_states, idx_mapping, idx_mapping_np = _select_states(states, idx_mapping)
req_ids = _prepare_req_ids(selected_states)
if _same_composition(cached_batch, req_ids, idx_mapping_np):
assert cached_batch is not None
cached_batch._repack_dynamic_fields(selected_states)
return cached_batch
if cached_batch is not None:
return cached_batch._rebuild(
selected_states,
idx_mapping,
idx_mapping_np,
req_ids,
)
prompt_embeds, prompt_embeds_mask = _prepare_prompt_embeds(selected_states)
negative_prompt_embeds, negative_prompt_embeds_mask = _prepare_negative_prompt_embeds(selected_states)
do_true_cfg, true_cfg_scale, cfg_normalize = _prepare_cfg_scalars(selected_states)
return cls(
req_ids=req_ids,
num_reqs=len(selected_states),
num_reqs_after_padding=len(selected_states),
idx_mapping=idx_mapping,
idx_mapping_np=idx_mapping_np,
latents=_prepare_latents(selected_states),
timesteps=_prepare_timesteps(selected_states),
guidance=_prepare_guidance(selected_states),
do_true_cfg=do_true_cfg,
true_cfg_scale=true_cfg_scale,
cfg_normalize=cfg_normalize,
image_latents=_prepare_image_latents(selected_states),
prompt_embeds=prompt_embeds,
prompt_embeds_mask=prompt_embeds_mask,
negative_prompt_embeds=negative_prompt_embeds,
negative_prompt_embeds_mask=negative_prompt_embeds_mask,
img_shapes=_prepare_img_shapes(selected_states),
txt_seq_lens=_prepare_seq_lens(selected_states, "txt_seq_lens"),
negative_txt_seq_lens=_prepare_seq_lens(
selected_states,
"negative_txt_seq_lens",
),
)
def scatter_latents(
states: Sequence[DiffusionRequestState],
input_batch: InputBatch,
) -> None:
"""Scatter the step-updated latents back into persistent request states.
This is the CPU fallback of the vLLM-style post-update path. The mapping is
driven entirely by ``input_batch.idx_mapping_np`` so the runner remains free
to keep request states in its own persistent storage layout.
"""
_scatter_batch_tensor_by_mapping(
states,
input_batch.idx_mapping_np,
attr_name="latents",
value=input_batch.latents,
)
DiffusionInputBatch = InputBatch
vllm_omni/diffusion/worker/utils.py
+39 -1
View File
@@ -4,6 +4,7 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
@@ -109,8 +110,14 @@ class DiffusionRequestState:
return self.step_index == 0 or self.timesteps is None
class BaseRunnerOutput(ABC):
@abstractmethod
def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
pass
@dataclass
class RunnerOutput:
class RunnerOutput(BaseRunnerOutput):
"""Output of a single denoising step for a request.
NOTE: `latents` may be None when returned through IPC to avoid
@@ -122,3 +129,34 @@ class RunnerOutput:
step_index: int | None = None
finished: bool = False
result: DiffusionOutput | None = None
def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
return self if self.req_id == sched_req_id else None
@dataclass
class BatchRunnerOutput(BaseRunnerOutput):
runner_outputs: list[RunnerOutput]
_id_to_idx: dict[str, int] = field(init=False, repr=False)
def __post_init__(self) -> None:
self._id_to_idx = {out.req_id: i for i, out in enumerate(self.runner_outputs)}
def __getitem__(self, req_id: str) -> RunnerOutput | None:
"""access single RunnerOutput by req_id"""
idx = self._id_to_idx.get(req_id)
return self.runner_outputs[idx] if idx is not None else None
def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
return self[sched_req_id]
@property
def req_ids(self) -> list[str]:
return list(self._id_to_idx.keys())
def __len__(self) -> int:
return len(self.runner_outputs)
@classmethod
def from_list(cls, runner_output_list: list[RunnerOutput]) -> BatchRunnerOutput:
return cls(runner_outputs=runner_output_list)
vllm_omni/engine/async_omni_engine.py
+1 -1
View File
@@ -1329,7 +1329,7 @@ class AsyncOmniEngine:
final_output_type = "audio" if model_class_name and supports_audio_output(model_class_name) else "image"
stage_engine_args = {
"max_num_seqs": 1,
"max_num_seqs": kwargs.get("max_num_seqs") or 1,
"parallel_config": parallel_config,
"model_class_name": kwargs.get("model_class_name", None),
"step_execution": kwargs.get("step_execution", False),