vllm-omni/docs/user_guide/diffusion/attention_backends/dense_backends.md

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.

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.

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.

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

FlashAttention 4 on Blackwell

Install the optional CUDA 13 extra:

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.

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.

Use FLASHINFER_ATTN or TORCH_SDPA as a workaround:

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.

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:

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.