[feat] Add opt-in FlashInfer attention backend for Wan

[klhhhhh]

Summary

This PR adds an opt-in FlashInfer dense-attention backend for the standard
Wan inference path.

Enable it before constructing the generator:

FASTVIDEO_ATTENTION_BACKEND=FLASHINFER fastvideo generate --config <config>

The implementation calls:

flashinfer.prefill.single_prefill_with_kv_cache(
    ...,
    kv_layout="NHD",
    backend="auto",
)

Backend selection remains explicit and model-scoped. This PR adds
FLASHINFER only to Wan's declared supported_attention_backends; it does
not implicitly treat every layer that supports FLASH_ATTN as
FlashInfer-compatible.

Motivation

FastVideo already depends on flashinfer-python for other kernel paths, but
does not expose FlashInfer's dense prefill attention through the common
attention backend interface.

Wan's video DiT recomputes full-sequence Q/K/V at each denoising step, so
FlashInfer's prefill attention API is the relevant dense-attention path rather
than its autoregressive decode API.

Implementation

The new backend:

• Preserves FastVideo's BSHD input/output contract.
• Invokes FlashInfer with NHD tensors after removing the batch dimension.
• Supports dense self-attention and cross-attention.
• Supports MHA and GQA.
• Supports causal attention.
• Converts tokenizer-style padding and additive masks to FlashInfer's boolean
  [Q, K] custom-mask contract.
• Combines padding and causal masks when both are present.
• Supports FP16 and BF16 kernel inputs.
• Temporarily casts other floating-point inputs to BF16 and restores the
  original output dtype, with a warning.
• Rejects grad-enabled Q/K/V because this initial backend is inference-only.
• Restricts the initial validated head dimensions to 64, 128, and 256.
• Requires NVIDIA compute capability sm80 or newer.
• Warns when batch size is greater than one because the initial implementation
  launches one single-request prefill kernel per batch item.

Model scope

The initial validated model scope is:

• Wan-AI/Wan2.1-T2V-1.3B-Diffusers
• Standard Wan dense-attention inference path

This PR intentionally does not add FlashInfer support to every existing model
that declares FLASH_ATTN.

Other model families, causal Wan variants, encoders, VAEs, and additional
attention shapes will be validated and enabled in follow-up PRs by explicitly
adding FLASHINFER to their own supported_attention_backends declarations.

The FlashInfer cuDNN batched dense-attention path
(cudnn_batch_prefill_with_kv_cache) is also out of scope for this PR and will
be evaluated separately.

Failure behavior

Selection follows FastVideo's existing model-level backend contract:

• A layer must explicitly declare FLASHINFER support.
• Unvalidated layers follow the existing selector fallback behavior.
• Missing flashinfer-python, pre-sm80 hardware, and unsupported head
  dimensions fail with explicit errors after a layer has selected FlashInfer.
• Training fails on the first grad-enabled attention forward with a clear
  inference-only error.

Tests

This PR adds coverage for:

Backend behavior

• FastVideo BSHD to FlashInfer NHD adaptation.
• Padding-mask expansion.
• Causal and padding-mask composition.
• Cross-attention and GQA argument handling.
• Rejection of unsupported per-head [B, H, Q, K] masks.
• Preservation of the output dtype.

Selector behavior

• Layers that do not explicitly declare FLASHINFER fall back normally.
• Layers that explicitly declare FLASHINFER resolve to the new backend.

CUDA platform resolution

• Successful FlashInfer backend resolution.
• sm80 capability gate.
• Missing flashinfer-python.
• Unsupported head dimension.

Real CUDA parity

The GPU test compares the real FlashInfer CUDA kernel against Torch SDPA for:

• FP16 and BF16.
• Head dimensions 64, 128, and 256.
• Dense self-attention.
• Cross-attention and GQA.
• Causal attention.
• Causal attention combined with a padding mask.

Example:

CUDA_VISIBLE_DEVICES=0 python -m pytest \
  fastvideo/tests/attention/test_flashinfer_backend.py \
  -k real_cuda -vs

Preliminary end-to-end result

Hardware and workload:

Item	Value
GPU	NVIDIA GB10, sm121
Model	Wan-AI/Wan2.1-T2V-1.3B-Diffusers
Resolution	832 × 480
Frames	77
Denoising steps	50
Guidance scale	6.0
Seed	1024
Warmup runs	1 per backend
Measured runs	3 per backend
Output saving	Enabled

Measured full-pipeline wall times:

Backend	Measured runs	Median
FastVideo FLASH_ATTN	673.527s, 658.267s, 660.769s	660.769s
FlashInfer single-prefill	649.441s, 649.768s, 654.594s	649.768s

This produced:

FlashInfer end-to-end speedup: 1.017x
Median wall-time reduction: 11.001s
Relative wall-time reduction: 1.66%

All three measured FlashInfer runs were faster than all three measured
FlashAttention runs in this experiment. The FlashInfer measurements also had
lower run-to-run variance:

FLASH_ATTN coefficient of variation: 1.23%
FLASHINFER coefficient of variation: 0.44%

This result should be treated as preliminary rather than a final kernel-level
performance claim:

• The sample count is small.
• Video encoding and output writing were included in the measured wall time.
• FlashAttention was run first, so filesystem/model-cache state may have
  favored the second benchmark arm.
• The full pipeline includes text encoding, linear/MLP layers, VAE decoding,
  scheduling, and other work outside attention.
• FlashInfer used backend="auto"; the final internal kernel was not pinned to
  FA2 or FA3.

A follow-up benchmark will reverse the backend order, disable output saving,
increase the number of measured runs, and profile the attention kernels
separately.

Limitations

• Inference-only.
• Initially enabled only for the validated standard Wan path.
• Uses the single-request prefill API once per batch item.
• Does not yet use FlashInfer's batched cuDNN prefill API.
• Does not implement persistent paged KV-cache or decode attention.
• Does not support arbitrary floating-point attention bias.
• Does not support per-head 4D custom masks.
• No automatic backend selection is introduced by this PR.

Follow-ups

Separate PRs will cover:

1. Validation and explicit opt-in for additional FastVideo model families.
2. Causal Wan validation.
3. FlashInfer cuDNN batched dense attention.
4. Larger-batch performance.
5. Kernel-level profiling and backend-pinned FA2/FA3 comparisons.
6. Additional mask and head-dimension coverage.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[klhhhhh]

The benchmark script.

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Compare complete FastVideo inference with FlashAttention and FlashInfer.

Each backend runs in a fresh subprocess because attention selection is resolved
when model components are constructed and CUDA state/model memory must not leak
between benchmark arms.

Example:
    CUDA_VISIBLE_DEVICES=0 python examples/inference/benchmark_attention_backends.py \
        --config scripts/inference/inference_wan.yaml \
        --override request.prompt="A fox running through snow" \
        --override request.inputs.prompt_path=null \
        --warmups 1 --repeats 3 --save-outputs
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import subprocess
import sys
import time
from copy import deepcopy
from dataclasses import asdict
from pathlib import Path
from typing import Any

_BACKENDS = ("FLASH_ATTN", "FLASHINFER")
_RESULT_PREFIX = "FASTVIDEO_BACKEND_BENCHMARK_RESULT="


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", required=True, help="Nested FastVideo inference YAML/JSON config")
    parser.add_argument("--override",
                        action="append",
                        default=[],
                        help="Dotted generator/request override; repeat as needed")
    parser.add_argument("--output-dir", default="outputs/attention_backend_comparison")
    parser.add_argument("--warmups", type=int, default=1)
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument("--save-outputs",
                        action=argparse.BooleanOptionalAction,
                        default=True,
                        help="Save measured videos for visual comparison (included in wall time)")
    parser.add_argument("--backend-order", nargs=2, choices=_BACKENDS, default=list(_BACKENDS))
    parser.add_argument("--worker-backend", choices=_BACKENDS, help=argparse.SUPPRESS)
    return parser.parse_args()


def _synchronize_cuda(torch_module: Any) -> None:
    if torch_module.cuda.is_available():
        torch_module.cuda.synchronize()


def _run_worker(args: argparse.Namespace) -> None:
    # This must happen before importing FastVideo: backend selection is folded
    # into component construction and must stay fixed for the worker lifetime.
    backend = args.worker_backend
    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = backend

    import torch

    from fastvideo import VideoGenerator
    from fastvideo.entrypoints.cli.inference_config import build_generate_run_config

    if not torch.cuda.is_available():
        raise RuntimeError("This benchmark requires an NVIDIA CUDA GPU")
    if backend == "FLASHINFER" and torch.cuda.get_device_capability() < (8, 0):
        raise RuntimeError("FLASHINFER requires an NVIDIA GPU with compute capability sm80 or newer")

    config_args = argparse.Namespace(config=args.config)
    config_overrides = [item if item.startswith("--") else f"--{item}" for item in args.override]
    run_config = build_generate_run_config(config_args, overrides=config_overrides)
    if run_config.generator.engine.num_gpus != 1:
        raise ValueError("This comparison script currently requires generator.engine.num_gpus=1")

    backend_dir = Path(args.output_dir).resolve() / backend.lower()
    backend_dir.mkdir(parents=True, exist_ok=True)

    load_started = time.perf_counter()
    generator = VideoGenerator.from_config(run_config.generator)
    load_seconds = time.perf_counter() - load_started

    def run_once(phase: str, index: int, *, measured: bool) -> float:
        request = deepcopy(run_config.request)
        request.output.output_path = str(backend_dir)
        request.output.output_video_name = f"{backend.lower()}_{phase}_{index:02d}"
        request.output.save_video = args.save_outputs if measured else False
        request.output.return_frames = False

        _synchronize_cuda(torch)
        started = time.perf_counter()
        generator.generate(request)
        _synchronize_cuda(torch)
        return time.perf_counter() - started

    warmup_seconds = [run_once("warmup", index, measured=False) for index in range(args.warmups)]
    measured_seconds = [run_once("run", index, measured=True) for index in range(args.repeats)]
    result = {
        "backend": backend,
        "device": torch.cuda.get_device_name(torch.cuda.current_device()),
        "device_capability": list(torch.cuda.get_device_capability()),
        "torch_version": torch.__version__,
        "cuda_version": torch.version.cuda,
        "config": str(Path(args.config).resolve()),
        "generator": asdict(run_config.generator),
        "request": asdict(run_config.request),
        "load_seconds": load_seconds,
        "warmup_seconds": warmup_seconds,
        "measured_seconds": measured_seconds,
        "median_seconds": statistics.median(measured_seconds),
        "mean_seconds": statistics.mean(measured_seconds),
        "save_outputs": args.save_outputs,
        "output_dir": str(backend_dir),
    }
    print(f"{_RESULT_PREFIX}{json.dumps(result, default=str)}", flush=True)


def _run_backend_subprocess(args: argparse.Namespace, backend: str) -> dict[str, Any]:
    command = [
        sys.executable,
        str(Path(__file__).resolve()),
        "--config",
        args.config,
        "--output-dir",
        args.output_dir,
        "--warmups",
        str(args.warmups),
        "--repeats",
        str(args.repeats),
        "--worker-backend",
        backend,
        "--save-outputs" if args.save_outputs else "--no-save-outputs",
    ]
    for override in args.override:
        command.extend(("--override", override))

    print(f"\n===== {backend} =====", flush=True)
    process = subprocess.Popen(command,
                               cwd=Path.cwd(),
                               env=os.environ.copy(),
                               stdout=subprocess.PIPE,
                               stderr=subprocess.STDOUT,
                               text=True,
                               bufsize=1)
    result: dict[str, Any] | None = None
    assert process.stdout is not None
    for line in process.stdout:
        print(line, end="", flush=True)
        if line.startswith(_RESULT_PREFIX):
            result = json.loads(line[len(_RESULT_PREFIX):])
    return_code = process.wait()
    if return_code != 0:
        raise RuntimeError(f"{backend} worker failed with exit code {return_code}")
    if result is None:
        raise RuntimeError(f"{backend} worker exited without a benchmark result")
    return result


def _write_summary(args: argparse.Namespace, results: list[dict[str, Any]]) -> Path:
    by_backend = {result["backend"]: result for result in results}
    flash_seconds = by_backend["FLASH_ATTN"]["median_seconds"]
    flashinfer_seconds = by_backend["FLASHINFER"]["median_seconds"]
    summary = {
        "results": by_backend,
        "comparison": {
            "flash_attn_median_seconds": flash_seconds,
            "flashinfer_median_seconds": flashinfer_seconds,
            "flashinfer_speedup": flash_seconds / flashinfer_seconds,
            "flashinfer_time_change_percent": (flashinfer_seconds / flash_seconds - 1.0) * 100.0,
        },
        "timing_scope": "VideoGenerator.generate wall time with CUDA synchronization",
        "notes": [
            "Warmup runs are excluded from statistics.",
            "Saved-video encoding is included when --save-outputs is enabled.",
            "A speedup greater than 1.0 means FLASHINFER was faster.",
        ],
    }
    output_dir = Path(args.output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    summary_path = output_dir / "summary.json"
    summary_path.write_text(json.dumps(summary, indent=2, default=str) + "\n", encoding="utf-8")
    return summary_path


def main() -> None:
    args = _parse_args()
    if args.warmups < 1:
        raise ValueError("--warmups must be at least 1 so JIT/caches are excluded")
    if args.repeats < 1:
        raise ValueError("--repeats must be at least 1")
    if set(args.backend_order) != set(_BACKENDS):
        raise ValueError("--backend-order must contain FLASH_ATTN and FLASHINFER exactly once")
    if args.worker_backend is not None:
        _run_worker(args)
        return

    results = [_run_backend_subprocess(args, backend) for backend in args.backend_order]
    summary_path = _write_summary(args, results)
    comparison = json.loads(summary_path.read_text(encoding="utf-8"))["comparison"]
    print("\n===== Comparison =====")
    print(f"FLASH_ATTN median: {comparison['flash_attn_median_seconds']:.3f} s")
    print(f"FLASHINFER median:  {comparison['flashinfer_median_seconds']:.3f} s")
    print(f"FLASHINFER speedup: {comparison['flashinfer_speedup']:.3f}x")
    print(f"Summary: {summary_path}")


if __name__ == "__main__":
    main()