diff --git a/benchmarks/diffusion/bench_attention_backends.py b/benchmarks/diffusion/bench_attention_backends.py new file mode 100644 index 000000000..0d34b2192 --- /dev/null +++ b/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() diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 4a19ad8c4..75915377f 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/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. diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py new file mode 100644 index 000000000..f27fe1870 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/flashinfer_attn.py b/vllm_omni/diffusion/attention/backends/flashinfer_attn.py new file mode 100644 index 000000000..1d2fa893b --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index a77223d16..356285654 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/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). diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index c13bd3c0c..8fd85fd9e 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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. diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index b5c4f9df9..b4fa1209f 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 64bb67e22..a1d72cf66 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/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) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index c92fd687f..29059a6d6 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index f3bc7c928..c16f74eef 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/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, *, diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index a3cc32c52..a5057739f 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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()