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

Dense Attention Backends

Use these backends to run dense attention. Optional quantized or sparse modes remain inactive unless they are explicitly configured. 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 is FA4 only (flash_attn.cute / vllm-omni[fa4]). If FA4 is unavailable, explicit FLASH_ATTN is rejected and automatic selection continues to another compatible Blackwell backend. Hopper-only FA2/FA3 wheels are not used, even if they import. 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. On Blackwell, FLASH_ATTN is FA4 only: Hopper-only FA2/FA3 wheels are not used, even if they import. If FA4 is missing, explicit FLASH_ATTN is rejected and automatic selection continues to other Blackwell backends.

TRTLLM_ATTN

TRTLLM_ATTN runs FlashInfer's trtllm-gen FMHA kernels and is the platform default on datacenter Blackwell for models that declare a compatible attention path. Selected without a quant or skip_softmax block, it runs dense BF16 attention at FA4-level performance.

vllm serve <model> --omni \
  --diffusion-attention-backend TRTLLM_ATTN

See TRTLLM Attention for its requirements and for the optional SAGE quantization and Skip-Softmax modes.

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.

On Blackwell, auto resolves to FlashInfer cute-dsl, which cannot run a nontrivial custom mask. Automatic selection may fall back to SDPA for those masks. An explicit FLASHINFER_ATTN selection does not: sequence-parallel auto-padding that would create a padding mask is rejected at capability preflight. Use TORCH_SDPA, pin quant.flashinfer_backend to fa2/fa3 where those kernels exist, or choose a mask-capable backend.

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.