diff --git a/benchmarks/diffusion/README.md b/benchmarks/diffusion/README.md index 06bf01726..16a86c88f 100644 --- a/benchmarks/diffusion/README.md +++ b/benchmarks/diffusion/README.md @@ -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). diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 77b36b3d9..f2b77ff8d 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -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)) diff --git a/benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md b/benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md index ce022f1a8..172a8ed05 100644 --- a/benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md +++ b/benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md @@ -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 --- diff --git a/docs/.nav.yml b/docs/.nav.yml index ddd2b0308..94de4dc94 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -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 diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 6d5782a6e..1c08b8bef 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -737,6 +737,16 @@ step-scoped phases. The reference implementation is Do not enable `step_execution=True` until those four methods are implemented and validated against the request-level path. +If you want the pipeline to work with the experimental batched step-wise path +(`max_num_seqs > 1`), also see: +[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md). + +If you expose this in example scripts or recipes, keep it opt-in. Surface +runtime features like `step_execution` as optional flags instead of silently +turning them on. For Qwen-Image-style serving examples, document +`--step-execution` as the feature gate and `--max-num-seqs N` as the +companion batching knob. + ### Cache Acceleration #### TeaCache diff --git a/docs/design/feature/diffusion_continuous_batching.md b/docs/design/feature/diffusion_continuous_batching.md new file mode 100644 index 000000000..ab8ff69fe --- /dev/null +++ b/docs/design/feature/diffusion_continuous_batching.md @@ -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) diff --git a/docs/design/feature/diffusion_step_execution.md b/docs/design/feature/diffusion_step_execution.md index b8c81f04f..797a1cb48 100644 --- a/docs/design/feature/diffusion_step_execution.md +++ b/docs/design/feature/diffusion_step_execution.md @@ -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. diff --git a/docs/design/index.md b/docs/design/index.md index 31420550f..61aa50483 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -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 diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index b0c7e9fc7..e0fe00299 100644 --- a/docs/design/module/dit_module.md +++ b/docs/design/module/dit_module.md @@ -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. --- diff --git a/docs/user_guide/diffusion/step_execution.md b/docs/user_guide/diffusion/step_execution.md index f8c9fa8dd..189510916 100644 --- a/docs/user_guide/diffusion/step_execution.md +++ b/docs/user_guide/diffusion/step_execution.md @@ -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). diff --git a/recipes/Qwen/Qwen-Image.md b/recipes/Qwen/Qwen-Image.md new file mode 100644 index 000000000..de0d73c07 --- /dev/null +++ b/recipes/Qwen/Qwen-Image.md @@ -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. diff --git a/recipes/README.md b/recipes/README.md index ac493e75f..4e9536b92 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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 diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py new file mode 100644 index 000000000..dca9a001a --- /dev/null +++ b/tests/diffusion/test_diffusion_engine.py @@ -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 diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 6c76d0a46..c242a0572 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -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) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 06f8cd14d..9417d540b 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -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( diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 9ec06e810..230838fbb 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -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): diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index f1cf4f9b7..36ceb6cfc 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -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 diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index c4c62dd59..f2c3f4eed 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -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(): diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index 81eba172c..f8d2b364a 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/vllm_omni/diffusion/executor/abstract.py @@ -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 diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 72436c79b..9f00ced38 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -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. diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 1c251867a..1817e7fdc 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -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] = {} diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index c569b2330..c786f82f9 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -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, ) diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index 6a7ee3d3e..d2c0f1e75 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -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: diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index 4db6f4135..8eef211ff 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -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] diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py index f641648e9..5029e3593 100644 --- a/vllm_omni/diffusion/sched/request_scheduler.py +++ b/vllm_omni/diffusion/sched/request_scheduler.py @@ -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" diff --git a/vllm_omni/diffusion/sched/step_scheduler.py b/vllm_omni/diffusion/sched/step_scheduler.py index 4d995dcf4..58449281b 100644 --- a/vllm_omni/diffusion/sched/step_scheduler.py +++ b/vllm_omni/diffusion/sched/step_scheduler.py @@ -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: diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index d36c5c644..2f3ad5594 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -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 = [] diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 535f053c3..cc2d00f1a 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -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) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 927bbeb1a..887ad27d3 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -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) diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py new file mode 100644 index 000000000..48c5bb57b --- /dev/null +++ b/vllm_omni/diffusion/worker/input_batch.py @@ -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 diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index af88913e7..1e98a7784 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/utils.py @@ -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) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 66096de7c..a5a17c719 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -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),