[Blackwell] Add CUDNN_ATTN and FLASHINFER_ATTN backends for diffusion (auto-route) (#3079)

Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: Zhang <jianmusings@gmail.com>
Signed-off-by: Zhang Jian <jianmusings@gmail.com>
Co-authored-by: frauttauteffasu <247219357+frauttauteffasu@users.noreply.github.com>
Co-authored-by: Zhang Jian <jianmusings@gmail.com>

benchmarks/diffusion/bench_attention_backends.py

@@ -0,0 +1,328 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Standalone diagnostic for the diffusion attention backends added in #3079.
+
+Exercises the same synthetic attention shape through each backend and each
+torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
+CUDNN_ATTN underperforms on some SKUs. Two typical causes:
+
+  (a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
+      and silently falls back to MATH.
+  (b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
+      walking CUDNN -> FLASH -> MATH.
+
+Run:
+    python benchmarks/diffusion/bench_attention_backends.py --preset hv15
+    python benchmarks/diffusion/bench_attention_backends.py --preset wan22
+    python benchmarks/diffusion/bench_attention_backends.py \
+        --batch 1 --heads 24 --seq 14336 --head-dim 128
+
+Optional extras:
+    pip install --pre flash-attn-4        # FA4 is currently pre-release only
+    pip install -U flashinfer              # latest FlashInfer (0.6.9)
+
+The table at the end is the data we want on the PR — surface the row where a
+backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+# Presets mirror the DiT hot-path attention call (single-stream block) at the
+# shapes we run in the PR's validation configs.
+_PRESETS = {
+    # HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
+    "hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
+    # Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
+    "wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
+    # Smaller image-gen shape for quick smoke tests.
+    "flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
+}
+
+_SDPA_BACKENDS = [
+    ("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
+    ("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
+    ("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
+    ("MATH", [SDPBackend.MATH]),
+    # The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
+    # what torch actually picks when multiple are allowed.
+    ("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
+]
+
+
+def _env_report() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
+    cc = torch.cuda.get_device_capability()
+    print("=" * 72)
+    print(f"GPU           : {torch.cuda.get_device_name()}")
+    print(f"SM capability : sm_{cc[0]}{cc[1]}")
+    print(f"torch         : {torch.__version__}")
+    print(f"cuDNN version : {torch.backends.cudnn.version()}")
+    print("flashinfer    : ", end="")
+    try:
+        import flashinfer
+
+        print(getattr(flashinfer, "__version__", "present"))
+    except Exception as e:
+        print(f"not installed ({type(e).__name__})")
+    print("=" * 72)
+
+
+def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
+    # Layout: (B, S, H, D) — what the PR's backends accept as input.
+    q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    return q, k, v
+
+
+def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
+    # Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
+    mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
+    if pad_tokens > 0:
+        mask[..., -pad_tokens:] = float("-inf")
+    return mask
+
+
+def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
+    """Return (median ms, error string). Skips on any exception so that a
+    backend that rejects our args (wrong dtype, missing JIT module, unsupported
+    kwarg value) doesn't abort the whole sweep."""
+    try:
+        for _ in range(warmup):
+            fn(*args)
+        torch.accelerator.synchronize()
+        times = []
+        for _ in range(iters):
+            torch.accelerator.synchronize()
+            t0 = time.perf_counter()
+            fn(*args)
+            torch.accelerator.synchronize()
+            times.append((time.perf_counter() - t0) * 1000.0)
+        times.sort()
+        return times[len(times) // 2], ""
+    except Exception as e:  # noqa: BLE001 — probe script, keep the whole table going
+        msg = str(e).split("\n", 1)[0][:60]
+        return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
+
+
+def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
+    # q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
+    q_t = q.permute(0, 2, 1, 3)
+    k_t = k.permute(0, 2, 1, 3)
+    v_t = v.permute(0, 2, 1, 3)
+    rows: list[tuple[str, float, str]] = []
+    for name, backends in _SDPA_BACKENDS:
+
+        def _call(mask=attn_mask):
+            with sdpa_kernel(backends):
+                return torch.nn.functional.scaled_dot_product_attention(
+                    q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
+                )
+
+        ms, err = _time_call(_call)
+        rows.append((name, ms, err))
+    return rows
+
+
+def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
+    """Call FlashInfer's dense single-prefill.
+
+    ``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
+    that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
+    we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
+    Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
+    """
+    try:
+        import inspect
+
+        from flashinfer.prefill import single_prefill_with_kv_cache
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
+    if backend is not None:
+        sig = inspect.signature(single_prefill_with_kv_cache)
+        if "backend" not in sig.parameters:
+            return float("nan"), "no-backend-kwarg"
+        kwargs["backend"] = backend
+
+    if attn_mask is not None:
+        # Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
+        # FlashInfer expects True = keep, False = masked out.
+        mask_2d = attn_mask[0, 0]
+        kwargs["custom_mask"] = mask_2d != float("-inf")
+
+    def _call():
+        out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
+        return out.unsqueeze(0)
+
+    return _time_call(_call)
+
+
+def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
+    """Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
+    ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
+    cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
+    ``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
+    try:
+        from flash_attn.cute import flash_attn_func
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    def _call():
+        # FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
+        return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
+
+    return _time_call(_call)
+
+
+def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
+    """Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
+    (which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
+    straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
+    try:
+        from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    b, s, h, d = q.shape
+    qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
+    kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
+
+    def _call():
+        return cudnn_batch_prefill_with_kv_cache(
+            q.reshape(b * s, h, d),
+            k.reshape(b * s, h, d),
+            v.reshape(b * s, h, d),
+            qo_indptr=qo_indptr,
+            kv_indptr=kv_indptr,
+            max_qo_len=s,
+            max_kv_len=s,
+            sm_scale=scale,
+            causal=False,
+        )
+
+    return _time_call(_call)
+
+
+def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
+    print(f"\n{title}")
+    print("-" * 88)
+    header = f"{'backend':<24} {'median (ms)':>14}    {'vs baseline':>12}    status"
+    print(header)
+    print("-" * 88)
+    baseline_ms = None
+    if baseline_name is not None:
+        for name, ms, _ in rows:
+            if name == baseline_name and ms == ms:  # not NaN
+                baseline_ms = ms
+                break
+    for name, ms, err in rows:
+        ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
+        if baseline_ms is not None and ms == ms:
+            ratio = baseline_ms / ms
+            ratio_str = f"{ratio:>11.2f}x"
+        else:
+            ratio_str = f"{'—':>12}"
+        status = "ok" if not err else f"FAILED ({err})"
+        print(f"{name:<24} {ms_str}    {ratio_str}    {status}")
+
+
+def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
+    """Return (backend, ms) of the fastest non-failing row, or None."""
+    ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
+    if not ok_rows:
+        return None
+    return min(ok_rows, key=lambda x: x[1])
+
+
+def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
+    torch.manual_seed(0)
+    q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
+    scale = 1.0 / (shape["head_dim"] ** 0.5)
+
+    rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
+    rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
+    # `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
+    for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
+        rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
+    rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
+    rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
+
+    mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
+    rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
+    rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
+    rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
+
+    return rows_nomask, rows_mask
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
+    parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
+    parser.add_argument("--batch", type=int, default=None)
+    parser.add_argument("--heads", type=int, default=None)
+    parser.add_argument("--seq", type=int, default=None)
+    parser.add_argument("--head-dim", type=int, default=None)
+    parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
+    parser.add_argument("--device", default="cuda:0")
+    args = parser.parse_args()
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+    _env_report()
+    print(f"dtype         : {args.dtype}")
+
+    presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
+    summary: dict[str, tuple[str, float] | None] = {}
+
+    for preset_name in presets:
+        shape = dict(_PRESETS[preset_name])
+        if not args.sweep:
+            for k in ("batch", "heads", "seq", "head_dim"):
+                v = getattr(args, k if k != "head_dim" else "head_dim")
+                if v is not None:
+                    shape[k] = v
+
+        print("\n" + "=" * 88)
+        print(
+            f"Preset: {preset_name}  |  batch={shape['batch']} heads={shape['heads']} "
+            f"seq={shape['seq']} head_dim={shape['head_dim']}"
+        )
+        print("=" * 88)
+
+        rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
+        _print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
+        _print_table(
+            f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
+        )
+
+        summary[preset_name] = _pick_winner(rows_nomask)
+
+    if args.sweep or len(presets) > 1:
+        print("\n" + "=" * 88)
+        print("Winners per preset (no-mask path)")
+        print("=" * 88)
+        for preset_name, winner in summary.items():
+            if winner is None:
+                print(f"  {preset_name:<10} — no successful backend")
+            else:
+                name, ms = winner
+                print(f"  {preset_name:<10} {name:<24} {ms:>8.3f} ms")
+
+    print("\nNotes:")
+    print("  * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
+    print("  * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
+    print("  * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
+    print("  * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
+
+
+if __name__ == "__main__":
+    main()

docs/user_guide/diffusion/attention_backends.md

@@ -1,28 +1,24 @@
 # Diffusion Attention Backends
 
-This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, and how to use SageAttention.
+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.
 
 ## Overview
 
 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.
 
-This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
+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.
 
-On CUDA, the practical choices today are:
-
-- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed.
-- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
-- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.
-
-If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
+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.
 
 ## Backend Options
 
 | Value | Notes |
 |---|---|
-| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. |
-| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
-| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |
+| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` 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` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
+| `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`. |
 
 ## Configuration
 
@@ -95,6 +91,80 @@ config = OmniDiffusionConfig(
 
 A plain dict is also accepted and normalized to `AttentionConfig`.
 
+## Platform Defaults
+
+### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
+
+Auto-route preference, in order:
+
+1. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
+2. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
+3. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
+4. `TORCH_SDPA` — last resort
+
+The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
+
+**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.
+
+### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
+
+Auto-route preference:
+
+1. `FLASH_ATTN` — when `flash-attn` is installed
+2. `TORCH_SDPA` — fallback
+
+`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.
+
+## 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 |
+|---|---|---|---|---|
+| 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.
+
+## Known Limitations
+
+### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
+
+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).
+
+**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Lightricks/LTX-2 ...
+```
+
+### FA4 not yet integrated
+
+FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
+
+## 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.
+
 ## SageAttention Installation
 
 vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
@@ -115,11 +185,33 @@ Quick check:
 python -c "import sageattention; print(sageattention.__file__)"
 ```
 
-## Usage
+## Usage Examples
 
-### Enable SageAttention
+### Default (auto-route)
 
-Example: HunyuanVideo-1.5 text-to-video
+```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 \
@@ -128,7 +220,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
     --height 480 --width 832 --num-frames 33 \
     --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
     --tensor-parallel-size 2 \
-    --output ../tmp/hv15_modelopt_sage.mp4
+    --output hv15_sage.mp4
 ```
 
 Example: Wan2.2 TI2V 5B
@@ -173,13 +265,13 @@ python examples/offline_inference/text_to_video/text_to_video.py \
 
 ## Validation Guidance
 
-Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
+Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
 
 Always compare:
 
 - End-to-end runtime
-- DiT / diffusion stage runtime
-- Output quality against a known-good baseline
+- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
+- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
 
 At minimum, keep the same:
 
@@ -187,6 +279,19 @@ At minimum, keep the same:
 - prompt
 - seed
 - resolution
-- frame count
-- inference steps
-- parallel config
+- frame count / step count
+- parallel config (TP / CFG-parallel / Ulysses degrees)
+
+## Reproducing the Benchmark Table
+
+The end-to-end numbers above were collected by running `text_to_video.py` /
+`text_to_image.py` with the same prompt and seed while varying
+`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
+backends without loading a model:
+
+```bash
+python benchmarks/diffusion/bench_attention_backends.py --preset hv15
+```
+
+It runs all three BF16 backends on representative DiT attention shapes and
+prints a ranking table at the end.

vllm_omni/diffusion/attention/backends/cudnn_attn.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
+
+logger = init_logger(__name__)
+
+
+class CuDNNAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # cuDNN 9.5+ FMHA on Blackwell supports any head_dim divisible by 8
+        # up to 256 for BF16/FP16. Empty list = "accept any"; the sdpa_kernel
+        # fallback chain handles shapes cuDNN can't take.
+        return []
+
+    @staticmethod
+    def get_name() -> str:
+        return "CUDNN_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["CuDNNAttentionImpl"]:
+        return CuDNNAttentionImpl
+
+
+class CuDNNAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        attention_mask = None
+        if attn_metadata:
+            attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode="broadcast_k")
+
+        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+        # Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a
+        # PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask,
+        # cuDNN gets runtime-disabled in the same call and the dispatcher falls
+        # through to MATH even though cuDNN alone handles the mask fine
+        # (~11 ms vs ~215 ms for MATH on sm_120 HV-1.5 shapes).
+        #
+        # Fall back to the default SDPA dispatcher if cuDNN rejects the shape,
+        # e.g. under torch.compile where Dynamo sees a symbolic head_dim and
+        # cuDNN's kernel selection fails (observed in LTX-2 audio attention).
+        # The unpinned dispatcher then picks EFFICIENT/MATH instead of raising.
+        try:
+            with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    query,
+                    key,
+                    value,
+                    attn_mask=attention_mask,
+                    dropout_p=0.0,
+                    is_causal=self.causal,
+                    scale=self.softmax_scale,
+                )
+        except RuntimeError as e:
+            if "No available kernel" not in str(e):
+                raise
+            logger.warning_once(
+                "cuDNN SDPA rejected this shape; falling back to default SDPA dispatcher. "
+                "Set DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA for the full dispatcher path on every call."
+            )
+            output = torch.nn.functional.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                is_causal=self.causal,
+                scale=self.softmax_scale,
+            )
+        return output.permute(0, 2, 1, 3)

vllm_omni/diffusion/attention/backends/flashinfer_attn.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from flashinfer.prefill import single_prefill_with_kv_cache
+
+    HAS_FLASHINFER = True
+except Exception as e:
+    HAS_FLASHINFER = False
+    logger.warning(
+        "FlashInfer is unavailable; FLASHINFER_ATTN backend will not work. Reason: %s",
+        e,
+    )
+
+
+class FlashInferAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        # FlashInfer single_prefill_with_kv_cache accepts a ``custom_mask`` kwarg
+        # (2D boolean, ``True`` = keep) for non-causal attention. See reviewer
+        # comment on #3079 and the flashinfer docs:
+        #   https://docs.flashinfer.ai/generated/flashinfer.prefill.single_prefill_with_kv_cache.html
+        return True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # FlashInfer dense prefill is well-tested for these head_dims on
+        # Ampere/Hopper/Blackwell. Covers the dominant diffusion DiT shapes
+        # (SD3 = 64, Flux/HV/Wan = 128, joint-attn = 256).
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashInferAttentionImpl"]:
+        return FlashInferAttentionImpl
+
+
+class FlashInferAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    @staticmethod
+    def _pack_mask_for_flashinfer(
+        attn_mask: torch.Tensor, batch_size: int, qo_len: int, kv_len: int
+    ) -> torch.Tensor | None:
+        """Convert a diffusion-style attn_mask into the boolean form
+        FlashInfer's ``custom_mask`` expects (``True`` = keep).
+
+        Returns either ``(qo_len, kv_len)`` (shared across the batch) or
+        ``(batch_size, qo_len, kv_len)`` (per-sample), or ``None`` when the
+        mask is all-keep (elide). Only boolean masks are handled here;
+        additive/float masks raise ``ValueError`` so the caller falls back to
+        SDPA, which applies them with the correct softmax semantics. Shape
+        mismatches also raise ``ValueError``.
+        """
+        mask = attn_mask
+        if mask.dtype != torch.bool:
+            # Additive masks (0 / -inf / -1e4 / finfo.min) cannot be faithfully
+            # reduced to a boolean keep-mask here; SDPA handles them correctly.
+            raise ValueError(
+                f"non-boolean attn_mask (dtype={mask.dtype}); FlashInfer custom_mask "
+                "is boolean-only — deferring to SDPA"
+            )
+        # Diffusion masks arrive as (qo,kv), (1,1,kv), (B,1,1,kv), (B,1,qo,kv)
+        # or (B,H,qo,kv). The mask is identical across heads, so collapse the
+        # head dim, but keep a real batch dim — indexing mask[0] would reuse
+        # sample 0's padding for every sample (wrong under CFG / mixed lengths).
+        if mask.dim() == 4:
+            mask = mask[:, 0]  # (B, qo|1, kv)
+        if mask.dim() == 3 and mask.shape[0] == 1:
+            mask = mask[0]  # (qo|1, kv) — shared across the batch
+        try:
+            if mask.dim() >= 3:
+                mask = mask.broadcast_to((batch_size, qo_len, kv_len))
+            else:
+                mask = mask.broadcast_to((qo_len, kv_len))
+        except RuntimeError as e:
+            raise ValueError(
+                f"attn_mask shape {tuple(attn_mask.shape)} cannot broadcast to "
+                f"(batch={batch_size}, qo_len={qo_len}, kv_len={kv_len})"
+            ) from e
+        if mask.all():
+            return None
+        # ``broadcast_to`` returns a non-contiguous view; materialize for the
+        # kernel, which reads from GPU memory directly.
+        return mask.contiguous()
+
+    def _sdpa_fallback(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+
+        impl = SDPAImpl(
+            num_heads=query.shape[2],
+            head_size=query.shape[3],
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+        )
+        return impl.forward_cuda(query, key, value, attn_metadata)
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        if not HAS_FLASHINFER:
+            raise ImportError(
+                "FLASHINFER_ATTN backend requires flashinfer. "
+                "Install it or set DIFFUSION_ATTENTION_BACKEND to another backend."
+            )
+
+        # Try the custom_mask path; if the mask can't be packed into the
+        # (qo_len, kv_len) layout FlashInfer expects, fall back to SDPA
+        # instead of risking an illegal-memory-access crash in the kernel.
+        # Input layout is (B, S, H, D); FlashInfer dense prefill takes (S, H, D).
+        batch_size = query.shape[0]
+
+        custom_mask = None
+        if attn_metadata is not None and attn_metadata.attn_mask is not None:
+            try:
+                custom_mask = self._pack_mask_for_flashinfer(
+                    attn_metadata.attn_mask,
+                    batch_size=batch_size,
+                    qo_len=query.shape[1],
+                    kv_len=key.shape[1],
+                )
+            except ValueError as e:
+                logger.debug("Falling back to SDPA for mask path: %s", e)
+                return self._sdpa_fallback(query, key, value, attn_metadata)
+            # FlashInfer cannot combine causal masking with a custom_mask; rather
+            # than silently dropping the explicit mask (diverging from SDPA), let
+            # SDPA handle the causal+mask case correctly.
+            if custom_mask is not None and self.causal:
+                logger.debug("causal=True with explicit attn_mask; deferring to SDPA")
+                return self._sdpa_fallback(query, key, value, attn_metadata)
+
+        outputs = []
+        for b in range(batch_size):
+            kwargs: dict = {
+                "sm_scale": self.softmax_scale,
+                "causal": self.causal,
+                "return_lse": False,
+            }
+            if custom_mask is not None:
+                kwargs["custom_mask"] = custom_mask if custom_mask.dim() == 2 else custom_mask[b]
+            out = single_prefill_with_kv_cache(query[b], key[b], value[b], **kwargs)
+            outputs.append(out)
+
+        if batch_size == 1:
+            return outputs[0].unsqueeze(0)
+        return torch.stack(outputs, dim=0)

vllm_omni/diffusion/attention/backends/registry.py

@@ -59,6 +59,8 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"
     TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"
     SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"
+    CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
+    FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
 
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).

vllm_omni/diffusion/diffusion_engine.py

@@ -691,6 +691,10 @@ class DiffusionEngine:
             dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
             prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
 
+        # Audio pipelines round audio token count from num_frames; the default
+        # of 1 yields seq_len=1 K/V which cuDNN SDPA refuses under torch.compile.
+        # 2 is the minimum that produces audio_num_frames > 1.
+        num_frames = 2 if supports_audio_input or supports_audio_output(self.od_config.model_class_name) else 1
         req = OmniDiffusionRequest(
             prompts=[prompt],
             request_ids=["dummy_req_id"],
@@ -698,6 +702,7 @@ class DiffusionEngine:
                 height=height,
                 width=width,
                 num_inference_steps=num_inference_steps,
+                num_frames=num_frames,
                 # Keep warmup path minimal and robust across text encoders.
                 # Some models may fail when warmup implicitly triggers
                 # classifier-free guidance with an empty negative prompt.

vllm_omni/diffusion/layers/rope.py

@@ -94,10 +94,18 @@ class RotaryEmbedding(CustomOp):
         self.interleaved = not is_neox_style
         self.apply_rotary_emb_flash_attn = None
         self.has_mindie = False
+        # ``find_spec("flash_attn")`` is True as long as *any* package publishes
+        # the ``flash_attn`` namespace — including ``flash-attn-4``, which ships
+        # only ``flash_attn.cute`` and no ``flash_attn.ops``. Guard the import
+        # so a partial namespace doesn't crash the RoPE init; the CUDA forward
+        # path uses ``vllm.vllm_flash_attn.layers.rotary`` anyway.
         if find_spec("flash_attn") is not None:
-            from flash_attn.ops.triton.rotary import apply_rotary
+            try:
+                from flash_attn.ops.triton.rotary import apply_rotary
 
-            self.apply_rotary_emb_flash_attn = apply_rotary
+                self.apply_rotary_emb_flash_attn = apply_rotary
+            except ImportError:
+                pass
         if find_spec("mindiesd") is not None:
             self.has_mindie = True
 

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -67,13 +67,20 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) ->
 
 def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
     cos, sin = freqs
-    x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1)  # [B, S, C // 2]
+    # Concrete pair count instead of -1 keeps SDPA shape static under torch.compile.
+    x_real, x_imag = x.unflatten(2, (x.shape[2] // 2, 2)).unbind(-1)  # [B, S, C // 2]
     x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
     out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
     return out
 
 
-def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+def apply_split_rotary_emb(
+    x: torch.Tensor,
+    freqs: tuple[torch.Tensor, torch.Tensor],
+    *,
+    head_dim: int,
+) -> torch.Tensor:
+    # `head_dim` is plumbed in (not inferred via `-1`) so SDPA shape stays static under torch.compile.
     cos, sin = freqs
 
     x_dtype = x.dtype
@@ -83,7 +90,7 @@ def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Ten
         # The cos/sin batch dim may only be broadcastable, so take batch size from x
         b = x.shape[0]
         _, h, t, _ = cos.shape
-        x = x.reshape(b, t, h, -1).swapaxes(1, 2)
+        x = x.reshape(b, t, h, head_dim).swapaxes(1, 2)
         needs_reshape = True
 
     # Split last dim (2*r) into (d=2, r)
@@ -110,7 +117,7 @@ def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Ten
     out = out.reshape(*out.shape[:-2], last)
 
     if needs_reshape:
-        out = out.swapaxes(1, 2).reshape(b, t, -1)
+        out = out.swapaxes(1, 2).reshape(b, t, h * head_dim)
 
     out = out.to(dtype=x_dtype)
     return out
@@ -488,12 +495,14 @@ class LTX2AudioVideoAttnProcessor:
                     key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
                 )
             elif attn.rope_type == "split":
-                query = apply_split_rotary_emb(query, query_rotary_emb)
-                key = apply_split_rotary_emb(key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb)
+                query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim)
+                key = apply_split_rotary_emb(
+                    key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, head_dim=attn.head_dim
+                )
 
-        query = query.unflatten(2, (attn.heads, -1))
-        key = key.unflatten(2, (attn.heads, -1))
-        value = value.unflatten(2, (attn.heads, -1))
+        query = query.unflatten(2, (attn.heads, attn.head_dim))
+        key = key.unflatten(2, (attn.heads, attn.head_dim))
+        value = value.unflatten(2, (attn.heads, attn.head_dim))
 
         attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
         hidden_states = attn.attn(query, key, value, attn_metadata)
@@ -502,7 +511,7 @@ class LTX2AudioVideoAttnProcessor:
 
         # LTX-2.3: per-head gated attention
         if attn.to_gate_logits is not None:
-            hidden_states = hidden_states.unflatten(2, (attn.heads, -1))  # [B, T, H, D]
+            hidden_states = hidden_states.unflatten(2, (attn.heads, attn.head_dim))  # [B, T, H, D]
             # 2.0 * sigmoid so zero-init gates produce 1.0 (identity)
             gates = 2.0 * torch.sigmoid(gate_logits)  # [B, T, H]
             hidden_states = hidden_states * gates.unsqueeze(-1)
@@ -1391,9 +1400,11 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
             # Reshape freqs to be compatible with multi-head attention
             b = cos_freq.shape[0]
             t = cos_freq.shape[1]
+            # Concrete per-head dim instead of -1 — see apply_split_rotary_emb
+            r = self.dim // self.num_attention_heads // 2
 
-            cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1)
-            sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1)
+            cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, r)
+            sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, r)
 
             cos_freqs = torch.swapaxes(cos_freq, 1, 2)  # (B,H,T,D//2)
             sin_freqs = torch.swapaxes(sin_freq, 1, 2)  # (B,H,T,D//2)

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -157,6 +157,9 @@ class _VideoAudioScheduler:
 
 
 class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
+    # Audio is diffused jointly with video; warmup must size audio tokens.
+    support_audio_output = True
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py

@@ -115,6 +115,9 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
     - CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU
     """
 
+    # Audio is diffused jointly with video; warmup must size audio tokens.
+    support_audio_output = True
+
     def __init__(
         self,
         *,

vllm_omni/platforms/cuda/platform.py

@@ -58,14 +58,32 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
     ) -> str:
         from vllm_omni.diffusion.envs import PACKAGES_CHECKER
 
-        # Check compute capability for Flash Attention support
-        # Flash Attention requires compute capability >= 8.0 and < 10.0
+        # Check compute capability for Flash Attention support.
+        # FA requires sm_80+. Blackwell (sm_10x/sm_12x) only works with FA builds
+        # that include the Blackwell CUTE kernel — plain FA2 will crash there.
+        #
+        # Known Blackwell SKUs:
+        #   sm_100 = B200 / GB200 (datacenter)
+        #   sm_103 = B300 / GB300 (Blackwell Ultra)
+        #   sm_120 = RTX Pro 6000, RTX 50-series (consumer)
+        #   sm_121 = consumer Blackwell refresh
+        _known_blackwell_sms = {(10, 0), (10, 3), (12, 0), (12, 1)}
         compute_capability = cls.get_device_capability()
         compute_supported = False
+        is_blackwell = False
+        sm_str = ""
         if compute_capability is not None:
             major, minor = compute_capability
             capability = major * 10 + minor
-            compute_supported = 80 <= capability < 100
+            compute_supported = capability >= 80
+            sm_str = f"sm_{major}{minor}"
+            # Accept major in {10, 11, 12} to cover future Blackwell refreshes.
+            is_blackwell = major in (10, 11, 12)
+            if is_blackwell and (major, minor) not in _known_blackwell_sms:
+                logger.info(
+                    "Detected Blackwell-class GPU %s (untested variant); routing to CUDNN_ATTN with SDPA fallback.",
+                    sm_str,
+                )
 
         # Check if FA packages are available
         packages_info = PACKAGES_CHECKER.get_packages_info()
@@ -74,13 +92,38 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         # Both compute capability and packages must be available for FA
         flash_attn_supported = compute_supported and packages_available
 
+        # cuDNN 9.5+ ships Blackwell FMHA kernels. If the runtime is older,
+        # the CUDNN_ATTN default would still work via internal fallback but
+        # without the tuned Blackwell path, so we skip routing there.
+        cudnn_version = torch.backends.cudnn.version() or 0
+        cudnn_blackwell_ready = cudnn_version >= 90500
+
+        # FlashInfer edges cuDNN by ~4% at the kernel level on sm_120 but
+        # regresses ~2x at e2e on HV-1.5 because its dense-prefill path can't
+        # take 2D attn_masks and the SDPA fallback dispatches to
+        # EFFICIENT_ATTENTION (~25 ms) instead of the cuDNN mask path (~11 ms).
+        # CUDNN_ATTN pins sdpa_kernel([CUDNN_ATTENTION]) directly so masked
+        # calls keep the cuDNN path. Blackwell default prefers CUDNN_ATTN;
+        # users can opt into FLASHINFER_ATTN explicitly for no-mask workloads.
+        flashinfer_available = False
+        try:
+            import flashinfer  # noqa: F401
+
+            flashinfer_available = True
+        except Exception as e:
+            # A partially installed / ABI-mismatched wheel can raise OSError or
+            # RuntimeError from extension loading, not just ImportError. This
+            # runs during default backend selection, so a probe failure must
+            # not abort startup — just treat FlashInfer as unavailable.
+            logger.debug("FlashInfer probe failed (%s); treating as unavailable", e)
+
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
                 if not compute_supported:
                     logger.warning(
-                        "Flash Attention requires GPU with compute capability >= 8.0 "
-                        "and < 10.0. Falling back to TORCH_SDPA backend."
+                        "Flash Attention requires GPU with compute capability >= 8.0. "
+                        "Falling back to TORCH_SDPA backend."
                     )
                 elif not packages_available:
                     logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.")
@@ -90,6 +133,29 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
+        if is_blackwell and cudnn_blackwell_ready:
+            logger.info(
+                "Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell %s, cuDNN %d)",
+                sm_str,
+                cudnn_version,
+            )
+            return DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
+
+        if is_blackwell and flashinfer_available:
+            logger.info(
+                "Defaulting to diffusion attention backend FLASHINFER_ATTN (Blackwell %s, cuDNN unavailable)",
+                sm_str,
+            )
+            return DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
+
+        if is_blackwell and not cudnn_blackwell_ready:
+            logger.warning(
+                "Detected Blackwell %s but cuDNN %d < 9.5 — no tuned Blackwell FMHA. "
+                "Falling through to FLASH_ATTN / SDPA.",
+                sm_str,
+                cudnn_version,
+            )
+
         if flash_attn_supported:
             logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
             return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()