Diffusion Attention Backends¶

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.

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.

For the internal selector, registry, and platform contract, see Attention Backend Selection.

Choose a guide¶

Backend options¶

Configuration¶

Backend selection is resolved in this order:

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. Do not pass it together with an explicit default.backend in the structured configuration.

Global default¶

vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN

# Backwards-compatible environment variable
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN

Per-role configuration¶

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. Category per_role[role_category] match.
3. default.
4. Platform default.

# Dotted flags
vllm-omni serve <model> \
  --diffusion-attention-config.default.backend FLASH_ATTN \
  --diffusion-attention-config.per_role.cross.backend TORCH_SDPA

# Equivalent JSON
vllm-omni serve <model> \
  --diffusion-attention-config \
  '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'

Python API¶

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")},
    ),
    ...,
)

Backend-specific typed blocks are documented with their consumers:

• quant: FlashInfer and TRTLLM SAGE.
• skip_softmax: TRTLLM Skip-Softmax.
• block_sparse: RainFusion.
• fastvideo_vsa_topk: FastVideo VSA.

Platform defaults¶

Blackwell (sm_100 / sm_103 / sm_120 / sm_121)¶

The CUDA auto-route preference is:

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 not auto-selected on workstation Blackwell (sm_120/sm_121), for other head dimensions, or for paths that require an unsupported mask.

Hopper, Ada, and Ampere¶

The CUDA auto-route uses FLASH_ATTN when available and otherwise falls back to TORCH_SDPA. CUDNN_ATTN and FLASHINFER_ATTN remain explicit options.

Other platforms validate an explicit backend and choose their own default through the platform implementation. Check the startup log to confirm the resolved backend.

Choosing a backend manually¶

Override the platform default when you need:

• a correctness reference (TORCH_SDPA);
• a backend-specific workaround;
• training/rollout kernel alignment (Hub backends); or
• an explicitly validated sparse or quantized speedup.

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.

Reference benchmark¶

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.

Mask-heavy DiTs favored CUDNN_ATTN; lighter-mask or TP-saturated workloads were close enough that users should benchmark their exact model and shape.

Compatibility anchors¶

The following headings preserve links to sections that moved into dedicated guides.

TRTLLM_ATTN Backend and Skip-Softmax¶

See TRTLLM Attention.

TRTLLM_ATTN SAGE Quantization¶

See TRTLLM Attention.

RAINFUSION_ATTN Backend and Block-Sparse Video Attention¶

See RainFusion.

SageAttention Installation¶

See SageAttention.

SageAttention3 Installation¶

See SageAttention.

HuggingFace Kernels Hub Backends¶

See Hugging Face Hub Backends.