[bugfix]: preserve FlashInfer correction under autocast

fastvideo-kernel/benchmarks/bench_attn_kernels.py benches all four
attention kernels on the real logged Wan2.1-1.3B shapes plus an L sweep
(1k/4k/16k/32760, H=12 D=128 B=1):

  bf16_flash_attn|bf16_sdpa  high-precision baseline (labels which ran)
  fp8_sageattn               SageAttention qk-int8/pv-fp8 (PR #1496 path)
  nvfp4_local                attn_qat_infer sageattn_blackwell (sm_120a)
  nvfp4_flashinfer           flashinfer nvfp4_attention_sm120 (SM 12.0/12.1)

Per kernel x shape: CUDA-event latency (mean/std), TFLOPS
(4*B*H*L^2*D), and accuracy (cos / max abs err / mean rel err) vs a
query-chunked fp32 SDPA reference. Inputs are biased (+1.0 mean shift
on q/k) by default because zero-mean randn is blind to the flashinfer
per-block-mean qk-correction head-broadcast bug; --include-zeromean
adds the zero-mean rows for contrast. Lanes probe availability + arch
gates and skip with a printed reason; oversized shapes are preflighted
(chunked reference, qk_correction workspace check) instead of OOMing.

Replaces bench_nvfp4_fp8_attention.py (logic folded in; docs pointer
updated, CPU unit tests renamed to test_attn_kernels_bench.py and
extended for the new --shapes/--bias surface).

docs/inference/optimizations.md

@@ -27,6 +27,7 @@ This page describes the various options for speeding up generation times in Fast
 - Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`
 - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
 - Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE`
+- FlashInfer SageAttention3 (SM120 NVFP4): `FASTVIDEO_ATTENTION_BACKEND=FLASHINFER_SAGE_ATTN3`
 - Attn-QAT inference (modified SageAttention3 FP4, sm_120/RTX 5090): `FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER`
 - Video MoBA Attention: `FASTVIDEO_ATTENTION_BACKEND=VMOBA_ATTN`
 - Sparse Linear Attention: `FASTVIDEO_ATTENTION_BACKEND=SLA_ATTN`
@@ -237,6 +238,60 @@ If you are using `uv` and `torch==2.8.0`, make sure that
 
 To use Sage Attention 3 in FastVideo, follow the `README.md` in the linked repository to install the package from source.
 
+### FlashInfer SageAttention3
+
+**`FLASHINFER_SAGE_ATTN3`**
+
+Wraps [FlashInfer](https://github.com/flashinfer-ai/flashinfer)'s dense SM120
+NVFP4 attention (`nvfp4_attention_sm120_quantize_qkv` +
+`nvfp4_attention_sm120_fwd`, added in
+[flashinfer#3640](https://github.com/flashinfer-ai/flashinfer/pull/3640)).
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=FLASHINFER_SAGE_ATTN3 python example.py
+```
+
+#### Kernel requirements
+
+- Compute capability **12.0 only** (consumer Blackwell, e.g. RTX 5090), CUDA 12.8+
+- `head_dim` 64 or 128, fp16/bf16, equal-shape q/k/v (self-attention)
+- `flashinfer-python` installed (the kernel is JIT-built on first use)
+
+Where any of these do not hold — including when flashinfer is not installed at
+all — the backend logs one warning and transparently falls back to torch SDPA,
+so it can run anywhere.
+
+#### Attention shape logging
+
+The fallback exists on purpose: this backend doubles as a shape-collection
+vehicle for tuning the SageAttention3 kernel on SM 12.1 (DGX Spark). Set
+
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=FLASHINFER_SAGE_ATTN3   # or SAGE_ATTN_THREE
+export FASTVIDEO_ATTN_SHAPE_LOG=/tmp/attn_shapes.jsonl
+```
+
+and run your normal workload (no GPU/kernel required). Each process writes
+`/tmp/attn_shapes.jsonl.rank<N>` (or `.pid<N>` when not distributed) with one
+JSON line per unique attention shape:
+
+```json
+{"backend": "FLASHINFER_SAGE_ATTN3", "batch": 1, "num_q_heads": 12,
+ "num_kv_heads": 12, "seq_len_q": 75600, "seq_len_kv": 75600, "head_dim": 128,
+ "dtype": "torch.bfloat16", "causal": false, "sm_scale": 0.0884, "count": 1200,
+ "torch": "2.8.0", "device": "NVIDIA GB10", "sm": "12.1"}
+```
+
+Counts accumulate per shape; the file is flushed periodically and at exit.
+Hand the JSONL files to the kernel folks. The `SAGE_ATTN_THREE` backend records
+to the same log, so either implementation can be used for collection.
+
+A collected log can directly drive the unified attention kernel benchmark
+(per-shape latency, TFLOPS, and accuracy for the bf16 baseline
+(flash-attn/SDPA), SageAttention FP8, the local SageAttention3 port, and
+FlashInfer NVFP4):
+`python fastvideo-kernel/benchmarks/bench_attn_kernels.py --shapes-jsonl /tmp/attn_shapes.jsonl.pid<N>`.
+
 ### V-MoBA / SLA / SageSLA
 
 These backends are model-specific and require the corresponding kernels and

examples/inference/optimizations/FastWan_QAD_TAEHV.py

@@ -46,8 +46,11 @@ OUTPUT_PATH = "video_samples"
 # guidance 1.0). Loaded on top of the base Wan2.1 pipeline; pass
 # ``--distilled_model ''`` to run the base weights instead.
 DEFAULT_DISTILLED_MODEL = "FastVideo/FastWan-QAD-1.3B"
-DISTILLED_WEIGHTS_FILE = (
-    "generator_inference_transformer/diffusion_pytorch_model.safetensors"
+# The HF repo moved the transformer weights; try current layout first,
+# then the legacy path (pre-2026-07 uploads).
+DISTILLED_WEIGHTS_FILES = (
+    "transformer/diffusion_pytorch_model.safetensors",
+    "generator_inference_transformer/diffusion_pytorch_model.safetensors",
 )
 
 # TAEHV checkpoint for Wan2.1. Clone https://github.com/madebyollin/taehv to get
@@ -101,7 +104,14 @@ def resolve_distilled_weights(hf_id: str) -> str:
     if os.path.exists(hf_id):
         return hf_id
     from huggingface_hub import hf_hub_download
-    return hf_hub_download(repo_id=hf_id, filename=DISTILLED_WEIGHTS_FILE)
+    from huggingface_hub.errors import EntryNotFoundError
+    for filename in DISTILLED_WEIGHTS_FILES:
+        try:
+            return hf_hub_download(repo_id=hf_id, filename=filename)
+        except EntryNotFoundError:
+            continue
+    raise FileNotFoundError(
+        f"None of {DISTILLED_WEIGHTS_FILES} found in {hf_id!r}")
 
 
 @contextlib.contextmanager

fastvideo-kernel/benchmarks/bench_attn_kernels.py

@@ -0,0 +1,542 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Benchmark the four attention kernels: speed, accuracy, TFLOPS.
+
+Lanes (each probes availability + arch gate and skips cleanly, with a
+printed reason, where it cannot run):
+
+  bf16_flash_attn / bf16_sdpa
+                    High-precision baseline: flash-attn (flash_attn_func)
+                    when importable, else torch SDPA — the lane name says
+                    which one is running. (Prefix follows --dtype.)
+  fp8_sageattn      ``sageattention.sageattn`` — the attention op activated by
+                    examples/inference/optimizations/fp8_wan2_1_1_3b.py
+                    (PR #1496, FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN →
+                    fastvideo/attention/backends/sage_attn.py). sageattn
+                    auto-dispatches per-arch; on sm89+ it selects the
+                    qk-int8 / pv-fp8 CUDA kernels.
+  nvfp4_local       fastvideo-kernel/attn_qat_infer (local SageAttention3
+                    port, ``sageattn_blackwell``). Built for sm_120a only.
+  nvfp4_flashinfer  flashinfer.nvfp4_attention_sm120_quantize_qkv +
+                    nvfp4_attention_sm120_fwd. Compute capability 12.0/12.1,
+                    head_dim 64/128.
+
+Accuracy for every lane is measured against an fp32 torch-SDPA ground truth
+(chunked over query length so the score matrix never materializes whole)
+computed on the same fixed-seed inputs. Reported: cosine similarity,
+max abs err, mean rel err.
+
+Inputs are BIASED by default: q and k get a nonzero mean shift (--bias, v
+stays zero-mean so a constant output offset doesn't inflate cosine
+similarity). Zero-mean randn is blind to mean-dependent quantization errors —
+biased inputs caught flashinfer's former per-block-mean qk-correction bug
+(fixed by flashinfer#3838). Use
+--include-zeromean to also run the zero-mean rows for contrast.
+
+All lanes run at the [B, H, L, D] interface.
+
+Usage:
+  python bench_attn_kernels.py                          # default shapes
+  python bench_attn_kernels.py --shapes 1x12x32760x128,1x24x4096x64
+  python bench_attn_kernels.py --causal --json out.json
+  python bench_attn_kernels.py --shapes-jsonl /tmp/attn_shapes.jsonl.pid1234
+      # shapes collected via FASTVIDEO_ATTN_SHAPE_LOG (fastvideo/attention/
+      # shape_logger.py); recorded causal/dtype/sm_scale are honored.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import inspect
+import json
+import statistics
+import sys
+from collections.abc import Callable
+from pathlib import Path
+
+import torch
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+
+# ---------------------------------------------------------------------------
+# Pure helpers (CPU-unit-tested in fastvideo/tests/attention/
+# test_attn_kernels_bench.py) — keep them torch-free where possible.
+# ---------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class Shape:
+    batch: int
+    num_q_heads: int
+    num_kv_heads: int
+    seq_len_q: int
+    seq_len_kv: int
+    head_dim: int
+    causal: bool
+    dtype: str = "torch.bfloat16"
+    sm_scale: float | None = None
+
+    def label(self) -> str:
+        base = f"{self.batch}x{self.num_q_heads}x{self.seq_len_q}x{self.head_dim}"
+        if self.seq_len_kv != self.seq_len_q or self.num_kv_heads != self.num_q_heads:
+            base += f" (kv {self.num_kv_heads}x{self.seq_len_kv})"
+        if self.causal:
+            base += " causal"
+        return base
+
+
+def attn_tflops(shape: Shape, ms: float) -> float:
+    """TFLOPS of one dense attention forward.
+
+    FLOPs = 4 * B * H * Lq * Lk * D:
+      QK^T is B*H*Lq*Lk*D multiply-adds = 2*B*H*Lq*Lk*D FLOPs, and P@V is
+      another 2*B*H*Lq*Lk*D. Causal masking computes ~half the score matrix,
+      so FLOPs are halved. TFLOPS = FLOPs / (ms * 1e9).
+    """
+    flops = 4.0 * shape.batch * shape.num_q_heads * shape.seq_len_q * shape.seq_len_kv * shape.head_dim
+    if shape.causal:
+        flops *= 0.5
+    return flops / (ms * 1e9)
+
+
+def parse_shapes_cli(spec: str, causal: bool, dtype: str) -> list[Shape]:
+    """Parse --shapes 'BxHxLxD,BxHxLxD,...' (self-attention shapes)."""
+    shapes = []
+    for part in spec.split(","):
+        try:
+            b, h, length, d = (int(x) for x in part.lower().split("x"))
+        except ValueError as e:
+            raise argparse.ArgumentTypeError(f"bad shape {part!r}, expected BxHxLxD") from e
+        shapes.append(Shape(batch=b, num_q_heads=h, num_kv_heads=h, seq_len_q=length,
+                            seq_len_kv=length, head_dim=d, causal=causal, dtype=dtype))
+    return shapes
+
+
+def parse_shapes_jsonl(paths: list[str], default_dtype: str) -> list[Shape]:
+    """Parse fastvideo/attention/shape_logger.py output (one JSON object per
+    line; keys: batch, num_q_heads, num_kv_heads, seq_len_q, seq_len_kv,
+    head_dim, dtype, causal, sm_scale, count, ...). Dedupes across files,
+    preserving first-seen order."""
+    shapes: list[Shape] = []
+    seen: set[Shape] = set()
+    for path in paths:
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                rec = json.loads(line)
+                shape = Shape(
+                    batch=int(rec["batch"]),
+                    num_q_heads=int(rec["num_q_heads"]),
+                    num_kv_heads=int(rec.get("num_kv_heads", rec["num_q_heads"])),
+                    seq_len_q=int(rec["seq_len_q"]),
+                    seq_len_kv=int(rec.get("seq_len_kv", rec["seq_len_q"])),
+                    head_dim=int(rec["head_dim"]),
+                    causal=bool(rec.get("causal", False)),
+                    dtype=str(rec.get("dtype") or default_dtype),
+                    sm_scale=rec.get("sm_scale"),
+                )
+                if shape not in seen:
+                    seen.add(shape)
+                    shapes.append(shape)
+    return shapes
+
+
+@dataclasses.dataclass(frozen=True)
+class LaneSpec:
+    name: str
+    exact_cap: tuple[int, int] | None = None  # kernel runs on exactly this SM
+    caps: tuple[tuple[int, int], ...] | None = None  # kernel runs only on these SMs
+    min_cap: tuple[int, int] | None = None
+    head_dims: tuple[int, ...] | None = None  # None = any
+    same_shape: bool = False  # requires q/k/v equal shapes (self-attention)
+    same_heads: bool = False  # requires num_q_heads == num_kv_heads
+
+
+def skip_reason(spec: LaneSpec, shape: Shape, cap: tuple[int, int] | None,
+                unavailable: str | None) -> str | None:
+    """Why this lane cannot run this shape (None = it can)."""
+    if unavailable:
+        return unavailable
+    if cap is None:
+        return "no CUDA device"
+    if spec.exact_cap is not None and cap != spec.exact_cap:
+        return (f"requires SM {spec.exact_cap[0]}.{spec.exact_cap[1]}, "
+                f"device is SM {cap[0]}.{cap[1]}")
+    if spec.caps is not None and cap not in spec.caps:
+        allowed = "/".join(f"{a}.{b}" for a, b in spec.caps)
+        return f"requires SM {allowed}, device is SM {cap[0]}.{cap[1]}"
+    if spec.min_cap is not None and cap < spec.min_cap:
+        return (f"requires SM >= {spec.min_cap[0]}.{spec.min_cap[1]}, "
+                f"device is SM {cap[0]}.{cap[1]}")
+    if spec.head_dims is not None and shape.head_dim not in spec.head_dims:
+        return f"head_dim {shape.head_dim} not in {list(spec.head_dims)}"
+    if spec.same_shape and (shape.seq_len_q != shape.seq_len_kv or shape.num_q_heads != shape.num_kv_heads):
+        return "requires equal q/k/v shapes (self-attention)"
+    if spec.same_heads and shape.num_q_heads != shape.num_kv_heads:
+        return "requires num_q_heads == num_kv_heads"
+    return None
+
+
+def shape_workspace_bytes(shape: Shape, dtype_bytes: int, ref_chunk: int) -> int:
+    """Rough peak device memory to bench one shape: q/k/v in the bench dtype,
+    their fp32 copies + fp32 reference output for the ground truth, and the
+    fp32 score-matrix chunk (x2 for the math-backend softmax temporary)."""
+    numel_q = shape.batch * shape.num_q_heads * shape.seq_len_q * shape.head_dim
+    numel_kv = shape.batch * shape.num_kv_heads * shape.seq_len_kv * shape.head_dim
+    lq = shape.seq_len_q if shape.causal else min(ref_chunk, shape.seq_len_q)
+    scores = shape.batch * shape.num_q_heads * lq * shape.seq_len_kv * 4
+    return (numel_q + 2 * numel_kv) * (dtype_bytes + 4) + numel_q * 4 + 2 * scores
+
+
+def format_table(headers: list[str], rows: list[list[str]]) -> str:
+    """Aligned plain-text table (left-justified columns)."""
+    cols = list(zip(headers, *rows)) if rows else [(h,) for h in headers]
+    widths = [max(len(str(c)) for c in col) for col in cols]
+    lines = [
+        "  ".join(str(h).ljust(w) for h, w in zip(headers, widths)),
+        "  ".join("-" * w for w in widths),
+    ]
+    for row in rows:
+        lines.append("  ".join(str(c).ljust(w) for c, w in zip(row, widths)))
+    return "\n".join(lines)
+
+
+# ---------------------------------------------------------------------------
+# Lanes (GPU)
+# ---------------------------------------------------------------------------
+
+AttnFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor, bool, float], torch.Tensor]
+
+_REF_CHUNK = 2048
+
+
+@dataclasses.dataclass
+class Lane:
+    spec: LaneSpec
+    fn: AttnFn | None  # None when unavailable
+    unavailable: str | None = None
+
+
+def _run_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
+    kwargs: dict = {"is_causal": causal, "scale": scale}
+    if q.shape[1] != k.shape[1]:
+        kwargs["enable_gqa"] = True
+    return torch.nn.functional.scaled_dot_product_attention(q, k, v, **kwargs)
+
+
+def _ref_sdpa_fp32(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
+    """fp32 SDPA ground truth, chunked over query length.
+
+    Where flash/mem-efficient SDPA reject fp32 (e.g. SM121), torch falls back
+    to the math backend, which materializes the full [B, H, Lq, Lkv] score
+    matrix -- ~48 GiB for the 32760-token default shape (host OOM-kill on
+    unified-memory devices like GB10). Non-causal attention is exact under
+    query chunking; causal shapes stay unchunked (per-chunk is_causal would
+    mis-align the mask), so they keep the old memory ceiling.
+    """
+    if causal or q.shape[2] <= _REF_CHUNK:
+        return _run_sdpa(q, k, v, causal, scale)
+    return torch.cat([_run_sdpa(q[:, :, i:i + _REF_CHUNK], k, v, False, scale)
+                      for i in range(0, q.shape[2], _REF_CHUNK)], dim=2)
+
+
+def _load_baseline(dtype_tag: str) -> Lane:
+    """High-precision baseline: flash-attn when importable, else torch SDPA.
+
+    The lane name records which implementation is actually timed.
+    """
+    try:
+        from flash_attn import flash_attn_func  # noqa: PLC0415
+    except ImportError:
+        return Lane(LaneSpec(f"{dtype_tag}_sdpa"), _run_sdpa)
+
+    def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
+        # flash_attn_func takes [B, L, H, D]; only the head dim must be
+        # contiguous, so plain transposes are accepted without a copy.
+        out = flash_attn_func(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
+                              softmax_scale=scale, causal=causal)
+        return out.transpose(1, 2)
+
+    # flash-attn's fp16/bf16 kernels need sm80+.
+    return Lane(LaneSpec(f"{dtype_tag}_flash_attn", min_cap=(8, 0)), fn)
+
+
+def _load_fp8_sageattn() -> Lane:
+    # sageattn's fp8 (qk-int8 / pv-fp8) kernels need sm89+; the package
+    # dispatches per-arch internally.
+    spec = LaneSpec("fp8_sageattn", min_cap=(8, 9), same_heads=True)
+    try:
+        from sageattention import sageattn  # noqa: PLC0415
+    except ImportError as e:
+        return Lane(spec, None, f"sageattention not installed ({e})")
+
+    def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
+        return sageattn(q, k, v, tensor_layout="HND", is_causal=causal, sm_scale=scale)
+
+    return Lane(spec, fn)
+
+
+def _load_nvfp4_local(single_level_p_quant: bool) -> Lane:
+    spec = LaneSpec("nvfp4_local", exact_cap=(12, 0), head_dims=(64, 128), same_shape=True)
+    for p in (_REPO_ROOT, _REPO_ROOT / "fastvideo-kernel", _REPO_ROOT / "fastvideo-kernel" / "python"):
+        if str(p) not in sys.path:
+            sys.path.insert(0, str(p))
+    try:
+        from attn_qat_infer import sageattn_blackwell  # noqa: PLC0415
+    except ImportError as e:
+        return Lane(spec, None, f"attn_qat_infer import failed ({e}); build fastvideo-kernel with sm_120a")
+
+    def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
+        # sageattn_blackwell pads L to a multiple of 128 internally and trims.
+        return sageattn_blackwell(q, k, v, is_causal=causal,
+                                  single_level_p_quant=single_level_p_quant, sm_scale=scale)
+
+    return Lane(spec, fn)
+
+
+def _load_nvfp4_flashinfer() -> Lane:
+    # SM121 (GB10 / DGX Spark) is enabled by flashinfer#3897.
+    spec = LaneSpec("nvfp4_flashinfer", caps=((12, 0), (12, 1)), head_dims=(64, 128), same_shape=True)
+    try:
+        from flashinfer import (  # noqa: PLC0415
+            nvfp4_attention_sm120_fwd, nvfp4_attention_sm120_quantize_qkv)
+    except ImportError as e:
+        return Lane(spec, None, f"flashinfer nvfp4_attention_sm120 not installed ({e})")
+
+    quant_kwargs = {}
+    if "per_block_mean" in inspect.signature(nvfp4_attention_sm120_quantize_qkv).parameters:
+        quant_kwargs["per_block_mean"] = True
+
+    def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
+        seq_len = q.shape[2]
+        packed = nvfp4_attention_sm120_quantize_qkv(q, k, v, **quant_kwargs)
+        out, _lse = nvfp4_attention_sm120_fwd(*packed, sm_scale=scale, causal=causal, out_dtype=q.dtype)
+        # flashinfer pads seq_len to a multiple of 128; trim it back off.
+        return out[:, :, :seq_len, :]
+
+    return Lane(spec, fn)
+
+
+def build_lanes(args: argparse.Namespace) -> list[Lane]:
+    return [
+        _load_baseline(args.dtype),
+        _load_fp8_sageattn(),
+        _load_nvfp4_local(single_level_p_quant=args.p_quant == "single"),
+        _load_nvfp4_flashinfer(),
+    ]
+
+
+# ---------------------------------------------------------------------------
+# Measurement
+# ---------------------------------------------------------------------------
+
+_DTYPES = {"torch.bfloat16": torch.bfloat16, "torch.float16": torch.float16,
+           "bf16": torch.bfloat16, "fp16": torch.float16}
+
+
+def make_qkv(shape: Shape, seed: int, bias: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Fixed-seed q/k/v; q and k get a +bias mean shift. v stays zero-mean so
+    a constant output offset doesn't inflate the cosine-similarity column."""
+    gen = torch.Generator(device="cuda").manual_seed(seed)  # fixed seed per shape
+    dtype = _DTYPES.get(shape.dtype, torch.bfloat16)
+
+    def rand(heads: int, seq: int, shift: float) -> torch.Tensor:
+        t = torch.randn(shape.batch, heads, seq, shape.head_dim,
+                        generator=gen, device="cuda", dtype=dtype)
+        return t + shift if shift else t
+
+    return (rand(shape.num_q_heads, shape.seq_len_q, bias),
+            rand(shape.num_kv_heads, shape.seq_len_kv, bias),
+            rand(shape.num_kv_heads, shape.seq_len_kv, 0.0))
+
+
+def accuracy_stats(out: torch.Tensor, ref: torch.Tensor) -> dict[str, float]:
+    out = out.float()
+    diff = (out - ref).abs()
+    cos = torch.nn.functional.cosine_similarity(out.flatten(), ref.flatten(), dim=0)
+    return {
+        "cos_sim": cos.item(),
+        "max_abs_err": diff.max().item(),
+        # mean relative error: mean |out-ref| normalized by the reference's
+        # mean magnitude (elementwise |out-ref|/|ref| blows up near ref~0).
+        "mean_rel_err": (diff.mean() / ref.abs().mean().clamp_min(1e-12)).item(),
+    }
+
+
+def time_fn(fn: Callable[[], torch.Tensor], warmup: int, iters: int) -> dict[str, float]:
+    """CUDA-event latency stats (ms) over `iters` timed calls."""
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    times = []
+    for _ in range(iters):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return {"mean_ms": statistics.fmean(times), "std_ms": statistics.pstdev(times),
+            "median_ms": statistics.median(times)}
+
+
+def bench_shape(shape: Shape, bias: float, lanes: list[Lane], cap: tuple[int, int] | None,
+                args: argparse.Namespace) -> list[dict]:
+    dist = f"biased(+{bias:g})" if bias else "zeromean"
+    results = []
+    runnable = [(lane, skip_reason(lane.spec, shape, cap, lane.unavailable)) for lane in lanes]
+
+    oom_reason = None
+    if any(reason is None for _, reason in runnable):
+        dtype_bytes = 2  # bf16/fp16
+        need = shape_workspace_bytes(shape, dtype_bytes, _REF_CHUNK)
+        free = torch.cuda.mem_get_info()[0]
+        if need > 0.9 * free:
+            oom_reason = (f"would OOM: needs ~{need / 2**30:.1f} GiB workspace, "
+                          f"{free / 2**30:.1f} GiB free")
+            print(f"  SKIP shape: {oom_reason}")
+
+    q = k = v = ref = None
+    if oom_reason is None and any(reason is None for _, reason in runnable):
+        q, k, v = make_qkv(shape, args.seed, bias)
+        scale = shape.sm_scale or shape.head_dim**-0.5
+        # fp32 SDPA ground truth (accuracy baseline only; never timed).
+        ref = _ref_sdpa_fp32(q.float(), k.float(), v.float(), shape.causal, scale)
+
+    for lane, reason in runnable:
+        rec: dict = {"shape": shape.label(), "dist": dist, "lane": lane.spec.name,
+                     **dataclasses.asdict(shape)}
+        reason = reason or oom_reason
+        if reason is not None:
+            rec.update(status="skip", reason=reason)
+            results.append(rec)
+            continue
+        scale = shape.sm_scale or shape.head_dim**-0.5
+        try:
+            out = lane.fn(q, k, v, shape.causal, scale)
+            rec.update(accuracy_stats(out, ref))
+            del out
+            rec.update(status="ok", **time_fn(lambda: lane.fn(q, k, v, shape.causal, scale),
+                                              args.warmup, args.iters))
+            rec["tflops"] = attn_tflops(shape, rec["mean_ms"])
+        except Exception as e:  # lane failed at runtime: report, keep going
+            rec.update(status="fail", reason=f"{type(e).__name__}: {e}")
+        results.append(rec)
+    del q, k, v, ref
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return results
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+# The real logged workload (FASTVIDEO_ATTN_SHAPE_LOG on 5090/GB10/GB200 runs of
+# Wan2.1-1.3B 480p, 81 frames: latent grid 21x30x52 = 32760 tokens, 12 heads x
+# head_dim 128, bf16, non-causal) plus a seq-len sweep at the same head
+# geometry. 32760 appears once (deduped).
+DEFAULT_SHAPE_DIMS = [
+    (1, 12, 1024, 128),
+    (1, 12, 4096, 128),
+    (1, 12, 16384, 128),
+    (1, 12, 32760, 128),  # Wan2.1-1.3B 480p self-attention (the logged shape)
+]
+
+
+def default_shapes(causal: bool, dtype: str) -> list[Shape]:
+    return [
+        Shape(batch=b, num_q_heads=h, num_kv_heads=h, seq_len_q=length, seq_len_kv=length,
+              head_dim=d, causal=causal, dtype=dtype) for b, h, length, d in DEFAULT_SHAPE_DIMS
+    ]
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0],
+                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    p.add_argument("--shapes", metavar="BxHxLxD,...",
+                   help="comma-separated self-attention shapes; replaces the default suite")
+    p.add_argument("--shapes-jsonl", nargs="+", metavar="PATH",
+                   help="attention shape log(s) written via FASTVIDEO_ATTN_SHAPE_LOG "
+                   "(fastvideo/attention/shape_logger.py); replaces the default shapes")
+    p.add_argument("--causal", action="store_true",
+                   help="run --shapes/default shapes causally (jsonl shapes keep their recorded flag)")
+    p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16",
+                   help="input dtype for default shapes / jsonl records without one")
+    p.add_argument("--bias", type=float, default=1.0,
+                   help="mean shift added to q and k (biased inputs expose mean-dependent "
+                   "quantization bugs that zero-mean randn hides)")
+    p.add_argument("--include-zeromean", action="store_true",
+                   help="additionally run every shape with zero-mean inputs for contrast")
+    p.add_argument("--p-quant", choices=["single", "two"], default="single",
+                   help="nvfp4_local P-matrix quantization: single-level or two-level "
+                   "(sageattn_blackwell single_level_p_quant knob)")
+    p.add_argument("--warmup", type=int, default=10, help="warmup iterations per lane")
+    p.add_argument("--iters", type=int, default=50, help="timed iterations per lane")
+    p.add_argument("--seed", type=int, default=0, help="input seed (fixed per shape)")
+    p.add_argument("--json", metavar="PATH", help="also dump results as JSON")
+    return p.parse_args(argv)
+
+
+def _fmt(rec: dict, key: str, spec: str) -> str:
+    return format(rec[key], spec) if key in rec else "-"
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = parse_args(argv)
+    dtype_name = f"torch.{'bfloat16' if args.dtype == 'bf16' else 'float16'}"
+    if args.shapes:
+        shapes = parse_shapes_cli(args.shapes, args.causal, dtype_name)
+    elif args.shapes_jsonl:
+        shapes = parse_shapes_jsonl(args.shapes_jsonl, dtype_name)
+    else:
+        shapes = default_shapes(args.causal, dtype_name)
+
+    cap = None
+    meta: dict = {"torch": torch.__version__, "warmup": args.warmup, "iters": args.iters,
+                  "seed": args.seed, "bias": args.bias}
+    if torch.cuda.is_available():
+        cap = torch.cuda.get_device_capability(0)
+        meta.update(device=torch.cuda.get_device_name(0), sm=f"{cap[0]}.{cap[1]}")
+        print(f"Device: {meta['device']} (SM {meta['sm']}), torch {torch.__version__}")
+    else:
+        print("No CUDA device: all lanes will be skipped (shape/skip dry run).")
+
+    lanes = build_lanes(args)
+    biases = [args.bias] + ([0.0] if args.include_zeromean and args.bias else [])
+    results = []
+    for shape in shapes:
+        for bias in biases:
+            print(f"\n== {shape.label()}  dtype={shape.dtype}  bias={bias:g} ==")
+            results.extend(bench_shape(shape, bias, lanes, cap, args))
+
+    headers = ["shape", "dist", "lane", "status", "mean_ms", "std_ms", "TFLOPS",
+               "cos_sim", "max_abs_err", "mean_rel_err"]
+    rows = []
+    for rec in results:
+        status = rec["status"] if rec["status"] == "ok" else f"{rec['status']}: {rec['reason']}"
+        rows.append([
+            rec["shape"], rec["dist"], rec["lane"], status,
+            _fmt(rec, "mean_ms", ".3f"),
+            _fmt(rec, "std_ms", ".3f"),
+            _fmt(rec, "tflops", ".1f"),
+            _fmt(rec, "cos_sim", ".6f"),
+            _fmt(rec, "max_abs_err", ".4e"),
+            _fmt(rec, "mean_rel_err", ".4e"),
+        ])
+    print("\n" + format_table(headers, rows))
+    print("\nAccuracy is vs chunked fp32 torch-SDPA ground truth on identical inputs "
+          "(the fp32 reference itself is never timed).")
+
+    if args.json:
+        with open(args.json, "w") as f:
+            json.dump({"meta": meta, "results": results}, f, indent=2)
+        print(f"Wrote {args.json}")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/attention/backends/flashinfer_sageattn3.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FlashInfer SageAttention3 (SM120 dense NVFP4 attention) backend.
+
+Wraps ``flashinfer.nvfp4_attention_sm120_quantize_qkv`` +
+``flashinfer.nvfp4_attention_sm120_fwd`` (flashinfer-ai/flashinfer#3640).
+The kernel is gated to compute capability 12.0 (consumer Blackwell, e.g.
+RTX 5090), head_dim 64/128, fp16/bf16, and self-attention with equal
+q/k/v shapes. flashinfer pads seq_len up to a multiple of 128 internally;
+this backend trims the padding back off the output. Note flashinfer's own
+accuracy tests only cover seq_len multiples of 128 — padded-key softmax
+mass for other lengths is flashinfer's semantics, unverified here.
+
+Anywhere the kernel cannot run (flashinfer missing, non-SM120 device,
+unsupported head size/dtype/shape) the backend logs one warning and falls
+back to torch SDPA, so it stays usable as a shape-collection vehicle
+(``FASTVIDEO_ATTN_SHAPE_LOG``) on machines without the kernel — e.g.
+DGX Spark (SM121), which the collected shape data is being gathered for.
+
+Works with both ``LocalAttention`` and ``DistributedAttention`` (sequence
+parallel): inputs arrive as [B, L, H, D] and are transposed internally.
+"""
+
+import torch
+
+from fastvideo.attention import shape_logger
+from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+try:
+    from flashinfer import (nvfp4_attention_sm120_fwd, nvfp4_attention_sm120_quantize_qkv)
+except ImportError:
+    nvfp4_attention_sm120_fwd = None  # type: ignore[assignment]
+    nvfp4_attention_sm120_quantize_qkv = None  # type: ignore[assignment]
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
+_SUPPORTED_HEAD_DIMS = (64, 128)
+
+_fallback_warned = False
+
+
+def _fall_back(reason: str) -> bool:
+    """Log the first fallback reason once per process; returns False."""
+    global _fallback_warned
+    if not _fallback_warned:
+        logger.warning("FLASHINFER_SAGE_ATTN3: %s. Falling back to torch SDPA "
+                       "(subsequent fallbacks are silent).", reason)
+        _fallback_warned = True
+    return False
+
+
+def _can_use_flashinfer(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, dropout: float) -> bool:
+    """Whether the flashinfer SM120 kernel can serve this call ([B, H, L, D] inputs)."""
+    if nvfp4_attention_sm120_fwd is None:
+        return _fall_back("flashinfer with nvfp4_attention_sm120 is not installed")
+    if not q.is_cuda:
+        return _fall_back(f"device {q.device} is not CUDA")
+    cap = torch.cuda.get_device_capability(q.device)
+    # SM121 (GB10 / DGX Spark) is enabled by flashinfer#3897.
+    if cap not in ((12, 0), (12, 1)):
+        return _fall_back(f"compute capability {cap[0]}.{cap[1]} is not 12.0/12.1 (kernel is SM120/SM121-only)")
+    if q.dtype not in _SUPPORTED_DTYPES:
+        return _fall_back(f"dtype {q.dtype} is not fp16/bf16")
+    if q.shape[-1] not in _SUPPORTED_HEAD_DIMS:
+        return _fall_back(f"head_dim {q.shape[-1]} not in {_SUPPORTED_HEAD_DIMS}")
+    if not (q.shape == k.shape == v.shape):
+        return _fall_back(f"q/k/v shapes differ (q={tuple(q.shape)}, k={tuple(k.shape)}, "
+                          f"v={tuple(v.shape)}); kernel requires equal-shape self-attention")
+    if dropout > 0:
+        return _fall_back(f"dropout_p={dropout} is not supported by the kernel")
+    return True
+
+
+class FlashInferSageAttention3Backend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER_SAGE_ATTN3"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashInferSageAttention3Impl"]:
+        return FlashInferSageAttention3Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+
+class FlashInferSageAttention3Impl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """Input is [B, L, H, D] (FastVideo impl convention); output matches."""
+        if shape_logger.enabled:
+            shape_logger.record(FlashInferSageAttention3Backend.get_name(),
+                                query,
+                                key,
+                                value,
+                                causal=self.causal,
+                                sm_scale=self.softmax_scale)
+
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+
+        if _can_use_flashinfer(q, k, v, self.dropout):
+            seq_len = q.shape[2]
+            # FlashInfer computes qk_correction with an fp32 matmul and its
+            # forward requires that tensor to remain fp32. Ambient pipeline
+            # autocast would otherwise downcast the matmul result to bf16.
+            with torch.autocast(device_type=q.device.type, enabled=False):
+                packed = nvfp4_attention_sm120_quantize_qkv(q.contiguous(), k.contiguous(), v.contiguous())
+            output, _ = nvfp4_attention_sm120_fwd(*packed,
+                                                  sm_scale=self.softmax_scale,
+                                                  causal=self.causal,
+                                                  out_dtype=query.dtype)
+            # flashinfer pads seq_len to a multiple of 128; trim it back off.
+            return output[:, :, :seq_len, :].transpose(1, 2).contiguous()
+
+        attn_kwargs = {
+            "dropout_p": self.dropout,
+            "is_causal": self.causal,
+            "scale": self.softmax_scale,
+        }
+        if q.shape[1] != k.shape[1]:
+            attn_kwargs["enable_gqa"] = True
+        output = torch.nn.functional.scaled_dot_product_attention(q, k, v, **attn_kwargs)
+        return output.transpose(1, 2)

fastvideo/attention/backends/sage_attn3.py

@@ -3,6 +3,7 @@
 import torch
 from sageattn3 import sageattn3_blackwell
 
+from fastvideo.attention import shape_logger
 from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
                                                    AttentionMetadataBuilder)
 from fastvideo.logger import init_logger
@@ -78,6 +79,14 @@ class SageAttention3Impl(AttentionImpl):
     ) -> torch.Tensor:
         """Call sageattn3_blackwell directly. Input is already [B, H, L, D]
         and contiguous from preprocess_qkv."""
+        if shape_logger.enabled:
+            shape_logger.record("SAGE_ATTN_THREE",
+                                query,
+                                key,
+                                value,
+                                causal=self.causal,
+                                sm_scale=self.softmax_scale,
+                                layout="bhsd")
         output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
         return output
 

fastvideo/attention/shape_logger.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Env-gated attention shape logger.
+
+Set ``FASTVIDEO_ATTN_SHAPE_LOG=/path/to/shapes.jsonl`` to record every unique
+attention shape executed by the backends that call :func:`record`. Each unique
+(backend, batch, heads, seq lens, head_dim, dtype, causal) combination becomes
+one JSONL line with an occurrence count plus first-seen metadata (torch
+version, device name, SM arch).
+
+Purpose: collect the real attention workloads FastVideo runs so kernel authors
+(e.g. FlashInfer SageAttention3 on DGX Spark / SM121) can tune against actual
+shapes instead of guesses.
+
+Workers are separate processes, so each process writes its own file: the
+configured path is suffixed with ``.rank<N>`` when ``WORLD_SIZE>1``, else
+``.pid<N>``. The file is rewritten on flush (counts mutate, so append-only
+JSONL would go stale); flushes happen every ``_FLUSH_EVERY`` new unique shapes
+and at interpreter exit.
+
+Unset env var = zero overhead: call sites guard on the module-level
+``enabled`` bool and never enter this module's code.
+"""
+
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import threading
+
+import torch
+
+_LOG_PATH: str | None = os.environ.get("FASTVIDEO_ATTN_SHAPE_LOG")
+enabled: bool = _LOG_PATH is not None
+
+_FLUSH_EVERY = 1  # rewrite the file every N new unique shapes (crash-safety)
+
+_lock = threading.Lock()
+_entries: dict[tuple, dict] = {}
+_new_since_flush = 0
+_atexit_registered = False
+
+
+def _log_path() -> str:
+    assert _LOG_PATH is not None
+    if int(os.environ.get("WORLD_SIZE", "1")) > 1:
+        return f"{_LOG_PATH}.rank{os.environ.get('RANK', '0')}"
+    return f"{_LOG_PATH}.pid{os.getpid()}"
+
+
+def _device_meta(device: torch.device) -> dict:
+    meta: dict = {"torch": torch.__version__, "device": str(device), "sm": None}
+    if device.type == "cuda" and torch.cuda.is_available():
+        cap = torch.cuda.get_device_capability(device)
+        meta["device"] = torch.cuda.get_device_name(device)
+        meta["sm"] = f"{cap[0]}.{cap[1]}"
+    return meta
+
+
+def _flush_locked() -> None:
+    """Rewrite the whole snapshot; atomic via os.replace."""
+    path = _log_path()
+    tmp = path + ".tmp"
+    with open(tmp, "w") as f:
+        for entry in _entries.values():
+            f.write(json.dumps(entry) + "\n")
+    os.replace(tmp, path)
+
+
+def _flush_at_exit() -> None:
+    with _lock:
+        if _entries:
+            _flush_locked()
+
+
+def record(backend: str,
+           q: torch.Tensor,
+           k: torch.Tensor,
+           v: torch.Tensor,
+           *,
+           causal: bool,
+           sm_scale: float | None,
+           layout: str = "bshd") -> None:
+    """Record one attention call.
+
+    ``layout`` describes what the calling backend sees: ``"bshd"`` for
+    [batch, seq, heads, head_dim] (the FastVideo impl.forward convention) or
+    ``"bhsd"`` for [batch, heads, seq, head_dim] (backends that transpose in
+    ``preprocess_qkv``, e.g. SAGE_ATTN_THREE).
+    """
+    global _new_since_flush, _atexit_registered
+    if not enabled:
+        return
+    if layout == "bhsd":
+        batch, num_q_heads, seq_len_q, head_dim = q.shape
+        num_kv_heads, seq_len_kv = k.shape[1], k.shape[2]
+    else:
+        batch, seq_len_q, num_q_heads, head_dim = q.shape
+        num_kv_heads, seq_len_kv = k.shape[2], k.shape[1]
+    key = (backend, batch, num_q_heads, num_kv_heads, seq_len_q, seq_len_kv, head_dim, str(q.dtype), bool(causal))
+    with _lock:
+        entry = _entries.get(key)
+        if entry is not None:
+            entry["count"] += 1
+            return
+        _entries[key] = {
+            "backend": backend,
+            "batch": batch,
+            "num_q_heads": num_q_heads,
+            "num_kv_heads": num_kv_heads,
+            "seq_len_q": seq_len_q,
+            "seq_len_kv": seq_len_kv,
+            "head_dim": head_dim,
+            "dtype": str(q.dtype),
+            "causal": bool(causal),
+            "sm_scale": float(sm_scale) if sm_scale is not None else None,
+            "count": 1,
+            **_device_meta(q.device),
+        }
+        if not _atexit_registered:
+            atexit.register(_flush_at_exit)
+            _atexit_registered = True
+        _new_since_flush += 1
+        if _new_since_flush >= _FLUSH_EVERY:
+            _flush_locked()
+            _new_since_flush = 0

fastvideo/configs/models/dits/base.py

@@ -24,6 +24,7 @@ class DiTArchConfig(ArchConfig):
                                                  AttentionBackendEnum.TORCH_SDPA,
                                                  AttentionBackendEnum.VIDEO_SPARSE_ATTN,
                                                  AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE,
+                                                 AttentionBackendEnum.FLASHINFER_SAGE_ATTN3,
                                                  AttentionBackendEnum.ATTN_QAT_INFER,
                                                  AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.SLA_ATTN,
                                                  AttentionBackendEnum.SAGE_SLA_ATTN)

fastvideo/envs.py

@@ -208,11 +208,22 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention
     # - "SAGE_ATTN": use Sage Attention
     # - "SAGE_ATTN_THREE": use Sage Attention 3
+    # - "FLASHINFER_SAGE_ATTN3": use FlashInfer SageAttention3 (SM120 NVFP4);
+    #   falls back to torch SDPA where the kernel cannot run
     # FLASH_ATTN uses FlashAttention-3/2; to run FlashAttention-4 set
     # FASTVIDEO_FA4=1 as well (see below).
     "FASTVIDEO_ATTENTION_BACKEND":
     lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
 
+    # If set to a file path, attention backends that support it (currently
+    # FLASHINFER_SAGE_ATTN3 and SAGE_ATTN_THREE) record every unique attention
+    # shape they run to <path>.rank<N>/<path>.pid<N> as JSONL, with occurrence
+    # counts and device metadata. Used to hand real workload shapes to kernel
+    # authors (e.g. SageAttention3 tuning for DGX Spark / SM121).
+    # See fastvideo/attention/shape_logger.py. Read once at import time.
+    "FASTVIDEO_ATTN_SHAPE_LOG":
+    lambda: os.getenv("FASTVIDEO_ATTN_SHAPE_LOG", None),
+
     # If set (=1), the FLASH_ATTN backend uses FlashAttention-4
     # (flash_attn.cute). FA4 is opt-in and never auto-selected just because it
     # is installed. Below sm90, grad-enabled and GQA calls are routed to FA2

fastvideo/pipelines/stages/denoising.py

@@ -64,9 +64,10 @@ class DenoisingStage(PipelineStage):
         self.attn_backend = get_attn_backend(
             head_size=attn_head_size,
             dtype=torch.float16,  # TODO(will): hack
-            supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN,
-                                          AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN,
-                                          AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE)  # hack
+            supported_attention_backends=(
+                AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN, AttentionBackendEnum.VMOBA_ATTN,
+                AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE,
+                AttentionBackendEnum.FLASHINFER_SAGE_ATTN3, AttentionBackendEnum.ATTN_QAT_INFER)  # hack
         )
 
     def forward(

fastvideo/pipelines/stages/sr_denoising.py

@@ -61,9 +61,10 @@ class SRDenoisingStage(PipelineStage):
         self.attn_backend = get_attn_backend(
             head_size=attn_head_size,
             dtype=torch.float16,  # TODO(will): hack
-            supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN,
-                                          AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
-                                          AttentionBackendEnum.SAGE_ATTN_THREE)  # hack
+            supported_attention_backends=(
+                AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN,
+                AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE,
+                AttentionBackendEnum.FLASHINFER_SAGE_ATTN3, AttentionBackendEnum.ATTN_QAT_INFER)  # hack
         )
 
     def add_noise_to_lq(self, lq_latents: torch.Tensor, strength: float = 0.7) -> torch.Tensor:

fastvideo/platforms/cuda.py

@@ -140,6 +140,12 @@ class CudaPlatformBase(Platform):
             except ImportError as e:
                 logger.info(e)
                 logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.")
+        elif selected_backend == AttentionBackendEnum.FLASHINFER_SAGE_ATTN3:
+            # No import guard needed: the backend imports without flashinfer
+            # and falls back to torch SDPA internally (it stays usable as a
+            # shape-collection vehicle via FASTVIDEO_ATTN_SHAPE_LOG).
+            logger.info("Using FlashInfer SageAttention3 backend.")
+            return "fastvideo.attention.backends.flashinfer_sageattn3.FlashInferSageAttention3Backend"
         elif selected_backend == AttentionBackendEnum.ATTN_QAT_INFER:
             from fastvideo.attention.backends.attn_qat_infer import (  # noqa: F401
                 AttnQatInferBackend, is_attn_qat_infer_available)

fastvideo/platforms/interface.py

@@ -15,6 +15,7 @@ class AttentionBackendEnum(enum.Enum):
     TORCH_SDPA = enum.auto()
     SAGE_ATTN = enum.auto()
     SAGE_ATTN_THREE = enum.auto()
+    FLASHINFER_SAGE_ATTN3 = enum.auto()
     ATTN_QAT_INFER = enum.auto()
     ATTN_QAT_TRAIN = enum.auto()
     VIDEO_SPARSE_ATTN = enum.auto()

fastvideo/tests/attention/test_attn_kernels_bench.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only tests for the pure parts of
+fastvideo-kernel/benchmarks/bench_attn_kernels.py (flops formula, shape
+parsing, lane-skip logic, table formatting). GPU lanes are exercised
+manually on real hardware."""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+from pathlib import Path
+
+import pytest
+
+_SCRIPT = Path(__file__).resolve().parents[3] / "fastvideo-kernel" / "benchmarks" / "bench_attn_kernels.py"
+
+spec = importlib.util.spec_from_file_location("bench_attn_kernels", _SCRIPT)
+bench = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = bench  # dataclasses resolve annotations via sys.modules
+spec.loader.exec_module(bench)
+
+
+def _shape(**overrides):
+    base = dict(batch=1, num_q_heads=12, num_kv_heads=12, seq_len_q=1024,
+                seq_len_kv=1024, head_dim=128, causal=False)
+    base.update(overrides)
+    return bench.Shape(**base)
+
+
+def test_attn_tflops_formula():
+    # 4*B*H*Lq*Lk*D FLOPs at 1 ms.
+    shape = _shape(batch=2, num_q_heads=3, seq_len_q=100, seq_len_kv=200, head_dim=64)
+    expected = 4 * 2 * 3 * 100 * 200 * 64 / 1e9  # TFLOPS at 1 ms
+    assert bench.attn_tflops(shape, 1.0) == pytest.approx(expected)
+    # Causal halves the FLOPs; 2 ms halves the rate again.
+    causal = _shape(batch=2, num_q_heads=3, seq_len_q=100, seq_len_kv=200, head_dim=64, causal=True)
+    assert bench.attn_tflops(causal, 2.0) == pytest.approx(expected / 4)
+
+
+def test_parse_shapes_jsonl_dedupe_and_defaults(tmp_path):
+    rec = {"backend": "FLASHINFER_SAGE_ATTN3", "batch": 1, "num_q_heads": 12,
+           "num_kv_heads": 12, "seq_len_q": 75600, "seq_len_kv": 75600,
+           "head_dim": 128, "dtype": "torch.bfloat16", "causal": False,
+           "sm_scale": 0.0884, "count": 1200, "torch": "2.8.0",
+           "device": "NVIDIA GB10", "sm": "12.1"}
+    minimal = {"batch": 2, "num_q_heads": 8, "seq_len_q": 512, "head_dim": 64}
+    path = tmp_path / "shapes.jsonl"
+    path.write_text(json.dumps(rec) + "\n\n" + json.dumps(rec) + "\n" + json.dumps(minimal) + "\n")
+
+    shapes = bench.parse_shapes_jsonl([str(path)], default_dtype="torch.float16")
+    assert len(shapes) == 2  # duplicate line deduped
+    full, mini = shapes
+    assert (full.batch, full.num_q_heads, full.seq_len_q, full.head_dim) == (1, 12, 75600, 128)
+    assert full.sm_scale == pytest.approx(0.0884)
+    assert full.dtype == "torch.bfloat16"
+    # Missing keys fall back: kv defaults to q, dtype to the CLI default.
+    assert (mini.num_kv_heads, mini.seq_len_kv) == (8, 512)
+    assert mini.dtype == "torch.float16"
+    assert mini.causal is False and mini.sm_scale is None
+
+
+def test_skip_reason():
+    sm120 = bench.LaneSpec("nvfp4", exact_cap=(12, 0), head_dims=(64, 128), same_shape=True)
+    ok = _shape()
+    assert bench.skip_reason(sm120, ok, (12, 0), None) is None
+    assert "not installed" in bench.skip_reason(sm120, ok, (12, 0), "not installed")
+    assert "no CUDA" in bench.skip_reason(sm120, ok, None, None)
+    assert "SM 12.0" in bench.skip_reason(sm120, ok, (10, 0), None)  # GB200
+    assert "SM 12.0" in bench.skip_reason(sm120, ok, (12, 1), None)  # Spark
+    assert "head_dim 96" in bench.skip_reason(sm120, _shape(head_dim=96), (12, 0), None)
+    cross = _shape(seq_len_kv=256)
+    assert "equal q/k/v" in bench.skip_reason(sm120, cross, (12, 0), None)
+
+    fp8 = bench.LaneSpec("fp8", min_cap=(8, 9), same_heads=True)
+    assert bench.skip_reason(fp8, ok, (9, 0), None) is None
+    assert bench.skip_reason(fp8, cross, (12, 0), None) is None  # cross-attn ok
+    assert "SM >= 8.9" in bench.skip_reason(fp8, ok, (8, 6), None)
+    assert "num_q_heads" in bench.skip_reason(fp8, _shape(num_kv_heads=4), (9, 0), None)
+
+    anywhere = bench.LaneSpec("sdpa_bf16")
+    assert bench.skip_reason(anywhere, cross, (7, 5), None) is None
+
+
+def test_format_table_alignment():
+    headers = ["shape", "lane", "TFLOPS"]
+    rows = [["1x12x75600x128", "nvfp4_local", "812.3"],
+            ["1x8x512x64", "sdpa_bf16", "9.1"]]
+    table = bench.format_table(headers, rows)
+    lines = table.splitlines()
+    assert len(lines) == 4  # header, rule, 2 rows
+    assert len({len(line) for line in lines}) == 1  # columns aligned
+    assert "nvfp4_local" in lines[2] and "sdpa_bf16" in lines[3]
+    assert bench.format_table(headers, []).splitlines()[0].startswith("shape")
+
+
+def test_cli_defaults_and_shapes():
+    args = bench.parse_args([])
+    assert (args.warmup, args.iters, args.p_quant) == (10, 50, "single")
+    assert args.bias == 1.0 and not args.include_zeromean  # biased inputs by default
+    shapes = bench.default_shapes(causal=False, dtype="torch.bfloat16")
+    assert bench.Shape(1, 12, 12, 32760, 32760, 128, False) in shapes  # Wan2.1 480p
+    assert len(shapes) == len(set(shapes))  # sweep + logged shapes deduped
+    with pytest.raises(SystemExit) as exc:
+        bench.parse_args(["--help"])
+    assert exc.value.code == 0
+
+
+def test_parse_shapes_cli():
+    shapes = bench.parse_shapes_cli("1x12x32760x128,2x24x512x64", causal=True,
+                                    dtype="torch.float16")
+    assert shapes == [
+        bench.Shape(1, 12, 12, 32760, 32760, 128, True, "torch.float16"),
+        bench.Shape(2, 24, 24, 512, 512, 64, True, "torch.float16"),
+    ]
+    with pytest.raises(Exception, match="BxHxLxD"):
+        bench.parse_shapes_cli("1x12x128", causal=False, dtype="torch.bfloat16")
+
+
+def test_shape_workspace_bytes_scales_with_seq():
+    small = bench.shape_workspace_bytes(_shape(seq_len_q=1024, seq_len_kv=1024), 2, 2048)
+    big = bench.shape_workspace_bytes(_shape(seq_len_q=32760, seq_len_kv=32760), 2, 2048)
+    assert big > small > 0

fastvideo/tests/attention/test_attn_shape_logger.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only tests for fastvideo.attention.shape_logger.
+
+The logger reads FASTVIDEO_ATTN_SHAPE_LOG once at module import, so each test
+sets the env var and reloads the module to get a fresh, enabled instance.
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+import os
+
+import pytest
+import torch
+
+import fastvideo.attention.shape_logger as shape_logger_mod
+
+
+@pytest.fixture(autouse=True)
+def _reset_logger_module():
+    """Reload the module disabled after each test so the atexit flush never
+    targets a deleted tmp_path."""
+    yield
+    os.environ.pop("FASTVIDEO_ATTN_SHAPE_LOG", None)
+    importlib.reload(shape_logger_mod)
+
+
+def _fresh_logger(monkeypatch, path: str):
+    monkeypatch.setenv("FASTVIDEO_ATTN_SHAPE_LOG", str(path))
+    monkeypatch.delenv("WORLD_SIZE", raising=False)
+    return importlib.reload(shape_logger_mod)
+
+
+def _read_jsonl(path: str) -> list[dict]:
+    with open(path) as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def test_disabled_by_default(monkeypatch, tmp_path):
+    monkeypatch.delenv("FASTVIDEO_ATTN_SHAPE_LOG", raising=False)
+    mod = importlib.reload(shape_logger_mod)
+    assert mod.enabled is False
+    q = torch.zeros(1, 4, 2, 8)
+    mod.record("X", q, q, q, causal=False, sm_scale=None)  # no-op, no crash
+    assert not mod._entries
+
+
+def test_dedupe_count_and_dump_roundtrip(monkeypatch, tmp_path):
+    log_base = str(tmp_path / "shapes.jsonl")
+    mod = _fresh_logger(monkeypatch, log_base)
+    assert mod.enabled is True
+
+    q = torch.zeros(2, 16, 4, 64, dtype=torch.float32)  # [B, L, H, D]
+    k = torch.zeros(2, 32, 4, 64, dtype=torch.float32)
+    for _ in range(3):
+        mod.record("FLASHINFER_SAGE_ATTN3", q, k, k, causal=False, sm_scale=0.125)
+    # different causal flag -> new entry
+    mod.record("FLASHINFER_SAGE_ATTN3", q, k, k, causal=True, sm_scale=0.125)
+    # bhsd layout: [B, H, L, D] with same logical dims dedupes onto entry 1
+    q_bhsd = q.transpose(1, 2)
+    k_bhsd = k.transpose(1, 2)
+    mod.record("FLASHINFER_SAGE_ATTN3", q_bhsd, k_bhsd, k_bhsd, causal=False, sm_scale=0.125, layout="bhsd")
+    # different backend name -> new entry
+    mod.record("SAGE_ATTN_THREE", q, k, k, causal=False, sm_scale=0.125)
+
+    assert len(mod._entries) == 3
+    mod._flush_at_exit()
+
+    out_path = f"{log_base}.pid{os.getpid()}"
+    entries = _read_jsonl(out_path)
+    assert len(entries) == 3
+
+    main = next(e for e in entries if e["backend"] == "FLASHINFER_SAGE_ATTN3" and e["causal"] is False)
+    assert main["count"] == 4
+    assert main["batch"] == 2
+    assert main["num_q_heads"] == 4
+    assert main["num_kv_heads"] == 4
+    assert main["seq_len_q"] == 16
+    assert main["seq_len_kv"] == 32
+    assert main["head_dim"] == 64
+    assert main["dtype"] == "torch.float32"
+    assert main["sm_scale"] == 0.125
+    assert main["torch"] == torch.__version__
+    assert "device" in main and "sm" in main
+
+
+def test_periodic_flush_on_new_unique_shapes(monkeypatch, tmp_path):
+    log_base = str(tmp_path / "flush.jsonl")
+    mod = _fresh_logger(monkeypatch, log_base)
+    monkeypatch.setattr(mod, "_FLUSH_EVERY", 4)
+
+    out_path = f"{log_base}.pid{os.getpid()}"
+    for i in range(1, 5):
+        q = torch.zeros(1, i, 2, 8)
+        mod.record("B", q, q, q, causal=False, sm_scale=None)
+    # 4th unique shape triggered a flush without any atexit involvement
+    assert len(_read_jsonl(out_path)) == 4
+
+
+def test_rank_suffix_when_distributed(monkeypatch, tmp_path):
+    log_base = str(tmp_path / "dist.jsonl")
+    mod = _fresh_logger(monkeypatch, log_base)
+    monkeypatch.setenv("WORLD_SIZE", "4")
+    monkeypatch.setenv("RANK", "2")
+    q = torch.zeros(1, 4, 2, 8)
+    mod.record("B", q, q, q, causal=False, sm_scale=None)
+    mod._flush_at_exit()
+    assert os.path.exists(f"{log_base}.rank2")

fastvideo/tests/attention/test_flashinfer_sageattn3.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only tests for the FLASHINFER_SAGE_ATTN3 backend.
+
+flashinfer is deliberately NOT installed in this environment: these tests
+exercise the registration wiring and the torch-SDPA fallback path, which is
+exactly what runs on machines without the SM120 kernel (e.g. shape collection
+on DGX Spark). No GPU required.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+import torch
+
+from fastvideo.attention.backends import flashinfer_sageattn3 as flashinfer_backend
+from fastvideo.attention.backends.flashinfer_sageattn3 import (FlashInferSageAttention3Backend,
+                                                               FlashInferSageAttention3Impl)
+from fastvideo.attention.selector import backend_name_to_enum
+from fastvideo.platforms.cuda import CudaPlatform
+from fastvideo.platforms.interface import AttentionBackendEnum
+
+
+def test_enum_and_name_roundtrip():
+    enum_val = backend_name_to_enum("FLASHINFER_SAGE_ATTN3")
+    assert enum_val is AttentionBackendEnum.FLASHINFER_SAGE_ATTN3
+    assert FlashInferSageAttention3Backend.get_name() == "FLASHINFER_SAGE_ATTN3"
+    assert FlashInferSageAttention3Backend.get_impl_cls() is FlashInferSageAttention3Impl
+    assert FlashInferSageAttention3Backend.get_supported_head_sizes() == [64, 128]
+
+
+def test_cuda_platform_pin_returns_backend_cls():
+    cls_path = CudaPlatform.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER_SAGE_ATTN3,
+                                                 head_size=128,
+                                                 dtype=torch.bfloat16)
+    assert cls_path == ("fastvideo.attention.backends.flashinfer_sageattn3.FlashInferSageAttention3Backend")
+
+
+def _manual_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: float,
+                      causal: bool) -> torch.Tensor:
+    """Reference attention on [B, H, L, D] tensors."""
+    scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * scale
+    if causal:
+        seq_len = q.shape[2]
+        mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1)
+        scores.masked_fill_(mask, float("-inf"))
+    probs = torch.softmax(scores, dim=-1)
+    return torch.matmul(probs, v.float())
+
+
+@pytest.mark.parametrize("causal", [False, True])
+@pytest.mark.parametrize("head_dim", [64, 128])
+def test_fallback_forward_matches_reference(causal: bool, head_dim: int):
+    torch.manual_seed(0)
+    batch, seq_len, num_heads = 2, 17, 4
+    scale = 1.0 / math.sqrt(head_dim)
+    impl = FlashInferSageAttention3Impl(num_heads=num_heads,
+                                        head_size=head_dim,
+                                        causal=causal,
+                                        softmax_scale=scale)
+
+    # FastVideo impl convention: [B, L, H, D]
+    q = torch.randn(batch, seq_len, num_heads, head_dim, dtype=torch.float32)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+
+    out = impl.forward(q, k, v, attn_metadata=None)
+    assert out.shape == q.shape
+
+    ref = _manual_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), scale,
+                            causal).transpose(1, 2)
+    torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4)
+
+
+def test_fallback_forward_gqa():
+    torch.manual_seed(1)
+    batch, seq_len, head_dim = 1, 9, 64
+    scale = 1.0 / math.sqrt(head_dim)
+    impl = FlashInferSageAttention3Impl(num_heads=4, head_size=head_dim, causal=False, softmax_scale=scale)
+
+    q = torch.randn(batch, seq_len, 4, head_dim)
+    k = torch.randn(batch, seq_len, 2, head_dim)
+    v = torch.randn_like(k)
+
+    out = impl.forward(q, k, v, attn_metadata=None)
+
+    k_rep = k.repeat_interleave(2, dim=2)
+    v_rep = v.repeat_interleave(2, dim=2)
+    ref = _manual_attention(q.transpose(1, 2), k_rep.transpose(1, 2), v_rep.transpose(1, 2), scale,
+                            causal=False).transpose(1, 2)
+    torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4)
+
+
+def test_flashinfer_quantize_disables_ambient_autocast(monkeypatch):
+    """Keep FlashInfer's qk_correction fp32 inside the bf16 pipeline context."""
+    correction_dtype = None
+
+    def fake_quantize(q, k, v):
+        nonlocal correction_dtype
+        correction = torch.matmul(q.float(), k.float().transpose(-2, -1))
+        correction_dtype = correction.dtype
+        return q, k, v, q, k, v, correction
+
+    def fake_forward(q, _k, _v, _qs, _ks, _vs, correction, **_kwargs):
+        assert correction.dtype == torch.float32
+        return q, None
+
+    monkeypatch.setattr(flashinfer_backend, "_can_use_flashinfer", lambda *_args: True)
+    monkeypatch.setattr(flashinfer_backend, "nvfp4_attention_sm120_quantize_qkv", fake_quantize)
+    monkeypatch.setattr(flashinfer_backend, "nvfp4_attention_sm120_fwd", fake_forward)
+
+    impl = FlashInferSageAttention3Impl(num_heads=2,
+                                        head_size=64,
+                                        causal=False,
+                                        softmax_scale=1 / math.sqrt(64))
+    q = torch.randn(1, 8, 2, 64)
+    with torch.autocast("cpu", dtype=torch.bfloat16):
+        out = impl.forward(q, q, q, attn_metadata=None)
+
+    assert correction_dtype == torch.float32
+    assert out.shape == q.shape