docs: split diffusion feature guides
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed
1 parent
8b0a6779d4
commit
7872663bf0
21 files changed
+1138
-801
No files matched your search
@@ -1,494 +1,192 @@
|
||||
# Diffusion Attention Backends
|
||||
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
|
||||
Use this page to select and configure a diffusion attention backend. Backend-
|
||||
specific installation, tuning, and compatibility details live in separate
|
||||
guides so that the selection contract stays easy to scan.
|
||||
|
||||
## Overview
|
||||
Diffusion attention backend selection applies to DiT attention in image and
|
||||
video generation models. It does **not** change autoregressive LLM attention,
|
||||
which uses vLLM's attention selector.
|
||||
|
||||
Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
|
||||
For the internal selector, registry, and platform contract, see
|
||||
[Attention Backend Selection](../../design/feature/attention_backend_selection.md).
|
||||
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
|
||||
## Choose a guide
|
||||
|
||||
The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
|
||||
| Need | Guide |
|
||||
| --- | --- |
|
||||
| Select a conservative or platform-native dense kernel | [Dense Backends](attention_backends/dense_backends.md) |
|
||||
| Use FlashInfer trtllm-gen FMHA, Skip-Softmax, or TRTLLM SAGE quantization | [TRTLLM Attention](attention_backends/trtllm.md) |
|
||||
| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
|
||||
| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
|
||||
|
||||
## Backend Options
|
||||
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
|
||||
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
|
||||
| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
|
||||
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
|
||||
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
|
||||
| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
|
||||
| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
|
||||
| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
|
||||
| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |
|
||||
## Backend options
|
||||
|
||||
| Value | Family | Primary use | Detail |
|
||||
| --- | --- | --- | --- |
|
||||
| `TORCH_SDPA` | Dense | Conservative reference; always available | [Dense Backends](attention_backends/dense_backends.md#torch_sdpa) |
|
||||
| `FLASH_ATTN` | Dense | FlashAttention 4/3/2 depending on the installed package and GPU | [Dense Backends](attention_backends/dense_backends.md#flash_attn) |
|
||||
| `CUDNN_ATTN` | Dense | Mask-heavy DiTs on Blackwell with cuDNN 9.5 or newer | [Dense Backends](attention_backends/dense_backends.md#cudnn_attn) |
|
||||
| `FLASHINFER_ATTN` | Dense or quantized | FlashInfer batch prefill; optional mixed Q/K and V dtypes | [Dense Backends](attention_backends/dense_backends.md#flashinfer_attn) |
|
||||
| `TRTLLM_ATTN` | Dense, sparse, or quantized | Datacenter Blackwell with `head_dim=128` and compatible packed paths | [TRTLLM Attention](attention_backends/trtllm.md) |
|
||||
| `SAGE_ATTN` | Quantized | SageAttention 2.2 INT8 attention | [SageAttention](attention_backends/sage.md#sage_attn) |
|
||||
| `SAGE_ATTN_3` | Quantized | SageAttention3 on Blackwell | [SageAttention](attention_backends/sage.md#sage_attn_3) |
|
||||
| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
|
||||
|
||||
## Configuration
|
||||
|
||||
Diffusion attention backends can be configured three ways, in priority order:
|
||||
Backend selection is resolved in this order:
|
||||
|
||||
1. **`--diffusion-attention-config`** — structured per-role config (highest priority).
|
||||
2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend.
|
||||
3. **Platform default** — used when nothing is configured.
|
||||
1. `--diffusion-attention-config` per-role configuration.
|
||||
2. `--diffusion-attention-backend` or `DIFFUSION_ATTENTION_BACKEND` as a
|
||||
global default.
|
||||
3. The current platform's default.
|
||||
|
||||
`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`.
|
||||
`--diffusion-attention-backend` is shorthand for
|
||||
`--diffusion-attention-config.default.backend`. Do not pass it together with an
|
||||
explicit `default.backend` in the structured configuration.
|
||||
|
||||
### Global default
|
||||
|
||||
Set the default backend for every diffusion attention layer:
|
||||
|
||||
```bash
|
||||
# CLI flag
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
|
||||
|
||||
# Environment variable (also recognized for backwards compatibility)
|
||||
export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
|
||||
# Backwards-compatible environment variable
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
```
|
||||
|
||||
### Per-role configuration
|
||||
|
||||
Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order:
|
||||
Roles are declared by each diffusion model. Common categories are `self` and
|
||||
`cross`; a model may also use a more specific role such as
|
||||
`ltx2.audio_to_video`. Resolution order is:
|
||||
|
||||
1. Exact `per_role[role]` match
|
||||
2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`)
|
||||
3. `default`
|
||||
4. Platform default
|
||||
|
||||
Use vLLM-style dotted flags or one JSON blob:
|
||||
1. Exact `per_role[role]` match.
|
||||
2. Category `per_role[role_category]` match.
|
||||
3. `default`.
|
||||
4. Platform default.
|
||||
|
||||
```bash
|
||||
# Dotted flags
|
||||
vllm-omni serve <model> \
|
||||
--diffusion-attention-config.default.backend FLASH_ATTN \
|
||||
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
|
||||
--diffusion-attention-config.default.backend FLASH_ATTN \
|
||||
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
|
||||
|
||||
# JSON
|
||||
# Equivalent JSON
|
||||
vllm-omni serve <model> \
|
||||
--diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
|
||||
--diffusion-attention-config \
|
||||
'{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
|
||||
```
|
||||
|
||||
A backend that needs configuration exposes it as a typed field on the spec:
|
||||
`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
|
||||
`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
|
||||
|
||||
```bash
|
||||
--diffusion-attention-config.default.backend TRTLLM_ATTN \
|
||||
--diffusion-attention-config.default.skip_softmax.target_sparsity 0.5
|
||||
```
|
||||
|
||||
### Programmatic API
|
||||
|
||||
When constructing `OmniDiffusionConfig` directly:
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.data import (
|
||||
AttentionConfig,
|
||||
AttentionSpec,
|
||||
OmniDiffusionConfig,
|
||||
)
|
||||
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(backend="FLASH_ATTN"),
|
||||
per_role={
|
||||
"cross": AttentionSpec(backend="TORCH_SDPA"),
|
||||
},
|
||||
per_role={"cross": AttentionSpec(backend="TORCH_SDPA")},
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
A plain dict is also accepted and normalized to `AttentionConfig`.
|
||||
Backend-specific typed blocks are documented with their consumers:
|
||||
|
||||
#### FlashInfer QK16/V8 quantized attention
|
||||
- `quant`: [FlashInfer](attention_backends/dense_backends.md#flashinfer-quantized-attention)
|
||||
and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
|
||||
- `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
|
||||
- `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
|
||||
|
||||
`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
|
||||
|
||||
```python
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="FLASHINFER_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="bfloat16",
|
||||
dtype_vo="fp8_e4m3",
|
||||
),
|
||||
),
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
|
||||
|
||||
## Platform Defaults
|
||||
## Platform defaults
|
||||
|
||||
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
|
||||
|
||||
Auto-route preference, in order:
|
||||
The CUDA auto-route preference is:
|
||||
|
||||
1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path**
|
||||
2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
|
||||
3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
|
||||
4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
|
||||
5. `TORCH_SDPA` — last resort
|
||||
1. `TRTLLM_ATTN` on datacenter Blackwell (sm_100/sm_103) when FlashInfer is
|
||||
available, `head_dim=128`, and the model declares a compatible packed or
|
||||
mask-free path.
|
||||
2. `CUDNN_ATTN` when cuDNN 9.5 or newer is available.
|
||||
3. `FLASHINFER_ATTN` when FlashInfer is available but cuDNN is too old.
|
||||
4. `FLASH_ATTN` when a compatible package is installed.
|
||||
5. `TORCH_SDPA`.
|
||||
|
||||
`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.
|
||||
`TRTLLM_ATTN` is not auto-selected on workstation Blackwell
|
||||
(sm_120/sm_121), for other head dimensions, or for paths that require an
|
||||
unsupported mask.
|
||||
|
||||
`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback.
|
||||
### Hopper, Ada, and Ampere
|
||||
|
||||
The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
|
||||
The CUDA auto-route uses `FLASH_ATTN` when available and otherwise falls back
|
||||
to `TORCH_SDPA`. `CUDNN_ATTN` and `FLASHINFER_ATTN` remain explicit options.
|
||||
|
||||
**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
|
||||
Other platforms validate an explicit backend and choose their own default
|
||||
through the platform implementation. Check the startup log to confirm the
|
||||
resolved backend.
|
||||
|
||||
### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
|
||||
## Choosing a backend manually
|
||||
|
||||
Auto-route preference:
|
||||
Override the platform default when you need:
|
||||
|
||||
1. `FLASH_ATTN` — when `flash-attn` is installed
|
||||
2. `TORCH_SDPA` — fallback
|
||||
- a correctness reference (`TORCH_SDPA`);
|
||||
- a backend-specific workaround;
|
||||
- training/rollout kernel alignment (Hub backends); or
|
||||
- an explicitly validated sparse or quantized speedup.
|
||||
|
||||
`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
|
||||
The startup log prints the resolved backend and whether it came from explicit
|
||||
configuration or platform defaulting. If no resolution message appears, check
|
||||
earlier logs for diffusion-stage initialization failures.
|
||||
|
||||
## TRTLLM_ATTN Backend and Skip-Softmax
|
||||
## Reference benchmark
|
||||
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of
|
||||
that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed
|
||||
(algorithm: [design doc](../../design/feature/skip_softmax.md)).
|
||||
The following BF16 results were measured on an sm_120 RTX Pro 6000 Blackwell
|
||||
with the same prompt and seed across runs. Treat them as reference results, not
|
||||
portable guarantees.
|
||||
|
||||
Enable it through the typed `skip_softmax` block on the attention spec:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. |
|
||||
| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
|
||||
| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
|
||||
"skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}'
|
||||
```
|
||||
|
||||
Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction):
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only
|
||||
as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with
|
||||
sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no
|
||||
`skip_softmax` block.
|
||||
|
||||
Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
|
||||
`TRTLLM_ATTN` raises rather than silently degrading.
|
||||
|
||||
## TRTLLM_ATTN SAGE Quantization
|
||||
|
||||
On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
|
||||
FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
|
||||
in low precision. Enable it through the typed `quant` block on the attention spec:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
|
||||
| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
|
||||
| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
|
||||
|
||||
V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
|
||||
neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
|
||||
|
||||
The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
|
||||
fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
|
||||
`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
|
||||
`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
|
||||
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
|
||||
availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
|
||||
(B300); `int8` QK kernels are compiled for **SM100 only**.
|
||||
|
||||
## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
|
||||
|
||||
`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
|
||||
128-token block of keys, ranks them per query block, and computes attention against only the
|
||||
top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
|
||||
spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
|
||||
|
||||
Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
|
||||
first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
|
||||
nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
|
||||
exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
|
||||
delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
|
||||
|
||||
Enable it through the typed `block_sparse` block on the attention spec, shared by every
|
||||
block-sparse backend:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
|
||||
| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
|
||||
| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
|
||||
|
||||
```bash
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
|
||||
"block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
|
||||
```
|
||||
|
||||
Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
|
||||
cheapest way to recover structure, because the early high-noise steps decide global layout while the
|
||||
later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
|
||||
`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
|
||||
|
||||
Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
|
||||
incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
|
||||
blocks — use Ulysses SP (`ring_degree=1`).
|
||||
|
||||
### Which geometries run sparse
|
||||
|
||||
MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
|
||||
necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
|
||||
does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
|
||||
promoted rows remain visible to every query and key.
|
||||
|
||||
For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
|
||||
128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
|
||||
MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
|
||||
and kernel tiling aligned, at the cost of retaining more blocks.
|
||||
|
||||
8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
|
||||
creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
|
||||
grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
|
||||
that implements protected video tails for rf_v2.
|
||||
|
||||
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
|
||||
|
||||
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
|
||||
|
||||
| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
|
||||
|---|---|---|---|---|
|
||||
| Model | Shape | `TORCH_SDPA` | `CUDNN_ATTN` | `FLASHINFER_ATTN` |
|
||||
| --- | --- | ---: | ---: | ---: |
|
||||
| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
|
||||
| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
|
||||
| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
|
||||
| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
|
||||
|
||||
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
|
||||
Mask-heavy DiTs favored `CUDNN_ATTN`; lighter-mask or TP-saturated workloads
|
||||
were close enough that users should benchmark their exact model and shape.
|
||||
|
||||
## Known Limitations
|
||||
## Compatibility anchors
|
||||
|
||||
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
|
||||
The following headings preserve links to sections that moved into dedicated
|
||||
guides.
|
||||
|
||||
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
## TRTLLM_ATTN Backend and Skip-Softmax
|
||||
|
||||
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
|
||||
See [TRTLLM Attention](attention_backends/trtllm.md#skip-softmax).
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
## TRTLLM_ATTN SAGE Quantization
|
||||
|
||||
### FlashAttention-4 on Blackwell
|
||||
See [TRTLLM Attention](attention_backends/trtllm.md#sage-quantization).
|
||||
|
||||
Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
|
||||
## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[fa4]'
|
||||
```
|
||||
|
||||
`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
|
||||
when it is unavailable. Version `4.0.0b18` is required; earlier beta
|
||||
wheels had known JIT failures on Blackwell.
|
||||
|
||||
## Choosing a Backend Manually
|
||||
|
||||
### When to override the default
|
||||
|
||||
- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
|
||||
- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
|
||||
- **Workaround for known issues**: see Known Limitations above.
|
||||
|
||||
### Verifying which backend is in use
|
||||
|
||||
The startup log prints one of:
|
||||
|
||||
```
|
||||
Using diffusion attention backend 'CUDNN_ATTN' # explicit override
|
||||
Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
|
||||
Defaulting to diffusion attention backend SDPA # nothing else available
|
||||
```
|
||||
|
||||
If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
|
||||
See [RainFusion](attention_backends/rainfusion.md).
|
||||
|
||||
## SageAttention Installation
|
||||
|
||||
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
|
||||
|
||||
Build from source:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention
|
||||
|
||||
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
|
||||
pip install . --no-build-isolation
|
||||
```
|
||||
|
||||
Quick check:
|
||||
|
||||
```bash
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
See [SageAttention](attention_backends/sage.md#installation).
|
||||
|
||||
## SageAttention3 Installation
|
||||
|
||||
vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
|
||||
|
||||
Build from source:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention/sageattention3_blackwell
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
Quick check:
|
||||
|
||||
```bash
|
||||
python -c "import sageattn3; print(sageattn3.__file__)"
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
|
||||
- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
|
||||
See [SageAttention](attention_backends/sage.md#sageattention3-installation).
|
||||
|
||||
## HuggingFace Kernels Hub Backends
|
||||
|
||||
To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
|
||||
|
||||
The following backend options are supported:
|
||||
- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
|
||||
- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
|
||||
|
||||
### Installation
|
||||
|
||||
To use these backends, you must install the `kernels` library:
|
||||
|
||||
```bash
|
||||
pip install kernels==0.14.1
|
||||
```
|
||||
|
||||
If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
|
||||
|
||||
### Usage
|
||||
|
||||
Select a Hub backend using the global CLI flag or environment variables:
|
||||
|
||||
```bash
|
||||
# Environment variable
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
|
||||
|
||||
# CLI flag
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Default (auto-route)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
|
||||
--output hv15.mp4
|
||||
```
|
||||
|
||||
On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
|
||||
|
||||
### Explicit backend selection
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
|
||||
--output ltx2.mp4
|
||||
```
|
||||
|
||||
### SageAttention (lossy)
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output hv15_sage.mp4
|
||||
```
|
||||
|
||||
Example: Wan2.2 TI2V 5B
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 704 --width 1280 --num-frames 49 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_sage.mp4
|
||||
```
|
||||
|
||||
### Enable SageAttention3
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/hv15_sage3.mp4
|
||||
```
|
||||
See [Hugging Face Hub Backends](attention_backends/huggingface_hub.md).
|
||||
@@ -0,0 +1,110 @@
|
||||
# Dense Attention Backends
|
||||
|
||||
Use these backends when you want dense attention without backend-specific
|
||||
sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark
|
||||
the fastest compatible kernel for your model, shape, and hardware.
|
||||
|
||||
For selection precedence and per-role configuration, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
|
||||
## `TORCH_SDPA`
|
||||
|
||||
`TORCH_SDPA` calls PyTorch `scaled_dot_product_attention` and lets PyTorch's
|
||||
dispatcher choose the implementation. It is always available and is the most
|
||||
conservative reference when validating another backend.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend TORCH_SDPA
|
||||
```
|
||||
|
||||
## `FLASH_ATTN`
|
||||
|
||||
`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it
|
||||
prefers the CuTe-based FlashAttention 4 path and falls back to
|
||||
FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
|
||||
Ampere it is the preferred automatic route when a compatible package is
|
||||
installed.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
### FlashAttention 4 on Blackwell
|
||||
|
||||
Install the optional CUDA 13 extra:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[fa4]'
|
||||
```
|
||||
|
||||
Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on
|
||||
Blackwell. If the CuTe path is unavailable, the backend falls back to the
|
||||
compatible FlashAttention 3 or 2 path.
|
||||
|
||||
## `CUDNN_ATTN`
|
||||
|
||||
`CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful
|
||||
for mask-heavy DiTs and is automatically preferred on Blackwell when cuDNN
|
||||
9.5 or newer is available and the higher-priority TRTLLM route is not
|
||||
compatible.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend CUDNN_ATTN
|
||||
```
|
||||
|
||||
### LTX-2.0 limitation
|
||||
|
||||
LTX-2 audio attention has a symbolic head dimension during `torch.compile`
|
||||
tracing. The cuDNN SDPA selector rejects that symbolic dimension and Dynamo
|
||||
aborts compilation. This is tracked in
|
||||
[issue #3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
|
||||
Use `FLASHINFER_ATTN` or `TORCH_SDPA` as a workaround:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
|
||||
## `FLASHINFER_ATTN`
|
||||
|
||||
`FLASHINFER_ATTN` uses FlashInfer's batch-prefill wrapper. It is an explicit
|
||||
option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
|
||||
is installed but cuDNN is too old for `CUDNN_ATTN`.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN
|
||||
```
|
||||
|
||||
### FlashInfer quantized attention
|
||||
|
||||
The backend accepts an `AttentionSpec.quant` block. For QK16/V8, keep Q and K
|
||||
in FP16 or BF16 and use FP8 E4M3 for V:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import (
|
||||
AttentionConfig,
|
||||
AttentionSpec,
|
||||
AttnQuantSpec,
|
||||
OmniDiffusionConfig,
|
||||
)
|
||||
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="FLASHINFER_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="bfloat16",
|
||||
dtype_vo="fp8_e4m3",
|
||||
),
|
||||
),
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
`dtype_qk` controls Q and K; `dtype_vo` controls V. Mixed-dtype
|
||||
configurations require FlashInfer 0.6.16rc1 or newer. The shared quantization
|
||||
schema is also consumed by TRTLLM, but each backend validates its own allowed
|
||||
fields and values; see [TRTLLM SAGE quantization](trtllm.md#sage-quantization).
|
||||
@@ -0,0 +1,32 @@
|
||||
# Hugging Face Hub Attention Backends
|
||||
|
||||
Hub backends load the same published kernels commonly used with Hugging Face
|
||||
Diffusers. They are useful when training and serving must use matching kernel
|
||||
implementations to reduce numerical drift and sampling divergence.
|
||||
|
||||
| Backend | Kernel | Platform |
|
||||
| --- | --- | --- |
|
||||
| `FLASH_ATTN_HUB` | `kernels-community/flash-attn2` | Compatible CUDA GPUs |
|
||||
| `FLASH_ATTN_3_HUB` | `kernels-community/flash-attn3` | Hopper sm_90 or newer |
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install kernels==0.14.1
|
||||
```
|
||||
|
||||
If `kernels` is unavailable, vLLM-Omni warns and falls back to the
|
||||
corresponding local FlashAttention implementation. On CUDA devices older than
|
||||
Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
|
||||
|
||||
# Equivalent CLI selection
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
|
||||
```
|
||||
|
||||
For per-role kernel alignment, configure the Hub backend on only the relevant
|
||||
role; see the [attention backend overview](../attention_backends.md#per-role-configuration).
|
||||
@@ -0,0 +1,72 @@
|
||||
# RainFusion Attention
|
||||
|
||||
`RAINFUSION_ATTN` runs MindIE-SD RainFusion (`rf_v2`) block-sparse video
|
||||
attention on Ascend NPU. It pools 128-token key blocks, ranks them per query
|
||||
block, and attends to the highest-scoring blocks after arranging video tokens
|
||||
in `(t, h, w)` order.
|
||||
|
||||
Only the video segment is sparse. Prefix rows such as text, visual conditions,
|
||||
and audio, plus first-frame blocks, remain dense. Unsupported calls—including
|
||||
warmup steps, skipped layers, missing video geometry, and video segments below
|
||||
32 blocks—delegate to `FLASH_ATTN`, so compatible models can select RainFusion
|
||||
globally.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity |
|
||||
| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense |
|
||||
| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense |
|
||||
|
||||
```bash
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN",\
|
||||
"block_sparse":{"sparsity":0.8,"start_step":0}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse=BlockSparseSpec(
|
||||
sparsity=0.8,
|
||||
start_step=0,
|
||||
skip_layers="0-1",
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
Tune `start_step`, then `sparsity`, then `skip_layers`. Increasing the dense
|
||||
early-step window is usually the cheapest way to recover global structure;
|
||||
use layer exclusions only after a same-seed dense comparison identifies
|
||||
sensitive blocks.
|
||||
|
||||
## Requirements and compatibility
|
||||
|
||||
RainFusion requires Ascend NPU and `mindiesd`. Selecting it on another
|
||||
platform raises. It is incompatible with ring sequence parallelism because
|
||||
the kernel needs the complete key sequence for ranking; use Ulysses sequence
|
||||
parallelism with `ring_degree=1`.
|
||||
|
||||
## Geometry handling
|
||||
|
||||
RainFusion handles arbitrary video grids by rearranging video spatially and,
|
||||
when necessary, promoting a real-video suffix to the always-kept prefix. It
|
||||
does not pad because `rf_v2` does not consume a padding attention mask.
|
||||
|
||||
For MiniMax-H3 at 1344x768, grid `(62, 24, 42)` produces 62,496 video rows.
|
||||
The implementation promotes 2,976 rows and leaves 59,520 sparse rows
|
||||
(`465 x 128`) so the mask and kernel tiling remain aligned.
|
||||
|
||||
Spatial grids aligned to 8x8 are still preferable. A latent height or width
|
||||
not divisible by 8 creates an always-kept suffix and reduces realized
|
||||
sparsity. At the image level, that means width and height divisible by 256.
|
||||
Protected video tails require a compatible MindIE-SD release.
|
||||
|
||||
For common configuration and selector behavior, see the
|
||||
[attention backend overview](../attention_backends.md) and the
|
||||
[backend selection design](../../../design/feature/attention_backend_selection.md).
|
||||
@@ -0,0 +1,63 @@
|
||||
# SageAttention
|
||||
|
||||
SageAttention backends provide lossy low-precision attention for diffusion
|
||||
models. Validate output quality against `TORCH_SDPA` at the same seed before
|
||||
using either backend in production.
|
||||
|
||||
## `SAGE_ATTN`
|
||||
|
||||
`SAGE_ATTN` uses SageAttention 2.2 with INT8-quantized attention and FP16
|
||||
accumulation.
|
||||
|
||||
### Installation
|
||||
|
||||
Install SageAttention into the same environment as vLLM-Omni:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention
|
||||
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
|
||||
pip install . --no-build-isolation
|
||||
```
|
||||
|
||||
Verify the installation:
|
||||
|
||||
```bash
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
|
||||
Select it globally:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
|
||||
```
|
||||
|
||||
## `SAGE_ATTN_3`
|
||||
|
||||
`SAGE_ATTN_3` uses the SageAttention3 Blackwell implementation.
|
||||
|
||||
### SageAttention3 installation
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention/sageattention3_blackwell
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
Verify the installation:
|
||||
|
||||
```bash
|
||||
python -c "import sageattn3; print(sageattn3.__file__)"
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN_3
|
||||
```
|
||||
|
||||
`SAGE_ATTN_3` requires CUDA, an importable `sageattn3`, and a Blackwell-class
|
||||
GPU. Its kernel assumes the query-head count equals the key/value-head count.
|
||||
GQA and MQA diffusion calls therefore fall back to PyTorch SDPA for
|
||||
correctness.
|
||||
|
||||
For common configuration and platform routing, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
@@ -0,0 +1,103 @@
|
||||
# TRTLLM Attention
|
||||
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
|
||||
BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
|
||||
|
||||
For common selection and per-role configuration, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
|
||||
## Requirements
|
||||
|
||||
Selecting `TRTLLM_ATTN` requires:
|
||||
|
||||
- datacenter Blackwell, sm_100 or sm_103;
|
||||
- `head_dim=128`;
|
||||
- FlashInfer; and
|
||||
- a model-declared packed or mask-free compatible path.
|
||||
|
||||
It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
|
||||
explicit selection raises instead of silently changing the backend.
|
||||
|
||||
## Dense mode
|
||||
|
||||
Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend TRTLLM_ATTN
|
||||
```
|
||||
|
||||
## Skip-Softmax
|
||||
|
||||
Skip-Softmax trades some fidelity for speed by skipping low-contribution
|
||||
attention work. See the [feature design](../../../design/feature/skip_softmax.md)
|
||||
for the algorithm and implementation contract.
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve |
|
||||
| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` |
|
||||
| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` |
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
|
||||
"skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
skip_softmax=SkipSoftmaxSpec(
|
||||
target_sparsity=0.65,
|
||||
disabled_until_timestep=0.86,
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase
|
||||
sparsity only after comparing quality against dense output at the same seed.
|
||||
The benefit grows with sequence length because only attention is accelerated.
|
||||
|
||||
## SAGE quantization
|
||||
|
||||
TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
|
||||
matrix multiplications use low precision.
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense |
|
||||
| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` |
|
||||
| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` |
|
||||
|
||||
V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
|
||||
Only block sizes with compiled kernels are accepted.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
|
||||
"quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="fp8_e4m3",
|
||||
q_block_size=1,
|
||||
k_block_size=16,
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are
|
||||
available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only.
|
||||
The shared `AttnQuantSpec` is backend typed: values intended for
|
||||
`FLASHINFER_ATTN` are rejected here and vice versa.
|
||||
@@ -1,389 +1,75 @@
|
||||
# CPU Offloading for Diffusion Models
|
||||
|
||||
## Overview
|
||||
vLLM-Omni provides three CPU-offload strategies for diffusion models. Choose
|
||||
the coarsest strategy that meets your memory target: finer-grained streaming
|
||||
saves more device memory but adds more transfer and setup overhead.
|
||||
|
||||
vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:
|
||||
For the shared factory, discovery, and lifecycle contract, see the
|
||||
[CPU Offloading design](../../design/feature/cpu_offloading.md).
|
||||
|
||||
1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
|
||||
2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap.
|
||||
## Choose a strategy
|
||||
|
||||
Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority.
|
||||
| Strategy | Device residency | Parallel scope | Guide |
|
||||
| --- | --- | --- | --- |
|
||||
| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](cpu_offload/model_level.md) |
|
||||
| Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](cpu_offload/layerwise.md) |
|
||||
| Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](cpu_offload/distributed_layerwise.md) |
|
||||
|
||||
All strategies use pinned host memory for faster transfers where applicable.
|
||||
Configuration priority is:
|
||||
|
||||
1. Distributed layerwise offloading.
|
||||
2. Layerwise offloading.
|
||||
3. Model-level offloading.
|
||||
|
||||
Treat the flags as mutually exclusive. If more than one is enabled, the
|
||||
higher-priority strategy is selected.
|
||||
|
||||
## Quick selection
|
||||
|
||||
- Use [model-level offloading](cpu_offload/model_level.md) when swapping whole
|
||||
encoders and DiTs is enough to fit and phase-boundary transfers are
|
||||
acceptable.
|
||||
- Use [layerwise offloading](cpu_offload/layerwise.md) for compute-heavy video
|
||||
DiTs where block transfers can overlap computation.
|
||||
- Use [distributed layerwise offloading](cpu_offload/distributed_layerwise.md)
|
||||
when a multi-rank deployment also needs bounded device residency and,
|
||||
optionally, sharded host weights.
|
||||
|
||||
## Supported models
|
||||
|
||||
| Architecture | Example models | DiT class | Model-level | Layerwise | Distributed layerwise | Layerwise block attributes |
|
||||
| --- | --- | --- | :---: | :---: | :---: | --- |
|
||||
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` |
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | yes | yes | — | `down_blocks`, `up_blocks` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules |
|
||||
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` |
|
||||
|
||||
Model-level support requires discoverable DiT and encoder components.
|
||||
Layerwise support requires transformer block topology. Distributed support
|
||||
reuses that topology but still requires validation for the model, checkpoint,
|
||||
and parallel configuration.
|
||||
|
||||
## Compatibility anchors
|
||||
|
||||
These headings preserve links to sections that moved into dedicated guides.
|
||||
|
||||
## Model-level (Sequential) Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
|
||||
|
||||
- **When encoders run**: DiT transformer is offloaded to CPU
|
||||
- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU.
|
||||
- **VAE**: Stays resident on GPU
|
||||
|
||||
Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
|
||||
|
||||
### Usage
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
|
||||
```
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
|
||||
```
|
||||
|
||||
### To Support a Model
|
||||
|
||||
Implement the `SupportsComponentDiscovery` protocol to declare which
|
||||
submodules serve as pipeline components (used by offloading, HSDP
|
||||
sharding, and other framework features):
|
||||
|
||||
```python
|
||||
from typing import ClassVar
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
|
||||
class MyPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = [] # optional
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = ... # DiT — stays on GPU during denoising
|
||||
self.text_encoder = ... # Encoder — offloaded to CPU during denoising
|
||||
self.vision_model = ... # Encoder — offloaded to CPU during denoising
|
||||
self.vae = ... # VAE — always on GPU
|
||||
```
|
||||
|
||||
- `_dit_modules`: attribute names of denoising submodules (kept on GPU
|
||||
during the diffusion loop).
|
||||
- `_encoder_modules`: attribute names of encoder/vision submodules
|
||||
(offloaded to CPU during the diffusion loop).
|
||||
- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
|
||||
part of the mutual exclusion hooks).
|
||||
- `_resident_modules`: attribute names of small submodules that must
|
||||
stay on GPU during layerwise offloading (e.g. embedders, connectors).
|
||||
Optional — defaults to `[]`.
|
||||
|
||||
All attribute names support dotted paths for nested submodules
|
||||
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
|
||||
|
||||
Both DiT and encoder lists are needed because the offload hooks use
|
||||
mutual exclusion: when one group runs, the other moves to CPU.
|
||||
|
||||
### Limitations
|
||||
- Cold start latency increases
|
||||
- Adds overhead from CPU-GPU transfers between encoder and denoising phases
|
||||
- Support single GPU only for now
|
||||
|
||||
|
||||
### Component offloading for split models (e.g. Cosmos3)
|
||||
|
||||
Some models split their transformer into mutually-exclusive *components* that run
|
||||
in different phases of a single forward pass rather than as separate pipeline
|
||||
components -- e.g. Cosmos3's understanding (reasoner) component runs once per
|
||||
generation while the generation (generator) component runs every denoising step.
|
||||
Such models have no separate text encoder to swap against, so the transformer
|
||||
owns a small model-local offload path and wraps each phase with
|
||||
`with self._offload_context(name):`
|
||||
|
||||
```python
|
||||
class Cosmos3VFMTransformer(nn.Module):
|
||||
def forward(self, ...):
|
||||
with self._offload_context("reasoner"):
|
||||
... # understanding pass, runs once
|
||||
with self._offload_context("generator"):
|
||||
... # denoising pass, runs every step
|
||||
```
|
||||
|
||||
Model-level offloading then keeps exactly one component GPU-resident at a time
|
||||
(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The
|
||||
pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the
|
||||
transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading
|
||||
works for these models too -- each component declares its own block container via
|
||||
`_layerwise_offload_blocks_attrs`.
|
||||
|
||||
See [Model-Level Offloading](cpu_offload/model_level.md).
|
||||
|
||||
## Layerwise (Blockwise) Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Layerwise offloading keeps only one transformer block on GPU at a time.
|
||||
|
||||
As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.
|
||||
|
||||
**Execution Flow:**
|
||||
|
||||
| Block | Pre-forward Hook | Forward | Post-forward Hook |
|
||||
|-------|------------------|---------|-------------------|
|
||||
| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 |
|
||||
| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 |
|
||||
| ... | ... | ... | ... |
|
||||
| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) |
|
||||
|
||||
Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward.
|
||||
|
||||
Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.
|
||||
|
||||
### Usage
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
# Text-to-video
|
||||
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)
|
||||
|
||||
# Or image-to-video
|
||||
m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)
|
||||
```
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
# Text-to-video
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
|
||||
|
||||
# Or image-to-video
|
||||
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
|
||||
```
|
||||
|
||||
### To Support a Model
|
||||
|
||||
Models must define the blocks attribute name for layerwise offloading:
|
||||
|
||||
```python
|
||||
class WanTransformer3DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"] # Attribute names containing transformer blocks
|
||||
|
||||
def __init__(self):
|
||||
self.blocks = nn.ModuleList([...]) # Transformer blocks
|
||||
```
|
||||
|
||||
For models with multiple block types:
|
||||
|
||||
```python
|
||||
class Flux2Transformer2DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
|
||||
```
|
||||
|
||||
### Limitations
|
||||
- Cold start latency increases because offloaded components must be moved to CPU
|
||||
during setup; layerwise offload may add extra weight consolidation and pinning
|
||||
work.
|
||||
- Performance depends on compute cost and H2D bandwidth as well
|
||||
- Support single GPU only for now
|
||||
|
||||
See [Layerwise Offloading](cpu_offload/layerwise.md).
|
||||
|
||||
## Distributed Layerwise Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Distributed layerwise offloading extends single-GPU layerwise offloading to
|
||||
multi-device deployments. In the default mode, each rank stores only
|
||||
**1/dp_size** of the model weights in host memory; full layer weights are
|
||||
reconstructed at runtime via **AllGather** on a dedicated communication stream,
|
||||
overlapped with computation via a fixed double-buffer scheme. The
|
||||
`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
|
||||
weights independently and does not shard weights across ranks.
|
||||
|
||||
For the implementation design and compatibility matrix, see
|
||||
[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
|
||||
|
||||
**Key features:**
|
||||
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
|
||||
reconstructed per-layer via `all_gather_into_tensor`
|
||||
- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
|
||||
time, regardless of model size
|
||||
- **DP multi-concurrency**: N concurrent requests processed in parallel
|
||||
(AllGather only gathers weight shards, which are request-independent)
|
||||
- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
|
||||
page cache, eliminating O(dp_size × model_size) RSS during model creation
|
||||
- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
|
||||
(CANN/HCCL) via vLLM-Omni's platform abstraction
|
||||
|
||||
**Execution Flow (per device):**
|
||||
|
||||
```
|
||||
Compute Stream: [Layer N] [Layer N+1] [Layer N+2]
|
||||
H2D Stream: [H2D Shard N+1] [H2D Shard N+2]
|
||||
AllGather: [AG N+1] [AG N+2]
|
||||
Slot Usage: Slot0: Layer N Slot1: Layer N+1
|
||||
```
|
||||
|
||||
### Usage
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
# 4× GPU/NPU with AllGather (recommended)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
|
||||
# Without DLO AllGather (each rank streams standard-loader rank-local weights)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
|
||||
# With SP instead of DP (long sequences)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4
|
||||
```
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
m = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True, # default
|
||||
)
|
||||
```
|
||||
|
||||
### CLI Flags
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
||||
| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
|
||||
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
|
||||
| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
|
||||
|
||||
### How model weights are loaded (mmap path)
|
||||
|
||||
When DLO + AllGather is active, the offloader:
|
||||
|
||||
1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
|
||||
from the normally-created transformer
|
||||
2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
|
||||
initialization weights
|
||||
3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
|
||||
which return views into the OS page cache (shared across all ranks, 0 RSS)
|
||||
4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
|
||||
(e.g. Cosmos3's `time_embedder` → FP32)
|
||||
5. **Restores non-persistent buffers** from saved copies
|
||||
|
||||
This approach requires **zero model-specific code changes** — no pipeline or
|
||||
transformer modifications are needed.
|
||||
|
||||
### OffloadPlan (declarative topology metadata)
|
||||
|
||||
Models can optionally declare an `OffloadPlan` class variable to provide
|
||||
topology metadata (block attribute names, submodules to offload) without
|
||||
any offload-specific logic:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.offloader import OffloadPlan
|
||||
|
||||
class MyPipeline(nn.Module):
|
||||
_dit_modules = ["transformer"]
|
||||
_offload_plan = OffloadPlan(
|
||||
block_attrs={"transformer": ("blocks",)},
|
||||
offload_submodules={"context_encoder": "layers"},
|
||||
)
|
||||
```
|
||||
|
||||
When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
|
||||
and heuristic attribute search. This is backward-compatible — existing models
|
||||
work without any changes.
|
||||
|
||||
### DP Multi-concurrency
|
||||
|
||||
When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
|
||||
up to `dp_size` requests per denoise step. Each DP rank processes a different
|
||||
request while AllGather synchronizes only weight shards (request-independent).
|
||||
|
||||
**Requirements for concurrent requests:**
|
||||
- `num_inference_steps` must be specified explicitly (None is not allowed)
|
||||
- All concurrent requests must have the same `num_inference_steps` value
|
||||
- These constraints exist because AllGather is a collective that requires
|
||||
every rank to participate at each denoise step
|
||||
|
||||
### Limitations
|
||||
- Online quantization (FP8) is incompatible with mmap loading — falls back to
|
||||
regular loading only when the model is outside the DLO+AllGather mmap path;
|
||||
online quantization with DLO+AllGather is rejected. Use
|
||||
`--dlo-no-use-allgather` or disable online quantization.
|
||||
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
|
||||
weight-loader callbacks are bypassed. The no-AllGather path retains the
|
||||
standard TP loader and is experimental.
|
||||
- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
|
||||
with no-AllGather is accepted at configuration level, but has limited
|
||||
end-to-end validation.
|
||||
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
|
||||
|
||||
**Module Discovery**
|
||||
|
||||
The offloader discovers pipeline components in two ways:
|
||||
|
||||
1. **Protocol-based** (preferred): If the pipeline implements
|
||||
`SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`,
|
||||
`_vae_modules`, and `_resident_modules` class variables are used
|
||||
directly. All attribute names support dotted paths (e.g.
|
||||
`"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
|
||||
|
||||
2. **Fallback attribute scan**: Otherwise, the offloader scans for
|
||||
well-known attribute names:
|
||||
- **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
|
||||
- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
|
||||
- **VAE**: `vae`, `audio_vae`
|
||||
|
||||
**Hook System**
|
||||
|
||||
Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.
|
||||
|
||||
**Backend Architecture**
|
||||
|
||||
```
|
||||
OffloadBackend (base class)
|
||||
├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
|
||||
│ (delegates to a pipeline's enable_omni_model_cpu_offload
|
||||
│ for split models like Cosmos3)
|
||||
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
|
||||
│ (single-GPU, full weights on host)
|
||||
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
|
||||
(multi-GPU, sharded AllGather or rank-local mode)
|
||||
```
|
||||
|
||||
Factory function `get_offload_backend()` selects the appropriate backend based on
|
||||
configuration.
|
||||
|
||||
For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's
|
||||
`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its
|
||||
reasoner/generator components inside the model forward pass.
|
||||
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
|
||||
|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
|
||||
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
|
||||
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
|
||||
|
||||
**Notes:**
|
||||
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
|
||||
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
|
||||
- Distributed Layerwise Offloading reuses the layerwise block topology, but the
|
||||
model and parallelism combination still needs validation. AllGather mode
|
||||
additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
|
||||
retains the standard loader's rank-local tensors. See the
|
||||
[DLO design](../../design/feature/distributed_layerwise_offload.md) and
|
||||
[Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
|
||||
for usage examples.
|
||||
See [Distributed Layerwise Offloading](cpu_offload/distributed_layerwise.md).
|
||||
@@ -0,0 +1,118 @@
|
||||
# Distributed Layerwise Offloading
|
||||
|
||||
Distributed layerwise offloading (DLO) extends block streaming to multi-device
|
||||
deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
|
||||
the host weights and reconstructs each layer at runtime. Without AllGather,
|
||||
each rank streams its standard-loader rank-local weights independently.
|
||||
|
||||
See the [DLO feature design](../../../design/feature/distributed_layerwise_offload.md)
|
||||
for the implementation contract and compatibility matrix.
|
||||
|
||||
## Execution model
|
||||
|
||||
DLO overlaps three operations with a fixed two-block device buffer:
|
||||
|
||||
```text
|
||||
Compute stream: [Layer N] [Layer N+1] [Layer N+2]
|
||||
H2D stream: [H2D shard N+1] [H2D shard N+2]
|
||||
AllGather: [AG N+1] [AG N+2]
|
||||
Slots: slot 0: Layer N slot 1: Layer N+1
|
||||
```
|
||||
|
||||
AllGather communicates only request-independent weight shards, so data-
|
||||
parallel ranks may process different requests concurrently.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
# Four ranks with sharded host weights and AllGather
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
|
||||
# Standard-loader rank-local weights, without DLO AllGather
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
|
||||
# Sequence parallel deployment
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
```
|
||||
|
||||
## Flags
|
||||
|
||||
| Flag | Meaning | Default |
|
||||
| --- | --- | --- |
|
||||
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
||||
| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
|
||||
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
|
||||
| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` |
|
||||
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
||||
|
||||
## mmap weight loading
|
||||
|
||||
The DLO plus AllGather path:
|
||||
|
||||
1. saves non-persistent buffers such as RoPE frequencies;
|
||||
2. moves the normally created transformer to the meta device;
|
||||
3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
|
||||
4. calls model-specific `post_load_weights()` conversions; and
|
||||
5. restores the saved non-persistent buffers.
|
||||
|
||||
This avoids one full checkpoint RSS copy per rank and does not require
|
||||
model-specific loading code.
|
||||
|
||||
## Declarative topology
|
||||
|
||||
Models may declare an `OffloadPlan` instead of embedding offload logic:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.offloader import OffloadPlan
|
||||
|
||||
|
||||
class MyPipeline(nn.Module):
|
||||
_dit_modules = ["transformer"]
|
||||
_offload_plan = OffloadPlan(
|
||||
block_attrs={"transformer": ("blocks",)},
|
||||
offload_submodules={"context_encoder": "layers"},
|
||||
)
|
||||
```
|
||||
|
||||
When no plan exists, discovery falls back to
|
||||
`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
|
||||
|
||||
## Data-parallel concurrency
|
||||
|
||||
With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
|
||||
up to `dp_size` requests per denoising step. Every concurrent request must set
|
||||
the same explicit `num_inference_steps`; `None` is rejected because every rank
|
||||
must enter each collective.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Online FP8 quantization is rejected with the DLO plus AllGather mmap path.
|
||||
Use `--dlo-no-use-allgather` or disable online quantization.
|
||||
- Tensor parallel size greater than one is rejected in the mmap path because
|
||||
it bypasses TP-aware loader callbacks. The no-AllGather path retains the
|
||||
standard loader and remains experimental with TP.
|
||||
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
|
||||
AllGather has limited end-to-end validation.
|
||||
- Resident leading layers require `--dlo-no-use-allgather` and a model
|
||||
`OffloadPlan` that declares eligible `resident_dit_paths`.
|
||||
- DP concurrency requires an explicit, identical inference-step count.
|
||||
|
||||
See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
|
||||
for an end-to-end example.
|
||||
@@ -0,0 +1,65 @@
|
||||
# Layerwise Offloading
|
||||
|
||||
Layerwise, or blockwise, offloading keeps one transformer block on the
|
||||
accelerator and prefetches the next block while the current block computes.
|
||||
It is best suited to compute-heavy video DiTs whose block execution can hide
|
||||
host-to-device transfers.
|
||||
|
||||
## Execution flow
|
||||
|
||||
Each block has a pre-forward and post-forward hook. Parameters are consolidated
|
||||
in pinned host tensors and rematerialized for execution on a dedicated copy
|
||||
stream.
|
||||
|
||||
| Block | Pre-forward hook | Forward | Post-forward hook |
|
||||
| --- | --- | --- | --- |
|
||||
| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 |
|
||||
| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 |
|
||||
| ... | ... | ... | ... |
|
||||
| last block | Prefetch block 0 | Compute last block | Free last block |
|
||||
|
||||
Encoders, VAE modules, and non-block DiT modules such as embeddings and norms
|
||||
remain device resident.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
enable_layerwise_offload=True,
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--omni --enable-layerwise-offload
|
||||
```
|
||||
|
||||
## Model integration
|
||||
|
||||
Transformer classes declare containers of executable blocks:
|
||||
|
||||
```python
|
||||
class WanTransformer3DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||
|
||||
|
||||
class Flux2Transformer2DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = [
|
||||
"transformer_blocks",
|
||||
"single_transformer_blocks",
|
||||
]
|
||||
```
|
||||
|
||||
See the [layerwise design](../../../design/feature/layerwise_offload.md) for
|
||||
the discovery and hook invariants. `OffloadPlan` is a separate declarative
|
||||
topology path for distributed layerwise offload.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Single device only.
|
||||
- Setup consolidates and pins block parameters, increasing cold-start time.
|
||||
- Performance depends on block compute time and host-to-device bandwidth;
|
||||
lightweight blocks may not hide transfers.
|
||||
@@ -0,0 +1,81 @@
|
||||
# Model-Level Offloading
|
||||
|
||||
Model-level, or sequential, offloading keeps only the pipeline component group
|
||||
currently executing on the accelerator. It is the simplest offload strategy
|
||||
and is selected with `--enable-cpu-offload`.
|
||||
|
||||
## How it works
|
||||
|
||||
Pre-forward hooks enforce mutual exclusion between DiT and encoder modules:
|
||||
|
||||
- before an encoder runs, the DiT moves to CPU;
|
||||
- before a DiT runs, encoders and other DiTs move to CPU; and
|
||||
- VAE modules remain on the accelerator.
|
||||
|
||||
Pinned host memory reduces transfer overhead. Transfers occur at phase
|
||||
boundaries, so cold-start and encoder-to-denoiser transitions become slower.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
enable_cpu_offload=True,
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--omni --enable-cpu-offload
|
||||
```
|
||||
|
||||
## Model integration
|
||||
|
||||
Pipelines should implement `SupportsComponentDiscovery`:
|
||||
|
||||
```python
|
||||
from typing import ClassVar
|
||||
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
|
||||
|
||||
class MyPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = []
|
||||
```
|
||||
|
||||
All entries may be dotted paths. DiT and encoder lists are both required for
|
||||
mutual exclusion. VAE modules are pinned but not swapped; resident modules are
|
||||
small modules that must stay on the accelerator for layerwise paths.
|
||||
|
||||
## Split-model components
|
||||
|
||||
Some models, such as Cosmos3, split one transformer into mutually exclusive
|
||||
components that run in different phases. The pipeline exposes
|
||||
`enable_omni_model_cpu_offload`, and the backend delegates to the model-local
|
||||
contexts:
|
||||
|
||||
```python
|
||||
class Cosmos3VFMTransformer(nn.Module):
|
||||
def forward(self, ...):
|
||||
with self._offload_context("reasoner"):
|
||||
...
|
||||
with self._offload_context("generator"):
|
||||
...
|
||||
```
|
||||
|
||||
This preserves the same invariant—exactly one component is device resident—
|
||||
while reusing sequential `.to()` movers.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Single device only.
|
||||
- Higher cold-start latency.
|
||||
- Transfers between encoder and denoising phases add latency.
|
||||
|
||||
See the [model-level design](../../../design/feature/model_level_offload.md)
|
||||
for lifecycle and extension invariants.
|
||||
Reference in new issue
Block a user