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

docs: split diffusion feature guides

Browse Source 
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-12 08:55:41 +08:00
1 parent 8b0a6779d4
commit 7872663bf0
21 files changed
+1138 -801
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/diffusion/attention_backends.md
+111 -413
View File
@@ -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).