[perf]FlashInfer cuDNN batched prefill attention

docs/inference/optimizations.md

@@ -1,4 +1,3 @@
-
 # Optimizations
 
 This page describes the various options for speeding up generation times in FastVideo.
@@ -7,8 +6,7 @@ This page describes the various options for speeding up generation times in Fast
     Several options on this page behave differently on the GB10's unified-memory
     hardware — some give little or nothing there. See
     [DGX Spark: Performance & Tuning](../getting_started/installation/spark_performance.md)
-    for what actually helps on that platform and why. Two Sparks, one clip:
-    [Pair two NVIDIA DGX Sparks](../getting_started/installation/spark_pair.md).
+    for what actually helps on that platform and why.
 
 ## Table of Contents
 
@@ -31,6 +29,7 @@ This page describes the various options for speeding up generation times in Fast
 
 - Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
 - Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
+- FlashInfer prefill: `FASTVIDEO_ATTENTION_BACKEND=FLASHINFER`
 - 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`
@@ -61,6 +60,40 @@ You can also set the environment variable on the command line:
 FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
 ```
 
+### FlashInfer cuDNN prefill
+
+**`FLASHINFER`**
+
+Wan attention layers can use FlashInfer's dense prefill kernels. Install a
+FlashInfer release that provides `cudnn_batch_prefill_with_kv_cache`, then
+select either the established per-sample kernel or the batched cuDNN SDPA path:
+
+```bash
+uv pip install flashinfer-python
+
+# FlashInfer single_prefill_with_kv_cache (default)
+FASTVIDEO_ATTENTION_BACKEND=FLASHINFER \
+FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single python example.py
+
+# FlashInfer batched cuDNN SDPA
+FASTVIDEO_ATTENTION_BACKEND=FLASHINFER \
+FASTVIDEO_FLASHINFER_PREFILL_BACKEND=cudnn python example.py
+```
+
+The cuDNN path is inference-only, currently requires head size 128, and does
+not accept arbitrary custom attention masks. It supports dense self-attention,
+cross-attention, GQA, and causal attention. FastVideo raises an error instead
+of silently selecting another kernel when these constraints are not met.
+
+For a kernel-level comparison against FlashAttention and FlashInfer's single
+prefill path, run:
+
+```bash
+python examples/inference/benchmark_flashinfer_cudnn.py \
+  --batch-size 1 2 --sequence-length 1024 4096 8192 \
+  --warmups 10 --repeats 50 --output results/flashinfer_cudnn.json
+```
+
 ### Flash Attention
 
 **`FLASH_ATTN`**
@@ -145,11 +178,11 @@ PyTorch 2.12.0 and CUDA 13 environment for this kernel.
 Branch-to-`nvidia-cutlass-dsl` compatibility (the fork tracks the CuTe DSL API
 surface closely):
 
-| fork branch | cutlass-dsl | notes |
-|---|---|---|
-| `fp4` | `==4.4.2` (+ `nvidia-cutlass-dsl-libs-base==4.4.2`) | validated set on GB200: `quack-kernels==0.4.1`, `flashinfer-python==0.6.8`, `CUTE_DSL_ENABLE_TVM_FFI=1`, `FASTVIDEO_FA4=1` |
-| `fix/cutlass-dsl-4.5` | `>=4.5.2` | carries the `cute.core.ThrMma` -> `cute.ThrMma` fix |
-| any | 4.6-era | unsupported: `cute.make_fragment` was removed at module level; fails at CuTe JIT trace |
+| fork branch           | cutlass-dsl                                         | notes                                                        |
+| --------------------- | --------------------------------------------------- | ------------------------------------------------------------ |
+| `fp4`                 | `==4.4.2` (+ `nvidia-cutlass-dsl-libs-base==4.4.2`) | validated set on GB200: `quack-kernels==0.4.1`, `flashinfer-python==0.6.8`, `CUTE_DSL_ENABLE_TVM_FFI=1`, `FASTVIDEO_FA4=1` |
+| `fix/cutlass-dsl-4.5` | `>=4.5.2`                                           | carries the `cute.core.ThrMma` -> `cute.ThrMma` fix          |
+| any                   | 4.6-era                                             | unsupported: `cute.make_fragment` was removed at module level; fails at CuTe JIT trace |
 
 `FASTVIDEO_FA4=1` is required alongside the fork: it ships no compiled
 FlashAttention-2, so dense attention paths raise ImportError without the FA4
@@ -391,8 +424,8 @@ The Wan result below measures the existing generic
 but it is **not** a benchmark or numerical gate for the stricter regional
 fullgraph path above.
 
-| Config | Effect |
-|---|---|
+| Config                                            | Effect                                                       |
+| ------------------------------------------------- | ------------------------------------------------------------ |
 | Wan2.1-T2V-1.3B, A100-80GB, 480×832×81f, 50 steps | end-to-end **259.7s → 198.1s (−23.7%)**; per-step **4.91 → 3.78 s/it** |
 
 The speedup is **configuration-dependent** — it varies with model,
@@ -518,11 +551,11 @@ remaining steps. The technique is the LinearAG variant of Adaptive Guidance
 
 Set the `FASTVIDEO_CFG_GATE_STEP` environment variable to a float in `[0, 1]`:
 
-| Value | Behavior |
-|-------|----------|
-| `1.0` (default) | Disabled — legacy two-pass CFG every step. |
-| `0.5` | Cache the delta after `len(timesteps) * 0.5` steps; reuse for the rest. |
-| `0.0` | Cache from the very first step (most aggressive). |
+| Value           | Behavior                                                     |
+| --------------- | ------------------------------------------------------------ |
+| `1.0` (default) | Disabled — legacy two-pass CFG every step.                   |
+| `0.5`           | Cache the delta after `len(timesteps) * 0.5` steps; reuse for the rest. |
+| `0.0`           | Cache from the very first step (most aggressive).            |
 
 ```bash
 export FASTVIDEO_CFG_GATE_STEP=0.5

fastvideo/attention/backends/flashinfer.py

@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FlashInfer dense prefill attention backend.
+
+This backend uses either FlashInfer's single-request NHD kernel once per batch
+item or its batched cuDNN SDPA kernel. Select the implementation with
+``FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single|cudnn`` (default: ``single``).
+Both paths preserve FastVideo's BSHD/SP contract and support self-attention,
+cross-attention, GQA, and causal attention. Arbitrary masks remain on the
+single-request path because FlashInfer's cuDNN entry point has no custom-mask
+argument.
+FlashInfer's prefill API is inference-only here; training must use FLASH_ATTN or
+TORCH_SDPA.
+
+``forward`` reads ``attn_mask``/``is_causal`` off ``attn_metadata`` by
+attribute, not by ``isinstance(FlashInferMetadata)``: several model forwards
+(e.g. HYWorld) build ``SDPAMetadata`` directly and hand it to whichever
+backend the selector resolved. Any metadata dataclass with the same field
+names satisfies this backend; do not add a strict type check here without
+also updating those call sites.
+
+A layer must list FLASHINFER explicitly in its ``supported_attention_backends``
+to use this backend — there is no implicit bridge from FLASH_ATTN support, since
+this kernel's masking/GQA conventions have not been vetted per-model.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+
+
+@dataclass
+class FlashInferMetadata(AttentionMetadata):
+    current_timestep: int
+    attn_mask: torch.Tensor | None = None
+    is_causal: bool = False
+
+
+class FlashInferMetadataBuilder(AttentionMetadataBuilder):
+
+    def prepare(self):
+        pass
+
+    def build(self, current_timestep: int, attn_mask: torch.Tensor | None = None) -> FlashInferMetadata:  # type: ignore
+        return FlashInferMetadata(current_timestep=current_timestep, attn_mask=attn_mask)
+
+
+class FlashInferBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # FlashInfer 0.6.x's FA2 fallback is unsafe for some other dimensions.
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashInferImpl"]:
+        return FlashInferImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type[FlashInferMetadata]:
+        return FlashInferMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type[FlashInferMetadataBuilder]:
+        return FlashInferMetadataBuilder
+
+
+def _mask_for_sample(attn_mask: torch.Tensor, sample: int, query_len: int, key_len: int) -> torch.Tensor:
+    """Convert FastVideo's padding/additive mask to FlashInfer's [Q, K] bool mask."""
+    mask = attn_mask.to(dtype=torch.bool) if not attn_mask.dtype.is_floating_point else attn_mask >= 0
+    if mask.dim() == 2:
+        if mask.shape[-1] > key_len:
+            raise ValueError(f"Invalid FLASHINFER mask length: expected at most {key_len}, got {mask.shape[-1]}")
+        key_mask = mask[sample]
+        if key_mask.shape[0] < key_len:
+            key_mask = torch.nn.functional.pad(key_mask, (key_len - key_mask.shape[0], 0), value=True)
+        return key_mask.unsqueeze(0).expand(query_len, -1)
+    if mask.dim() == 3:
+        mask = mask[sample]
+    elif mask.dim() == 4:
+        mask = mask[sample, 0]
+    else:
+        raise ValueError(f"Unsupported FLASHINFER attention mask shape: {attn_mask.shape}")
+    if mask.shape[-2:] != (query_len, key_len):
+        if mask.shape[-2] == 1 and mask.shape[-1] == key_len:
+            mask = mask.expand(query_len, -1)
+        else:
+            raise ValueError(f"FLASHINFER mask must broadcast to [{query_len}, {key_len}], got {mask.shape}")
+    return mask
+
+
+class FlashInferImpl(AttentionImpl):
+
+    _CUDNN_WORKSPACE_BYTES = 128 * 1024 * 1024
+    # FlashInfer documents 128 MiB as sufficient for typical prefill. Share one
+    # allocation per device instead of reserving it once per transformer layer.
+    _cudnn_workspaces: dict[torch.device, torch.Tensor] = {}
+
+    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:
+        del num_heads, num_kv_heads, prefix, extra_impl_args
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.head_size = head_size
+        self.prefill_backend = envs.FASTVIDEO_FLASHINFER_PREFILL_BACKEND
+        if self.prefill_backend not in ("single", "cudnn"):
+            raise ValueError("FASTVIDEO_FLASHINFER_PREFILL_BACKEND must be 'single' or 'cudnn'; "
+                             f"got {self.prefill_backend!r}")
+        if self.prefill_backend == "cudnn" and head_size != 128:
+            raise ValueError("FlashInfer cuDNN prefill requires head size 128 in FastVideo; "
+                             f"got {head_size}. Use FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single instead.")
+
+    def _forward_cudnn(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, causal: bool) -> torch.Tensor:
+        from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
+
+        batch_size, query_len = query.shape[:2]
+        key_len = key.shape[1]
+        workspace = self._cudnn_workspaces.get(query.device)
+        if workspace is None:
+            workspace = torch.empty(self._CUDNN_WORKSPACE_BYTES, dtype=torch.uint8, device=query.device)
+            self._cudnn_workspaces[query.device] = workspace
+
+        query_offsets = torch.arange(batch_size + 1, dtype=torch.int32, device=query.device) * query_len
+        key_offsets = torch.arange(batch_size + 1, dtype=torch.int32, device=query.device) * key_len
+        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=key_offsets,
+            batch_offsets_v=key_offsets,
+            batch_offsets_units="tokens",
+            causal=causal,
+            return_lse=False,
+        )
+        return output.view_as(query)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
+                attn_metadata: FlashInferMetadata | None) -> torch.Tensor:
+        if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad):
+            raise RuntimeError("FLASHINFER backend is inference-only; use FLASH_ATTN or TORCH_SDPA for training.")
+
+        original_dtype = query.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            query = query.to(torch.bfloat16)
+            key = key.to(torch.bfloat16)
+            value = value.to(torch.bfloat16)
+
+        mask = attn_metadata.attn_mask if attn_metadata is not None else None
+        causal = self.causal or bool(attn_metadata is not None and getattr(attn_metadata, "is_causal", False))
+        if self.prefill_backend == "cudnn":
+            if mask is not None:
+                raise ValueError("FlashInfer cuDNN prefill does not support arbitrary attention masks; "
+                                 "use FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single instead.")
+            output = self._forward_cudnn(query, key, value, causal)
+            return output.to(original_dtype) if output.dtype != original_dtype else output
+
+        from flashinfer.prefill import single_prefill_with_kv_cache
+
+        outputs = []
+        for sample in range(query.shape[0]):
+            custom_mask = None
+            if mask is not None:
+                custom_mask = _mask_for_sample(mask, sample, query.shape[1], key.shape[1]).to(query.device)
+                if causal:
+                    causal_mask = torch.ones((query.shape[1], key.shape[1]), dtype=torch.bool,
+                                             device=query.device).tril(key.shape[1] - query.shape[1])
+                    custom_mask = custom_mask & causal_mask
+            outputs.append(
+                single_prefill_with_kv_cache(query[sample],
+                                             key[sample],
+                                             value[sample],
+                                             custom_mask=custom_mask,
+                                             causal=causal and custom_mask is None,
+                                             kv_layout="NHD",
+                                             sm_scale=self.softmax_scale))
+        output = torch.stack(outputs)
+        return output.to(original_dtype) if output.dtype != original_dtype else output

fastvideo/envs.py

@@ -20,6 +20,7 @@
     FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None
     FASTVIDEO_TRACE_FUNCTION: int = 0
     FASTVIDEO_ATTENTION_BACKEND: str | None = None
+    FASTVIDEO_FLASHINFER_PREFILL_BACKEND: str = "single"
     FASTVIDEO_FA4: bool = False
     FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN: bool = False
     FASTVIDEO_INFERENCE_TORCH_COMPILE: bool = False
@@ -218,6 +219,12 @@ def maybe_convert_int(value: str | None) -> int | None:
     "FASTVIDEO_ATTENTION_BACKEND":
     lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
 
+    # Select the concrete FlashInfer prefill implementation. ``single`` uses
+    # single_prefill_with_kv_cache once per sample; ``cudnn`` uses FlashInfer's
+    # batched cuDNN SDPA entry point.
+    "FASTVIDEO_FLASHINFER_PREFILL_BACKEND":
+    lambda: os.getenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "single").strip().lower(),
+
     # 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/platforms/cuda.py

@@ -173,7 +173,36 @@ def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, hea
 
         logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND)
         logger.info("Selected backend: %s", selected_backend)
-        if selected_backend == AttentionBackendEnum.SAGE_ATTN:
+        if selected_backend == AttentionBackendEnum.FLASHINFER:
+            if not cls.has_device_capability(80):
+                raise RuntimeError("FLASHINFER requires an NVIDIA GPU with compute capability sm80 or newer.")
+            if dtype not in (torch.float16, torch.bfloat16):
+                logger.warning("FLASHINFER will cast %s inputs to bfloat16 for the kernel.", dtype)
+            try:
+                from flashinfer.prefill import single_prefill_with_kv_cache  # noqa: F401
+
+                if envs.FASTVIDEO_FLASHINFER_PREFILL_BACKEND == "cudnn":
+                    from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache  # noqa: F401
+
+                from fastvideo.attention.backends.flashinfer import (  # noqa: F401
+                    FlashInferBackend)
+            except ImportError as e:
+                raise ImportError("FLASHINFER selected but flashinfer-python is not importable. "
+                                  "Install the FastVideo Linux dependencies or "
+                                  "`uv pip install flashinfer-python`.") from e
+            if head_size not in FlashInferBackend.get_supported_head_sizes():
+                raise ValueError(f"FLASHINFER does not safely support head size {head_size}; "
+                                 f"supported sizes are {FlashInferBackend.get_supported_head_sizes()}.")
+            # The cudnn prefill arm narrows FlashInfer's general head-size support to
+            # 128 only (see FlashInferImpl.__init__). Check it here too so a bad
+            # combination fails at backend selection instead of per-layer, deep into
+            # model construction.
+            if envs.FASTVIDEO_FLASHINFER_PREFILL_BACKEND == "cudnn" and head_size != 128:
+                raise ValueError(f"FLASHINFER cuDNN prefill requires head size 128; got {head_size}. "
+                                 "Use FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single instead.")
+            logger.info("Using FlashInfer attention backend.")
+            return "fastvideo.attention.backends.flashinfer.FlashInferBackend"
+        elif selected_backend == AttentionBackendEnum.SAGE_ATTN:
             try:
                 from sageattention import sageattn  # noqa: F401
 

fastvideo/platforms/interface.py

@@ -12,6 +12,7 @@
 
 class AttentionBackendEnum(enum.Enum):
     FLASH_ATTN = enum.auto()
+    FLASHINFER = enum.auto()
     TORCH_SDPA = enum.auto()
     SAGE_ATTN = enum.auto()
     SAGE_ATTN_THREE = enum.auto()