[perf]FlashInfer cuDNN batched prefill attention

[klhhhhh]

Dependency

This PR depends on:

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

PR #1799 introduces the opt-in FlashInfer attention backend for Wan using
single_prefill_with_kv_cache.

This PR extends that backend with FlashInfer's batched cuDNN prefill kernel and
should be reviewed and merged after #1799.

Summary

This PR adds an opt-in FlashInfer cuDNN prefill path using:

flashinfer.prefill.cudnn_batch_prefill_with_kv_cache

The implementation processes the complete batch through one packed cuDNN
attention call instead of invoking the single-request kernel once per batch
item.

Enable it before constructing the model:

FASTVIDEO_ATTENTION_BACKEND=FLASHINFER
FASTVIDEO_FLASHINFER_PREFILL_BACKEND=cudnn

The existing single-request implementation remains the default:

FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single

Therefore, this PR does not change the default behavior introduced by #1799.

Motivation

The FlashInfer backend introduced in #1799 invokes
single_prefill_with_kv_cache separately for every item in the batch:

sample 0 -> single_prefill_with_kv_cache
sample 1 -> single_prefill_with_kv_cache
...
sample N -> single_prefill_with_kv_cache

This is effective for batch size 1, but larger batches require multiple kernel
launches followed by output stacking.

FlashInfer's cuDNN prefill API accepts packed Q/K/V tensors and batch offsets,
allowing the complete batch to be processed by one batched attention call.

This is particularly beneficial for large-batch GQA cross-attention with short
or medium query lengths and long KV sequences.

Implementation

FlashInfer implementation selection

This PR adds:

FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single|cudnn

The default value is single.

The variable only selects the implementation within the FLASHINFER backend.
The existing top-level selection remains unchanged:

FASTVIDEO_ATTENTION_BACKEND=FLASHINFER

Batched Q/K/V packing

FastVideo provides attention tensors in BSHD layout:

query: [B, Sq, Hq, D]
key:   [B, Sk, Hkv, D]
value: [B, Sk, Hkv, D]

The cuDNN path packs them as:

query: [B * Sq, Hq, D]
key:   [B * Sk, Hkv, D]
value: [B * Sk, Hkv, D]

It constructs token-unit batch offsets:

query_offsets = (
    torch.arange(
        batch_size + 1,
        dtype=torch.int32,
        device=query.device,
    )
    * query_len
)

kv_offsets = (
    torch.arange(
        batch_size + 1,
        dtype=torch.int32,
        device=query.device,
    )
    * key_len
)

The packed tensors are passed to:

output, _ = cudnn_batch_prefill_with_kv_cache(
    query.flatten(0, 1),
    key.flatten(0, 1),
    value.flatten(0, 1),
    self.softmax_scale,
    workspace,
    max_token_per_sequence=query_len,
    max_sequence_kv=key_len,
    batch_offsets_q=query_offsets,
    batch_offsets_o=query_offsets,
    batch_offsets_k=kv_offsets,
    batch_offsets_v=kv_offsets,
    batch_offsets_units="tokens",
    causal=causal,
    return_lse=False,
)

The output is then restored to FastVideo's BSHD layout.

Workspace reuse

The cuDNN API requires a workspace buffer. This PR allocates a 128 MiB
workspace and shares it across attention layers on the same CUDA device:

_cudnn_workspaces: dict[torch.device, torch.Tensor] = {}

This prevents every Transformer attention layer from allocating its own
workspace.

Availability checks

When cuDNN is selected, CUDA backend initialization verifies that:

cudnn_batch_prefill_with_kv_cache

is available.

This makes an incompatible FlashInfer installation fail during backend
initialization rather than during the first attention forward.

Supported behavior

The cuDNN path currently supports:

• Dense self-attention
• Cross-attention with different query and KV lengths
• MHA
• GQA
• Causal attention
• FP16 and BF16 inputs
• FastVideo's BSHD input/output contract

Current restrictions:

• Inference-only
• NVIDIA sm80 or newer
• Head dimension 128
• No arbitrary custom attention masks

Unsupported configurations raise explicit errors instead of silently falling
back to another kernel.

Tests

This PR adds CPU/mock coverage for:

• Selecting the cuDNN prefill path
• BSHD-to-packed-Q/K/V conversion
• Query and KV token offsets
• Batched GQA tensor shapes
• Output shape restoration
• Rejection of unsupported head dimensions
• Rejection of arbitrary custom masks
• Preservation of the existing single-prefill path

It also adds a real CUDA parity test against Torch SDPA covering:

• Batch sizes 1 and 2
• BF16
• Head dimension 128
• Cross-attention
• GQA
• Different query and KV sequence lengths

Test command:

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

Kernel benchmark

The benchmark compares:

• FlashAttention
• FlashInfer single-request prefill
• FlashInfer cuDNN batched prefill

The cuDNN workspace and batch offsets are created outside the timed region.
Latency is measured with CUDA events.

The benchmark script is included below for reproducibility only. It is not
added to the repository by this PR.

Benchmark script

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Benchmark FlashAttention and FlashInfer prefill paths.

Defaults target cuDNN's intended use case: batched GQA with unequal Q/KV
lengths. Static workspace and offsets are prepared outside the timed region.
Both the raw cuDNN path and the FastVideo wrapper are reported.

Example:
    python examples/inference/benchmark_flashinfer_cudnn.py \
      --batch-size 1 2 4 8 --query-length 256 1024 --kv-length 4096 \
      --num-heads 12 --num-kv-heads 4 --warmups 20 --repeats 100
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
from pathlib import Path
from typing import Callable


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--batch-size", type=int, nargs="+", default=[1, 2, 4, 8])
    parser.add_argument("--query-length", type=int, nargs="+", default=[256, 1024])
    parser.add_argument("--kv-length", type=int, nargs="+", default=[4096])
    parser.add_argument("--num-heads", type=int, default=12)
    parser.add_argument("--num-kv-heads", type=int, default=4)
    parser.add_argument("--head-size", type=int, default=128)
    parser.add_argument("--dtype", choices=("float16", "bfloat16"), default="bfloat16")
    parser.add_argument("--causal", action=argparse.BooleanOptionalAction, default=False)
    parser.add_argument("--warmups", type=int, default=20)
    parser.add_argument("--repeats", type=int, default=100)
    parser.add_argument("--output", type=Path, default=Path("results/flashinfer_cudnn.json"))
    return parser.parse_args()


def percentile(samples: list[float], quantile: float) -> float:
    ordered = sorted(samples)
    return ordered[round((len(ordered) - 1) * quantile)]


def benchmark_cuda(fn: Callable[[], object], warmups: int, repeats: int) -> list[float]:
    import torch

    for _ in range(warmups):
        fn()
    torch.cuda.synchronize()
    samples = []
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    for _ in range(repeats):
        start.record()
        fn()
        end.record()
        end.synchronize()
        samples.append(start.elapsed_time(end))
    return samples


def main() -> None:
    args = parse_args()
    if args.warmups < 1 or args.repeats < 1:
        raise ValueError("--warmups and --repeats must be positive")
    if args.head_size != 128:
        raise ValueError("The FlashInfer cuDNN arm currently requires --head-size 128")
    if args.num_heads % args.num_kv_heads:
        raise ValueError("--num-heads must be divisible by --num-kv-heads")
    if args.causal and set(args.query_length) != set(args.kv_length):
        raise ValueError("Use matching --query-length and --kv-length values with --causal")

    import torch
    from flash_attn import flash_attn_func
    from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache, single_prefill_with_kv_cache

    from fastvideo.attention.backends.flashinfer import FlashInferImpl, FlashInferMetadata

    if not torch.cuda.is_available():
        raise RuntimeError("This benchmark requires an NVIDIA CUDA GPU")
    if torch.cuda.get_device_capability() < (8, 0):
        raise RuntimeError("FlashInfer requires sm80 or newer")

    dtype = getattr(torch, args.dtype)
    device = torch.device("cuda", torch.cuda.current_device())
    scale = args.head_size**-0.5
    workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device)
    metadata = FlashInferMetadata(current_timestep=0)
    results: list[dict[str, object]] = []

    for batch_size in args.batch_size:
        for query_len in args.query_length:
            for kv_len in args.kv_length:
                if args.causal and query_len != kv_len:
                    continue
                torch.manual_seed(0)
                query = torch.randn((batch_size, query_len, args.num_heads, args.head_size),
                                    device=device,
                                    dtype=dtype)
                key = torch.randn((batch_size, kv_len, args.num_kv_heads, args.head_size),
                                  device=device,
                                  dtype=dtype)
                value = torch.randn_like(key)
                query_offsets = torch.arange(batch_size + 1, dtype=torch.int32, device=device) * query_len
                kv_offsets = torch.arange(batch_size + 1, dtype=torch.int32, device=device) * kv_len

                def run_flash_attn() -> torch.Tensor:
                    return flash_attn_func(query, key, value, softmax_scale=scale, causal=args.causal)

                def run_single_raw() -> torch.Tensor:
                    outputs = [
                        single_prefill_with_kv_cache(query[i],
                                                     key[i],
                                                     value[i],
                                                     causal=args.causal,
                                                     kv_layout="NHD",
                                                     sm_scale=scale) for i in range(batch_size)
                    ]
                    return torch.stack(outputs)

                def run_cudnn_raw() -> torch.Tensor:
                    output, _ = cudnn_batch_prefill_with_kv_cache(
                        query.flatten(0, 1),
                        key.flatten(0, 1),
                        value.flatten(0, 1),
                        scale,
                        workspace,
                        max_token_per_sequence=query_len,
                        max_sequence_kv=kv_len,
                        batch_offsets_q=query_offsets,
                        batch_offsets_o=query_offsets,
                        batch_offsets_k=kv_offsets,
                        batch_offsets_v=kv_offsets,
                        batch_offsets_units="tokens",
                        causal=args.causal,
                        return_lse=False,
                    )
                    return output.view_as(query)

                os.environ["FASTVIDEO_FLASHINFER_PREFILL_BACKEND"] = "cudnn"
                fastvideo_cudnn = FlashInferImpl(num_heads=args.num_heads,
                                                 num_kv_heads=args.num_kv_heads,
                                                 head_size=args.head_size,
                                                 causal=args.causal,
                                                 softmax_scale=scale)
                runners = {
                    "flash_attn": run_flash_attn,
                    "flashinfer_single_raw": run_single_raw,
                    "flashinfer_cudnn_raw": run_cudnn_raw,
                    "flashinfer_cudnn_fastvideo": lambda: fastvideo_cudnn.forward(query, key, value, metadata),
                }

                shape_results: dict[str, float] = {}
                for name, runner in runners.items():
                    samples = benchmark_cuda(runner, args.warmups, args.repeats)
                    median_ms = statistics.median(samples)
                    shape_results[name] = median_ms
                    results.append({
                        "backend": name,
                        "batch_size": batch_size,
                        "query_length": query_len,
                        "kv_length": kv_len,
                        "num_heads": args.num_heads,
                        "num_kv_heads": args.num_kv_heads,
                        "head_size": args.head_size,
                        "dtype": args.dtype,
                        "causal": args.causal,
                        "median_ms": median_ms,
                        "p20_ms": percentile(samples, 0.2),
                        "p80_ms": percentile(samples, 0.8),
                        "samples_ms": samples,
                    })
                    print(f"{name:30s} B={batch_size} Q={query_len} KV={kv_len}: {median_ms:.3f} ms")

                cudnn_ms = shape_results["flashinfer_cudnn_raw"]
                print(f"  cuDNN speedup vs FlashAttention: {shape_results['flash_attn'] / cudnn_ms:.3f}x")
                print(f"  cuDNN speedup vs single:         {shape_results['flashinfer_single_raw'] / cudnn_ms:.3f}x")
                wrapper_overhead = shape_results["flashinfer_cudnn_fastvideo"] / cudnn_ms - 1
                print(f"  FastVideo wrapper overhead:      {wrapper_overhead * 100:.2f}%")

    payload = {
        "environment": {
            "gpu": torch.cuda.get_device_name(device),
            "compute_capability": list(torch.cuda.get_device_capability(device)),
            "torch": torch.__version__,
            "cuda": torch.version.cuda,
        },
        "timing_scope": "CUDA events; static workspace and offsets prepared outside timed region",
        "warmups": args.warmups,
        "repeats": args.repeats,
        "results": results,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote {args.output.resolve()}")


if __name__ == "__main__":
    main()

Benchmark invocation

The reported large-batch results were produced with:

python examples/inference/benchmark_flashinfer_cudnn.py \
  --batch-size 4 8 16 32 \
  --query-length 128 256 512 1024 \
  --kv-length 4096 \
  --num-heads 12 \
  --num-kv-heads 4 \
  --warmups 30 \
  --repeats 200 \
  --output results/flashinfer_cudnn_large_batch.json

Environment

Item	Value
GPU	NVIDIA GB10
Compute capability	sm121
PyTorch	2.12.0+cu130
CUDA	13.0
Input dtype	BF16
Query heads	12
KV heads	4
Head dimension	128
KV length	4096
Warmup iterations	30
Measured iterations	200

Results

Median kernel latency:

Batch	Q	FlashAttention	FlashInfer single	FlashInfer cuDNN	cuDNN vs FA	cuDNN vs single
4	128	0.2172 ms	0.2265 ms	0.2340 ms	0.929x	0.968x
4	256	0.3096 ms	0.3441 ms	0.3404 ms	0.910x	1.011x
4	512	0.6059 ms	0.6484 ms	0.6306 ms	0.961x	1.028x
4	1024	1.1821 ms	1.2289 ms	1.2344 ms	0.958x	0.996x
8	128	0.3865 ms	0.4590 ms	0.4110 ms	0.940x	1.117x
8	256	0.6138 ms	0.6996 ms	0.6416 ms	0.957x	1.090x
8	512	1.2027 ms	1.3027 ms	1.2306 ms	0.977x	1.059x
8	1024	2.3740 ms	2.5174 ms	2.4344 ms	0.975x	1.034x
16	128	0.7027 ms	0.8717 ms	0.7601 ms	0.924x	1.147x
16	256	1.3673 ms	1.4093 ms	1.2517 ms	1.092x	1.126x
16	512	2.4999 ms	2.6679 ms	2.4354 ms	1.026x	1.095x
16	1024	4.7858 ms	5.0631 ms	4.8593 ms	0.985x	1.042x
32	128	2.3976 ms	1.7882 ms	1.4495 ms	1.654x	1.234x
32	256	3.8638 ms	2.8510 ms	2.4834 ms	1.556x	1.148x
32	512	5.1901 ms	5.3928 ms	4.8849 ms	1.062x	1.104x
32	1024	9.6248 ms	10.2427 ms	9.8312 ms	0.979x	1.042x

Observations

The cuDNN path performs best for large-batch GQA cross-attention with short or
medium query lengths.

At B=32, Q=128, KV=4096, cuDNN achieves:

• 1.654x speedup over FlashAttention
• 1.234x speedup over FlashInfer single prefill
• Approximately 39.5% lower latency than FlashAttention
• Approximately 19.0% lower latency than FlashInfer single prefill

At B=32, Q=256, KV=4096, cuDNN achieves:

• 1.556x speedup over FlashAttention
• 1.148x speedup over FlashInfer single prefill

The benefit increases with batch size and decreases as query length grows.
For small batches, the single-request path remains competitive. For
Q=1024, FlashAttention is still slightly faster.

These results support keeping the cuDNN implementation opt-in rather than
making it the universal FlashInfer default.

Limitations

• The cuDNN path is opt-in.
• The validated head dimension is currently limited to 128.
• Arbitrary custom attention masks are not supported.
• The implementation is inference-only.
• The benchmark measures attention kernels rather than complete video
  generation.
• Results are specific to NVIDIA GB10 and should not be generalized to every
  NVIDIA architecture.
• End-to-end Wan performance should be evaluated separately.

Follow-ups

1. Benchmark end-to-end Wan inference with the complete cuDNN dispatch path.
2. Evaluate automatic dispatch based on batch size and query length.
3. Validate the implementation on A100, H100, and datacenter Blackwell GPUs.
4. Evaluate additional supported head dimensions.
5. Add compatible attention-mask support where possible.

[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=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

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

[mergify]

Pre-commit checks failed

Hi @klhhhhh, the pre-commit checks have failed. To fix them locally:

# Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install

# Run all checks and auto-fix what's possible
pre-commit run --all-files

Common fixes:

• yapf: yapf -i <file> (formatting)
• ruff: ruff check --fix <file> (linting)
• codespell: codespell --write-changes <file> (spelling)

After fixing, commit and push the changes. The checks will re-run automatically.

For future commits, pre-commit will run automatically on changed files before each commit.