[New model] Qwen-Image-2.1 support (rebase + review follow-ups) (#8099)
Signed-off-by: NumberWan <wantszkin2003@gmail.com> Signed-off-by: SamitHuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: SamitHuang <285365963@qq.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
This commit is contained in:
52 files changed
+8058
-136
No files matched your search
@@ -32,6 +32,7 @@ th {
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImage21Pipeline` | Qwen-Image-2.1 (T2I + image-conditioned generation, up to 4 condition images) | `Qwen/Qwen-Image-2.1` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2.1.md) |
|
||||
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
|
||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tongyi-MAI/Z-Image-Turbo.md) |
|
||||
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) |
|
||||
|
||||
@@ -163,6 +163,7 @@ The following tables show which models support each feature:
|
||||
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Qwen-Image-2.1** | ❌ | ❌ | ❓ | ✅ | ✅ | ❌ | ❓ | ❓ | ✅ (decode) | ✅ | ✅ |
|
||||
| **SenseNova-U1** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **SenseNova-U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **SenseNova-U1-A3B** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
@@ -181,7 +182,8 @@ The following tables show which models support each feature:
|
||||
> 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
|
||||
> 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
|
||||
> 7. MammothModa2 runs its DiT stage on the diffusion runner (`StageExecutionType.DIFFUSION`); Cache-DiT is enabled through the standard diffusion-stage knobs on the stage entry of the deploy YAML (`cache_backend: cache_dit`, plus optional `cache_config` / `enable_cache_dit_summary`). The runner installs the backend at startup and the pipeline adopts it per request. Only the repeated main-layer stack is cached; requests with `text_guidance_scale = 1.0` bypass the cache hooks.
|
||||
> 8. SenseNova-U1-A3B: `--cfg-parallel-size 2` and TP=2 both complete. CFG-P covers t2i 2-branch (`cond`+`uncond`), img2img 2-branch (`cond`+`img_cond` when `img_cfg_scale=1`), and img2img 3-branch (`cond`+`img_cond`+`uncond` when `img_cfg_scale` is neither 1 nor equal to `cfg_scale`; 3 branches on 2 ranks dispatch as `[[0, 2], [1]]`). Cache-DiT does not: the adapter is declared on `SenseNovaU1Model`, but discovery walks `_dit_modules` (`language_model` / `SenseNovaU1ForCausalLM`) and raises. Layerwise CPU offload is a no-op for the same reason — that class has no `_layerwise_offload_blocks_attrs`, so the 72.2 GiB expert weights stay resident.
|
||||
> 8. Qwen-Image-2.1 (`QwenImage21Pipeline`) handles both text-to-image and image-conditioned generation (up to 4 condition images) and supports LoRA through `QwenImageLoraLoaderMixin`. Ulysses sequence parallelism is wired and ring attention is not supported, but multi-GPU padding e2e and a parity number are not in this PR's H200 results, so SP stays ❓. Cache acceleration (TeaCache / Cache-DiT) is unavailable; its built-in KV cache for the text/condition prefix is a separate feature from those cache backends. Online FP8 is documented in the quantization guide with H200 numbers.
|
||||
> 9. SenseNova-U1-A3B: `--cfg-parallel-size 2` and TP=2 both complete. CFG-P covers t2i 2-branch (`cond`+`uncond`), img2img 2-branch (`cond`+`img_cond` when `img_cfg_scale=1`), and img2img 3-branch (`cond`+`img_cond`+`uncond` when `img_cfg_scale` is neither 1 nor equal to `cfg_scale`; 3 branches on 2 ranks dispatch as `[[0, 2], [1]]`). Cache-DiT does not: the adapter is declared on `SenseNovaU1Model`, but discovery walks `_dit_modules` (`language_model` / `SenseNovaU1ForCausalLM`) and raises. Layerwise CPU offload is a no-op for the same reason — that class has no `_layerwise_offload_blocks_attrs`, so the 72.2 GiB expert weights stay resident.
|
||||
|
||||
### VideoGen
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
# FP8 Quantization
|
||||
# FP8 Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
@@ -17,7 +17,7 @@ in deep DiT blocks.
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| -------- | --------- |
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ |
|
||||
@@ -93,8 +93,9 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
### Diffusion Models
|
||||
|
||||
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|
||||
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
|
||||
| ------- | ----------- | :-------: | :------: | ---------------- | ------------------ | ------------------ |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
|
||||
| Qwen-Image-2.1 | `Qwen/Qwen-Image-2.1` | Yes | Not validated | Skip image-stream MLPs for best fidelity; all-layer FP8 is loadable. On H200 (1024×1024, teapot prompt) DiT FP8-all is ~30.4 GB / 6.06 s vs BF16 ~36.9 GB / 7.34 s (~33 dB vs BF16). Text encoder FP8 covers the language model only — the vision tower and LM head stay BF16 | `img_mlp` | ✅︎ (language model only) |
|
||||
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
|
||||
| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
|
||||
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
|
||||
@@ -110,20 +111,20 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
| Model | Scope | Format | Status |
|
||||
|-------|-------|--------|--------|
|
||||
| Model | Scope | Format | Status |
|
||||
| ---------- | ---------------------------- | ---------------------------------------- | ----------------------------------- |
|
||||
| Qwen3-Omni | Thinker language-model stage | [ModelOpt](modelopt.md) `quant_algo=FP8` | Tested for thinker memory reduction |
|
||||
| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |
|
||||
| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |
|
||||
|
||||
Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
|
||||
model-specific guide says otherwise.
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| BAGEL | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling |
|
||||
| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline |
|
||||
| Model | Scope | Status | Notes |
|
||||
| --------- | ---------------------------------------- | ------------- | ----------------------------------------------- |
|
||||
| BAGEL | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling |
|
||||
| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline |
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -162,7 +163,7 @@ For a pipeline that exposes both a transformer and a quantization-aware text
|
||||
encoder, the scope is:
|
||||
|
||||
| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
|
||||
|---------------|-------------|--------------|-------------------------------------------------|
|
||||
| --------------- | ------------- | -------------- | ------------------------------------------------- |
|
||||
| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
|
||||
| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
|
||||
| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
|
||||
@@ -211,7 +212,7 @@ does not fit the workload.
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `method` | str | - | Quantization method (`"fp8"`) |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
|
||||
| `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling, or `"static"` when scales are available |
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""
|
||||
Example script for image editing with OmniGen2.
|
||||
@@ -209,14 +209,14 @@ def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
|
||||
parser = argparse.ArgumentParser(description="Edit images with vLLM-Omni.")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Qwen/Qwen-Image-Edit",
|
||||
help=(
|
||||
"Diffusion model name or local path. "
|
||||
"For multiple image inputs, use Qwen/Qwen-Image-Edit-2509 or Qwen/Qwen-Image-Edit-2511"
|
||||
"which supports QwenImageEditPlusPipeline."
|
||||
"For multiple image inputs, use Qwen/Qwen-Image-Edit-2509, "
|
||||
"Qwen/Qwen-Image-Edit-2511, or Qwen/Qwen-Image-2.1."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -368,7 +368,7 @@ def parse_args() -> argparse.Namespace:
|
||||
"--color-format",
|
||||
type=str,
|
||||
default="RGB",
|
||||
help="For Qwen-Image-Layered, set to RGBA.",
|
||||
help="For Qwen-Image-Layered or Qwen-Image-2.1, set to RGBA to preserve transparency.",
|
||||
)
|
||||
|
||||
# Cache-DiT specific parameters
|
||||
@@ -647,6 +647,7 @@ def main():
|
||||
|
||||
diffusion_params = OmniDiffusionSamplingParams(
|
||||
generator=generator,
|
||||
seed=args.seed,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_2,
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
python "$(dirname "${BASH_SOURCE[0]}")/image_edit.py" \
|
||||
--model Qwen/Qwen-Image-2.1 \
|
||||
--color-format RGBA \
|
||||
--seed 42 \
|
||||
--image qwen_bear.png \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
|
||||
--output qwen_image_21_edit.png \
|
||||
--num-inference-steps 50 \
|
||||
--cfg-scale 1.0 \
|
||||
"$@"
|
||||
@@ -108,7 +108,8 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
"--model",
|
||||
default="Qwen/Qwen-Image",
|
||||
help="Diffusion model name or local path. Supported models: "
|
||||
"Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
|
||||
"Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, Qwen/Qwen-Image-2.1, "
|
||||
"stepfun-ai/NextStep-1.1, "
|
||||
"black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
|
||||
"black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
|
||||
"HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
|
||||
|
||||
@@ -0,0 +1,418 @@
|
||||
# Qwen-Image-2.1
|
||||
|
||||
> Text-to-image and image-conditioned generation with Qwen-Image 2.1
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Qwen
|
||||
- Model: `Qwen/Qwen-Image-2.1`
|
||||
- Task: Text-to-image generation and image-conditioned generation (editing)
|
||||
- Mode: Offline inference and online serving with optional step-wise execution
|
||||
- Maintainer: Community
|
||||
|
||||
> **Model id:** `Qwen/Qwen-Image-2.1` follows the reference diffusers
|
||||
> documentation and may change before the official release — the final
|
||||
> checkpoint name is subject to the official release.
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe as a starting point for running `Qwen/Qwen-Image-2.1`. A single
|
||||
pipeline class, `QwenImage21Pipeline`, handles both pure text-to-image requests
|
||||
(no condition image) and image-conditioned requests with up to 4 condition
|
||||
images: prompt and condition images are encoded together by a Qwen3-VL text
|
||||
encoder, so editing-style requests go through the same serving path as
|
||||
text-to-image.
|
||||
|
||||
## References
|
||||
|
||||
- Related offline examples:
|
||||
[`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py),
|
||||
[`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py)
|
||||
- Supported models table:
|
||||
[`docs/models/supported_models.md`](../../docs/models/supported_models.md)
|
||||
- Feature compatibility matrix:
|
||||
[`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md)
|
||||
- Sibling recipes: [`Qwen-Image.md`](./Qwen-Image.md),
|
||||
[`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md)
|
||||
|
||||
## Offline Inference
|
||||
|
||||
### Text-to-image
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model Qwen/Qwen-Image-2.1 \
|
||||
--prompt "A ceramic teapot on a wooden table" \
|
||||
--output qwen_image_21_t2i.png \
|
||||
--num-inference-steps 50 \
|
||||
--cfg-scale 1.0
|
||||
```
|
||||
|
||||
### Image-conditioned generation (editing)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/image_to_image/image_edit.py \
|
||||
--model Qwen/Qwen-Image-2.1 \
|
||||
--color-format RGBA \
|
||||
--seed 42 \
|
||||
--image qwen_bear.png \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars" \
|
||||
--output qwen_image_21_edit.png \
|
||||
--num-inference-steps 50 \
|
||||
--cfg-scale 1.0
|
||||
```
|
||||
|
||||
or via the bundled launcher:
|
||||
|
||||
```bash
|
||||
bash examples/offline_inference/image_to_image/run_qwen_image_21.sh
|
||||
```
|
||||
|
||||
Up to 4 condition images can be passed via `--image`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/image_to_image/image_edit.py \
|
||||
--model Qwen/Qwen-Image-2.1 \
|
||||
--color-format RGBA \
|
||||
--seed 42 \
|
||||
--image input1.png input2.png \
|
||||
--prompt "Combine these images into a single scene" \
|
||||
--output qwen_image_21_multi.png \
|
||||
--num-inference-steps 50 \
|
||||
--cfg-scale 1.0
|
||||
```
|
||||
|
||||
Use `--color-format RGBA` to preserve transparency in condition images. The
|
||||
shared example otherwise loads inputs as RGB. Add `--width` and `--height`
|
||||
(multiples of 32) to set the output size; when omitted, Qwen-Image 2.1 derives
|
||||
it from the last condition image's aspect ratio at approximately 1024×1024.
|
||||
Use `--num-outputs-per-prompt` to generate multiple images and `--vae-use-tiling`
|
||||
to reduce VAE memory usage. The launcher accepts these options as well.
|
||||
|
||||
## Online Serving
|
||||
|
||||
Start the server:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image-2.1 --omni --port 8091
|
||||
```
|
||||
|
||||
To enable the step-wise runtime:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image-2.1 --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
```
|
||||
|
||||
Step execution batches compatible requests at the same KV-cache phase. The
|
||||
pipeline tags every request with a `batch_compatibility_key` (conditioning
|
||||
mode plus condition-image layout), so a newly arriving request that would
|
||||
otherwise join a batch already past its first denoising step is isolated by
|
||||
the scheduler and deferred until the in-flight batch drains — it no longer
|
||||
fails the whole batch. `--max-num-seqs` controls batch capacity.
|
||||
|
||||
### CUDA Graph decode
|
||||
|
||||
Qwen-Image-2.1 automatically uses CUDA Graph for supported fixed-shape denoising
|
||||
decode steps when `enable_cuda_graph_decode=True` (the default) and
|
||||
`enforce_eager=False`. Prefill remains eager. This works with both request
|
||||
execution and `--step-execution`.
|
||||
|
||||
To disable graph capture while keeping the rest of the compilation setup:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image-2.1 --omni --no-enable-cuda-graph-decode
|
||||
```
|
||||
|
||||
The same control is available as `Omni(..., enable_cuda_graph_decode=False)`
|
||||
and as `enable_cuda_graph_decode: false` on a stage in the deployment YAML.
|
||||
`--enforce-eager` / `enforce_eager: true` remains the broader switch: it
|
||||
disables both graph capture and automatic `torch.compile`. The two
|
||||
optimizations stack: the DiT blocks are regionally `torch.compile`'d and
|
||||
graph capture records the compiled (fused) kernels, while inductor's own
|
||||
cudagraphs stay disabled so the two graph layers never nest. Measured on a
|
||||
single H200 (1024×1024, 50 steps, seed 42, BF16, `true_cfg_scale=1.0`,
|
||||
warmup 1 + median of 3): compile+graph **5.93 s**, compile-only **6.88 s**,
|
||||
graph-only **9.52 s**, eager **8.28 s** end-to-end per image; combo and
|
||||
compile-only vs eager are **38.10 dB** PSNR (compile fusion drift), and
|
||||
graph-only vs eager is bit-identical.
|
||||
|
||||
The autoregressive engine's `compilation_config.cudagraph_mode` does not
|
||||
control this diffusion path; use `enable_cuda_graph_decode` (or
|
||||
`enforce_eager` for everything) to force eager decode execution.
|
||||
|
||||
TP/SP/ring parallelism, HSDP, cache hooks, quantized KV caches,
|
||||
dynamic LoRA, padded text masks, and a second in-flight request whose cache
|
||||
aliases an already-owned graph key fall back to eager decode.
|
||||
Model-level CPU offload does not disable decode graphs while
|
||||
`enable_cuda_graph_decode` stays on (the default). DiT weights stay on fixed
|
||||
device storage so captured pointers remain valid, and that storage is not
|
||||
released while the text encoder runs. Pass `--no-enable-cuda-graph-decode`
|
||||
(or `enable_cuda_graph_decode: false`) for the eager offload path, which
|
||||
moves DiT weights off the device and frees that VRAM. Peak memory for the
|
||||
default graph-plus-offload combination was not re-measured in this PR.
|
||||
Graphs keep separate entries for different image layouts and copy
|
||||
request-owned prefix K/V into static buffers before replay; each graph key
|
||||
has a single live owner at a time, so concurrent same-shape requests never
|
||||
overwrite each other's prefix. Those buffers require additional memory; graph
|
||||
entries are bounded by the transformer's cache limit.
|
||||
|
||||
### Verification
|
||||
|
||||
For a direct API smoke test:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Qwen/Qwen-Image-2.1",
|
||||
"prompt": "A ceramic teapot on a wooden table",
|
||||
"size": "1024x1024",
|
||||
"num_inference_steps": 50,
|
||||
"true_cfg_scale": 1.0,
|
||||
"seed": 42
|
||||
}'
|
||||
```
|
||||
|
||||
For image-conditioned generation, pass condition images through the
|
||||
OpenAI-compatible multimodal endpoint (`/v1/chat/completions` with
|
||||
`image_url` content parts, or `/v1/images/edits`), same as
|
||||
[`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md).
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Default | Notes |
|
||||
| --- | --- | --- |
|
||||
| `true_cfg_scale` | 1.0 | Disables the negative-prompt branch by default. Set above 1 with a `negative_prompt` to enable true CFG. Qwen-Image 2.1 supports true CFG only — there is no `guidance_scale` parameter. |
|
||||
| `num_inference_steps` | 50 | Number of denoising steps. |
|
||||
| `height` / `width` | 1024×1024 | Must be multiples of 32. For image-conditioned requests, leaving them unset derives the output size from the last condition image's aspect ratio at ~1024×1024. |
|
||||
| `negative_prompt` | None | Ignored at the default CFG scale of 1.0; used only when `true_cfg_scale > 1`. |
|
||||
| condition images | — | Up to 4 input images per request; more than 4 raises an error. |
|
||||
| `seed` | — | Fix for reproducible outputs. |
|
||||
|
||||
The default CFG scale of 1.0 runs only the positive-prompt branch at each
|
||||
denoising step. To enable negative-prompt guidance, explicitly set
|
||||
`--cfg-scale` (offline) or `true_cfg_scale` (API) above 1 and provide a
|
||||
negative prompt. This changes the generated output as well as compute cost.
|
||||
|
||||
### Prefix KV cache
|
||||
|
||||
The transformer caches the text/condition-image prefix across denoising steps
|
||||
(a prefill/decode split keyed by CFG branch). The cache is enabled
|
||||
automatically when the checkpoint declares `causal_condition: true` and is
|
||||
re-created per generation, so no state leaks between requests. It changes the
|
||||
compute path — cached prefix vs. full recompute at every step — so exact
|
||||
numerical reproduction against a no-cache reference (e.g. the diffusers
|
||||
pipeline) requires matching the cache setting on both sides. The prefix KV
|
||||
cache is orthogonal to weight quantization and works unchanged under FP8.
|
||||
|
||||
### FP8 quantization
|
||||
|
||||
Online FP8 quantization of the transformer is supported; the fused
|
||||
`to_qkv` projections load correctly, and the prefix KV cache path is
|
||||
unaffected. Only block-internal linears are eligible — the boundary
|
||||
projections (`img_in`/`txt_in`/`modulation`/`proj_out`/`time_text_embed`) stay
|
||||
in BF16 by construction:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model Qwen/Qwen-Image-2.1 \
|
||||
--prompt "A ceramic teapot on a wooden table" \
|
||||
--output qwen_image_21_fp8.png \
|
||||
--num-inference-steps 50 --cfg-scale 1.0 \
|
||||
--quantization fp8 --ignored-layers "img_mlp"
|
||||
```
|
||||
|
||||
`ignored_layers` entries are matched as name patterns for this model (e.g.
|
||||
`img_mlp` matches every block's `img_mlp.proj`/`img_mlp.gate_layer`/
|
||||
`img_mlp.out`); exact full prefixes like
|
||||
`transformer.transformer_blocks.3.attn.to_qkv` also work.
|
||||
|
||||
A per-module sensitivity sweep (fixed prompt embeds / latent / timestep, 5
|
||||
timesteps across the trajectory, relative L2 of the noise prediction vs BF16)
|
||||
ranks the groups: `img_mlp` is the most sensitive (3.8% max error on its own,
|
||||
driven by `img_mlp.proj` 2.8% and `img_mlp.out` 2.5%), then `attn.to_out`
|
||||
(2.7%), then `attn.to_qkv` (1.6%); errors are largest at the late (low-sigma)
|
||||
steps. End-to-end on H200, single GPU, same prompt as Quality Comparison
|
||||
(`"A ceramic teapot on a wooden table"`, 1024×1024, 50 steps, seed 42,
|
||||
eager):
|
||||
|
||||
| Config | Quantized block linears | PSNR vs BF16 | Steady-state time/image | Peak memory |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| BF16 | 0 / 160 | — | **7.34 s** | **36.9 GB** |
|
||||
| FP8 all layers | 160 / 160 | **33.0 dB** | **6.06 s** | **30.4 GB** |
|
||||
| FP8, `ignored_layers=["img_mlp"]` | 64 / 160 | **34.3 dB** | **6.98 s** | **34.9 GB** |
|
||||
|
||||
Keeping `img_mlp` in BF16 is still the higher-fidelity DiT setting vs
|
||||
all-layer FP8. On H200, all-layer FP8 is both leaner and faster than BF16
|
||||
at this shape.
|
||||
|
||||
## Quantization
|
||||
|
||||
Online FP8 is supported per component via `quantization_config`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image-2.1",
|
||||
quantization_config={
|
||||
# DiT: skip the sensitive image-stream MLPs if quality regresses.
|
||||
"transformer": {"method": "fp8", "ignored_layers": ["img_mlp"]},
|
||||
# Qwen3-VL text encoder: only the language-model linear layers are
|
||||
# quantized; the vision tower and the (unused) LM head stay BF16.
|
||||
"text_encoder": {"method": "fp8"},
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
Either component can be quantized on its own. Measured on H200 at 1024×1024
|
||||
(seed 42, 50 steps, same teapot prompt, eager): text-encoder FP8 lowers peak
|
||||
GPU memory from **36.9 GB** (BF16) to **30.4 GB** (−6.5 GB); T2I PSNR vs BF16
|
||||
**32.1 dB**. DiT FP8-all is **30.4 GB** / **6.06 s**. Combined DiT+text-encoder
|
||||
FP8 was not re-measured on H200. The edit path routes condition images through
|
||||
the BF16 vision tower, so it is unaffected by text-encoder FP8. See
|
||||
[`docs/user_guide/quantization/fp8.md`](../../docs/user_guide/quantization/fp8.md)
|
||||
for the scope rules.
|
||||
|
||||
### FP8 prefix KV storage
|
||||
|
||||
The cached prefix K/V can be stored in FP8 E4M3 (with per-token-per-head fp32
|
||||
scales), halving the prefix-cache memory — relevant for long prompts (up to
|
||||
8192 text tokens) and up to 4 condition images (~4096 latent tokens each).
|
||||
Prefill quantizes once; decode dequantizes per step, so attention compute is
|
||||
unchanged. Enable it through the model-specific `extras` config (e.g. in the
|
||||
deploy YAML's stage `extras` or as an `Omni(...)` engine kwarg):
|
||||
|
||||
```yaml
|
||||
extras:
|
||||
prefix_kv_cache_dtype: "fp8"
|
||||
```
|
||||
|
||||
Accepted values: `"fp8"` / `"fp8_e4m3"` (K and V quantized), `"fp8_v"` (V
|
||||
only — K stays in the native dtype), and `None` / `"auto"` (default, native
|
||||
dtype — behavior unchanged). This is independent of
|
||||
`diffusion_kv_cache_dtype`, which quantizes attention Q/K/V *compute* per
|
||||
forward pass on supported backends.
|
||||
|
||||
Measured on one NVIDIA H200, T2I 1024x1024 (seed 42, 50 steps, true CFG 4.0,
|
||||
eager, PSNR vs. the bf16 baseline on the same GPU): **`"fp8"` 34.2 dB**,
|
||||
**`"fp8_v"` 44.0 dB**. The error is dominated by
|
||||
K quantization: post-RoPE keys are the precision-sensitive half of the
|
||||
cache, so `"fp8_v"` buys back about 9.8 dB at 75% (instead of 50%) of the original
|
||||
cache size. The residual error is inherent e4m3 precision accumulated
|
||||
coherently over the denoising trajectory; finer scale granularity
|
||||
(per-tensor/per-head/per-token were compared) or Hadamard-rotated V did not
|
||||
improve it. The saving scales with prefix length. These sizes are calculated at the
|
||||
limit (8192 text tokens + 4 condition images, ~24.6k prefix tokens), not
|
||||
measured on the 1024 T2I run above: ~12.9 GB per CFG branch in bf16, ~6.6 GB
|
||||
in `"fp8"` and ~9.7 GB in `"fp8_v"`. That is where this option matters. Treat it as an opt-in for
|
||||
memory-bound long-prompt / multi-image workloads, not a free lunch;
|
||||
`"fp8_v"` is the better default trade-off when quality matters.
|
||||
|
||||
Note: a quantized prefix cache is not CUDA-graph capturable — requests using
|
||||
this option stay on the eager decode path (a warning is logged once).
|
||||
|
||||
### Why not the scheduler-managed paged KV (`DiffusionKVCacheMode.PAGED_SCHEDULER`)?
|
||||
|
||||
Evaluated and deliberately not adopted (see the `qwen21-p0-kv` commit message
|
||||
for the full analysis): the paged scheduler targets *cross-request* KV
|
||||
management — a statically sized page pool, block sharing/dedup, and eviction —
|
||||
whereas this prefix cache lives exactly as long as one generation and is
|
||||
sized once per request. Migrating would require routing decode attention
|
||||
through the worker paged adapter (writing fresh target K/V into pages every
|
||||
step), per-request page tables, and re-validating the block-causal
|
||||
piecewise-span masking and Ulysses/CFG constraints, for no memory benefit over
|
||||
the per-request dense tensors. The FP8 storage above captures the actual
|
||||
memory win at ~1% of the integration cost.
|
||||
|
||||
## Quality Comparison (head vs diffusers)
|
||||
|
||||
The reference is the diffusers Qwen-Image 2.1 pipeline
|
||||
([huggingface/diffusers#14804](https://github.com/huggingface/diffusers/pull/14804)).
|
||||
Both sides must pin the same checkpoint, inputs, seed, precision (BF16), and
|
||||
hardware.
|
||||
|
||||
vLLM-Omni (this recipe):
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model Qwen/Qwen-Image-2.1 \
|
||||
--prompt "A ceramic teapot on a wooden table" \
|
||||
--seed 42 \
|
||||
--output qwen_image_21_vllm_omni.png \
|
||||
--num-inference-steps 50 \
|
||||
--cfg-scale 1.0
|
||||
```
|
||||
|
||||
diffusers reference:
|
||||
|
||||
```python
|
||||
import torch
|
||||
from diffusers import QwenImage21Pipeline # diffusers PR #14804
|
||||
|
||||
pipe = QwenImage21Pipeline.from_pretrained(
|
||||
"Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
|
||||
).to("cuda")
|
||||
image = pipe(
|
||||
prompt="A ceramic teapot on a wooden table",
|
||||
num_inference_steps=50,
|
||||
true_cfg_scale=1.0,
|
||||
generator=torch.Generator(device="cuda").manual_seed(42),
|
||||
height=1024,
|
||||
width=1024,
|
||||
).images[0]
|
||||
image.save("qwen_image_21_diffusers.png")
|
||||
```
|
||||
|
||||
| Metric | vLLM-Omni | diffusers (#14804) | Status |
|
||||
| --- | --- | --- | --- |
|
||||
| Output image (1024×1024, 50 steps, seed 42, BF16, H200) | generated (eager + default CUDA-graph) | generated | measured |
|
||||
| PSNR / LPIPS vs reference | eager vs Diffusers **47.17 dB**; **default compile+graph** vs Diffusers **38.02 dB**; **default compile+graph** vs eager **38.10 dB** | — | measured (H200) |
|
||||
| Latency (time/image, same GPU, H200) | **default compile+graph** **6.04 s** (120.8 ms/step); eager **7.31 s** (146.2 ms/step) | **7.74 s** (154.9 ms/step) | measured |
|
||||
|
||||
The prefix KV cache changes the compute path (cached prefix vs. full recompute
|
||||
at every step), so exact numerical parity is not expected; match the cache
|
||||
setting on both sides or compare with PSNR/LPIPS rather than exact pixels.
|
||||
Eager stays numerically close to Diffusers (~47 dB). The **default** serving
|
||||
path stacks regional `torch.compile` with CUDA-graph decode; on H200 that
|
||||
combo lands at **~38.1 dB** vs eager (~1.21× on the Quality Comparison timed
|
||||
run). Graph-only vs eager is bit-identical; the ~38 dB drift is the compile
|
||||
fusion, not CUDA-graph replay. See "CUDA Graph decode" for the four-mode
|
||||
H200 breakdown (warmup 1 + median of 3).
|
||||
|
||||
Parity / feature measurements (same prompt/seed/steps/CFG, BF16, eager decode
|
||||
unless noted; 4× NVIDIA H200):
|
||||
|
||||
- [x] TP=2/4 output parity vs single-card — measured (1024×1024): TP1 vs TP2
|
||||
**39.32 dB**, TP1 vs TP4 **38.37 dB**, TP2 vs TP4 **45.77 dB**. Latency:
|
||||
TP1 **7.30 s**, TP2 **4.89 s**, TP4 **5.12 s** per image.
|
||||
- [x] VAE tiling vs non-tiled decode — measured: 1024×1024 nontile vs
|
||||
`--vae-use-tiling` **43.55 dB** (peak mem ~37.7 GB → ~32.7 GB); 2048×2048
|
||||
nontile vs tiled **44.14 dB** (peak mem ~57.9 GB → ~33.3 GB; gen ~36.8 s vs
|
||||
~37.1 s).
|
||||
- [x] CUDA-graph four-mode vs eager — measured on H200 (1024×1024, 50 steps,
|
||||
seed 42, BF16, warmup 1 + median of 3): compile+graph **5.93 s**,
|
||||
compile-only **6.88 s**, graph-only **9.52 s**, eager **8.28 s**. Combo /
|
||||
compile-only vs eager **38.10 dB**; graph-only vs eager bit-identical.
|
||||
Quality Comparison table uses a separate single-run A/B vs Diffusers
|
||||
(compile+graph **6.04 s**, eager **7.31 s**).
|
||||
|
||||
## Known Limitations
|
||||
|
||||
- Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported;
|
||||
`QwenImage21Pipeline` is registered in `_NO_CACHE_ACCELERATION`.
|
||||
- Ulysses sequence parallelism is wired; ring attention is not a supported path.
|
||||
A short H200 smoke (1024×1024, 8 steps, seed 42, true CFG 1.0, eager) finished:
|
||||
Ulysses SP=2 vs SP=1 is **46.5 dB**. That size's target length is divisible by 4,
|
||||
so SP padding was not exercised, and this is not the 50-step parity next to the
|
||||
TP tables. The features matrix therefore stays ❓. Ring SP=2 also ran, at
|
||||
**29.9 dB** vs SP=1, which is why it stays unsupported.
|
||||
- Only true CFG is exposed; there is no `guidance_scale` knob.
|
||||
- VAE tiling (`--vae-use-tiling`, also implied by `--vae-patch-parallel-size > 1`)
|
||||
decodes in 512px tiles with 384px stride (32 latent pixels per tile at 16x
|
||||
spatial compression). On CUDA OOM the tile size is automatically halved
|
||||
(keeping 25% overlap) and decoding retried, down to 128px tiles. Tile sizes
|
||||
are not yet exposed as engine arguments.
|
||||
@@ -87,6 +87,7 @@ recipes/
|
||||
| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
|
||||
| [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
|
||||
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4; offline text-to-image on XPU | 1x A100 80GB / 2x B200 / 1x Intel Arc Pro B70 32GB |
|
||||
| [`Qwen/Qwen-Image-2.1.md`](./Qwen/Qwen-Image-2.1.md) | Text-to-image and image-conditioned generation (editing) with prefix KV cache and CUDA graph decode | 1x GPU with CPU offload / 1x H200 |
|
||||
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
|
||||
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
|
||||
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat; offline end-to-end text-to-speech on XPU | 1x A100 80GB / 5x Intel Arc Pro B70 32GB |
|
||||
|
||||
@@ -12,9 +12,10 @@ import vllm_omni.diffusion.attention.backends.cudnn_attn as cudnn_backend
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend, CuDNNAttentionImpl
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
|
||||
"""Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V."""
|
||||
selected_backends = []
|
||||
@@ -40,6 +41,7 @@ def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
|
||||
assert selected_backends == [(SDPBackend.MATH,)]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_explicit_cudnn_rejects_kv_seq_len_one():
|
||||
impl = CuDNNAttentionImpl(
|
||||
num_heads=2,
|
||||
@@ -54,6 +56,7 @@ def test_explicit_cudnn_rejects_kv_seq_len_one():
|
||||
impl.forward_cuda(query, singleton_kv, singleton_kv)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
|
||||
selected_backends = []
|
||||
|
||||
@@ -77,6 +80,7 @@ def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
|
||||
assert selected_backends == [(SDPBackend.CUDNN_ATTENTION,)]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
|
||||
observed: dict[str, Any] = {}
|
||||
|
||||
@@ -112,6 +116,7 @@ def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
|
||||
assert observed["kwargs"]["attn_mask"] is None
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_cudnn_rejects_invalid_valid_kv_length():
|
||||
impl = CuDNNAttentionImpl(
|
||||
num_heads=2,
|
||||
@@ -130,12 +135,97 @@ def test_cudnn_rejects_invalid_valid_kv_length():
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_size", [8, 64, 128, 256])
|
||||
@pytest.mark.cpu
|
||||
def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size):
|
||||
assert CuDNNAttentionBackend.supports_head_size(head_size)
|
||||
assert head_size in CuDNNAttentionBackend.get_supported_head_sizes()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_size", [0, 7, 12, 320])
|
||||
@pytest.mark.cpu
|
||||
def test_cudnn_backend_rejects_incompatible_head_sizes(head_size):
|
||||
assert not CuDNNAttentionBackend.supports_head_size(head_size)
|
||||
assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2])
|
||||
@pytest.mark.parametrize("kv_heads", [1, 2])
|
||||
@pytest.mark.parametrize("query_offset", [0, 4])
|
||||
@pytest.mark.cpu
|
||||
def test_piecewise_matches_dense_block_causal_mask(monkeypatch, batch_size, kv_heads, query_offset):
|
||||
monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
|
||||
torch.manual_seed(42)
|
||||
query = torch.randn(batch_size, 12, 2, 8)
|
||||
key = torch.randn(batch_size, 12, kv_heads, 8)
|
||||
value = torch.randn_like(key)
|
||||
spans = [(2, 5), (7, 10)]
|
||||
positions = torch.arange(12)
|
||||
mask = positions[:, None] >= positions[None, :]
|
||||
for start, end in spans:
|
||||
mask[start:end, start:end] = True
|
||||
expected = torch.nn.functional.scaled_dot_product_attention(
|
||||
query[:, query_offset:].transpose(1, 2),
|
||||
key.transpose(1, 2),
|
||||
value.transpose(1, 2),
|
||||
attn_mask=mask[None, None, query_offset:],
|
||||
scale=0.5,
|
||||
enable_gqa=kv_heads == 1,
|
||||
).transpose(1, 2)
|
||||
impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
|
||||
actual = impl.forward_cuda(
|
||||
query[:, query_offset:], key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size)
|
||||
)
|
||||
torch.testing.assert_close(actual, expected)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_piecewise_preserves_discontinuous_query_ranges(monkeypatch):
|
||||
from vllm_omni.diffusion.attention.backends.abstract import QueryRange
|
||||
|
||||
monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
|
||||
query, key, value = [torch.randn(2, 12, 2, 8) for _ in range(3)]
|
||||
impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
|
||||
metadata = AttentionMetadata(full_attn_spans=[[(2, 5), (7, 12)]] * 2)
|
||||
expected = impl.forward_cuda(query, key, value, metadata)
|
||||
metadata.query_ranges = (QueryRange(0, 3, 0), QueryRange(3, 6, 8))
|
||||
actual = impl.forward_cuda(torch.cat([query[:, :3], query[:, 8:11]], dim=1), key, value, metadata)
|
||||
torch.testing.assert_close(actual, torch.cat([expected[:, :3], expected[:, 8:11]], dim=1))
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_piecewise_rejects_padding_mask(monkeypatch):
|
||||
impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
|
||||
query = torch.randn(1, 8, 2, 8)
|
||||
metadata = AttentionMetadata(full_attn_spans=[[(2, 8)]], attn_mask=torch.ones(1, 8, dtype=torch.bool))
|
||||
with pytest.raises(ValueError, match="padding mask"):
|
||||
impl.forward_cuda(query, query, query, metadata)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_dense_mask_takes_precedence_over_spans(monkeypatch):
|
||||
monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
|
||||
query, key, value = [torch.randn(2, 8, 2, 8) for _ in range(3)]
|
||||
mask = torch.ones(2, 1, 8, 8, dtype=torch.bool).tril()
|
||||
mask[1, :, :, 3] = False
|
||||
impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
|
||||
expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask))
|
||||
actual = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask, full_attn_spans=[[(2, 8)]] * 2))
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.parametrize("batch_size", [1, 2])
|
||||
def test_cudnn_piecewise_matches_dense_on_device(batch_size):
|
||||
generator = torch.Generator(device="cuda").manual_seed(42)
|
||||
query, key, value = [
|
||||
torch.randn(batch_size, 96, 2, 128, generator=generator, device="cuda", dtype=torch.bfloat16) for _ in range(3)
|
||||
]
|
||||
spans = [(8, 40), (48, 96)]
|
||||
mask = torch.ones(96, 96, device="cuda", dtype=torch.bool).tril()
|
||||
for start, end in spans:
|
||||
mask[start:end, start:end] = True
|
||||
impl = CuDNNAttentionImpl(num_heads=2, head_size=128, softmax_scale=128**-0.5)
|
||||
actual = impl.forward_cuda(query, key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size))
|
||||
expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask[None, None]))
|
||||
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.004)
|
||||
@@ -64,6 +64,7 @@ class E2EOperator:
|
||||
@dataclass
|
||||
class FakeWorldGroup:
|
||||
device_group: object
|
||||
cpu_group: object = None
|
||||
|
||||
|
||||
class DummyMixin(DistributedVaeMixin):
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from datetime import timedelta
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders import distributed_vae_executor as vae_executor
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class TileOperator:
|
||||
def split(self, tensor):
|
||||
return [vae_executor.TileTask(i, (i,), row[None]) for i, row in enumerate(tensor)], vae_executor.GridSpec(
|
||||
split_dims=(0,), grid_shape=(tensor.shape[0],)
|
||||
)
|
||||
|
||||
def exec(self, task):
|
||||
return task.tensor * 2
|
||||
|
||||
def merge(self, tiles, grid):
|
||||
return torch.cat([tiles[(i,)] for i in range(grid.grid_shape[0])])
|
||||
|
||||
|
||||
def run_recovery_worker(rank, init_file):
|
||||
torch.set_num_threads(1)
|
||||
dist.init_process_group(
|
||||
"gloo", init_method=f"file://{init_file}", rank=rank, world_size=2, timeout=timedelta(seconds=30)
|
||||
)
|
||||
try:
|
||||
group = SimpleNamespace(device_group=dist.group.WORLD, cpu_group=dist.group.WORLD)
|
||||
with pytest.MonkeyPatch.context() as patch:
|
||||
patch.setattr(vae_executor, "get_world_group", lambda: group)
|
||||
executor = vae_executor.DistributedVaeExecutor()
|
||||
executor.set_parallel_size(2)
|
||||
tensor = torch.arange(12, dtype=torch.float32).reshape(4, 3)
|
||||
operator = TileOperator()
|
||||
for stage, failing_rank in [("decode", 1), ("packing", 1), ("merge", 0), ("gather", 1)]:
|
||||
with pytest.MonkeyPatch.context() as patch:
|
||||
if rank == failing_rank:
|
||||
|
||||
def fail(*args, **kwargs):
|
||||
raise torch.OutOfMemoryError("injected OOM")
|
||||
|
||||
target, attribute = {
|
||||
"decode": (operator, "exec"),
|
||||
"packing": (executor, "_pack_local_tiles"),
|
||||
"merge": (operator, "merge"),
|
||||
"gather": (torch, "empty_like"),
|
||||
}[stage]
|
||||
patch.setattr(target, attribute, fail)
|
||||
with pytest.raises(torch.OutOfMemoryError, match="Distributed VAE"):
|
||||
executor.execute(tensor, operator)
|
||||
torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Distributed VAE mixed failures failed"):
|
||||
with executor._sync_errors("mixed failures"):
|
||||
if rank == 0:
|
||||
raise torch.OutOfMemoryError("retryable")
|
||||
raise ValueError("not retryable")
|
||||
torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2)
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def test_all_ranks_observe_failures_and_can_retry(tmp_path):
|
||||
torch.multiprocessing.spawn(run_recovery_worker, args=(str(tmp_path / "rendezvous"),), nprocs=2, join=True)
|
||||
@@ -0,0 +1,192 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
|
||||
get_qwen_image_21_pre_process_func,
|
||||
)
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler, StepScheduler
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def preprocess(tmp_path):
|
||||
vae_dir = tmp_path / "vae"
|
||||
vae_dir.mkdir()
|
||||
(vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4}))
|
||||
return get_qwen_image_21_pre_process_func(SimpleNamespace(model=str(tmp_path)))
|
||||
|
||||
|
||||
def make_request(request_id, sizes, color="red"):
|
||||
return OmniDiffusionRequest(
|
||||
request_id=request_id,
|
||||
prompt={
|
||||
"prompt": "A ceramic teapot",
|
||||
"multi_modal_data": {"image": [Image.new("RGB", size, color) for size in sizes]},
|
||||
},
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=1024, width=1024, num_inference_steps=2, true_cfg_scale=1.0, seed=42
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_class", [RequestScheduler, StepScheduler])
|
||||
@pytest.mark.parametrize(
|
||||
"first_sizes,second_sizes,compatible",
|
||||
[
|
||||
pytest.param([], [(64, 64)], False, id="text-and-edit"),
|
||||
pytest.param([(64, 64)], [], False, id="edit-and-text"),
|
||||
pytest.param([(64, 64)], [(64, 64), (64, 64)], False, id="image-count"),
|
||||
pytest.param([(64, 64)], [(128, 64)], False, id="image-aspect-ratio"),
|
||||
pytest.param([(64, 64), (128, 64)], [(128, 64), (64, 64)], False, id="slot-order"),
|
||||
pytest.param([], [], True, id="two-text-requests"),
|
||||
pytest.param([(64, 64)], [(64, 64)], True, id="different-image-content"),
|
||||
pytest.param([(64, 64)], [(128, 128)], True, id="same-resized-layout"),
|
||||
pytest.param([(64, 64), (128, 64)], [(128, 128), (256, 128)], True, id="same-multi-image-layout"),
|
||||
],
|
||||
)
|
||||
def test_scheduler_groups_by_condition_layout(preprocess, scheduler_class, first_sizes, second_sizes, compatible):
|
||||
first = preprocess(make_request("first", first_sizes))
|
||||
second = preprocess(make_request("second", second_sizes, color="blue"))
|
||||
scheduler = scheduler_class()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
|
||||
scheduler.add_request(first)
|
||||
scheduler.add_request(second)
|
||||
|
||||
scheduled = scheduler.schedule()
|
||||
expected_ids = ["first", "second"] if compatible else ["first"]
|
||||
assert scheduled.scheduled_request_ids == expected_ids
|
||||
|
||||
scheduler.finish_requests(expected_ids, DiffusionRequestStatus.FINISHED_COMPLETED)
|
||||
if not compatible:
|
||||
assert scheduler.schedule().scheduled_request_ids == ["second"]
|
||||
scheduler.finish_requests("second", DiffusionRequestStatus.FINISHED_COMPLETED)
|
||||
assert not scheduler.has_requests()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prompt", ["A teapot", {"prompt": "A teapot"}, {"prompt": "A teapot", "multi_modal_data": {"image": None}}]
|
||||
)
|
||||
def test_text_input_forms_share_compatibility_key(preprocess, prompt):
|
||||
request = make_request("text", [])
|
||||
request.prompt = prompt
|
||||
processed = preprocess(request)
|
||||
empty_list_request = preprocess(make_request("empty-list", []))
|
||||
assert processed.batch_compatibility_key is not None
|
||||
assert processed.batch_compatibility_key == empty_list_request.batch_compatibility_key
|
||||
|
||||
|
||||
@pytest.mark.parametrize("image_type", ["numpy", "tensor"])
|
||||
def test_preprocess_accepts_array_images_without_truth_testing(preprocess, image_type):
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
pixels = np.full((32, 32, 4), 255, dtype=np.uint8)
|
||||
image = torch.from_numpy(pixels).permute(2, 0, 1) if image_type == "tensor" else pixels
|
||||
request = make_request("array", [])
|
||||
request.prompt["multi_modal_data"]["image"] = image
|
||||
processed = preprocess(request)
|
||||
assert processed.batch_compatibility_key is not None
|
||||
assert processed.allow_mixed_step_phases is False
|
||||
|
||||
|
||||
def test_late_prefill_waits_until_all_decode_requests_finish(preprocess):
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
||||
|
||||
scheduler = StepScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=3))
|
||||
for request_id in ("a", "b"):
|
||||
scheduler.add_request(preprocess(make_request(request_id, [])))
|
||||
wave = scheduler.schedule()
|
||||
assert wave.scheduled_request_ids == ["a", "b"]
|
||||
scheduler.update_from_output(
|
||||
wave,
|
||||
BatchRunnerOutput.from_list(
|
||||
[RunnerOutput(request_id=request_id, step_index=1, finished=False) for request_id in ("a", "b")]
|
||||
),
|
||||
)
|
||||
scheduler.add_request(preprocess(make_request("late", [])))
|
||||
assert scheduler.schedule().scheduled_request_ids == ["a", "b"]
|
||||
scheduler.finish_requests("a", DiffusionRequestStatus.FINISHED_COMPLETED)
|
||||
assert scheduler.schedule().scheduled_request_ids == ["b"]
|
||||
scheduler.finish_requests("b", DiffusionRequestStatus.FINISHED_COMPLETED)
|
||||
assert scheduler.schedule().scheduled_request_ids == ["late"]
|
||||
|
||||
|
||||
def test_preprocess_normalizes_tuple_images_like_a_list(preprocess):
|
||||
image = Image.new("RGB", (64, 64), "red")
|
||||
tuple_request = make_request("tuple", [])
|
||||
tuple_request.prompt["multi_modal_data"]["image"] = (image,)
|
||||
|
||||
processed = preprocess(tuple_request)
|
||||
reference = preprocess(make_request("list", [(64, 64)]))
|
||||
|
||||
assert processed.batch_compatibility_key == reference.batch_compatibility_key
|
||||
assert processed.batch_compatibility_key[0] == "qwen_image_21"
|
||||
assert len(processed.batch_compatibility_key[1]) == 1
|
||||
info = processed.prompt["additional_information"]
|
||||
assert len(info["prompt_image"]) == 1
|
||||
assert len(info["vae_images"]) == 1
|
||||
assert len(info["input_image_sizes"]) == 1
|
||||
|
||||
|
||||
def test_preprocess_rejects_more_images_than_the_model_supports(preprocess):
|
||||
from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
|
||||
MAX_QWEN_IMAGE_21_INPUT_IMAGES,
|
||||
)
|
||||
|
||||
request = make_request("too-many", [(64, 64)] * (MAX_QWEN_IMAGE_21_INPUT_IMAGES + 1))
|
||||
with pytest.raises(ValueError, match="At most"):
|
||||
preprocess(request)
|
||||
|
||||
|
||||
def test_preprocess_defaults_true_cfg_scale(preprocess):
|
||||
request = make_request("cfg", [])
|
||||
request.sampling_params.true_cfg_scale = None
|
||||
processed = preprocess(request)
|
||||
assert processed.sampling_params.true_cfg_scale == 1.0
|
||||
|
||||
request = make_request("cfg-set", [])
|
||||
request.sampling_params.true_cfg_scale = 4.0
|
||||
assert preprocess(request).sampling_params.true_cfg_scale == 4.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size,expected_wider", [((128, 64), True), ((64, 128), False)])
|
||||
def test_preprocess_derives_output_size_from_condition_aspect(preprocess, size, expected_wider):
|
||||
request = make_request("aspect", [size])
|
||||
request.sampling_params.height = None
|
||||
request.sampling_params.width = None
|
||||
|
||||
processed = preprocess(request)
|
||||
|
||||
height = processed.sampling_params.height
|
||||
width = processed.sampling_params.width
|
||||
assert height % 32 == 0 and width % 32 == 0
|
||||
assert (width > height) is expected_wider
|
||||
info = processed.prompt["additional_information"]
|
||||
assert info["calculated_height"] == height
|
||||
assert info["calculated_width"] == width
|
||||
|
||||
|
||||
def test_preprocess_stages_rgba_vae_images_and_resized_prompt_images(preprocess):
|
||||
processed = preprocess(make_request("staged", [(96, 64)]))
|
||||
info = processed.prompt["additional_information"]
|
||||
|
||||
(vae_image,) = info["vae_images"]
|
||||
# RGB input gains an opaque alpha channel for the RGBA VAE; the frame dim is
|
||||
# inserted for the 3D (video-capable) VAE.
|
||||
assert vae_image.shape[1] == 4
|
||||
assert vae_image.shape[2] == 1
|
||||
(prompt_image,) = info["prompt_image"]
|
||||
assert isinstance(prompt_image, Image.Image)
|
||||
(width, height) = info["input_image_sizes"][0]
|
||||
assert prompt_image.size == (width, height)
|
||||
assert processed.batch_compatibility_key == ("qwen_image_21", ((width, height),))
|
||||
@@ -0,0 +1,318 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
|
||||
from vllm_omni.diffusion.models.qwen_image_21 import decode_graph
|
||||
from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class DecodeModel(torch.nn.Module):
|
||||
in_channels = 4
|
||||
|
||||
def _decode_graph_forward(self, entry):
|
||||
return (
|
||||
entry.hidden.sin()
|
||||
+ entry.k[0].mean(dim=(1, 2, 3))[:, None, None]
|
||||
+ entry.v[0].mean(dim=(1, 2, 3))[:, None, None]
|
||||
+ entry.freqs[None, :, :]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def graph_case(request, monkeypatch):
|
||||
device = "cuda" if request.node.get_closest_marker("cuda") else "cpu"
|
||||
manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), max_entries=2)
|
||||
monkeypatch.setattr(manager, "eligible", lambda: True)
|
||||
monkeypatch.setattr(manager, "_backend_name", lambda: "test")
|
||||
cleanup = []
|
||||
monkeypatch.setattr(decode_graph.current_omni_platform, "empty_cache", lambda: cleanup.append(True))
|
||||
return manager, device, cleanup
|
||||
|
||||
|
||||
def make_cache(device, value):
|
||||
return [
|
||||
{
|
||||
"cond": {
|
||||
"key": torch.full((1, 2, 1, 4), value, device=device),
|
||||
"value": torch.full((1, 2, 1, 4), value + 1, device=device),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def register(manager, cache, layout=(1, 2, 2)):
|
||||
device = cache[0]["cond"]["key"].device
|
||||
manager.register_prefill(
|
||||
kv_cache=cache,
|
||||
cache_branch="cond",
|
||||
prefix_len=2,
|
||||
img_shapes=[layout],
|
||||
target_freqs=torch.full((4, 4), float(layout[1]), device=device),
|
||||
joint_key_valid=None,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def decode(manager, cache, layout=(1, 2, 2)):
|
||||
device = cache[0]["cond"]["key"].device
|
||||
return manager.try_decode(
|
||||
hidden_states=torch.ones(1, 4, 4, device=device),
|
||||
timestep=torch.ones(1, device=device),
|
||||
kv_cache=cache,
|
||||
cache_branch="cond",
|
||||
img_shapes=[layout],
|
||||
img_mask=torch.tensor([[False, False, True]], device=device),
|
||||
encoder_hidden_states_mask=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_model_level_offload_without_staging_disables_decode_graphs():
|
||||
manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), model_level_offload=True)
|
||||
assert manager._offload_reason() is not None
|
||||
assert manager.eligible() is False
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_top_level_offload_hook_without_staging_disables_decode_graphs():
|
||||
model = DecodeModel()
|
||||
manager = decode_graph.QwenImage21DecodeGraphManager(model)
|
||||
registry = HookRegistry.get_or_create(model)
|
||||
# A present-but-empty registry (e.g. after hooks were removed) must not disqualify.
|
||||
assert manager._offload_reason() is None
|
||||
registry.register_hook("sequential_offload", ModelHook())
|
||||
assert manager._offload_reason() is not None
|
||||
assert manager.eligible() is False
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_persistent_staging_makes_model_level_offload_graph_eligible():
|
||||
model = DecodeModel()
|
||||
manager = decode_graph.QwenImage21DecodeGraphManager(model, model_level_offload=True)
|
||||
hook = SequentialOffloadHook(offload_targets=[], device=torch.device("cpu"))
|
||||
HookRegistry.get_or_create(model).register_hook(SequentialOffloadHook._HOOK_NAME, hook)
|
||||
# No staging storage yet: a swap would allocate fresh storage per generation.
|
||||
assert manager._offload_reason() is not None
|
||||
# Staging established: fixed device storage keeps captured pointers valid.
|
||||
hook._stager = object()
|
||||
assert manager._offload_reason() is None
|
||||
# A second, non-offload hook on the top-level module still disqualifies.
|
||||
model._hook_registry.register_hook("teacache", ModelHook())
|
||||
assert manager._offload_reason() is not None
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_allocation_oom_releases_partial_buffers_and_remembers_failure(graph_case, monkeypatch):
|
||||
manager, device, cleanup = graph_case
|
||||
cache = make_cache(device, 1.0)
|
||||
constructor = decode_graph.QwenImage21DecodeGraphEntry
|
||||
allocated = []
|
||||
|
||||
def fail_allocation(**kwargs):
|
||||
entry = constructor(**kwargs)
|
||||
allocated.append(weakref.ref(entry.k[0]))
|
||||
raise torch.OutOfMemoryError("injected allocation OOM")
|
||||
|
||||
monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", fail_allocation)
|
||||
register(manager, cache)
|
||||
register(manager, cache)
|
||||
assert len(allocated) == 1
|
||||
assert allocated[0]() is None
|
||||
assert list(manager.entries.values()) == [None]
|
||||
assert len(cleanup) == 1
|
||||
assert torch.equal(cache[0]["cond"]["key"], torch.ones(1, 2, 1, 4))
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_evict_before_allocating_replacement(graph_case, monkeypatch):
|
||||
manager, device, _ = graph_case
|
||||
manager.max_entries = 1
|
||||
cache = make_cache(device, 1.0)
|
||||
register(manager, cache)
|
||||
old_buffer = weakref.ref(next(iter(manager.entries.values())).k[0])
|
||||
constructor = decode_graph.QwenImage21DecodeGraphEntry
|
||||
|
||||
def allocate_after_release(**kwargs):
|
||||
assert old_buffer() is None
|
||||
return constructor(**kwargs)
|
||||
|
||||
monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", allocate_after_release)
|
||||
register(manager, cache, (1, 1, 4))
|
||||
assert len(manager.entries) == 1
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_registration_preserves_request_cache_and_separates_layouts(graph_case):
|
||||
manager, device, _ = graph_case
|
||||
first, second = make_cache(device, 1.0), make_cache(device, 5.0)
|
||||
original = first[0]["cond"]["key"]
|
||||
register(manager, first)
|
||||
register(manager, second)
|
||||
assert first[0]["cond"]["key"] is original
|
||||
assert torch.equal(original, torch.ones_like(original))
|
||||
register(manager, second, (1, 1, 4))
|
||||
assert len(manager.entries) == 2
|
||||
entries = list(manager.entries.values())
|
||||
assert not torch.equal(entries[0].freqs, entries[1].freqs)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.parametrize("failure_phase", ["warmup", "capture"])
|
||||
def test_capture_failure_releases_buffers_and_keeps_request_cache(graph_case, monkeypatch, failure_phase):
|
||||
manager, device, cleanup = graph_case
|
||||
cache = make_cache(device, 1.0)
|
||||
register(manager, cache)
|
||||
entry = next(iter(manager.entries.values()))
|
||||
entry_ref, buffer_ref = weakref.ref(entry), weakref.ref(entry.k[0])
|
||||
del entry
|
||||
|
||||
original = manager.model._decode_graph_forward
|
||||
|
||||
def fail_capture(entry):
|
||||
output = original(entry)
|
||||
if failure_phase == "warmup" or torch.cuda.is_current_stream_capturing():
|
||||
raise torch.OutOfMemoryError("injected capture OOM")
|
||||
return output
|
||||
|
||||
monkeypatch.setattr(manager.model, "_decode_graph_forward", fail_capture)
|
||||
assert decode(manager, cache) is None
|
||||
assert entry_ref() is None and buffer_ref() is None
|
||||
assert list(manager.entries.values()) == [None]
|
||||
register(manager, cache)
|
||||
assert decode(manager, cache) is None
|
||||
assert len(cleanup) == 1
|
||||
assert bool((cache[0]["cond"]["key"] == 1).all())
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.gpu
|
||||
def test_replay_uses_active_request_and_layout(graph_case):
|
||||
manager, device, _ = graph_case
|
||||
first, second = make_cache(device, 1.0), make_cache(device, 5.0)
|
||||
first_key = weakref.ref(first[0]["cond"]["key"])
|
||||
with torch.inference_mode():
|
||||
register(manager, first)
|
||||
# A second live request aliasing the same key is refused: it keeps its
|
||||
# own cache and decodes eagerly while the owner is in flight.
|
||||
register(manager, second)
|
||||
assert decode(manager, second) is None
|
||||
output = decode(manager, first)
|
||||
expected = torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2
|
||||
torch.testing.assert_close(output, expected)
|
||||
# Once the owner's cache is released, the key is free to adopt.
|
||||
del first
|
||||
assert first_key() is None
|
||||
register(manager, second)
|
||||
output = decode(manager, second)
|
||||
expected = (
|
||||
torch.ones_like(output).sin() + second[0]["cond"]["key"].mean() + second[0]["cond"]["value"].mean() + 2
|
||||
)
|
||||
torch.testing.assert_close(output, expected)
|
||||
original = decode(manager, second).clone()
|
||||
register(manager, second, (1, 1, 4))
|
||||
other = decode(manager, second, (1, 1, 4))
|
||||
torch.testing.assert_close(other, torch.ones_like(other).sin() + 5 + 6 + 1)
|
||||
torch.testing.assert_close(decode(manager, second), original)
|
||||
assert all(entry.captures == 1 for entry in manager.entries.values())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("inference", [False, True])
|
||||
def test_prefix_reuse_refreshes_every_layer_without_retaining_requests(graph_case, inference):
|
||||
manager, device, _ = graph_case
|
||||
with torch.inference_mode(inference):
|
||||
cache = make_cache(device, 1.0) + make_cache(device, 2.0)
|
||||
register(manager, cache)
|
||||
entry = next(iter(manager.entries.values()))
|
||||
entry.refresh_prefix(cache)
|
||||
versions = [tensor._version for tensor in entry.k + entry.v]
|
||||
entry.refresh_prefix(cache)
|
||||
assert versions == [tensor._version for tensor in entry.k + entry.v]
|
||||
cache[1]["cond"]["value"] = torch.full_like(cache[1]["cond"]["value"], 7.0)
|
||||
entry.refresh_prefix(cache)
|
||||
assert bool((entry.v[1] == 7).all())
|
||||
source = weakref.ref(cache[0]["cond"]["key"])
|
||||
del cache
|
||||
assert source() is None
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_prefix_reuse_detects_inplace_changes_and_prefill_registration(graph_case):
|
||||
manager, device, _ = graph_case
|
||||
cache = make_cache(device, 1.0)
|
||||
register(manager, cache)
|
||||
entry = next(iter(manager.entries.values()))
|
||||
entry.refresh_prefix(cache)
|
||||
cache[0]["cond"]["value"].add_(3)
|
||||
entry.refresh_prefix(cache)
|
||||
assert bool((entry.v[0] == 5).all())
|
||||
register(manager, cache)
|
||||
assert not entry.prefix_sources
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_all_valid_mask_shares_graph_entry_and_padding_falls_back(graph_case):
|
||||
manager, device, _ = graph_case
|
||||
cache = make_cache(device, 1.0)
|
||||
register(manager, cache)
|
||||
entry = next(iter(manager.entries.values()))
|
||||
for valid in (True, False):
|
||||
manager.register_prefill(
|
||||
kv_cache=cache,
|
||||
cache_branch="cond",
|
||||
prefix_len=2,
|
||||
img_shapes=[(1, 2, 2)],
|
||||
target_freqs=torch.ones(4, 4),
|
||||
joint_key_valid=torch.full((1, 6), valid),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
assert list(manager.entries.values()) == [entry]
|
||||
assert entry.attn_metadata is None
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.gpu
|
||||
def test_batch_changes_and_padding_do_not_reuse_stale_prefix(graph_case):
|
||||
manager, device, _ = graph_case
|
||||
first = make_cache(device, 1.0)
|
||||
second = make_cache(device, 4.0)
|
||||
merged = [
|
||||
{"cond": {part: torch.cat([first[0]["cond"][part], second[0]["cond"][part]]) for part in ("key", "value")}}
|
||||
]
|
||||
with torch.inference_mode():
|
||||
register(manager, first)
|
||||
single = decode(manager, first).clone()
|
||||
register(manager, merged)
|
||||
kwargs = dict(
|
||||
hidden_states=torch.ones(2, 4, 4, device=device),
|
||||
timestep=torch.ones(2, device=device),
|
||||
kv_cache=merged,
|
||||
cache_branch="cond",
|
||||
img_shapes=[(1, 2, 2)],
|
||||
img_mask=torch.tensor([[False, False, True]] * 2, device=device),
|
||||
encoder_hidden_states_mask=torch.ones(2, 3, dtype=torch.bool, device=device),
|
||||
)
|
||||
result = manager.try_decode(**kwargs)
|
||||
assert result is not None and result.shape == (2, 4, 4)
|
||||
expected = torch.ones_like(result).sin() + torch.tensor([5.0, 11.0], device=device)[:, None, None]
|
||||
torch.testing.assert_close(result, expected)
|
||||
# A rebuilt cache (e.g. reordered rows) is a different live cache
|
||||
# aliasing the same key: decode refuses the shared entry until the
|
||||
# previous owner's tensors are released.
|
||||
flipped = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}]
|
||||
kwargs["kv_cache"] = flipped
|
||||
assert manager.try_decode(**kwargs) is None
|
||||
del merged
|
||||
register(manager, flipped)
|
||||
torch.testing.assert_close(manager.try_decode(**kwargs), expected.flip(0))
|
||||
kwargs["encoder_hidden_states_mask"][1, 0] = False
|
||||
assert manager.try_decode(**kwargs) is None
|
||||
torch.testing.assert_close(decode(manager, first), single, rtol=0, atol=0)
|
||||
@@ -0,0 +1,402 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""CPU unit tests for QwenImage21Pipeline helpers: prepare_latents duplication
|
||||
order, KV-cache assembly across step phases, and the post-process func."""
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
|
||||
QwenImage21Pipeline,
|
||||
get_qwen_image_21_post_process_func,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _toy_pipeline() -> QwenImage21Pipeline:
|
||||
"""A pipeline shell with the attributes prepare_latents reads, no weights."""
|
||||
pipeline = QwenImage21Pipeline.__new__(QwenImage21Pipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.vae_scale_factor = 16
|
||||
pipeline.latent_channels = 4
|
||||
return pipeline
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# prepare_latents
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prepare_latents_duplicates_condition_latents_request_major():
|
||||
"""batch=2 prompts x n=2 outputs must order condition latents A,A,B,B.
|
||||
|
||||
``encode_prompt`` expands prompt embeddings request-major (A,A,B,B); the
|
||||
condition latents must follow the same order or generations 2 and 3 swap
|
||||
condition images. ``repeat`` would produce the wrong (A,B,A,B) tile-major
|
||||
order; ``repeat_interleave`` is correct.
|
||||
"""
|
||||
pipeline = _toy_pipeline()
|
||||
cond_a = torch.full((1, 4, 1, 8, 8), 1.0)
|
||||
cond_b = torch.full((1, 4, 1, 8, 8), 2.0)
|
||||
stacked = torch.cat([cond_a, cond_b], dim=0) # one row per request
|
||||
|
||||
latents, image_latents = pipeline.prepare_latents(
|
||||
[stacked],
|
||||
batch_size=4,
|
||||
num_channels_latents=4,
|
||||
height=64,
|
||||
width=64,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
generator=torch.Generator().manual_seed(0),
|
||||
)
|
||||
|
||||
assert latents.shape == (4, 16, 4) # packed: 4x4 latent tokens, 4 channels
|
||||
assert image_latents.shape == (4, 64, 4) # packed: 8x8 condition tokens
|
||||
row_means = image_latents.mean(dim=(1, 2))
|
||||
torch.testing.assert_close(row_means, torch.tensor([1.0, 1.0, 2.0, 2.0]))
|
||||
|
||||
|
||||
def test_prepare_latents_without_duplication_when_rows_match_batch():
|
||||
pipeline = _toy_pipeline()
|
||||
stacked = torch.arange(4, dtype=torch.float32).view(4, 1, 1, 1, 1).expand(4, 4, 1, 2, 2).contiguous()
|
||||
|
||||
_, image_latents = pipeline.prepare_latents(
|
||||
[stacked],
|
||||
batch_size=4,
|
||||
num_channels_latents=4,
|
||||
height=32,
|
||||
width=32,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
generator=torch.Generator().manual_seed(0),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(image_latents.mean(dim=(1, 2)), torch.arange(4, dtype=torch.float32))
|
||||
|
||||
|
||||
def test_prepare_latents_rejects_non_divisible_duplication():
|
||||
pipeline = _toy_pipeline()
|
||||
stacked = torch.zeros(2, 4, 1, 8, 8)
|
||||
|
||||
with pytest.raises(ValueError, match="Cannot duplicate"):
|
||||
pipeline.prepare_latents(
|
||||
[stacked],
|
||||
batch_size=3,
|
||||
num_channels_latents=4,
|
||||
height=64,
|
||||
width=64,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
generator=torch.Generator().manual_seed(0),
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_latents_rejects_generator_list_length_mismatch():
|
||||
pipeline = _toy_pipeline()
|
||||
|
||||
with pytest.raises(ValueError, match="list of generators"):
|
||||
pipeline.prepare_latents(
|
||||
None,
|
||||
batch_size=2,
|
||||
num_channels_latents=4,
|
||||
height=32,
|
||||
width=32,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
generator=[torch.Generator().manual_seed(0)],
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_latents_concatenates_multiple_condition_images():
|
||||
pipeline = _toy_pipeline()
|
||||
first = torch.full((2, 4, 1, 4, 4), 1.0)
|
||||
second = torch.full((2, 4, 1, 2, 2), 3.0)
|
||||
|
||||
_, image_latents = pipeline.prepare_latents(
|
||||
[first, second],
|
||||
batch_size=2,
|
||||
num_channels_latents=4,
|
||||
height=32,
|
||||
width=32,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
generator=torch.Generator().manual_seed(0),
|
||||
)
|
||||
|
||||
# Token dim concatenates the packed condition images: 16 + 4 per row.
|
||||
assert image_latents.shape == (2, 20, 4)
|
||||
torch.testing.assert_close(image_latents[:, :16].mean(), torch.tensor(1.0))
|
||||
torch.testing.assert_close(image_latents[:, 16:].mean(), torch.tensor(3.0))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# KV-cache assembly across step phases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _state(cache) -> SimpleNamespace:
|
||||
return SimpleNamespace(extra={"kv_cache": cache})
|
||||
|
||||
|
||||
def _decode_cache(rows: int, value: float) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"cond": {
|
||||
"key": torch.full((rows, 2), value),
|
||||
"value": torch.full((rows, 2), value + 0.5),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_kv_cache_phase_classification():
|
||||
assert QwenImage21Pipeline._kv_cache_phase(None) == "none"
|
||||
assert QwenImage21Pipeline._kv_cache_phase([{}]) == "prefill"
|
||||
assert QwenImage21Pipeline._kv_cache_phase(_decode_cache(1, 0.0)) == "decode"
|
||||
|
||||
|
||||
def test_assemble_kv_cache_all_none():
|
||||
assert QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(None)]) == (None, False)
|
||||
|
||||
|
||||
def test_assemble_kv_cache_prefill_returns_fresh_owned_cache():
|
||||
states = [_state([{} for _ in range(3)]), _state([{} for _ in range(3)])]
|
||||
|
||||
kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states)
|
||||
|
||||
assert take_ownership is True
|
||||
assert kv_cache == [{}, {}, {}]
|
||||
# Fresh blocks, not aliases of either request's cache: the batched prefill
|
||||
# writes into these and _scatter_kv_cache splits them back per request.
|
||||
for block in kv_cache:
|
||||
assert all(block is not request_block for state in states for request_block in state.extra["kv_cache"])
|
||||
|
||||
|
||||
def test_assemble_kv_cache_decode_concatenates_in_request_order():
|
||||
first, second = _decode_cache(2, 1.0), _decode_cache(1, 3.0)
|
||||
states = [_state(first), _state(second)]
|
||||
|
||||
kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states)
|
||||
|
||||
assert take_ownership is False
|
||||
assert len(kv_cache) == 1
|
||||
cond = kv_cache[0]["cond"]
|
||||
torch.testing.assert_close(cond["key"], torch.tensor([[1.0, 1.0], [1.0, 1.0], [3.0, 3.0]]))
|
||||
torch.testing.assert_close(cond["value"], torch.tensor([[1.5, 1.5], [1.5, 1.5], [3.5, 3.5]]))
|
||||
# Decode only reads the cache: per-request tensors must stay untouched.
|
||||
assert first[0]["cond"]["key"].shape == (2, 2)
|
||||
assert second[0]["cond"]["key"].shape == (1, 2)
|
||||
|
||||
|
||||
def test_assemble_kv_cache_decode_merges_every_part():
|
||||
first = [{"cond": {"key": torch.ones(1, 1), "value": torch.ones(1, 1), "key_scale": torch.ones(1)}}]
|
||||
second = [{"cond": {"key": torch.zeros(1, 1), "value": torch.zeros(1, 1), "key_scale": torch.zeros(1)}}]
|
||||
|
||||
kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache([_state(first), _state(second)])
|
||||
|
||||
assert take_ownership is False
|
||||
assert set(kv_cache[0]["cond"]) == {"key", "value", "key_scale"}
|
||||
torch.testing.assert_close(kv_cache[0]["cond"]["key_scale"], torch.tensor([1.0, 0.0]))
|
||||
|
||||
|
||||
def test_assemble_kv_cache_mixed_phases_raises():
|
||||
"""Defensive backstop: a late prefill must never join a decoding batch.
|
||||
|
||||
The preprocessor sets ``allow_mixed_step_phases = False`` so the
|
||||
StepScheduler defers late requests until the batch finishes decoding; this
|
||||
raise is the last line of defense if a mixed batch is ever assembled.
|
||||
"""
|
||||
states = [_state([{}]), _state(_decode_cache(1, 0.0))]
|
||||
|
||||
with pytest.raises(ValueError, match="mixed KV-cache phases"):
|
||||
QwenImage21Pipeline._assemble_kv_cache(states)
|
||||
|
||||
with pytest.raises(ValueError, match="mixed KV-cache phases"):
|
||||
QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(_decode_cache(1, 0.0))])
|
||||
|
||||
|
||||
def test_assemble_kv_cache_decode_rejects_mismatched_prefill_prompt_seq_len():
|
||||
"""Homogeneous merge invariant: denoise_step must sub-batch before assembling."""
|
||||
short = _state(_decode_cache(1, 1.0))
|
||||
short.extra["prefill_prompt_seq_len"] = 4
|
||||
long = _state(_decode_cache(1, 2.0))
|
||||
long.extra["prefill_prompt_seq_len"] = 8
|
||||
|
||||
with pytest.raises(ValueError, match="sub-batch by prefill_prompt_seq_len"):
|
||||
QwenImage21Pipeline._assemble_kv_cache([short, long])
|
||||
|
||||
# Same baked length remains batchable.
|
||||
same_a = _state(_decode_cache(1, 1.0))
|
||||
same_a.extra["prefill_prompt_seq_len"] = 4
|
||||
same_b = _state(_decode_cache(1, 2.0))
|
||||
same_b.extra["prefill_prompt_seq_len"] = 4
|
||||
_, take_ownership = QwenImage21Pipeline._assemble_kv_cache([same_a, same_b])
|
||||
assert take_ownership is False
|
||||
|
||||
|
||||
def test_split_decode_groups_by_prefill_prompt_seq_len():
|
||||
"""Routine serving: separate prefills → joint decode must split, not fail."""
|
||||
short = _state(_decode_cache(1, 1.0))
|
||||
short.request_id = "short"
|
||||
short.extra["prefill_prompt_seq_len"] = 4
|
||||
long = _state(_decode_cache(1, 2.0))
|
||||
long.request_id = "long"
|
||||
long.extra["prefill_prompt_seq_len"] = 8
|
||||
short2 = _state(_decode_cache(1, 3.0))
|
||||
short2.request_id = "short2"
|
||||
short2.extra["prefill_prompt_seq_len"] = 4
|
||||
|
||||
states = [short, long, short2]
|
||||
assert QwenImage21Pipeline._needs_decode_seq_len_split(states) is True
|
||||
groups = QwenImage21Pipeline._split_decode_groups(states)
|
||||
assert [[s.request_id for s in g] for g in groups] == [["short", "short2"], ["long"]]
|
||||
|
||||
same = [_state(_decode_cache(1, 1.0)), _state(_decode_cache(1, 2.0))]
|
||||
for s in same:
|
||||
s.extra["prefill_prompt_seq_len"] = 4
|
||||
assert QwenImage21Pipeline._needs_decode_seq_len_split(same) is False
|
||||
assert QwenImage21Pipeline._split_decode_groups(same) == [same]
|
||||
|
||||
|
||||
def test_scatter_kv_cache_splits_rows_back_per_request():
|
||||
states = [_state([{}]), _state([{}])]
|
||||
merged = [
|
||||
{
|
||||
"cond": {
|
||||
"key": torch.tensor([[1.0], [2.0], [3.0]]),
|
||||
"value": torch.tensor([[4.0], [5.0], [6.0]]),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
QwenImage21Pipeline._scatter_kv_cache(states, [1, 2], merged)
|
||||
|
||||
torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[1.0]]))
|
||||
torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["value"], torch.tensor([[4.0]]))
|
||||
torch.testing.assert_close(states[1].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[2.0], [3.0]]))
|
||||
scattered = states[1].extra["kv_cache"][0]["cond"]["key"]
|
||||
assert scattered.is_contiguous()
|
||||
# Must own storage (clone), not a view into the merged prefill tensor —
|
||||
# otherwise step-mode CUDA-graph decode never matches owner identity.
|
||||
assert scattered.data_ptr() != merged[0]["cond"]["key"][1:3].data_ptr()
|
||||
merged[0]["cond"]["key"].fill_(0.0)
|
||||
torch.testing.assert_close(scattered, torch.tensor([[2.0], [3.0]]))
|
||||
|
||||
|
||||
def test_encode_prompt_expands_masks_request_major():
|
||||
"""Two prompts × n=2 must keep embeds and masks in A,A,B,B order.
|
||||
|
||||
A 2D ``mask.repeat(1, n, 1).view(B*n, S)`` would wrongly produce A,B,A,B.
|
||||
"""
|
||||
pipeline = _toy_pipeline()
|
||||
embeds = torch.tensor([[[1.0], [1.0]], [[2.0], [0.0]]]) # [B=2, S=2, D=1]
|
||||
mask = torch.tensor([[True, True], [True, False]])
|
||||
image_pad = torch.tensor([[False, True], [False, False]])
|
||||
|
||||
def _fake_get(prompt, images_per_prompt, max_sequence_length=None, prompt_name="prompt"):
|
||||
del prompt, images_per_prompt, max_sequence_length, prompt_name
|
||||
return embeds.clone(), mask.clone(), image_pad.clone()
|
||||
|
||||
pipeline._get_qwen_prompt_embeds = _fake_get # type: ignore[method-assign]
|
||||
|
||||
out_embeds, out_mask, out_pad = pipeline.encode_prompt(
|
||||
["short", "loooooonger"],
|
||||
num_images_per_prompt=2,
|
||||
)
|
||||
|
||||
assert out_embeds.shape[0] == 4
|
||||
# embeds request-major: A,A,B,B
|
||||
torch.testing.assert_close(out_embeds[:, 0, 0], torch.tensor([1.0, 1.0, 2.0, 2.0]))
|
||||
# masks must follow the same order (not A,B,A,B)
|
||||
assert out_mask.tolist() == [
|
||||
[True, True],
|
||||
[True, True],
|
||||
[True, False],
|
||||
[True, False],
|
||||
]
|
||||
assert out_pad.tolist() == [
|
||||
[False, True],
|
||||
[False, True],
|
||||
[False, False],
|
||||
[False, False],
|
||||
]
|
||||
|
||||
|
||||
def test_resolve_prompt_seq_len_keeps_prefill_padding_after_batch_shrink():
|
||||
"""Longest-prompt cancel must not shrink the surviving request's KV layout."""
|
||||
short = _state([{}])
|
||||
short.extra["image_pad_mask"] = torch.tensor([[True, False]])
|
||||
short.extra["prefill_prompt_seq_len"] = 4 # baked in when batched with a longer prompt
|
||||
|
||||
# Current InputBatch only sees the short prompt's native length.
|
||||
assert QwenImage21Pipeline._resolve_prompt_seq_len([short], current_seq_len=2) == 4
|
||||
padded = QwenImage21Pipeline._pad_rows_to_seq_len([short.extra["image_pad_mask"]], 4)
|
||||
assert padded.tolist() == [[True, False, False, False]]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# post-process func
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def postprocess(tmp_path):
|
||||
vae_dir = tmp_path / "vae"
|
||||
vae_dir.mkdir()
|
||||
(vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4}))
|
||||
return get_qwen_image_21_post_process_func(SimpleNamespace(model=str(tmp_path)))
|
||||
|
||||
|
||||
def test_postprocess_tensor_returns_pil_images(postprocess):
|
||||
from PIL import Image
|
||||
|
||||
images = torch.rand(2, 4, 16, 16) * 2 - 1
|
||||
|
||||
result = postprocess(images)
|
||||
|
||||
assert len(result) == 2
|
||||
assert all(isinstance(image, Image.Image) for image in result)
|
||||
|
||||
|
||||
def test_postprocess_envelope_unwraps_payload_and_keeps_metadata(postprocess):
|
||||
from PIL import Image
|
||||
|
||||
envelope = {
|
||||
"payload": {"image": torch.rand(1, 4, 16, 16), "seed": 42},
|
||||
"metadata": {"request_id": "abc"},
|
||||
}
|
||||
|
||||
result = postprocess(envelope)
|
||||
|
||||
assert result["metadata"] == {"request_id": "abc"}
|
||||
assert result["payload"]["seed"] == 42
|
||||
assert len(result["payload"]["image"]) == 1
|
||||
assert isinstance(result["payload"]["image"][0], Image.Image)
|
||||
|
||||
|
||||
def test_postprocess_envelope_requires_image_payload(postprocess):
|
||||
with pytest.raises(ValueError, match=r"payload\['image'\]"):
|
||||
postprocess({"payload": {"seed": 42}, "metadata": {}})
|
||||
|
||||
|
||||
def test_release_captured_graphs_delegates_to_transformer() -> None:
|
||||
pipeline = QwenImage21Pipeline.__new__(QwenImage21Pipeline)
|
||||
released: list[bool] = []
|
||||
pipeline.transformer = SimpleNamespace(release_captured_graphs=lambda: released.append(True))
|
||||
|
||||
pipeline.release_captured_graphs()
|
||||
|
||||
assert released == [True]
|
||||
|
||||
|
||||
def test_postprocess_envelope_normalizes_non_dict_metadata(postprocess):
|
||||
result = postprocess({"payload": {"image": torch.rand(1, 4, 16, 16)}, "metadata": "bogus"})
|
||||
|
||||
assert result["metadata"] == {}
|
||||
@@ -0,0 +1,101 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""CPU lock for the Qwen-Image-2.1 CFG tail slice.
|
||||
|
||||
2.1 prepends condition latents, so the target noise is the tail. The shared
|
||||
mixin slices the head. These tests stub ``predict_noise`` and check the
|
||||
returned rows are the tail, including rank-0 / rank-1 gather order.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.qwen_image_21 import cfg_parallel as cfg_mod
|
||||
from vllm_omni.diffusion.models.qwen_image_21.cfg_parallel import QwenImage21CFGParallelMixin
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
PREFIX = 3
|
||||
TARGET = 2
|
||||
SCALE = 4.0
|
||||
|
||||
|
||||
class _Stub(QwenImage21CFGParallelMixin):
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
def predict_noise(self, **kwargs: Any) -> torch.Tensor:
|
||||
self.calls.append(kwargs)
|
||||
return kwargs["joint"]
|
||||
|
||||
|
||||
class _Gather:
|
||||
def __init__(self, rank_tails: list[torch.Tensor]) -> None:
|
||||
self.rank_tails = rank_tails
|
||||
self.seen: list[torch.Tensor] = []
|
||||
|
||||
def all_gather(self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False) -> list[torch.Tensor]:
|
||||
del dim
|
||||
assert separate_tensors
|
||||
self.seen.append(input_.clone())
|
||||
return [tensor.clone() for tensor in self.rank_tails]
|
||||
|
||||
|
||||
def _joint(fill: float) -> torch.Tensor:
|
||||
prefix = torch.full((1, PREFIX, 1), fill)
|
||||
target = torch.full((1, TARGET, 1), fill + 100.0)
|
||||
return torch.cat([prefix, target], dim=1)
|
||||
|
||||
|
||||
def _cfg(positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
return negative + SCALE * (positive - negative)
|
||||
|
||||
|
||||
def test_single_rank_cfg_uses_target_tail_not_condition_head() -> None:
|
||||
pipe = _Stub()
|
||||
positive = _joint(1.0)
|
||||
negative = _joint(2.0)
|
||||
actual = pipe.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=True,
|
||||
true_cfg_scale=SCALE,
|
||||
positive_kwargs={"joint": positive, "branch": "pos"},
|
||||
negative_kwargs={"joint": negative, "branch": "neg"},
|
||||
output_slice=TARGET,
|
||||
)
|
||||
tail = _cfg(positive[:, -TARGET:], negative[:, -TARGET:])
|
||||
head = _cfg(positive[:, :TARGET], negative[:, :TARGET])
|
||||
torch.testing.assert_close(actual, tail)
|
||||
assert not torch.allclose(actual, head)
|
||||
assert [call["branch"] for call in pipe.calls] == ["pos", "neg"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
def test_cfg_parallel_gather_order_uses_target_tails(monkeypatch: pytest.MonkeyPatch, rank: int) -> None:
|
||||
pipe = _Stub()
|
||||
positive = _joint(1.0)
|
||||
negative = _joint(2.0)
|
||||
pos_tail = positive[:, -TARGET:]
|
||||
neg_tail = negative[:, -TARGET:]
|
||||
group = _Gather([pos_tail, neg_tail])
|
||||
monkeypatch.setattr(cfg_mod, "_get_cfg_world_size_or_one", lambda: 2)
|
||||
monkeypatch.setattr(cfg_mod, "get_classifier_free_guidance_rank", lambda: rank)
|
||||
monkeypatch.setattr(cfg_mod, "get_cfg_group", lambda: group)
|
||||
|
||||
actual = pipe.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=True,
|
||||
true_cfg_scale=SCALE,
|
||||
positive_kwargs={"joint": positive, "branch": "pos"},
|
||||
negative_kwargs={"joint": negative, "branch": "neg"},
|
||||
output_slice=TARGET,
|
||||
)
|
||||
|
||||
expected_branch = "pos" if rank == 0 else "neg"
|
||||
expected_local = pos_tail if rank == 0 else neg_tail
|
||||
assert [call["branch"] for call in pipe.calls] == [expected_branch]
|
||||
torch.testing.assert_close(group.seen[0], expected_local)
|
||||
torch.testing.assert_close(actual, _cfg(pos_tail, neg_tail))
|
||||
head = _cfg(positive[:, :TARGET], negative[:, :TARGET])
|
||||
assert not torch.allclose(actual, head)
|
||||
@@ -0,0 +1,64 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""L1 tests for checkpoint-to-pipeline routing of Qwen-Image 2.1."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def test_registry_entry_loads_the_pipeline_class():
|
||||
from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline
|
||||
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
||||
|
||||
assert DiffusionModelRegistry._try_load_model_cls("QwenImage21Pipeline") is QwenImage21Pipeline
|
||||
|
||||
|
||||
def test_pipeline_declares_request_batching_and_step_execution():
|
||||
from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline
|
||||
|
||||
assert QwenImage21Pipeline.supports_request_batch
|
||||
assert QwenImage21Pipeline.supports_step_execution
|
||||
|
||||
|
||||
def test_process_funcs_resolve_to_the_pipeline_module():
|
||||
from vllm_omni.diffusion.models.qwen_image_21 import pipeline_qwen_image_21
|
||||
from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS
|
||||
|
||||
pre_name = _DIFFUSION_PRE_PROCESS_FUNCS["QwenImage21Pipeline"]
|
||||
post_name = _DIFFUSION_POST_PROCESS_FUNCS["QwenImage21Pipeline"]
|
||||
assert pre_name == "get_qwen_image_21_pre_process_func"
|
||||
assert post_name == "get_qwen_image_21_post_process_func"
|
||||
assert callable(getattr(pipeline_qwen_image_21, pre_name))
|
||||
assert callable(getattr(pipeline_qwen_image_21, post_name))
|
||||
|
||||
|
||||
def test_multimodal_metadata_caps_condition_images_at_four():
|
||||
from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_21_MAX_INPUT_IMAGES, get_diffusion_model_metadata
|
||||
|
||||
metadata = get_diffusion_model_metadata("QwenImage21Pipeline")
|
||||
|
||||
assert metadata.supports_multimodal_inputs
|
||||
assert metadata.max_multimodal_image_inputs == 4
|
||||
assert metadata.max_multimodal_image_inputs == QWEN_IMAGE_21_MAX_INPUT_IMAGES
|
||||
|
||||
|
||||
def test_model_index_resolves_and_enriches(tmp_path):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
|
||||
|
||||
(tmp_path / "model_index.json").write_text(
|
||||
json.dumps({"_class_name": "QwenImage21Pipeline"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert resolve_model_class_name(str(tmp_path)) == "QwenImage21Pipeline"
|
||||
|
||||
config = OmniDiffusionConfig(model=str(tmp_path))
|
||||
config.enrich_config()
|
||||
|
||||
assert config.model_class_name == "QwenImage21Pipeline"
|
||||
assert config.supports_multimodal_inputs
|
||||
assert config.max_multimodal_image_inputs == 4
|
||||
@@ -0,0 +1,46 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""CPU lock: adjacent condition images stay separate token blocks.
|
||||
|
||||
Two condition images with no text between them are one run of True in
|
||||
``image_pad_mask``. Block ids come from ``img_shapes`` token counts, so they
|
||||
must not collapse into one bidirectional attention span.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import (
|
||||
QwenImage21Transformer2DModel,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def test_adjacent_condition_images_stay_separate_blocks() -> None:
|
||||
# text(2) | cond0(2) | cond1(3) | text(1) | target(2)
|
||||
# cond0 and cond1 are one contiguous True run.
|
||||
image_pad_mask = torch.tensor([False, False, True, True, True, True, True, False, True, True])
|
||||
img_shapes = [(1, 1, 2), (1, 1, 3), (1, 1, 2)]
|
||||
|
||||
image_ids, target_token_mask = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes)
|
||||
|
||||
assert image_ids.tolist() == [-1, -1, 0, 0, 1, 1, 1, -1, 2, 2]
|
||||
assert target_token_mask.tolist() == [
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
True,
|
||||
True,
|
||||
]
|
||||
assert QwenImage21Transformer2DModel._build_full_attn_spans(image_ids) == [
|
||||
(2, 4),
|
||||
(4, 7),
|
||||
(8, 10),
|
||||
]
|
||||
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.qwen_image_21 import qwen_image_21_transformer as transformer
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"device",
|
||||
[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu])],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_qk_norm_matches_diffusers_rounding(monkeypatch, device, dtype):
|
||||
monkeypatch.setattr(
|
||||
"vllm.distributed.parallel_state.get_tp_group", lambda: SimpleNamespace(rank_in_group=0, world_size=1)
|
||||
)
|
||||
monkeypatch.setattr(transformer, "Attention", lambda **kwargs: torch.nn.Identity())
|
||||
attention = transformer.QwenImage21Attention(dim=256, heads=2, dim_head=128).to(device=device, dtype=dtype)
|
||||
values = torch.randn(2, 17, 2, 128, generator=torch.Generator().manual_seed(42)).to(device=device, dtype=dtype)
|
||||
for norm in (attention.norm_q, attention.norm_k):
|
||||
with torch.no_grad():
|
||||
norm.weight.copy_(torch.linspace(0.5, 1.5, 128, device=device, dtype=dtype))
|
||||
variance = values.float().square().mean(-1, keepdim=True)
|
||||
expected = (values * torch.rsqrt(variance + norm.eps)).to(dtype) * norm.weight
|
||||
torch.testing.assert_close(norm(values), expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"device",
|
||||
[pytest.param("meta", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu])],
|
||||
)
|
||||
def test_rope_frequencies_do_not_depend_on_loader_device(device):
|
||||
reference = transformer.QwenImage21Rope(theta=10000, axes_dim=[16, 56, 56])
|
||||
with torch.device(device):
|
||||
loaded = transformer.QwenImage21Rope(theta=10000, axes_dim=[16, 56, 56])
|
||||
for expected, actual in zip(reference.freqs, loaded.freqs):
|
||||
assert actual.device.type == "cpu"
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.gpu
|
||||
def test_scheduler_matches_diffusers_on_latent_device():
|
||||
import numpy as np
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
|
||||
from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import QwenImage21Pipeline
|
||||
|
||||
pipeline = object.__new__(QwenImage21Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cuda", 0)
|
||||
pipeline.scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True)
|
||||
timesteps, _ = pipeline.prepare_timesteps(50, None, 4096)
|
||||
reference = FlowMatchEulerDiscreteScheduler.from_config(pipeline.scheduler.config)
|
||||
reference.set_timesteps(sigmas=np.linspace(1.0, 1 / 50, 50), device=pipeline.device, mu=1.15)
|
||||
sample = torch.randn(1, 16, 64, generator=torch.Generator().manual_seed(42)).to(
|
||||
device=pipeline.device, dtype=torch.bfloat16
|
||||
)
|
||||
noise = sample.cos()
|
||||
actual = pipeline.scheduler.step(noise, timesteps[0], sample, return_dict=False)[0]
|
||||
expected = reference.step(noise, reference.timesteps[0], sample, return_dict=False)[0]
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
assert pipeline.scheduler.sigmas.device == sample.device
|
||||
assert timesteps.device == sample.device
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.gpu
|
||||
def test_timestep_frequencies_match_cpu_initialization():
|
||||
reference = transformer.QwenImage21TemporalTimesteps(timestep_dim=256)
|
||||
with torch.device("cuda"):
|
||||
loaded = transformer.QwenImage21TemporalTimesteps(timestep_dim=256)
|
||||
assert loaded.freqs.device.type == "cuda"
|
||||
torch.testing.assert_close(loaded.freqs.cpu(), reference.freqs, rtol=0, atol=0)
|
||||
timesteps = torch.tensor([0.0, 0.984375, 1.0], device="cuda")
|
||||
torch.testing.assert_close(loaded(timesteps), reference.to("cuda")(timesteps), rtol=0, atol=0)
|
||||
@@ -0,0 +1,266 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage21 import (
|
||||
DistributedAutoencoderKLQwenImage21,
|
||||
)
|
||||
from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
|
||||
AutoencoderKLQwenImage21,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def single_thread():
|
||||
previous = torch.get_num_threads()
|
||||
torch.set_num_threads(1)
|
||||
yield
|
||||
torch.set_num_threads(previous)
|
||||
|
||||
|
||||
def _make_vae(cls=AutoencoderKLQwenImage21, **overrides):
|
||||
"""Toy VAE with the 2.1 residual architecture; 4x spatial compression."""
|
||||
kwargs = dict(
|
||||
base_dim=16,
|
||||
decoder_base_dim=16,
|
||||
z_dim=4,
|
||||
dim_mult=[1, 2, 2],
|
||||
num_res_blocks=1,
|
||||
temperal_downsample=[False, False],
|
||||
is_residual=True,
|
||||
in_channels=4,
|
||||
out_channels=4,
|
||||
latents_mean=[0.0] * 4,
|
||||
latents_std=[1.0] * 4,
|
||||
scale_factor_spatial=None,
|
||||
scale_factor_temporal=None,
|
||||
)
|
||||
kwargs.update(overrides)
|
||||
return cls(**kwargs)
|
||||
|
||||
|
||||
def test_spatial_compression_ratio_derived_from_architecture():
|
||||
# One 2x downsampling stage per `temperal_downsample` entry.
|
||||
vae = _make_vae()
|
||||
assert vae.spatial_compression_ratio == 4
|
||||
|
||||
# The real Qwen-Image-2.1 checkpoint ships scale_factor_spatial=8, but its
|
||||
# four downsampling stages compress 16x; the architecture wins.
|
||||
with torch.device("meta"):
|
||||
vae = AutoencoderKLQwenImage21(
|
||||
base_dim=96,
|
||||
decoder_base_dim=144,
|
||||
z_dim=64,
|
||||
dim_mult=[1, 2, 4, 8, 8],
|
||||
num_res_blocks=2,
|
||||
temperal_downsample=[False, True, True, True],
|
||||
is_residual=True,
|
||||
in_channels=4,
|
||||
out_channels=4,
|
||||
scale_factor_spatial=8,
|
||||
)
|
||||
assert vae.spatial_compression_ratio == 16
|
||||
|
||||
|
||||
def test_tiled_decode_output_shape_and_closeness():
|
||||
torch.manual_seed(0)
|
||||
vae = _make_vae()
|
||||
z = torch.randn(1, 4, 1, 16, 16)
|
||||
ratio = vae.spatial_compression_ratio
|
||||
|
||||
with torch.no_grad():
|
||||
ref = vae._decode(z, return_dict=False)[0]
|
||||
vae.enable_tiling(
|
||||
tile_sample_min_height=32,
|
||||
tile_sample_min_width=32,
|
||||
tile_sample_stride_height=24,
|
||||
tile_sample_stride_width=24,
|
||||
)
|
||||
tiled = vae.decode(z, return_dict=False)[0]
|
||||
|
||||
# Regression: with the wrong (config-derived) ratio the tiled output was
|
||||
# cropped to a fraction of the sample size.
|
||||
assert tiled.shape == (1, 4, 1, 16 * ratio, 16 * ratio)
|
||||
assert tiled.shape == ref.shape
|
||||
|
||||
# Tiled decode blends overlapping tiles, so it is not bitwise-identical to
|
||||
# untiled decode; on this toy model the seam error stays small.
|
||||
mse = ((ref - tiled) ** 2).mean().item()
|
||||
psnr = 10 * math.log10(4.0 / mse)
|
||||
assert psnr > 20.0
|
||||
|
||||
# Tiled output is clamped to [-1, 1] like the untiled path.
|
||||
assert tiled.min() >= -1.0 and tiled.max() <= 1.0
|
||||
|
||||
|
||||
def test_distributed_tile_hooks_match_tiled_decode():
|
||||
"""tile_split -> tile_exec -> tile_merge must reproduce tiled_decode exactly."""
|
||||
torch.manual_seed(0)
|
||||
vae = _make_vae()
|
||||
vae.enable_tiling(
|
||||
tile_sample_min_height=32,
|
||||
tile_sample_min_width=32,
|
||||
tile_sample_stride_height=24,
|
||||
tile_sample_stride_width=24,
|
||||
)
|
||||
distributed = _make_vae(DistributedAutoencoderKLQwenImage21)
|
||||
distributed.enable_tiling(
|
||||
tile_sample_min_height=32,
|
||||
tile_sample_min_width=32,
|
||||
tile_sample_stride_height=24,
|
||||
tile_sample_stride_width=24,
|
||||
)
|
||||
distributed.load_state_dict(vae.state_dict())
|
||||
|
||||
z = torch.randn(1, 4, 1, 16, 16)
|
||||
with torch.no_grad():
|
||||
tiled = vae.tiled_decode(z, return_dict=False)[0]
|
||||
|
||||
tasks, grid_spec = DistributedAutoencoderKLQwenImage21.tile_split(distributed, z)
|
||||
# 16x16 latent, 8x8 latent tiles, stride 6 -> range(0, 16, 6) = 3 stops/axis.
|
||||
assert len(tasks) == 9
|
||||
assert grid_spec.grid_shape == (3, 3)
|
||||
assert grid_spec.split_dims == (3, 4)
|
||||
# Every tile carries all frames (T=1 for images).
|
||||
assert all(len(task.tensor) == z.shape[2] for task in tasks)
|
||||
|
||||
coord_tensor_map = {task.grid_coord: distributed.tile_exec(task) for task in tasks}
|
||||
merged = distributed.tile_merge(coord_tensor_map, grid_spec)
|
||||
|
||||
assert merged.shape == tiled.shape
|
||||
torch.testing.assert_close(merged, tiled)
|
||||
|
||||
|
||||
def test_tiled_encode_output_shape_and_closeness():
|
||||
"""tiled_encode is engaged for image-conditioned generation when tiling is on."""
|
||||
torch.manual_seed(0)
|
||||
vae = _make_vae()
|
||||
x = torch.randn(1, 4, 1, 64, 64)
|
||||
|
||||
with torch.no_grad():
|
||||
ref = vae._encode(x)
|
||||
vae.enable_tiling(
|
||||
tile_sample_min_height=32,
|
||||
tile_sample_min_width=32,
|
||||
tile_sample_stride_height=24,
|
||||
tile_sample_stride_width=24,
|
||||
)
|
||||
tiled = vae._encode(x)
|
||||
|
||||
ratio = vae.spatial_compression_ratio
|
||||
assert tiled.shape == (1, 8, 1, 64 // ratio, 64 // ratio)
|
||||
assert tiled.shape == ref.shape
|
||||
# Blended tile borders deviate from untiled encode but stay small.
|
||||
assert (ref - tiled).abs().max().item() < 1.0
|
||||
|
||||
|
||||
def test_encode_decode_round_trip_shape():
|
||||
torch.manual_seed(0)
|
||||
vae = _make_vae()
|
||||
x = torch.randn(1, 4, 1, 32, 32)
|
||||
with torch.no_grad():
|
||||
latents = vae.encode(x, return_dict=False)[0].mode()
|
||||
decoded = vae.decode(latents, return_dict=False)[0]
|
||||
ratio = vae.spatial_compression_ratio
|
||||
assert latents.shape == (1, 4, 1, 32 // ratio, 32 // ratio)
|
||||
assert decoded.shape == x.shape
|
||||
|
||||
|
||||
@pytest.mark.parametrize("persistent", [False, True])
|
||||
def test_adaptive_oom_restores_tile_configuration(monkeypatch, persistent):
|
||||
vae = _make_vae()
|
||||
vae.enable_tiling(
|
||||
tile_sample_min_height=512,
|
||||
tile_sample_min_width=512,
|
||||
tile_sample_stride_height=384,
|
||||
tile_sample_stride_width=384,
|
||||
)
|
||||
attempts = []
|
||||
z = torch.zeros(1, 4, 1, 4, 4)
|
||||
|
||||
def decode_tile(z, return_dict):
|
||||
attempts.append(vae.tile_sample_min_height)
|
||||
if persistent or vae.tile_sample_min_height > 128:
|
||||
raise torch.OutOfMemoryError("injected tile OOM")
|
||||
return (z,)
|
||||
|
||||
monkeypatch.setattr(vae, "tiled_decode", decode_tile)
|
||||
monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
|
||||
if persistent:
|
||||
with pytest.raises(torch.OutOfMemoryError):
|
||||
vae._tiled_decode_adaptive(z, return_dict=False)
|
||||
else:
|
||||
assert vae._tiled_decode_adaptive(z, return_dict=False)[0] is z
|
||||
assert attempts == [512, 256, 128]
|
||||
assert (
|
||||
vae.tile_sample_min_height,
|
||||
vae.tile_sample_min_width,
|
||||
vae.tile_sample_stride_height,
|
||||
vae.tile_sample_stride_width,
|
||||
) == (512, 512, 384, 384)
|
||||
|
||||
|
||||
def test_adaptive_oom_clears_feat_map_when_retries_exhausted(monkeypatch):
|
||||
"""Final min-tile OOM must not leave decoder activations on the device."""
|
||||
vae = _make_vae()
|
||||
vae.enable_tiling(
|
||||
tile_sample_min_height=512,
|
||||
tile_sample_min_width=512,
|
||||
tile_sample_stride_height=384,
|
||||
tile_sample_stride_width=384,
|
||||
)
|
||||
vae.clear_cache()
|
||||
z = torch.zeros(1, 4, 1, 4, 4)
|
||||
leaked = torch.ones(2, 2, device="cpu") # stand-in for a cached activation
|
||||
|
||||
def decode_tile(z, return_dict):
|
||||
# Mimic a partial tiled decode that already wrote into _feat_map before OOM.
|
||||
vae._feat_map[0] = leaked
|
||||
raise torch.OutOfMemoryError("injected tile OOM")
|
||||
|
||||
monkeypatch.setattr(vae, "tiled_decode", decode_tile)
|
||||
monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
|
||||
|
||||
with pytest.raises(torch.OutOfMemoryError):
|
||||
vae._tiled_decode_adaptive(z, return_dict=False)
|
||||
|
||||
assert all(slot is None for slot in vae._feat_map)
|
||||
assert all(slot is None for slot in vae._enc_feat_map)
|
||||
assert (
|
||||
vae.tile_sample_min_height,
|
||||
vae.tile_sample_min_width,
|
||||
vae.tile_sample_stride_height,
|
||||
vae.tile_sample_stride_width,
|
||||
) == (512, 512, 384, 384)
|
||||
|
||||
|
||||
def test_tiled_encode_oom_clears_enc_feat_map(monkeypatch):
|
||||
"""Tiled _encode must clear encoder activations if tiled_encode OOMs."""
|
||||
vae = _make_vae()
|
||||
vae.enable_tiling(
|
||||
tile_sample_min_height=32,
|
||||
tile_sample_min_width=32,
|
||||
tile_sample_stride_height=24,
|
||||
tile_sample_stride_width=24,
|
||||
)
|
||||
vae.clear_cache()
|
||||
leaked = torch.ones(2, 2, device="cpu")
|
||||
x = torch.zeros(1, 4, 1, 64, 64)
|
||||
|
||||
def encode_tile(x_in):
|
||||
vae._enc_feat_map[0] = leaked
|
||||
raise torch.OutOfMemoryError("injected tiled encode OOM")
|
||||
|
||||
monkeypatch.setattr(vae, "tiled_encode", encode_tile)
|
||||
|
||||
with pytest.raises(torch.OutOfMemoryError):
|
||||
vae._encode(x)
|
||||
|
||||
assert all(slot is None for slot in vae._enc_feat_map)
|
||||
assert all(slot is None for slot in vae._feat_map)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Unit tests for SequentialOffloadBackend."""
|
||||
|
||||
@@ -284,6 +284,102 @@ def test_component_selective_model_offload_requires_swap_counterpart(component,
|
||||
backend.enable(pipeline)
|
||||
|
||||
|
||||
def test_persistent_dit_staging_reuses_fixed_device_storage(accelerator_device) -> None:
|
||||
dit = _create_simple_module().to(accelerator_device)
|
||||
encoder = _create_simple_module().to(accelerator_device)
|
||||
original_weight = dit.linear.weight.detach().clone()
|
||||
|
||||
apply_sequential_offload(
|
||||
dit_modules=[dit],
|
||||
encoder_modules=[encoder],
|
||||
device=accelerator_device,
|
||||
pin_memory=False,
|
||||
persistent_dit_staging=True,
|
||||
)
|
||||
|
||||
dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
|
||||
encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
|
||||
assert dit_hook._stager is not None
|
||||
assert encoder_hook._stager is None # encoders keep plain move semantics
|
||||
|
||||
# Registration snapshots the CPU master and rebinds to it.
|
||||
assert dit.linear.weight.device.type == "cpu"
|
||||
|
||||
dit_hook._to_gpu(dit)
|
||||
first_ptr = dit.linear.weight.data_ptr()
|
||||
assert dit.linear.weight.device == accelerator_device
|
||||
torch.testing.assert_close(dit.linear.weight.data, original_weight)
|
||||
retained_storage = dit_hook._stager._device_storages[0]
|
||||
retained_ptr = retained_storage.data_ptr()
|
||||
|
||||
# Offload via a *different* module's hook (the cross-module swap path).
|
||||
# The module parameters move back to the CPU master, but the fixed device
|
||||
# storage stays allocated so CUDA-graph-captured weight pointers stay valid.
|
||||
encoder_hook._to_cpu(dit)
|
||||
assert dit.linear.weight.device.type == "cpu"
|
||||
assert dit_hook._stager._device_storages
|
||||
assert dit_hook._stager._device_storages[0].data_ptr() == retained_ptr
|
||||
dit_hook._to_gpu(dit)
|
||||
assert dit.linear.weight.data_ptr() == first_ptr
|
||||
torch.testing.assert_close(dit.linear.weight.data, original_weight)
|
||||
|
||||
# A second offload/load cycle through the DiT hook itself stays put too.
|
||||
dit_hook._to_cpu(dit)
|
||||
dit_hook._to_gpu(dit)
|
||||
assert dit.linear.weight.data_ptr() == first_ptr
|
||||
|
||||
remove_sequential_offload([dit, encoder])
|
||||
|
||||
|
||||
def test_default_sequential_offload_releases_device_storage(accelerator_device) -> None:
|
||||
"""Without persistent staging, offloading keeps plain move semantics."""
|
||||
dit = _create_simple_module().to(accelerator_device)
|
||||
encoder = _create_simple_module().to(accelerator_device)
|
||||
|
||||
apply_sequential_offload(
|
||||
dit_modules=[dit],
|
||||
encoder_modules=[encoder],
|
||||
device=accelerator_device,
|
||||
pin_memory=False,
|
||||
)
|
||||
|
||||
dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
|
||||
assert dit_hook._stager is None # no retained staging: GPU storage is freed on offload
|
||||
|
||||
dit_hook._to_gpu(dit)
|
||||
assert dit.linear.weight.device == accelerator_device
|
||||
dit_hook._to_cpu(dit)
|
||||
assert dit.linear.weight.device.type == "cpu"
|
||||
assert all(parameter.device.type == "cpu" for parameter in dit.parameters())
|
||||
|
||||
remove_sequential_offload([dit, encoder])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("decode_graphs", "expect_stager"),
|
||||
[(True, True), (False, False)],
|
||||
)
|
||||
def test_model_level_persistent_staging_gated_on_decode_graph(
|
||||
accelerator_device, decode_graphs: bool, expect_stager: bool
|
||||
) -> None:
|
||||
"""ModelLevelOffloadBackend only retains DiT staging for decode-graph models."""
|
||||
pipeline = nn.Module()
|
||||
pipeline.transformer = _create_simple_module()
|
||||
if decode_graphs:
|
||||
pipeline.transformer.enable_cuda_graph_decode = True
|
||||
pipeline.text_encoder = _create_simple_module()
|
||||
backend = ModelLevelOffloadBackend(
|
||||
OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
|
||||
accelerator_device,
|
||||
)
|
||||
backend.enable(pipeline)
|
||||
|
||||
hook = pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
|
||||
assert (hook._stager is not None) is expect_stager
|
||||
|
||||
backend.disable()
|
||||
|
||||
|
||||
def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
dit = _create_simple_module()
|
||||
encoder = _create_simple_module()
|
||||
|
||||
@@ -46,6 +46,21 @@ def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
|
||||
return OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None})
|
||||
|
||||
|
||||
def test_direct_config_extras_are_independent_dicts():
|
||||
first = OmniDiffusionConfig(model="test-model")
|
||||
second = OmniDiffusionConfig(model="test-model")
|
||||
|
||||
assert first.extras.get("prefix_kv_cache_dtype") is None
|
||||
first.extras["prefix_kv_cache_dtype"] = "fp8"
|
||||
assert second.extras == {}
|
||||
|
||||
|
||||
def test_direct_config_preserves_explicit_extras():
|
||||
config = OmniDiffusionConfig(model="test-model", extras={"prefix_kv_cache_dtype": "fp8"})
|
||||
|
||||
assert config.extras["prefix_kv_cache_dtype"] == "fp8"
|
||||
|
||||
|
||||
class TestParallelConfigPropagation:
|
||||
"""Core regression tests: parallel_config must survive serialization."""
|
||||
|
||||
|
||||
@@ -521,6 +521,35 @@ def test_compile_transformer_regionally_compiles_blocks(monkeypatch, use_hsdp):
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_compile_transformer_scopes_inductor_cudagraphs_to_decode_graph(monkeypatch):
|
||||
model = SimpleNamespace(enable_cuda_graph_decode=True)
|
||||
runner = _make_compile_runner(model)
|
||||
compile_calls = []
|
||||
|
||||
def _regionally_compile(target, *args, **kwargs):
|
||||
compile_calls.append(kwargs)
|
||||
return target
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile)
|
||||
import torch._inductor.config as inductor_config
|
||||
|
||||
before = inductor_config.triton.cudagraphs
|
||||
DiffusionModelRunner._compile_transformer(runner, "transformer")
|
||||
|
||||
assert inductor_config.triton.cudagraphs == before
|
||||
assert compile_calls == [
|
||||
{
|
||||
"dynamic": True,
|
||||
"options": {
|
||||
"triton.cudagraphs": False,
|
||||
"triton.cudagraph_trees": False,
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_compile_transformer_uses_regional_dynamic_false_config(monkeypatch):
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.media import generate_synthetic_image
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
|
||||
|
||||
MODEL = os.environ.get("QWEN_IMAGE_21_TEST_MODEL", "Qwen/Qwen-Image-2.1")
|
||||
SERVER_CASES = [
|
||||
pytest.param(
|
||||
OmniServerParams(model=MODEL, server_args=["--vae-use-tiling", "--max-num-seqs", "2"] + extra),
|
||||
id=name,
|
||||
marks=hardware_marks(res={"cuda": "H100"}),
|
||||
)
|
||||
for name, extra in [("request", []), ("step", ["--step-execution"])]
|
||||
]
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.advanced_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.skipif(
|
||||
not os.environ.get("QWEN_IMAGE_21_TEST_MODEL"),
|
||||
reason="Set QWEN_IMAGE_21_TEST_MODEL to a provisioned checkpoint until public weights are available.",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
|
||||
def test_text_to_image(omni_server: OmniServer, online_client: OnlineOmniClient):
|
||||
online_client.send_diffusion_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"messages": dummy_messages_from_mix_data(content_text="A red ceramic teapot on a wooden table."),
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 4,
|
||||
"seed": 42,
|
||||
"negative_prompt": "blurry, low quality",
|
||||
"true_cfg_scale": 4.0,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
|
||||
def test_image_edit(omni_server: OmniServer, online_client: OnlineOmniClient):
|
||||
image = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
online_client.send_diffusion_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"messages": dummy_messages_from_mix_data(content_text="Make the image blue.", image_data_url=image),
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 4,
|
||||
"seed": 42,
|
||||
"negative_prompt": "blurry, low quality",
|
||||
"true_cfg_scale": 4.0,
|
||||
},
|
||||
}
|
||||
)
|
||||
@@ -933,7 +933,7 @@ def test_image_edits_streaming_returns_ar_delta_then_image(streaming_image_edit_
|
||||
assert payloads[0]["index"] == 0
|
||||
assert payloads[1]["delta"] == " done"
|
||||
assert payloads[2]["output_format"] == "png"
|
||||
assert payloads[2]["size"] == "16x16"
|
||||
assert payloads[2]["size"] == "32x24"
|
||||
|
||||
image_payload = payloads[2]["data"][0]
|
||||
img = Image.open(io.BytesIO(base64.b64decode(image_payload["b64_json"])))
|
||||
@@ -2290,8 +2290,8 @@ def test_extract_images_from_result():
|
||||
class EmptyResult:
|
||||
pass
|
||||
|
||||
result = EmptyResult()
|
||||
images = _extract_images_from_result(result)
|
||||
empty_result = EmptyResult()
|
||||
images = _extract_images_from_result(empty_result)
|
||||
assert images == []
|
||||
|
||||
# Test nested batch: [np.array(shape=(3, 64, 64, 3))]
|
||||
@@ -2301,8 +2301,8 @@ def test_extract_images_from_result():
|
||||
def __init__(self):
|
||||
self.images = [batch]
|
||||
|
||||
result = BatchResult()
|
||||
images = _extract_images_from_result(result)
|
||||
batch_result = BatchResult()
|
||||
images = _extract_images_from_result(batch_result)
|
||||
assert len(images) == 3
|
||||
assert all(isinstance(img, Image.Image) for img in images)
|
||||
assert all(img.size == (64, 64) for img in images)
|
||||
@@ -2312,8 +2312,8 @@ def test_extract_images_from_result():
|
||||
def __init__(self):
|
||||
self.images = [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]
|
||||
|
||||
result = DictRequestOutput()
|
||||
images = _extract_images_from_result(result)
|
||||
dict_result = DictRequestOutput()
|
||||
images = _extract_images_from_result(dict_result)
|
||||
assert len(images) == 1
|
||||
assert isinstance(images[0], Image.Image)
|
||||
|
||||
@@ -2322,8 +2322,8 @@ def test_extract_images_from_result():
|
||||
def __init__(self):
|
||||
self.images = [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]
|
||||
|
||||
result = AttrRequestOutput()
|
||||
images = _extract_images_from_result(result)
|
||||
attr_result = AttrRequestOutput()
|
||||
images = _extract_images_from_result(attr_result)
|
||||
assert len(images) == 1
|
||||
assert isinstance(images[0], Image.Image)
|
||||
assert images[0].size == (32, 32)
|
||||
|
||||
@@ -271,7 +271,7 @@ def test_ltx_reference_image_size_policy(tmp_path, model_version: str, expected:
|
||||
|
||||
|
||||
def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured = {}
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
def fake_policy(*, model, revision=None):
|
||||
captured.update(model=model, revision=revision)
|
||||
@@ -495,6 +495,27 @@ def test_bagel_image_to_image_prompt_builder() -> None:
|
||||
assert result["negative_prompt"] == "ugly"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_qwen_image_21_image_to_image_prompt_builder() -> None:
|
||||
dummy_image = Image.new("RGBA", (64, 64))
|
||||
result = build_image_to_image_prompt(
|
||||
"QwenImage21Pipeline",
|
||||
prompt="make it dance",
|
||||
negative_prompt="blurry",
|
||||
input_image=dummy_image,
|
||||
height=1024,
|
||||
width=768,
|
||||
)
|
||||
assert result == {
|
||||
"prompt": "make it dance",
|
||||
"multi_modal_data": {"image": dummy_image},
|
||||
"height": 1024,
|
||||
"width": 768,
|
||||
"negative_prompt": "blurry",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_unknown_pipeline_preserves_canonical_text_to_image_prompt() -> None:
|
||||
@@ -785,7 +806,7 @@ def _patch_hunyuan_autotokenizer(monkeypatch: pytest.MonkeyPatch, tokenizer: obj
|
||||
return tokenizer
|
||||
|
||||
fake_transformers = type(sys)("transformers")
|
||||
fake_transformers.AutoTokenizer = _FakeAutoTokenizer
|
||||
setattr(fake_transformers, "AutoTokenizer", _FakeAutoTokenizer)
|
||||
monkeypatch.setitem(sys.modules, "transformers", fake_transformers)
|
||||
return calls
|
||||
|
||||
|
||||
@@ -25,6 +25,9 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
|
||||
EXCLUDED_MODELS = [
|
||||
"QwenImageLayeredPipeline",
|
||||
# Dedicated CPU/unit coverage in tests/diffusion/models/qwen_image_21/.
|
||||
# Tiny-model offline suite is not wired yet (2.1 transformer/VAE configs).
|
||||
"QwenImage21Pipeline",
|
||||
"GlmImagePipeline",
|
||||
"ZImagePipeline",
|
||||
"OvisImagePipeline",
|
||||
|
||||
@@ -845,6 +845,10 @@ class _DiffusionConfigProjection:
|
||||
pin_cpu_memory: bool = True
|
||||
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
|
||||
diffusion_compile_dynamic: bool = Field(default=True, strict=True)
|
||||
# Mirrors OmniDiffusionConfig.enable_cuda_graph_decode; enforced by the
|
||||
# model (Qwen-Image-2.1) together with enforce_eager, which lives in the
|
||||
# shared engine config.
|
||||
enable_cuda_graph_decode: bool = True
|
||||
fa_deterministic: bool = False
|
||||
vae_use_slicing: bool = False
|
||||
vae_use_tiling: bool = False
|
||||
|
||||
@@ -480,6 +480,10 @@ class StageDeployConfig:
|
||||
# Diffusion execution, cache, and VAE behavior.
|
||||
diffusion_compile_granularity: str | None = None
|
||||
diffusion_compile_dynamic: bool | None = None
|
||||
# CUDA graph capture of fixed-shape KV-cache decode steps (Qwen-Image-2.1
|
||||
# today). Independent of compilation_config.cudagraph_mode;
|
||||
# enforce_eager=True also disables it.
|
||||
enable_cuda_graph_decode: bool | None = None
|
||||
fa_deterministic: bool | None = None
|
||||
cache_backend: str | None = None
|
||||
cache_config: dict[str, Any] | None = None
|
||||
|
||||
@@ -10,11 +10,13 @@ from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
AttentionMetadata,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
|
||||
from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import piecewise_attn
|
||||
|
||||
|
||||
class CuDNNAttentionBackend(AttentionBackend):
|
||||
accept_output_buffer: bool = True
|
||||
supports_prefix_kv_slicing: bool = True
|
||||
supports_piecewise_spans: bool = True
|
||||
|
||||
# cuDNN 9.5+ FMHA on Blackwell: head_dim divisible by 8 and at most 256
|
||||
# for BF16/FP16. Used by automatic platform selection; explicit CUDNN_ATTN
|
||||
@@ -74,6 +76,23 @@ class CuDNNAttentionImpl(AttentionImpl):
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata | None = None,
|
||||
) -> torch.Tensor:
|
||||
if (
|
||||
attn_metadata is not None
|
||||
and attn_metadata.full_attn_spans is not None
|
||||
and (attn_metadata.attn_mask is None or attn_metadata.attn_mask.ndim != 4)
|
||||
):
|
||||
if attn_metadata.attn_mask is not None or "valid_kv_length" in attn_metadata.extra:
|
||||
raise ValueError("Piecewise cuDNN attention cannot combine spans with a padding mask.")
|
||||
return piecewise_attn(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_metadata.full_attn_spans,
|
||||
self.softmax_scale,
|
||||
self._piecewise_attention,
|
||||
query_ranges=attn_metadata.query_ranges,
|
||||
)
|
||||
|
||||
attention_mask = None
|
||||
if attn_metadata:
|
||||
valid_kv_length = attn_metadata.extra.get("valid_kv_length")
|
||||
@@ -97,6 +116,19 @@ class CuDNNAttentionImpl(AttentionImpl):
|
||||
mask_mode="broadcast_k",
|
||||
)
|
||||
|
||||
return self._attention(query, key, value, attention_mask, self.causal, self.softmax_scale)
|
||||
|
||||
def _piecewise_attention(self, query, key, value, *, causal, softmax_scale):
|
||||
mask = None
|
||||
if causal:
|
||||
# SDPA's causal flag aligns at the top left. Text segments need
|
||||
# bottom-right alignment so every query can also see its prefix.
|
||||
q_len, kv_len = query.shape[1], key.shape[1]
|
||||
mask = torch.ones(q_len, kv_len, dtype=torch.bool, device=query.device).tril(kv_len - q_len)
|
||||
mask = mask[None, None]
|
||||
return self._attention(query, key, value, mask, False, softmax_scale)
|
||||
|
||||
def _attention(self, query, key, value, attention_mask, causal, softmax_scale):
|
||||
enable_gqa = query.shape[2] != key.shape[2]
|
||||
kv_seq_len = key.shape[1]
|
||||
query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
|
||||
@@ -130,8 +162,8 @@ class CuDNNAttentionImpl(AttentionImpl):
|
||||
value,
|
||||
attn_mask=attention_mask,
|
||||
dropout_p=0.0,
|
||||
is_causal=self.causal,
|
||||
scale=self.softmax_scale,
|
||||
is_causal=causal,
|
||||
scale=softmax_scale,
|
||||
enable_gqa=enable_gqa,
|
||||
)
|
||||
return output.permute(0, 2, 1, 3)
|
||||
@@ -16,7 +16,7 @@ import diffusers
|
||||
import huggingface_hub
|
||||
import torch
|
||||
from PIL import Image
|
||||
from pydantic import Field, model_validator
|
||||
from pydantic import model_validator
|
||||
from typing_extensions import Self
|
||||
from vllm.config.utils import config
|
||||
from vllm.logger import init_logger
|
||||
@@ -981,6 +981,19 @@ class OmniDiffusionConfig:
|
||||
|
||||
# Compilation
|
||||
enforce_eager: bool = False
|
||||
# Capture fixed-shape KV-cache decode (denoising) steps into CUDA graphs.
|
||||
# Currently only Qwen-Image-2.1's transformer implements this; other models
|
||||
# ignore the flag. This is a per-model diffusion path, independent of the
|
||||
# AR engine's ``compilation_config.cudagraph_mode`` (which does not apply
|
||||
# to diffusion stages). It stacks with ``diffusion_compile_granularity``:
|
||||
# the DiT blocks are still torch.compile'd and graph capture records the
|
||||
# compiled (fused) kernels, while inductor's own cudagraphs stay off.
|
||||
# ``enforce_eager=True`` forces eager decode and
|
||||
# disables capture regardless of this flag; unsupported configurations
|
||||
# (SP/TP, ring, HSDP, dynamic LoRA, padded masks, quantized prefix KV, or
|
||||
# a second in-flight request aliasing the same graph key) log and fall
|
||||
# back to eager decode.
|
||||
enable_cuda_graph_decode: bool = True
|
||||
# Controls the generic compilation path used when a pipeline does not
|
||||
# provide its own setup_compile() implementation.
|
||||
diffusion_compile_granularity: str = "regional"
|
||||
@@ -1140,7 +1153,7 @@ class OmniDiffusionConfig:
|
||||
request_batch_max_wait_ms: float = 0.0
|
||||
|
||||
# Supplementary model specific parameters
|
||||
extras: dict[str, Any] = Field(default_factory=dict)
|
||||
extras: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def is_moe(self) -> bool:
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers.models.autoencoders.vae import DecoderOutput
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
||||
DistributedOperator,
|
||||
DistributedVaeMixin,
|
||||
GridSpec,
|
||||
TileTask,
|
||||
)
|
||||
from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
|
||||
AutoencoderKLQwenImage21,
|
||||
_cudnn_deterministic,
|
||||
_unpatchify,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DistributedAutoencoderKLQwenImage21(AutoencoderKLQwenImage21, DistributedVaeMixin):
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
||||
model = super().from_pretrained(*args, **kwargs)
|
||||
model.init_distributed()
|
||||
return model
|
||||
|
||||
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
# mostly copy from AutoencoderKLQwenImage21.tiled_decode
|
||||
_, _, num_frames, height, width = z.shape
|
||||
sample_height = height * self.spatial_compression_ratio
|
||||
sample_width = width * self.spatial_compression_ratio
|
||||
|
||||
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
||||
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
||||
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
||||
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
||||
|
||||
tile_sample_stride_height = self.tile_sample_stride_height
|
||||
tile_sample_stride_width = self.tile_sample_stride_width
|
||||
if self.config.patch_size is not None:
|
||||
sample_height = sample_height // self.config.patch_size
|
||||
sample_width = sample_width // self.config.patch_size
|
||||
tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
|
||||
tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
|
||||
blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height
|
||||
blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width
|
||||
else:
|
||||
blend_height = self.tile_sample_min_height - tile_sample_stride_height
|
||||
blend_width = self.tile_sample_min_width - tile_sample_stride_width
|
||||
|
||||
# Split z into overlapping tiles and decode them separately.
|
||||
# The tiles have an overlap to avoid seams between tiles.
|
||||
tiletask_list: list[TileTask] = []
|
||||
for i in range(0, height, tile_latent_stride_height):
|
||||
for j in range(0, width, tile_latent_stride_width):
|
||||
time_list = []
|
||||
for k in range(num_frames):
|
||||
tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
|
||||
time_list.append(tile)
|
||||
tiletask_list.append(
|
||||
TileTask(
|
||||
len(tiletask_list),
|
||||
(i // tile_latent_stride_height, j // tile_latent_stride_width),
|
||||
time_list,
|
||||
workload=time_list[0].shape[3] * time_list[0].shape[4],
|
||||
)
|
||||
)
|
||||
tile_spec = {
|
||||
"sample_height": sample_height,
|
||||
"sample_width": sample_width,
|
||||
"blend_height": blend_height,
|
||||
"blend_width": blend_width,
|
||||
"tile_sample_stride_height": tile_sample_stride_height,
|
||||
"tile_sample_stride_width": tile_sample_stride_width,
|
||||
}
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(3, 4),
|
||||
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
|
||||
tile_spec=tile_spec,
|
||||
output_dtype=self.dtype,
|
||||
)
|
||||
return tiletask_list, grid_spec
|
||||
|
||||
def tile_exec(self, task: TileTask) -> torch.Tensor:
|
||||
"""Decode a single latent tile into RGB space."""
|
||||
self.clear_cache()
|
||||
time = []
|
||||
with _cudnn_deterministic():
|
||||
for k in range(len(task.tensor)):
|
||||
self._conv_idx = [0]
|
||||
tile = self.post_quant_conv(task.tensor[k])
|
||||
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
|
||||
time.append(decoded)
|
||||
result = torch.cat(time, dim=2)
|
||||
return result
|
||||
|
||||
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
||||
"""Merge decoded tiles into a full image."""
|
||||
grid_h, grid_w = grid_spec.grid_shape
|
||||
self.clear_cache()
|
||||
|
||||
result_rows = []
|
||||
for i in range(grid_h):
|
||||
result_row = []
|
||||
for j in range(grid_w):
|
||||
tile = coord_tensor_map[(i, j)]
|
||||
if i > 0:
|
||||
tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
|
||||
if j > 0:
|
||||
tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
|
||||
result_row.append(
|
||||
tile[
|
||||
:,
|
||||
:,
|
||||
:,
|
||||
: grid_spec.tile_spec["tile_sample_stride_height"],
|
||||
: grid_spec.tile_spec["tile_sample_stride_width"],
|
||||
]
|
||||
)
|
||||
result_rows.append(torch.cat(result_row, dim=-1))
|
||||
dec = torch.cat(result_rows, dim=3)[
|
||||
:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
|
||||
]
|
||||
|
||||
if self.config.patch_size is not None:
|
||||
dec = _unpatchify(dec, patch_size=self.config.patch_size)
|
||||
|
||||
dec = torch.clamp(dec, min=-1.0, max=1.0)
|
||||
return dec
|
||||
|
||||
def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
|
||||
if not self.is_distributed_enabled():
|
||||
return super().tiled_decode(z, return_dict=return_dict)
|
||||
|
||||
logger.debug("Decode running with distributed executor")
|
||||
result = self.distributed_executor.execute(
|
||||
z,
|
||||
DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
|
||||
broadcast_result=True,
|
||||
)
|
||||
if not return_dict:
|
||||
return (result,)
|
||||
|
||||
return DecoderOutput(sample=result)
|
||||
@@ -3,6 +3,8 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
@@ -34,9 +36,9 @@ class TileTask:
|
||||
|
||||
@dataclass
|
||||
class DistributedOperator:
|
||||
split: callable
|
||||
exec: callable
|
||||
merge: callable
|
||||
split: Callable[[torch.Tensor], tuple[list[TileTask], GridSpec]]
|
||||
exec: Callable[[TileTask], torch.Tensor]
|
||||
merge: Callable[[dict[tuple[int, ...], torch.Tensor], GridSpec], torch.Tensor]
|
||||
|
||||
|
||||
class DistributedVaeExecutor:
|
||||
@@ -46,7 +48,11 @@ class DistributedVaeExecutor:
|
||||
|
||||
def __init__(self):
|
||||
# Use a dedicated process group spanning the complete worker WORLD.
|
||||
self.group = get_world_group().device_group
|
||||
world_group = get_world_group()
|
||||
self.group = world_group.device_group
|
||||
self._error_group = world_group.cpu_group
|
||||
# Reporting an OOM must not require another device allocation.
|
||||
self._error_status = torch.zeros((), dtype=torch.int32, device="cpu")
|
||||
self.world_size = dist.get_world_size(self.group)
|
||||
self.rank = dist.get_rank(self.group)
|
||||
self.parallel_size = 1
|
||||
@@ -56,8 +62,34 @@ class DistributedVaeExecutor:
|
||||
self.parallel_size = parallel_size
|
||||
self.parallel_mode = mode
|
||||
|
||||
@contextmanager
|
||||
def _sync_errors(self, stage: str):
|
||||
"""Agree on local-stage failures before entering the next collective."""
|
||||
status = 0
|
||||
message = ""
|
||||
try:
|
||||
yield
|
||||
except torch.OutOfMemoryError as error:
|
||||
status = 1
|
||||
message = str(error)
|
||||
except Exception as error:
|
||||
# A non-retryable error takes precedence over OOM on another rank.
|
||||
status = 2
|
||||
message = f"{type(error).__name__}: {error}"
|
||||
|
||||
self._error_status.fill_(status)
|
||||
dist.all_reduce(self._error_status, op=dist.ReduceOp.MAX, group=self._error_group)
|
||||
global_status = self._error_status.item()
|
||||
if global_status == 2:
|
||||
detail = message if status == 2 else "another rank failed"
|
||||
raise RuntimeError(f"Distributed VAE {stage} failed: {detail}")
|
||||
if global_status == 1:
|
||||
detail = message if status == 1 else "another rank ran out of memory"
|
||||
raise torch.OutOfMemoryError(f"Distributed VAE {stage} failed: {detail}")
|
||||
|
||||
def gather_tensors(self, tensor: torch.Tensor):
|
||||
gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
|
||||
with self._sync_errors("gather allocation"):
|
||||
gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
|
||||
dist.all_gather(gather_list, tensor, group=self.group)
|
||||
return gather_list if self.rank == 0 else None
|
||||
|
||||
@@ -66,11 +98,12 @@ class DistributedVaeExecutor:
|
||||
return tensor
|
||||
|
||||
def _compute_global_padding_shape(self, local_results, output_ndim: int, device):
|
||||
local_tile_max_dims = [0] * output_ndim
|
||||
for _, tile_tensor in local_results:
|
||||
for dim_idx, dim_size in enumerate(tile_tensor.shape):
|
||||
local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
|
||||
local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
|
||||
with self._sync_errors("shape preparation"):
|
||||
local_tile_max_dims = [0] * output_ndim
|
||||
for _, tile_tensor in local_results:
|
||||
for dim_idx, dim_size in enumerate(tile_tensor.shape):
|
||||
local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
|
||||
local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
|
||||
dist.all_reduce(
|
||||
local_shape_stat,
|
||||
op=dist.ReduceOp.MAX,
|
||||
@@ -100,17 +133,16 @@ class DistributedVaeExecutor:
|
||||
tid = int(meta_src[idx, 0])
|
||||
if tid < 0:
|
||||
continue
|
||||
slices = [slice(None)] * tiles_src[idx].ndim
|
||||
slices_list = [slice(None)] * tiles_src[idx].ndim
|
||||
for i, dim in enumerate(grid_spec.split_dims):
|
||||
slices[dim] = slice(0, int(meta_src[idx, i + 1]))
|
||||
slices = tuple(slices)
|
||||
coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][slices]
|
||||
slices_list[dim] = slice(0, int(meta_src[idx, i + 1]))
|
||||
coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][tuple(slices_list)]
|
||||
|
||||
return coord_tensor_map
|
||||
|
||||
def _balance_tasks(self, task_list, num_rank):
|
||||
workloads = [0] * num_rank
|
||||
assigned = [[] for _ in range(num_rank)]
|
||||
assigned: list[list[TileTask]] = [[] for _ in range(num_rank)]
|
||||
|
||||
for task in sorted(task_list, key=lambda t: t.workload, reverse=True):
|
||||
r = workloads.index(min(workloads))
|
||||
@@ -122,49 +154,51 @@ class DistributedVaeExecutor:
|
||||
def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
|
||||
pp_size = min(self.parallel_size, self.world_size)
|
||||
|
||||
# 1. Split into tiles
|
||||
tiletask_list, grid_spec = operator.split(z)
|
||||
tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
|
||||
|
||||
# 2. local decode
|
||||
assigned = self._balance_tasks(tiletask_list, pp_size)
|
||||
local_tasks = assigned[self.rank] if self.rank < pp_size else []
|
||||
local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
|
||||
with self._sync_errors("tile decode"):
|
||||
tiletask_list, grid_spec = operator.split(z)
|
||||
tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
|
||||
assigned = self._balance_tasks(tiletask_list, pp_size)
|
||||
local_tasks = assigned[self.rank] if self.rank < pp_size else []
|
||||
local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
|
||||
|
||||
# 3. compute shape per rank
|
||||
global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device)
|
||||
|
||||
# 5. prepare tile tensors
|
||||
output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
|
||||
local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
|
||||
local_results, global_padding_shape, grid_spec, z.device, output_dtype
|
||||
)
|
||||
with self._sync_errors("tile packing"):
|
||||
local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
|
||||
local_results, global_padding_shape, grid_spec, z.device, output_dtype
|
||||
)
|
||||
|
||||
# 6. gather tiles & meta
|
||||
meta_gather = self.gather_tensors(local_meta_tensor)
|
||||
tile_gather = self.gather_tensors(local_tile_tensor)
|
||||
|
||||
if self.rank != 0:
|
||||
result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks
|
||||
else:
|
||||
# 7. reconstruct full tensor (rank 0)
|
||||
coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
|
||||
result = operator.merge(coord_tensor_map, grid_spec)
|
||||
with self._sync_errors("tile merge"):
|
||||
if self.rank != 0:
|
||||
result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks
|
||||
else:
|
||||
# 7. reconstruct full tensor (rank 0)
|
||||
coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
|
||||
result = operator.merge(coord_tensor_map, grid_spec)
|
||||
|
||||
if broadcast_result:
|
||||
result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
|
||||
return result
|
||||
|
||||
def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype):
|
||||
shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
|
||||
if self.rank == 0:
|
||||
shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
|
||||
with self._sync_errors("output shape preparation"):
|
||||
shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
|
||||
if self.rank == 0:
|
||||
shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
|
||||
dist.broadcast(shape_tensor, src=0, group=self.group)
|
||||
|
||||
if self.rank != 0:
|
||||
sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
|
||||
else:
|
||||
sync_result = rank0_result
|
||||
with self._sync_errors("output allocation"):
|
||||
if self.rank != 0:
|
||||
sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
|
||||
else:
|
||||
sync_result = rank0_result
|
||||
dist.broadcast(sync_result, src=0, group=self.group)
|
||||
return sync_result
|
||||
|
||||
|
||||
@@ -23,6 +23,11 @@ lora_convert_mapping: dict[str, Callable] = {
|
||||
"QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
|
||||
"QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
|
||||
"QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
|
||||
# 2.1's single-stream blocks keep the same leaf names the converter's
|
||||
# protected n-grams cover (`attn.to_q/to_k/to_v`, `img_mlp`, ...), so the
|
||||
# 2.0 converter applies. Non-diffusers LoRAs that target the fused `to_qkv`
|
||||
# projection are not supported (the converter would split it into `to.qkv`).
|
||||
"QwenImage21Pipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
|
||||
"Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers,
|
||||
"Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers,
|
||||
}
|
||||
|
||||
@@ -26,6 +26,8 @@ class DiffusionModelMetadata:
|
||||
# FLUX.2 Klein supports up to four reference images.
|
||||
FLUX2_KLEIN_MAX_INPUT_IMAGES = 4
|
||||
QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
|
||||
# Qwen-Image 2.1 image-conditioned generation caps condition images at 4.
|
||||
QWEN_IMAGE_21_MAX_INPUT_IMAGES = 4
|
||||
# Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3.
|
||||
HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3
|
||||
# Boogu-Image editing (TI2I) supports a single reference image for now.
|
||||
@@ -42,6 +44,12 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
|
||||
supports_multimodal_inputs=True,
|
||||
max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
|
||||
),
|
||||
# Qwen-Image 2.1 handles both text-to-image (no image) and
|
||||
# image-conditioned requests through the same pipeline class.
|
||||
"QwenImage21Pipeline": DiffusionModelMetadata(
|
||||
supports_multimodal_inputs=True,
|
||||
max_multimodal_image_inputs=QWEN_IMAGE_21_MAX_INPUT_IMAGES,
|
||||
),
|
||||
"HunyuanImage3Pipeline": DiffusionModelMetadata(
|
||||
supports_multimodal_inputs=True,
|
||||
max_multimodal_image_inputs=HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Qwen Image 2.1 diffusion model components.
|
||||
|
||||
Lazy by necessity: the pipeline imports the distributed VAE
|
||||
(`vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage21`),
|
||||
which in turn imports the plain VAE module from this package. Eager imports
|
||||
here would make that a circular import whenever the distributed module is
|
||||
imported first.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
__all__ = [
|
||||
"AutoencoderKLQwenImage21",
|
||||
"QwenImage21Pipeline",
|
||||
"QwenImage21Transformer2DModel",
|
||||
"get_qwen_image_21_post_process_func",
|
||||
"get_qwen_image_21_pre_process_func",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name == "AutoencoderKLQwenImage21":
|
||||
from .autoencoder_kl_qwenimage21 import AutoencoderKLQwenImage21
|
||||
|
||||
return AutoencoderKLQwenImage21
|
||||
if name == "QwenImage21Transformer2DModel":
|
||||
from .qwen_image_21_transformer import QwenImage21Transformer2DModel
|
||||
|
||||
return QwenImage21Transformer2DModel
|
||||
if name in (
|
||||
"QwenImage21Pipeline",
|
||||
"get_qwen_image_21_post_process_func",
|
||||
"get_qwen_image_21_pre_process_func",
|
||||
):
|
||||
from .pipeline_qwen_image_21 import (
|
||||
QwenImage21Pipeline,
|
||||
get_qwen_image_21_post_process_func,
|
||||
get_qwen_image_21_pre_process_func,
|
||||
)
|
||||
|
||||
return {
|
||||
"QwenImage21Pipeline": QwenImage21Pipeline,
|
||||
"get_qwen_image_21_post_process_func": get_qwen_image_21_post_process_func,
|
||||
"get_qwen_image_21_pre_process_func": get_qwen_image_21_pre_process_func,
|
||||
}[name]
|
||||
raise AttributeError(name)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,190 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""CFG Parallel Mixin for Qwen-Image 2.1.
|
||||
|
||||
Differences from the 2.0 mixin (``qwen_image/cfg_parallel.py``):
|
||||
- The transformer owns a KV cache for the text/condition-image prefix, so the
|
||||
denoise loop passes a fresh per-generation ``kv_cache`` and names the CFG
|
||||
branch (``cache_branch="cond"/"uncond"``) on every forward.
|
||||
- The joint sequence is ``[condition images, ..., target image]`` — the target
|
||||
noise prediction is the *tail* of the output, not the head.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import (
|
||||
CFGParallelMixin,
|
||||
_get_cfg_world_size_or_one,
|
||||
_unwrap,
|
||||
_wrap,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
get_cfg_group,
|
||||
get_classifier_free_guidance_rank,
|
||||
)
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QwenImage21CFGParallelMixin(CFGParallelMixin, ProgressBarMixin):
|
||||
"""Base Mixin class for the Qwen-Image 2.1 pipeline providing shared CFG methods."""
|
||||
|
||||
def predict_noise_maybe_with_cfg(
|
||||
self,
|
||||
do_true_cfg: bool,
|
||||
true_cfg_scale: float,
|
||||
positive_kwargs: dict[str, Any],
|
||||
negative_kwargs: dict[str, Any] | None,
|
||||
cfg_normalize: bool = False,
|
||||
output_slice: int | None = None,
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, ...]:
|
||||
# The base implementation slices `[:, :output_slice]` per branch, which
|
||||
# matches the 2.0 condition-last layout. 2.1 concatenates condition
|
||||
# latents *before* the target, so the target prediction is the tail of
|
||||
# the joint output. Each branch must be tail-sliced *before* the CFG
|
||||
# combine: positive and negative prompts tokenize to different lengths,
|
||||
# so the full joint outputs do not align positionally.
|
||||
def tail_slice(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
|
||||
if output_slice is None:
|
||||
return _wrap(pred)
|
||||
return tuple(p[:, -output_slice:] for p in _wrap(pred))
|
||||
|
||||
if do_true_cfg:
|
||||
if _get_cfg_world_size_or_one() > 1:
|
||||
cfg_group = get_cfg_group()
|
||||
cfg_rank = get_classifier_free_guidance_rank()
|
||||
branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
|
||||
local_pred = tail_slice(self.predict_noise(**branch_kwargs))
|
||||
gathered = [cfg_group.all_gather(p, separate_tensors=True) for p in local_pred]
|
||||
positive_noise_pred = tuple(g[0] for g in gathered)
|
||||
negative_noise_pred = tuple(g[1] for g in gathered)
|
||||
else:
|
||||
positive_noise_pred = tail_slice(self.predict_noise(**positive_kwargs))
|
||||
negative_noise_pred = tail_slice(self.predict_noise(**negative_kwargs))
|
||||
return self.combine_cfg_noise(
|
||||
positive_noise_pred,
|
||||
negative_noise_pred,
|
||||
true_cfg_scale,
|
||||
cfg_normalize,
|
||||
**({} if kwargs is None else {"kwargs": kwargs}),
|
||||
)
|
||||
|
||||
pred = self.predict_noise(**positive_kwargs)
|
||||
if output_slice is not None:
|
||||
pred = _unwrap(tail_slice(pred))
|
||||
return pred
|
||||
|
||||
def diffuse(
|
||||
self,
|
||||
prompt_embeds: torch.Tensor,
|
||||
prompt_embeds_mask: torch.Tensor,
|
||||
negative_prompt_embeds: torch.Tensor | None,
|
||||
negative_prompt_embeds_mask: torch.Tensor | None,
|
||||
img_mask: torch.Tensor,
|
||||
negative_img_mask: torch.Tensor | None,
|
||||
latents: torch.Tensor,
|
||||
image_latents: torch.Tensor | None,
|
||||
img_shapes: list,
|
||||
timesteps: torch.Tensor,
|
||||
do_true_cfg: bool,
|
||||
true_cfg_scale: float,
|
||||
attention_kwargs: dict[str, Any] | None = None,
|
||||
additional_transformer_kwargs: dict[str, Any] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Diffusion loop with optional classifier-free guidance.
|
||||
|
||||
Args:
|
||||
prompt_embeds: Positive prompt embeddings
|
||||
prompt_embeds_mask: Mask for positive prompt
|
||||
negative_prompt_embeds: Negative prompt embeddings
|
||||
negative_prompt_embeds_mask: Mask for negative prompt
|
||||
img_mask: Joint-sequence image-slot mask for the positive branch
|
||||
(VLM image slots plus one slot per 2x2 group of target tokens)
|
||||
negative_img_mask: Same for the negative branch
|
||||
latents: Noise latents to denoise (target image only, packed)
|
||||
image_latents: Packed condition-image latents, prepended to the
|
||||
noise latents (default: None for pure text-to-image)
|
||||
img_shapes: Per-sample list of (frame, height, width) latent-token
|
||||
shapes, condition images first and the target last
|
||||
timesteps: Diffusion timesteps
|
||||
do_true_cfg: Whether to apply CFG
|
||||
true_cfg_scale: CFG scale factor
|
||||
attention_kwargs: Passed through to the transformer (attn_path hint)
|
||||
additional_transformer_kwargs: Extra kwargs for the transformer
|
||||
|
||||
Returns:
|
||||
Denoised latents
|
||||
"""
|
||||
# Fresh cache for every generation: the pipeline is a singleton shared
|
||||
# across requests, so cache state must never leak between them.
|
||||
cache_enabled = getattr(self.transformer, "causal_condition", False)
|
||||
kv_cache = [{} for _ in range(len(self.transformer.transformer_blocks))] if cache_enabled else None
|
||||
self.scheduler.set_begin_index(0)
|
||||
self.transformer.do_true_cfg = do_true_cfg
|
||||
additional_transformer_kwargs = additional_transformer_kwargs or {}
|
||||
|
||||
with self.progress_bar(total=len(timesteps)) as pbar:
|
||||
for i, t in enumerate(timesteps):
|
||||
if self.interrupt:
|
||||
continue
|
||||
self._current_timestep = t
|
||||
|
||||
# Broadcast timestep to match batch size
|
||||
timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
|
||||
|
||||
# Condition image latents go in front of the noise latents.
|
||||
latent_model_input = latents
|
||||
if image_latents is not None:
|
||||
latent_model_input = torch.cat([image_latents, latents], dim=1)
|
||||
|
||||
positive_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"timestep": timestep / 1000,
|
||||
"encoder_hidden_states_mask": prompt_embeds_mask,
|
||||
"encoder_hidden_states": prompt_embeds,
|
||||
"img_shapes": img_shapes,
|
||||
"img_mask": img_mask,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"kv_cache": kv_cache,
|
||||
"cache_branch": "cond",
|
||||
**additional_transformer_kwargs,
|
||||
}
|
||||
if do_true_cfg:
|
||||
negative_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"timestep": timestep / 1000,
|
||||
"encoder_hidden_states_mask": negative_prompt_embeds_mask,
|
||||
"encoder_hidden_states": negative_prompt_embeds,
|
||||
"img_shapes": img_shapes,
|
||||
"img_mask": negative_img_mask,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"kv_cache": kv_cache,
|
||||
"cache_branch": "uncond",
|
||||
**additional_transformer_kwargs,
|
||||
}
|
||||
else:
|
||||
negative_kwargs = None
|
||||
|
||||
# Prefill (first step) returns the full joint sequence and
|
||||
# caches the prefix; decode steps return target tokens only.
|
||||
# Slicing the target tail is correct for both.
|
||||
noise_pred = self.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg,
|
||||
true_cfg_scale,
|
||||
positive_kwargs,
|
||||
negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
output_slice=latents.size(1),
|
||||
)
|
||||
|
||||
latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
|
||||
|
||||
pbar.update()
|
||||
|
||||
return latents
|
||||
@@ -0,0 +1,457 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""CUDA Graph capture for Qwen-Image-2.1 KV-cache decode steps.
|
||||
|
||||
Decode (denoising steps 1..N) recomputes only the target image's tokens: the
|
||||
shape is identical every step and the timestep-independent prefix K/V come from
|
||||
the cross-step KV cache. That is the textbook CUDA-graph case, except the
|
||||
default cache stores prefix K/V in freshly allocated tensors, so a captured
|
||||
graph would bind addresses that the next request (or a step-mode cache merge)
|
||||
no longer owns.
|
||||
|
||||
Requests retain ownership of their prefix K/V. Decode refreshes separate
|
||||
fixed-address graph buffers when the active prefix tensors change. Graphs
|
||||
are keyed by branch, batch, prefix length, complete image layout,
|
||||
dtype, device and backend, so equal token counts with different RoPE layouts stay separate.
|
||||
|
||||
Each key admits a single live owner, tracked by a weak reference to the
|
||||
registering request's prefix K tensor. A second in-flight request whose cache
|
||||
aliases the same key is refused graph decode and stays eager until the
|
||||
owner's cache is released, so shared static buffers are never overwritten
|
||||
under a concurrent owner.
|
||||
|
||||
Memory note: graph entries own an additional copy of the prefix K/V. Decode
|
||||
attention concatenates that prefix with target K/V inside the captured region;
|
||||
transient tensors use the graph memory pool.
|
||||
|
||||
The model runner regionally torch.compile's the transformer blocks even when
|
||||
decode graphs are enabled, so capture records the compiled (fused) kernels;
|
||||
inductor's own cudagraphs stay off so the two graph layers never stack.
|
||||
|
||||
Fallbacks (all logged, all eager): CUDA unavailable, model-level CPU offload
|
||||
without persistent weight staging, sequence/ring/tensor parallelism,
|
||||
HSDP/offload hooks, KV-cache quantization, dynamic LoRA wrappers, padded text
|
||||
masks (the masked attention path branches on mask contents, which cannot be
|
||||
captured), unknown graph key at decode, a second live request aliasing an
|
||||
owned graph key, and capture failure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import weakref
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import (
|
||||
QwenImage21Transformer2DModel,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _dynamic_lora_wrappers_present(module: torch.nn.Module) -> bool:
|
||||
"""True once ``DiffusionLoRAManager`` has wrapped layers under ``module``.
|
||||
|
||||
A captured graph records the layers and branches that ran at capture time;
|
||||
binding or rescaling an adapter afterwards changes the math without
|
||||
changing the shapes the graph was keyed on, so graphs and dynamic LoRA
|
||||
must not mix. Same reasoning as ``sensenova_u1.paged_decode``.
|
||||
"""
|
||||
try:
|
||||
from vllm.lora.layers import BaseLayerWithLoRA
|
||||
except ImportError: # pragma: no cover - depends on the wheel
|
||||
return False
|
||||
return any(isinstance(m, BaseLayerWithLoRA) for m in module.modules())
|
||||
|
||||
|
||||
class QwenImage21DecodeGraphEntry:
|
||||
"""Static state for one captured decode graph.
|
||||
|
||||
Buffers are created outside inference mode so replay-side ``copy_`` is
|
||||
legal regardless of the caller's grad context. ``graph`` is None until the
|
||||
first decode step captures it. Failed entries are replaced by a ``None``
|
||||
marker in the manager so their buffers can be released.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
branch: str,
|
||||
batch_size: int,
|
||||
prefix_len: int,
|
||||
target_tokens: int,
|
||||
in_channels: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
kv_shapes: list[torch.Size],
|
||||
kv_dtype: torch.dtype,
|
||||
freqs: torch.Tensor,
|
||||
):
|
||||
self.branch = branch
|
||||
self.target_tokens = target_tokens
|
||||
self.prefix_len = prefix_len
|
||||
with torch.inference_mode(False):
|
||||
self.hidden = torch.zeros(batch_size, target_tokens, in_channels, dtype=dtype, device=device)
|
||||
self.timestep = torch.zeros(batch_size, dtype=dtype, device=device)
|
||||
self.k = [torch.zeros(shape, dtype=kv_dtype, device=device) for shape in kv_shapes]
|
||||
self.v = [torch.zeros(shape, dtype=kv_dtype, device=device) for shape in kv_shapes]
|
||||
self.target_token_mask = torch.ones(target_tokens, dtype=torch.bool, device=device)
|
||||
# RoPE frequencies are layout-dependent but step-independent. Clone
|
||||
# into a normal (non-inference) tensor: the prefill runs under
|
||||
# inference_mode, and the captured body takes views of this tensor
|
||||
# outside it, which is forbidden for inference tensors.
|
||||
self.freqs = freqs.clone()
|
||||
self.attn_metadata = None
|
||||
# Per-block cache dicts with the same shape as the legacy protocol.
|
||||
self.block_caches = [{branch: {"key": self.k[i], "value": self.v[i]}} for i in range(len(kv_shapes))]
|
||||
self.graph: torch.cuda.CUDAGraph | None = None
|
||||
self.output: torch.Tensor | None = None
|
||||
self.captures = 0
|
||||
self.prefix_sources: list[tuple[weakref.ReferenceType[torch.Tensor], int | None]] = []
|
||||
|
||||
def refresh_prefix(self, kv_cache: list[dict[str, dict[str, torch.Tensor]]]) -> None:
|
||||
sources = [block[self.branch][name] for block in kv_cache for name in ("key", "value")]
|
||||
versions = [None if tensor.is_inference() else tensor._version for tensor in sources]
|
||||
if len(sources) == len(self.prefix_sources) and all(
|
||||
previous() is tensor and version == current_version
|
||||
for (previous, version), tensor, current_version in zip(self.prefix_sources, sources, versions)
|
||||
):
|
||||
return
|
||||
# Prefill replaces inference-mode tensors; decode treats them as immutable.
|
||||
# Weak references avoid retaining completed requests or mistaking a reused
|
||||
# allocation for the same prefix. Normal tensors also track in-place writes.
|
||||
self.prefix_sources = []
|
||||
for i, block in enumerate(kv_cache):
|
||||
self.k[i].copy_(block[self.branch]["key"])
|
||||
self.v[i].copy_(block[self.branch]["value"])
|
||||
self.prefix_sources = [(weakref.ref(tensor), version) for tensor, version in zip(sources, versions)]
|
||||
|
||||
def capture(self, body) -> None:
|
||||
"""Warm up on a side stream, then capture one decode step.
|
||||
|
||||
Warm-up runs the same computation it will record, allocating cuBLAS
|
||||
workspaces and any lazy buffers outside the capture. Decode writes
|
||||
nothing to the cache (prefix K/V are read-only), so warm-up is
|
||||
idempotent.
|
||||
"""
|
||||
side = torch.cuda.Stream()
|
||||
side.wait_stream(torch.cuda.current_stream())
|
||||
try:
|
||||
with torch.cuda.stream(side), torch.inference_mode(False), torch.no_grad():
|
||||
for _ in range(2):
|
||||
body()
|
||||
finally:
|
||||
torch.cuda.current_stream().wait_stream(side)
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
# The platform-wide pool is shared with every other captured path in
|
||||
# the tree; entries only keep their *output* alive, whose tiny size
|
||||
# makes reuse safe because the caller clones it before any other graph
|
||||
# on the pool replays.
|
||||
with torch.inference_mode(False), torch.no_grad():
|
||||
with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
|
||||
self.output = body()
|
||||
self.graph = graph
|
||||
self.captures += 1
|
||||
|
||||
|
||||
class QwenImage21DecodeGraphManager:
|
||||
"""Owns the static KV buffers and captured graphs of one transformer."""
|
||||
|
||||
def __init__(self, model: QwenImage21Transformer2DModel, max_entries: int = 8, model_level_offload: bool = False):
|
||||
self.model = model
|
||||
self.max_entries = max_entries
|
||||
self.model_level_offload = model_level_offload
|
||||
self.entries: OrderedDict[tuple, QwenImage21DecodeGraphEntry | None] = OrderedDict()
|
||||
# Key -> weak reference to the prefix K tensor of the request that owns
|
||||
# the entry's static buffers. Dead references mean the owning request
|
||||
# finished and its cache was released, so the key is free to adopt.
|
||||
self._owners: dict[tuple, weakref.ReferenceType[torch.Tensor]] = {}
|
||||
self._static_eligible: bool | None = None
|
||||
|
||||
# ── eligibility ──
|
||||
|
||||
def _offload_reason(self) -> str | None:
|
||||
"""Reason decode graphs are incompatible with model-level offload, if any.
|
||||
|
||||
Model-level (sequential) offload registers its hook on the top-level
|
||||
transformer module, not on individual blocks. With persistent staging
|
||||
the hook keeps the DiT on fixed device storage and only rebinds
|
||||
``p.data`` around each generation, and the swap completes in
|
||||
``pre_forward`` before capture or replay begin, so captured weight
|
||||
pointers stay valid. Without staging every swap allocates fresh
|
||||
storage and replay would read stale pointers, so capture stays
|
||||
disabled and decode falls back to eager.
|
||||
"""
|
||||
registry = getattr(self.model, "_hook_registry", None)
|
||||
hooks = registry._hooks if registry is not None else {}
|
||||
if hooks:
|
||||
from vllm_omni.diffusion.offloader.sequential_backend import (
|
||||
SequentialOffloadHook,
|
||||
sequential_offload_staging_active,
|
||||
)
|
||||
|
||||
if set(hooks) <= {SequentialOffloadHook._HOOK_NAME} and sequential_offload_staging_active(self.model):
|
||||
return None
|
||||
return "model module carries offload/cache hooks"
|
||||
if self.model_level_offload:
|
||||
return "model-level CPU offload without persistent staging swaps weight storage after capture"
|
||||
return None
|
||||
|
||||
def _check_static_eligibility(self) -> bool:
|
||||
"""One-time checks that cannot change after weights are loaded."""
|
||||
model = self.model
|
||||
reason = self._offload_reason()
|
||||
param = next(model.parameters(), None)
|
||||
if reason is None and (param is None or param.device.type != "cuda"):
|
||||
reason = f"model is not on CUDA (device={None if param is None else param.device})"
|
||||
elif reason is None:
|
||||
parallel_config = getattr(model, "parallel_config", None)
|
||||
if parallel_config is not None:
|
||||
sp = getattr(parallel_config, "sequence_parallel_size", None) or 1
|
||||
if sp > 1 or getattr(parallel_config, "ring_degree", 1) > 1:
|
||||
reason = "sequence/ring parallelism changes decode shapes per rank"
|
||||
elif getattr(parallel_config, "use_hsdp", False):
|
||||
reason = "HSDP wraps blocks with gather hooks outside the captured region"
|
||||
if reason is None:
|
||||
from vllm.distributed import get_tensor_model_parallel_world_size
|
||||
|
||||
if get_tensor_model_parallel_world_size() > 1:
|
||||
reason = "tensor parallelism runs collectives inside the captured region"
|
||||
if reason is None:
|
||||
for block in model.transformer_blocks:
|
||||
if hasattr(block, "_hook_registry") or hasattr(block, "_omni_original_forward"):
|
||||
reason = "transformer blocks carry offload/cache hooks"
|
||||
break
|
||||
if reason is None and model.transformer_blocks:
|
||||
attn_layer = model.transformer_blocks[0].attn.attn
|
||||
if getattr(attn_layer, "_kv_cache_dtype", None) is not None:
|
||||
reason = "KV-cache quantization is enabled"
|
||||
|
||||
if reason is not None:
|
||||
logger.warning_once("Qwen-Image-2.1 CUDA graph decode disabled: %s. Falling back to eager decode.", reason)
|
||||
return False
|
||||
return True
|
||||
|
||||
def eligible(self) -> bool:
|
||||
if self._static_eligible is None:
|
||||
self._static_eligible = self._check_static_eligibility()
|
||||
return self._static_eligible
|
||||
|
||||
def _backend_name(self) -> str:
|
||||
backend = self.model.transformer_blocks[0].attn.attn.attn_backend
|
||||
return backend.get_name() if backend is not None else "custom"
|
||||
|
||||
# ── prefill registration ──
|
||||
|
||||
def register_prefill(
|
||||
self,
|
||||
*,
|
||||
kv_cache: list[dict[str, dict[str, torch.Tensor]]],
|
||||
cache_branch: str,
|
||||
prefix_len: int,
|
||||
img_shapes: list[tuple[int, int, int]],
|
||||
target_freqs: torch.Tensor,
|
||||
joint_key_valid: torch.Tensor | None,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Prepare graph buffers without changing the request's prefix K/V.
|
||||
|
||||
The first in-flight request to register a key owns its static buffers;
|
||||
registrations for the same key from another live request are refused
|
||||
so their decode stays eager until the owner completes.
|
||||
"""
|
||||
if not self.eligible() or torch.compiler.is_compiling():
|
||||
return
|
||||
if _dynamic_lora_wrappers_present(self.model):
|
||||
if self.entries:
|
||||
logger.warning_once(
|
||||
"Qwen-Image-2.1 CUDA graph decode: LoRA wrappers appeared; "
|
||||
"dropping captured graphs and falling back to eager decode."
|
||||
)
|
||||
self.entries.clear()
|
||||
return
|
||||
if joint_key_valid is not None and not bool(joint_key_valid.all()):
|
||||
# The masked decode path (varlen-unpack or 4D mask) branches on mask
|
||||
# contents and cannot be captured; this request stays eager.
|
||||
logger.warning_once(
|
||||
"Qwen-Image-2.1 CUDA graph decode: padded text mask present; this request falls back to eager decode."
|
||||
)
|
||||
return
|
||||
|
||||
first = kv_cache[0].get(cache_branch)
|
||||
if first is None or "key" not in first:
|
||||
return
|
||||
if any(name.endswith("_scale") for name in first):
|
||||
# Quantized (FP8) prefix storage: the static K/V buffers assume a single
|
||||
# dtype and the captured body would not see the dequant scales (a bf16 K +
|
||||
# fp8 V mix would even copy_ fp8 into bf16 buffers, silently dropping the
|
||||
# scales). Keep quantized-cache requests on the eager decode path.
|
||||
logger.warning_once(
|
||||
"Qwen-Image-2.1 CUDA graph decode: quantized prefix KV cache "
|
||||
"(prefix_kv_cache_dtype) is not graph-capturable; falling back to eager decode."
|
||||
)
|
||||
return
|
||||
batch_size = first["key"].shape[0]
|
||||
key = (
|
||||
cache_branch,
|
||||
batch_size,
|
||||
prefix_len,
|
||||
tuple(tuple(shape) for shape in img_shapes),
|
||||
dtype,
|
||||
first["key"].device.index,
|
||||
self._backend_name(),
|
||||
)
|
||||
|
||||
if key not in self.entries:
|
||||
while len(self.entries) >= self.max_entries:
|
||||
evicted_key, evicted_entry = self.entries.popitem(last=False)
|
||||
self._owners.pop(evicted_key, None)
|
||||
del evicted_entry
|
||||
logger.debug("Evicting decode graph entry %s (max_entries=%d)", evicted_key, self.max_entries)
|
||||
allocation_error = None
|
||||
try:
|
||||
self.entries[key] = QwenImage21DecodeGraphEntry(
|
||||
branch=cache_branch,
|
||||
batch_size=batch_size,
|
||||
prefix_len=prefix_len,
|
||||
target_tokens=target_freqs.shape[0],
|
||||
in_channels=self.model.in_channels,
|
||||
dtype=dtype,
|
||||
device=first["key"].device,
|
||||
kv_shapes=[block_cache[cache_branch]["key"].shape for block_cache in kv_cache],
|
||||
kv_dtype=first["key"].dtype,
|
||||
freqs=target_freqs,
|
||||
)
|
||||
except torch.OutOfMemoryError as exc:
|
||||
allocation_error = str(exc)
|
||||
if allocation_error is not None:
|
||||
# Keep a bounded failure marker without retaining tensors or the exception traceback.
|
||||
self.entries[key] = None
|
||||
current_omni_platform.empty_cache()
|
||||
logger.warning(
|
||||
"Qwen-Image-2.1 CUDA graph allocation failed for key=%s: %s. Falling back to eager decode.",
|
||||
key,
|
||||
allocation_error,
|
||||
)
|
||||
return
|
||||
else:
|
||||
self.entries.move_to_end(key)
|
||||
|
||||
entry = self.entries[key]
|
||||
if entry is not None:
|
||||
owner = self._owners.get(key)
|
||||
owner_prefix = owner() if owner is not None else None
|
||||
if owner_prefix is not None and owner_prefix is not first["key"]:
|
||||
# Another in-flight request still owns this entry's static
|
||||
# buffers; refreshing them would overwrite that request's
|
||||
# prefix mid-generation. This request decodes eagerly until
|
||||
# the owner's cache is released and the weak reference dies.
|
||||
logger.warning_once(
|
||||
"Qwen-Image-2.1 CUDA graph decode: entry key=%s is owned by another in-flight "
|
||||
"request; this request falls back to eager decode.",
|
||||
key,
|
||||
)
|
||||
return
|
||||
self._owners[key] = weakref.ref(first["key"])
|
||||
entry.prefix_sources = []
|
||||
|
||||
logger.debug(
|
||||
"Registered decode graph entry key=%s (entries=%d, captured=%d)",
|
||||
key,
|
||||
len(self.entries),
|
||||
sum(1 for e in self.entries.values() if e is not None and e.graph is not None),
|
||||
)
|
||||
|
||||
# ── decode replay ──
|
||||
|
||||
def try_decode(
|
||||
self,
|
||||
*,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
kv_cache: list[dict[str, dict[str, torch.Tensor]]],
|
||||
cache_branch: str,
|
||||
img_shapes: list[tuple[int, int, int]],
|
||||
img_mask: torch.Tensor,
|
||||
encoder_hidden_states_mask: torch.Tensor | None,
|
||||
) -> torch.Tensor | None:
|
||||
"""Replay the captured decode step, or return None to fall back to eager."""
|
||||
if not self.eligible() or torch.compiler.is_compiling():
|
||||
return None
|
||||
if _dynamic_lora_wrappers_present(self.model):
|
||||
self.clear()
|
||||
return None
|
||||
if encoder_hidden_states_mask is not None:
|
||||
text_positions = ~img_mask[0, : encoder_hidden_states_mask.shape[1]]
|
||||
if not bool(encoder_hidden_states_mask[:, text_positions].all()):
|
||||
return None
|
||||
first = kv_cache[0][cache_branch]
|
||||
cached_key = first["key"]
|
||||
if cached_key.device.type != "cuda":
|
||||
# Registration is device-agnostic (unit tests exercise the buffer
|
||||
# machinery on CPU), but capture/replay is CUDA-only: an empty
|
||||
# graph over CPU ops would replay as a stale no-op.
|
||||
return None
|
||||
key = (
|
||||
cache_branch,
|
||||
cached_key.shape[0],
|
||||
cached_key.shape[1],
|
||||
tuple(tuple(shape) for shape in img_shapes),
|
||||
hidden_states.dtype,
|
||||
cached_key.device.index,
|
||||
self._backend_name(),
|
||||
)
|
||||
entry = self.entries.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
self.entries.move_to_end(key)
|
||||
|
||||
owner = self._owners.get(key)
|
||||
owner_prefix = owner() if owner is not None else None
|
||||
if owner_prefix is not None and owner_prefix is not cached_key:
|
||||
# The entry is owned by a different in-flight request (see
|
||||
# register_prefill); this request stays on the eager decode path
|
||||
# until the owner's cache is released.
|
||||
return None
|
||||
|
||||
entry.hidden.copy_(hidden_states[:, -entry.target_tokens :])
|
||||
entry.timestep.copy_(timestep)
|
||||
# Graph buffers are scratch space; request caches must never alias them.
|
||||
entry.refresh_prefix(kv_cache)
|
||||
|
||||
if entry.graph is None:
|
||||
capture_error = None
|
||||
try:
|
||||
entry.capture(lambda: self.model._decode_graph_forward(entry))
|
||||
except Exception as exc:
|
||||
capture_error = str(exc)
|
||||
if capture_error is not None:
|
||||
self.entries[key] = None
|
||||
del entry
|
||||
current_omni_platform.empty_cache()
|
||||
logger.warning(
|
||||
"Qwen-Image-2.1 CUDA graph capture failed for key=%s: %s. "
|
||||
"Falling back to eager decode for this shape.",
|
||||
key,
|
||||
capture_error,
|
||||
)
|
||||
return None
|
||||
logger.info("Captured Qwen-Image-2.1 decode graph for key=%s", key)
|
||||
entry.graph.replay()
|
||||
assert entry.output is not None
|
||||
# Clone before returning: the output lives in the shared graph pool and
|
||||
# the next replay (of this or another graph on the pool) overwrites it.
|
||||
return entry.output.clone()
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Drop all entries, freeing buffers and captured graphs."""
|
||||
self.entries.clear()
|
||||
self._owners.clear()
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -98,6 +98,12 @@ class PinnedModuleStager:
|
||||
A module iterable is treated as one staging group. It uses one copy stream
|
||||
and one reusable completion event. Tensors sharing storage keep their
|
||||
shapes, strides, offsets, dtypes, and aliases across every transition.
|
||||
|
||||
With ``retain_device_storage`` the device storages survive ``offload``:
|
||||
``load`` then only rebinds to the same fixed storages (weights are
|
||||
immutable, so no H2D copy repeats). The stable storage addresses keep
|
||||
CUDA-graph-captured weight pointers valid across offload/load swaps, at
|
||||
the cost of keeping the device storage resident while offloaded.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -108,6 +114,7 @@ class PinnedModuleStager:
|
||||
pin_memory: bool = True,
|
||||
copy_stream: Any | None = None,
|
||||
cache_retention: BoundedAllocatorCache | None = None,
|
||||
retain_device_storage: bool = False,
|
||||
) -> None:
|
||||
modules = (module,) if isinstance(module, nn.Module) else tuple(module)
|
||||
if not modules or not all(isinstance(item, nn.Module) for item in modules):
|
||||
@@ -117,6 +124,7 @@ class PinnedModuleStager:
|
||||
self.copy_stream = copy_stream if copy_stream is not None else current_omni_platform.Stream()
|
||||
self._ready_event = current_omni_platform.Event()
|
||||
self.cache_retention = cache_retention
|
||||
self.retain_device_storage = retain_device_storage
|
||||
self.loaded = False
|
||||
self._groups = self._snapshot_groups(modules, pin_memory=pin_memory)
|
||||
self._device_storages: list[torch.Tensor] = []
|
||||
@@ -234,6 +242,12 @@ class PinnedModuleStager:
|
||||
def load(self) -> None:
|
||||
if self.loaded:
|
||||
return
|
||||
if self._device_storages:
|
||||
# Storage retained from a previous load still holds the immutable
|
||||
# weights; rebinding suffices, no H2D copy repeats.
|
||||
self._bind(self._device_storages)
|
||||
self.loaded = True
|
||||
return
|
||||
try:
|
||||
self._load_once()
|
||||
except torch.OutOfMemoryError:
|
||||
@@ -263,7 +277,8 @@ class PinnedModuleStager:
|
||||
self.loaded = False
|
||||
self._release_cache(force=True)
|
||||
raise
|
||||
self._device_storages.clear()
|
||||
if not self.retain_device_storage:
|
||||
self._device_storages.clear()
|
||||
self.loaded = False
|
||||
self._release_cache()
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload
|
||||
from .module_residency import PinnedModuleStager
|
||||
from .plan_resolver import resolve_offload_plan
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -60,6 +61,7 @@ class SequentialOffloadHook(ModelHook):
|
||||
pin_memory: bool = True,
|
||||
use_hsdp: bool = False,
|
||||
offload_after_context: bool = True,
|
||||
persistent_staging: bool = False,
|
||||
):
|
||||
# Modules to offload to CPU before this module runs
|
||||
self.offload_targets = offload_targets
|
||||
@@ -67,6 +69,21 @@ class SequentialOffloadHook(ModelHook):
|
||||
self.pin_memory = pin_memory
|
||||
self.use_hsdp = use_hsdp
|
||||
self.offload_after_context = offload_after_context
|
||||
self.persistent_staging = persistent_staging
|
||||
self._stager: PinnedModuleStager | None = None
|
||||
|
||||
def initialize_hook(self, module: nn.Module) -> nn.Module:
|
||||
if self.persistent_staging and not self.use_hsdp:
|
||||
# Fixed staging storage keeps CUDA-graph-captured weight pointers
|
||||
# valid across offload/load swaps. HSDP keeps the plain move
|
||||
# semantics (decode graphs refuse HSDP anyway).
|
||||
self._stager = PinnedModuleStager(
|
||||
module,
|
||||
self.device,
|
||||
pin_memory=self.pin_memory,
|
||||
retain_device_storage=True,
|
||||
)
|
||||
return module
|
||||
|
||||
@staticmethod
|
||||
def _move_params(
|
||||
@@ -101,7 +118,23 @@ class SequentialOffloadHook(ModelHook):
|
||||
moved = True
|
||||
return moved
|
||||
|
||||
@staticmethod
|
||||
def _module_stager(module: nn.Module) -> PinnedModuleStager | None:
|
||||
"""The persistent stager owned by ``module``'s own sequential hook, if any.
|
||||
|
||||
Cross-module swaps call ``_to_cpu(target)`` through the *running*
|
||||
module's hook, so staging state must be resolved on the target's hook
|
||||
to keep its ``loaded`` bookkeeping consistent.
|
||||
"""
|
||||
registry: HookRegistry | None = getattr(module, "_hook_registry", None)
|
||||
hook = registry.get_hook(SequentialOffloadHook._HOOK_NAME) if registry is not None else None
|
||||
return hook._stager if isinstance(hook, SequentialOffloadHook) else None
|
||||
|
||||
def _to_cpu(self, module: nn.Module) -> None:
|
||||
stager = self._module_stager(module)
|
||||
if stager is not None:
|
||||
stager.offload()
|
||||
return
|
||||
# XPU's allocator doesn't respect stream dependencies in empty_cache,
|
||||
# so non-blocking copies can race with cache eviction. Use blocking
|
||||
# copies on XPU to avoid NULL pointer errors during DMA.
|
||||
@@ -116,6 +149,10 @@ class SequentialOffloadHook(ModelHook):
|
||||
current_omni_platform.empty_cache()
|
||||
|
||||
def _to_gpu(self, module: nn.Module) -> None:
|
||||
stager = self._module_stager(module)
|
||||
if stager is not None:
|
||||
stager.load()
|
||||
return
|
||||
self._move_params(module, self.device, non_blocking=False)
|
||||
|
||||
def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]:
|
||||
@@ -147,6 +184,7 @@ def apply_sequential_offload(
|
||||
offload_initial_dits: bool = False,
|
||||
offload_dit_modules: Collection[nn.Module] | None = None,
|
||||
offload_encoder_modules: Collection[nn.Module] | None = None,
|
||||
persistent_dit_staging: bool = False,
|
||||
) -> None:
|
||||
"""Apply sequential offloading hooks to DiT and encoder modules.
|
||||
|
||||
@@ -165,6 +203,10 @@ def apply_sequential_offload(
|
||||
every DiT for backward compatibility.
|
||||
offload_encoder_modules: Encoder/stage modules allowed to move to CPU.
|
||||
None selects every supplied module for backward compatibility.
|
||||
persistent_dit_staging: Keep DiT weights on fixed device staging
|
||||
storage across swaps, so CUDA-graph-captured weight pointers stay
|
||||
valid. Encoders keep plain move semantics so their memory is
|
||||
actually freed while the DiT runs.
|
||||
|
||||
Example:
|
||||
>>> apply_sequential_offload(
|
||||
@@ -193,6 +235,7 @@ def apply_sequential_offload(
|
||||
pin_memory=pin_memory,
|
||||
use_hsdp=use_hsdp,
|
||||
offload_after_context=id(dit_mod) in selected_dit_ids,
|
||||
persistent_staging=persistent_dit_staging,
|
||||
)
|
||||
registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook)
|
||||
logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__)
|
||||
@@ -261,6 +304,15 @@ def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook:
|
||||
return hook
|
||||
|
||||
|
||||
def sequential_offload_staging_active(module: nn.Module) -> bool:
|
||||
"""True when the module's sequential offload hook serves fixed device storage.
|
||||
|
||||
Fixed staging storage keeps CUDA-graph-captured weight pointers valid
|
||||
across offload/load swaps.
|
||||
"""
|
||||
return SequentialOffloadHook._module_stager(module) is not None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def sequential_offload_component(module: nn.Module) -> Iterator[None]:
|
||||
"""Activate and release a hooked component called outside ``forward``."""
|
||||
@@ -362,6 +414,14 @@ class ModelLevelOffloadBackend(OffloadBackend):
|
||||
use_hsdp=self.config.use_hsdp,
|
||||
offload_dit_modules=selected_dits,
|
||||
offload_encoder_modules=selected_encoders,
|
||||
# Fixed DiT staging keeps decode-graph weight pointers valid;
|
||||
# only models that capture weight pointers (Qwen-Image-2.1's
|
||||
# CUDA-graph decode) pay for it — every other model keeps
|
||||
# plain move semantics so offloading actually frees DiT VRAM
|
||||
# while the encoders run.
|
||||
persistent_dit_staging=(
|
||||
self.device.type == "cuda" and any(getattr(dit, "enable_cuda_graph_decode", False) for dit in dits)
|
||||
),
|
||||
)
|
||||
except BaseException:
|
||||
try:
|
||||
|
||||
@@ -56,6 +56,11 @@ _DIFFUSION_MODELS = {
|
||||
"pipeline_qwen_image_layered",
|
||||
"QwenImageLayeredPipeline",
|
||||
),
|
||||
"QwenImage21Pipeline": (
|
||||
"qwen_image_21",
|
||||
"pipeline_qwen_image_21",
|
||||
"QwenImage21Pipeline",
|
||||
),
|
||||
"GlmImagePipeline": (
|
||||
"glm_image",
|
||||
"pipeline_glm_image",
|
||||
@@ -426,6 +431,10 @@ _NO_CACHE_ACCELERATION = {
|
||||
"Pi0Pipeline",
|
||||
"Pi05Pipeline",
|
||||
"LingBotWorldCausalDMDPipeline",
|
||||
# Qwen-Image 2.1's transformer carries its own prefix KV cache (keyed by CFG
|
||||
# branch) across denoising steps, which conflicts with cache_dit / tea_cache
|
||||
# step-skipping hooks.
|
||||
"QwenImage21Pipeline",
|
||||
}
|
||||
|
||||
|
||||
@@ -651,6 +660,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
|
||||
"AnimaPipeline": "get_anima_post_process_func",
|
||||
"QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
|
||||
"QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func",
|
||||
"QwenImage21Pipeline": "get_qwen_image_21_post_process_func",
|
||||
"GlmImagePipeline": "get_glm_image_post_process_func",
|
||||
"ZImagePipeline": "get_post_process_func",
|
||||
"OvisImagePipeline": "get_ovis_image_post_process_func",
|
||||
@@ -740,6 +750,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
|
||||
"BooguImageTurboPipeline": "get_boogu_image_pre_process_func",
|
||||
"QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
|
||||
"QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func",
|
||||
"QwenImage21Pipeline": "get_qwen_image_21_pre_process_func",
|
||||
"LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func",
|
||||
"LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func",
|
||||
"QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func",
|
||||
|
||||
@@ -76,6 +76,8 @@ class OmniDiffusionRequest:
|
||||
# A model preprocessor may keep selected requests on the legacy full-forward
|
||||
# path even when the engine is globally configured for step execution.
|
||||
use_step_execution: bool = True
|
||||
# Whether first and later denoising steps may share a step batch.
|
||||
allow_mixed_step_phases: bool = True
|
||||
# KV-recv wall-clock (ms), set by the runner's _prepare_request_for_forward
|
||||
# and carried to DiffusionOutput for the vllm_omni:diffusion_kv_load_s metric.
|
||||
kv_recv_ms: float = 0.0
|
||||
|
||||
@@ -13,6 +13,7 @@ from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
|
||||
from vllm_omni.diffusion.sched.interface import (
|
||||
DiffusionRequestStatus,
|
||||
DiffusionSchedulerOutput,
|
||||
SchedulerRequestState,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -62,6 +63,19 @@ class StepScheduler(BaseScheduler):
|
||||
)
|
||||
return request_id
|
||||
|
||||
def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool:
|
||||
if not super()._can_schedule_waiting(state):
|
||||
return False
|
||||
|
||||
is_first_step = self._request_progress[state.request_id].current_step == 0
|
||||
for request_id in self._running:
|
||||
running = self._request_states[request_id]
|
||||
if state.req.allow_mixed_step_phases and running.req.allow_mixed_step_phases:
|
||||
continue
|
||||
if (self._request_progress[request_id].current_step == 0) != is_first_step:
|
||||
return False
|
||||
return True
|
||||
|
||||
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: BaseRunnerOutput) -> set[str]:
|
||||
scheduled_request_ids = sched_output.scheduled_request_ids
|
||||
if not scheduled_request_ids and not sched_output.finished_req_ids:
|
||||
|
||||
@@ -206,7 +206,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
|
||||
self.kv_transfer_manager = (
|
||||
payload_transfer_manager if getattr(od_config, "kv_transfer_config", None) is None else None
|
||||
)
|
||||
self.init_omni_connectors(od_config, payload_transfer_manager, synchronous=True)
|
||||
self.init_omni_connectors(od_config, payload_transfer_manager, synchronous=True) # type: ignore[arg-type]
|
||||
self._kv_connector = None
|
||||
from vllm_omni.diffusion.diffusion_kv.kv_connector import KVReceiveProgress, native_prefetch_enabled
|
||||
|
||||
@@ -250,14 +250,25 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
|
||||
if model is None:
|
||||
return
|
||||
|
||||
compile_kwargs: dict[str, Any] = {"dynamic": self.od_config.diffusion_compile_dynamic}
|
||||
if getattr(model, "enable_cuda_graph_decode", False):
|
||||
# Decode-graph models still compile their blocks: graph capture
|
||||
# records the compiled (fused) kernels. Scope inductor cudagraphs
|
||||
# to this compile so the two graph layers never stack, without
|
||||
# changing torch.compile for every other model in the process.
|
||||
compile_kwargs["options"] = {
|
||||
"triton.cudagraphs": False,
|
||||
"triton.cudagraph_trees": False,
|
||||
}
|
||||
logger.info("Model runner: %s combines CUDA graph decode with torch.compile.", attr_name)
|
||||
|
||||
compile_granularity = self.od_config.diffusion_compile_granularity
|
||||
compile_dynamic = self.od_config.diffusion_compile_dynamic
|
||||
try:
|
||||
if compile_granularity == "full":
|
||||
model.compile(dynamic=compile_dynamic)
|
||||
model.compile(**compile_kwargs)
|
||||
compiled_model = model
|
||||
else:
|
||||
compiled_model = regionally_compile(model, dynamic=compile_dynamic)
|
||||
compiled_model = regionally_compile(model, **compile_kwargs)
|
||||
setattr(self.pipeline, attr_name, compiled_model)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
@@ -275,7 +286,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
|
||||
"compilation errors may surface on the first request.",
|
||||
attr_name,
|
||||
compile_granularity,
|
||||
compile_dynamic,
|
||||
compile_kwargs["dynamic"],
|
||||
)
|
||||
|
||||
def load_model(
|
||||
@@ -1454,7 +1465,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
|
||||
else req.denoise_completed
|
||||
)
|
||||
if finished and result is not None:
|
||||
self._maybe_send_stage_payload([req], [result])
|
||||
self._maybe_send_stage_payload([req], [result]) # type: ignore[list-item]
|
||||
runner_output_list.append(
|
||||
RunnerOutput(
|
||||
request_id=req.request_id,
|
||||
|
||||
@@ -585,6 +585,10 @@ class OrchestratorArgs:
|
||||
diffusion_attention_config: str | None = None
|
||||
diffusion_compile_granularity: str | None = None
|
||||
diffusion_compile_dynamic: bool | None = None
|
||||
# CUDA graph capture of fixed-shape KV-cache decode steps (Qwen-Image-2.1
|
||||
# today). None defers to the OmniDiffusionConfig default (enabled);
|
||||
# --enforce-eager also disables it.
|
||||
enable_cuda_graph_decode: bool | None = None
|
||||
cache_backend: str = "none"
|
||||
cache_config: str | None = None
|
||||
video_output_transport: dict[str, object] | None = None
|
||||
|
||||
@@ -614,6 +614,16 @@ class OmniServeCommand(CLISubcommand):
|
||||
"Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
|
||||
),
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
"--enable-cuda-graph-decode",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=None,
|
||||
help=(
|
||||
"Capture fixed-shape KV-cache decode steps into CUDA graphs on models that implement them "
|
||||
"(currently Qwen-Image-2.1). Independent of compilation_config.cudagraph_mode; "
|
||||
"disable with --no-enable-cuda-graph-decode. --enforce-eager also disables it."
|
||||
),
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
"--fa-deterministic",
|
||||
dest="fa_deterministic",
|
||||
|
||||
@@ -2279,6 +2279,7 @@ async def generate_images(
|
||||
|
||||
@router.post(
|
||||
"/v1/images/edits",
|
||||
response_model=None,
|
||||
responses={
|
||||
HTTPStatus.OK.value: {"model": ImageGenerationResponse},
|
||||
HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
|
||||
@@ -2286,6 +2287,7 @@ async def generate_images(
|
||||
HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
|
||||
},
|
||||
)
|
||||
@with_cancellation
|
||||
async def edit_images(
|
||||
raw_request: Request,
|
||||
image: list[UploadFile] | None = File(None),
|
||||
|
||||
@@ -473,7 +473,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
nested = getattr(request, "extra_body", None)
|
||||
explicit = getattr(request, "model_fields_set", None)
|
||||
if explicit is None:
|
||||
explicit = getattr(request, "__fields_set__", ())
|
||||
explicit = getattr(request, "__fields_set__", None) or ()
|
||||
explicit_root_args = {
|
||||
key: getattr(request, key) for key in explicit if key != "extra_body" and hasattr(request, key)
|
||||
}
|
||||
@@ -540,7 +540,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
def _truthy_extra_body_flag(request: Any, key: str) -> bool:
|
||||
if isinstance(request, dict):
|
||||
extra_body = request
|
||||
model_extra = {}
|
||||
model_extra: dict[str, Any] = {}
|
||||
else:
|
||||
extra_body = getattr(request, "extra_body", None) or {}
|
||||
model_extra = getattr(request, "model_extra", None) or {}
|
||||
@@ -778,7 +778,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
|
||||
# Override the prompts produced by chat-template preprocessing.
|
||||
is_img2img = engine_prompt_image is not None
|
||||
tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
|
||||
tprompt: OmniTextPrompt = OmniTextPrompt(prompt=extracted_prompt)
|
||||
if is_img2img:
|
||||
tprompt["modalities"] = ["img2img"]
|
||||
else:
|
||||
@@ -833,19 +833,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
if not extracted_prompt:
|
||||
return self.create_error_response("No text prompt found in messages")
|
||||
|
||||
tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
|
||||
tprompt["modalities"] = ["text"]
|
||||
text_prompt: OmniTextPrompt = OmniTextPrompt(prompt=extracted_prompt, modalities=["text"])
|
||||
|
||||
if reference_images:
|
||||
try:
|
||||
img_bytes = base64.b64decode(reference_images[0])
|
||||
img = Image.open(BytesIO(img_bytes))
|
||||
tprompt["multi_modal_data"] = {"image": img}
|
||||
tprompt["multi_modal_uuids"] = {"image": [f"{request_id}-image-0"]}
|
||||
text_prompt["multi_modal_data"] = {"image": img}
|
||||
text_prompt["multi_modal_uuids"] = {"image": [f"{request_id}-image-0"]}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
engine_prompts = [tprompt]
|
||||
engine_prompts = [text_prompt]
|
||||
except Exception as e:
|
||||
logger.warning("Failed to build text-output prompt for single-stage diffusion: %s", e)
|
||||
_image_gen_height = None
|
||||
@@ -1422,7 +1421,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
# Pydantic v2 uses `model_fields_set`; keep v1 fallback for compatibility.
|
||||
explicit_fields = getattr(request, "model_fields_set", None)
|
||||
if explicit_fields is None:
|
||||
explicit_fields = getattr(request, "__fields_set__", set())
|
||||
explicit_fields = getattr(request, "__fields_set__", None) or set()
|
||||
|
||||
for field_name in self._OPENAI_SAMPLING_FIELDS:
|
||||
if field_name not in explicit_fields:
|
||||
@@ -2175,6 +2174,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
# finish_reason is:
|
||||
# "tool_calls" for "auto" or "required" tool calls,
|
||||
# and "stop" for named tool calls.
|
||||
finish_reason_: str | None
|
||||
if (
|
||||
auto_tools_called
|
||||
or (tools_streamed[i] and not tool_choice_function_name)
|
||||
@@ -2626,7 +2626,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
if output_text:
|
||||
# Get the corresponding output token IDs
|
||||
output_token_ids = None
|
||||
if choice.index < len(final_res.outputs):
|
||||
if final_res is not None and choice.index < len(final_res.outputs):
|
||||
output_token_ids = final_res.outputs[choice.index].token_ids
|
||||
|
||||
self.request_logger.log_outputs(
|
||||
@@ -3234,7 +3234,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
resolve_stop_token_ids,
|
||||
)
|
||||
|
||||
build_kwargs: dict[str, Any] = {
|
||||
build_kwargs = {
|
||||
"task": "it2i" if reference_images else "t2i",
|
||||
"sys_type": use_system_prompt,
|
||||
"custom_system_prompt": custom_system_prompt,
|
||||
@@ -3282,7 +3282,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
**ar_stop_kwargs,
|
||||
)
|
||||
|
||||
engine_prompt: OmniTextPrompt = {"prompt": prompt}
|
||||
engine_prompt: OmniTextPrompt = OmniTextPrompt(prompt=prompt)
|
||||
if prompt_token_ids is not None:
|
||||
engine_prompt["prompt_token_ids"] = prompt_token_ids
|
||||
if system_prompt_type is not None:
|
||||
@@ -3327,7 +3327,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
# AR stop tokens: use stage_type=="llm" instead of comprehension_idx
|
||||
# (None for DictConfig where is_comprehension is nested in engine_args).
|
||||
if stage_type == "llm" and ar_stop_token_ids is not None:
|
||||
default_stage_params.stop_token_ids = ar_stop_token_ids
|
||||
setattr(default_stage_params, "stop_token_ids", ar_stop_token_ids)
|
||||
|
||||
if stage_type == "llm" and idx == 0 and ar_prompt_info is not None:
|
||||
if should_init_extra_args_for_non_diffusion_stages(model_class_name):
|
||||
@@ -3480,11 +3480,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
except Exception as e: # pragma: no cover - safeguard
|
||||
logger.warning("Failed to parse LoRA request: %s", e)
|
||||
|
||||
gen_prompt: OmniTextPrompt = {
|
||||
"prompt": prompt,
|
||||
"negative_prompt": negative_prompt,
|
||||
"modalities": ["image"],
|
||||
}
|
||||
gen_prompt: OmniTextPrompt = OmniTextPrompt(
|
||||
prompt=prompt,
|
||||
modalities=["image"],
|
||||
)
|
||||
if negative_prompt is not None:
|
||||
gen_prompt["negative_prompt"] = negative_prompt
|
||||
if pil_images:
|
||||
if len(pil_images) == 1:
|
||||
gen_prompt["multi_modal_data"] = {"image": pil_images[0]}
|
||||
@@ -3612,20 +3613,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
# ``images`` field, so fall back to the shared extractor used by the
|
||||
# offline example path before giving up.
|
||||
images = getattr(result, "images", None) or extract_images_from_outputs(result)
|
||||
stage_durations = result.stage_durations
|
||||
peak_memory_mb = result.peak_memory_mb
|
||||
stage_durations = getattr(result, "stage_durations", None)
|
||||
peak_memory_mb = getattr(result, "peak_memory_mb", None)
|
||||
response_metrics = getattr(result, "metrics", None) if return_stage_metrics else None
|
||||
cot_output = None
|
||||
|
||||
req_out = result
|
||||
if req_out:
|
||||
prompt_obj = getattr(req_out, "prompt", None)
|
||||
if isinstance(prompt_obj, dict):
|
||||
extra = prompt_obj.get("extra", {})
|
||||
if isinstance(extra, dict):
|
||||
ar_text = extra.get("ar_generated_text")
|
||||
if isinstance(ar_text, str) and ar_text.strip():
|
||||
cot_output = ar_text
|
||||
prompt_obj = getattr(result, "prompt", None)
|
||||
if isinstance(prompt_obj, dict):
|
||||
extra = prompt_obj.get("extra", {})
|
||||
if isinstance(extra, dict):
|
||||
ar_text = extra.get("ar_generated_text")
|
||||
if isinstance(ar_text, str) and ar_text.strip():
|
||||
cot_output = ar_text
|
||||
|
||||
return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output, response_metrics
|
||||
|
||||
@@ -3653,14 +3651,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
text = completion.text or ""
|
||||
if not text:
|
||||
continue
|
||||
chunk = ImageEditARDeltaChunk(
|
||||
delta_chunk = ImageEditARDeltaChunk(
|
||||
delta=text,
|
||||
index=completion.index,
|
||||
created=created,
|
||||
model=model,
|
||||
metrics=metrics,
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
yield f"data: {delta_chunk.model_dump_json()}\n\n"
|
||||
elif final_output_type == "image":
|
||||
images = self._flatten_diffusion_images(getattr(output, "images", []))
|
||||
if not images:
|
||||
@@ -3676,15 +3674,25 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
)
|
||||
for img in images
|
||||
]
|
||||
chunk = ImageEditImageChunk(
|
||||
actual_size = f"{images[0].width}x{images[0].height}"
|
||||
# Echo the request size only when it already matches the
|
||||
# decoded pixels. Use actual WxH for size=auto/None or when
|
||||
# the pipeline diverged (e.g. Qwen-Image-2.1 aspect-derived
|
||||
# outputs). Edit models that honor an explicit WxH keep the
|
||||
# same wire string because size == actual_size.
|
||||
if size in (None, "auto") or size != actual_size:
|
||||
chunk_size = actual_size
|
||||
else:
|
||||
chunk_size = size
|
||||
img_chunk = ImageEditImageChunk(
|
||||
data=image_data,
|
||||
output_format=output_format,
|
||||
size=size,
|
||||
size=chunk_size,
|
||||
created=created,
|
||||
model=model,
|
||||
metrics=metrics,
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
yield f"data: {img_chunk.model_dump_json()}\n\n"
|
||||
emitted_image = True
|
||||
if not emitted_image:
|
||||
raise RuntimeError("Streaming image edit completed without a final image output.")
|
||||
@@ -3704,7 +3712,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
)
|
||||
except OmniClientError as exc:
|
||||
logger.info("Client error during streaming image edit: %s", exc)
|
||||
chunk = ImageEditStreamError(
|
||||
client_error_chunk = ImageEditStreamError(
|
||||
created=created,
|
||||
model=model,
|
||||
error={
|
||||
@@ -3713,10 +3721,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
"code": exc.status_code,
|
||||
},
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
yield f"data: {client_error_chunk.model_dump_json()}\n\n"
|
||||
except Exception as exc:
|
||||
logger.exception("Streaming image edit failed: %s", exc)
|
||||
chunk = ImageEditStreamError(
|
||||
server_error_chunk = ImageEditStreamError(
|
||||
created=created,
|
||||
model=model,
|
||||
error={
|
||||
@@ -3725,7 +3733,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
"code": 500,
|
||||
},
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
yield f"data: {server_error_chunk.model_dump_json()}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
@staticmethod
|
||||
@@ -3832,11 +3840,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
logger.warning("Failed to decode reference image: %s", e)
|
||||
|
||||
# Build generation kwargs
|
||||
gen_prompt: OmniTextPrompt = {
|
||||
"prompt": prompt,
|
||||
"negative_prompt": negative_prompt,
|
||||
"modalities": ["image"],
|
||||
}
|
||||
gen_prompt: OmniTextPrompt = OmniTextPrompt(
|
||||
prompt=prompt,
|
||||
modalities=["image"],
|
||||
)
|
||||
if negative_prompt is not None:
|
||||
gen_prompt["negative_prompt"] = negative_prompt
|
||||
gen_params = OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
@@ -4062,6 +4071,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
)
|
||||
|
||||
# Build response
|
||||
content: str | list[dict[str, Any]]
|
||||
if not image_contents:
|
||||
content = "Image generation completed but no images were produced."
|
||||
else:
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def build_image_to_image_prompt(
|
||||
prompt: str,
|
||||
negative_prompt: str | None,
|
||||
input_image: Image.Image | list[Image.Image],
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
img_prompt: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": input_image},
|
||||
}
|
||||
if height is not None:
|
||||
img_prompt["height"] = height
|
||||
if width is not None:
|
||||
img_prompt["width"] = width
|
||||
if negative_prompt is not None:
|
||||
img_prompt["negative_prompt"] = negative_prompt
|
||||
return img_prompt
|
||||
@@ -91,6 +91,9 @@ from vllm_omni.model_extras.ming_image import (
|
||||
from vllm_omni.model_extras.ming_image import (
|
||||
build_text_to_image_prompt as build_ming_image_text_to_image_prompt,
|
||||
)
|
||||
from vllm_omni.model_extras.qwen_image_21 import (
|
||||
build_image_to_image_prompt as build_qwen_image_21_image_to_image_prompt,
|
||||
)
|
||||
from vllm_omni.model_extras.sana_video import SANA_VIDEO_EXTRA_BODY_PARAMS
|
||||
from vllm_omni.model_extras.sensenova_u1 import (
|
||||
SENSENOVA_U1_EXTRA_BODY_PARAMS,
|
||||
@@ -329,6 +332,9 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
|
||||
"text_to_image_prompt_builder": build_ming_image_text_to_image_prompt,
|
||||
"image_to_image_prompt_builder": build_ming_image_image_to_image_prompt,
|
||||
},
|
||||
"QwenImage21Pipeline": {
|
||||
"image_to_image_prompt_builder": build_qwen_image_21_image_to_image_prompt,
|
||||
},
|
||||
}
|
||||
_EXTRA_SPECS["MingImageLayeredDiffusionPipeline"] = _EXTRA_SPECS["MingImageDiffusionPipeline"]
|
||||
|
||||
@@ -473,10 +479,10 @@ def build_text_to_image_prompt(
|
||||
if builder is None:
|
||||
return prompt
|
||||
return builder(
|
||||
prompt=str(prompt["prompt"]),
|
||||
negative_prompt=prompt.get("negative_prompt"),
|
||||
height=height,
|
||||
width=width,
|
||||
str(prompt["prompt"]),
|
||||
prompt.get("negative_prompt"),
|
||||
height,
|
||||
width,
|
||||
)
|
||||
|
||||
|
||||
@@ -513,10 +519,10 @@ def build_image_to_video_prompt(
|
||||
if not isinstance(media_inputs, Mapping):
|
||||
raise TypeError("Canonical I2V prompt multi_modal_data must be a mapping.")
|
||||
return builder(
|
||||
prompt=str(prompt["prompt"]),
|
||||
negative_prompt=prompt.get("negative_prompt"),
|
||||
media_inputs=media_inputs,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
str(prompt["prompt"]),
|
||||
prompt.get("negative_prompt"),
|
||||
media_inputs,
|
||||
height,
|
||||
width,
|
||||
num_frames,
|
||||
)
|
||||
Reference in new issue
Block a user