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

fastvideo/attention/backends/flashinfer.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FlashInfer dense prefill attention backend.
+
+This backend uses FlashInfer's single-request NHD kernel once per batch item.
+That path preserves FastVideo's BSHD/SP contract and supports self-attention,
+cross-attention, GQA, causal attention, and tokenizer-style padding masks.
+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
+
+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):
+
+    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, head_size, num_kv_heads, prefix, extra_impl_args
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    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.")
+
+        from flashinfer.prefill import single_prefill_with_kv_cache
+
+        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))
+        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/models/dits/wanvideo.py

@@ -172,6 +172,7 @@ def __init__(self,
                                    softmax_scale=None,
                                    causal=False,
                                    supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
+                                                                 AttentionBackendEnum.FLASHINFER,
                                                                  AttentionBackendEnum.TORCH_SDPA))
 
     def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int):

fastvideo/platforms/cuda.py

@@ -173,7 +173,26 @@ 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
+
+                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()}.")
+            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()

fastvideo/tests/attention/test_flashinfer_backend.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fastvideo.attention.backends.flashinfer import FlashInferImpl, FlashInferMetadata, _mask_for_sample
+
+
+def test_padding_mask_is_front_padded_and_expanded() -> None:
+    mask = torch.tensor([[1, 0]], dtype=torch.int64)
+    actual = _mask_for_sample(mask, sample=0, query_len=3, key_len=4)
+    expected = torch.tensor([[1, 1, 1, 0]] * 3, dtype=torch.bool)
+    torch.testing.assert_close(actual, expected)
+
+
+def test_forward_preserves_bshd_contract_and_arguments(monkeypatch) -> None:
+    calls = []
+
+    def fake_kernel(q, k, v, **kwargs):
+        calls.append((q.shape, k.shape, v.shape, kwargs))
+        return q + 1
+
+    prefill = types.ModuleType("flashinfer.prefill")
+    prefill.single_prefill_with_kv_cache = fake_kernel
+    flashinfer = types.ModuleType("flashinfer")
+    flashinfer.prefill = prefill
+    monkeypatch.setitem(sys.modules, "flashinfer", flashinfer)
+    monkeypatch.setitem(sys.modules, "flashinfer.prefill", prefill)
+
+    impl = FlashInferImpl(num_heads=4, head_size=64, causal=True, softmax_scale=0.125)
+    query = torch.randn(2, 3, 4, 64, dtype=torch.bfloat16)
+    key = torch.randn(2, 5, 2, 64, dtype=torch.bfloat16)
+    value = torch.randn(2, 5, 2, 64, dtype=torch.bfloat16)
+    metadata = FlashInferMetadata(current_timestep=0, attn_mask=torch.ones(2, 5, dtype=torch.bool))
+
+    output = impl.forward(query, key, value, metadata)
+
+    assert output.shape == query.shape
+    assert len(calls) == 2
+    assert calls[0][0] == (3, 4, 64)
+    assert calls[0][1] == (5, 2, 64)
+    assert calls[0][3]["kv_layout"] == "NHD"
+    assert calls[0][3]["causal"] is False
+    assert calls[0][3]["custom_mask"].shape == (3, 5)
+    assert calls[0][3]["sm_scale"] == 0.125
+
+
+def _require_flashinfer_cuda() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("requires one NVIDIA CUDA GPU")
+    if torch.cuda.get_device_capability() < (8, 0):
+        pytest.skip("FlashInfer attention requires sm80 or newer")
+    pytest.importorskip("flashinfer.prefill", reason="flashinfer-python is not installed")
+
+
+def _sdpa_reference(query: torch.Tensor,
+                    key: torch.Tensor,
+                    value: torch.Tensor,
+                    *,
+                    scale: float,
+                    causal: bool = False,
+                    key_mask: torch.Tensor | None = None) -> torch.Tensor:
+    attn_mask = key_mask[:, None, None, :] if key_mask is not None else None
+    output = F.scaled_dot_product_attention(query.transpose(1, 2),
+                                            key.transpose(1, 2),
+                                            value.transpose(1, 2),
+                                            attn_mask=attn_mask,
+                                            is_causal=causal,
+                                            scale=scale,
+                                            enable_gqa=query.shape[2] != key.shape[2])
+    return output.transpose(1, 2)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("head_size", [64, 128, 256])
+def test_flashinfer_real_cuda_kernel_matches_sdpa(dtype: torch.dtype, head_size: int) -> None:
+    """Launch the real single-GPU FlashInfer kernel for every supported head size."""
+    _require_flashinfer_cuda()
+    torch.manual_seed(0)
+    device = torch.device("cuda", torch.cuda.current_device())
+    scale = head_size**-0.5
+    query = torch.randn(1, 128, 4, head_size, device=device, dtype=dtype)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    impl = FlashInferImpl(num_heads=4, head_size=head_size, causal=False, softmax_scale=scale)
+
+    actual = impl.forward(query, key, value, FlashInferMetadata(current_timestep=0))
+    expected = _sdpa_reference(query, key, value, scale=scale)
+
+    torch.cuda.synchronize(device)
+    tolerance = 3e-2 if dtype == torch.bfloat16 else 1e-2
+    torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance)
+
+
+@pytest.mark.parametrize("mode", ["causal", "cross_gqa", "causal_padding"])
+def test_flashinfer_real_cuda_kernel_attention_modes(mode: str) -> None:
+    """Exercise native causal, GQA/cross-attention, and combined custom masks."""
+    _require_flashinfer_cuda()
+    torch.manual_seed(1)
+    device = torch.device("cuda", torch.cuda.current_device())
+    dtype = torch.bfloat16
+    head_size = 128
+    query_len, key_len = (64, 96) if mode == "cross_gqa" else (128, 128)
+    query_heads, kv_heads = (4, 2) if mode == "cross_gqa" else (4, 4)
+    scale = head_size**-0.5
+    query = torch.randn(1, query_len, query_heads, head_size, device=device, dtype=dtype)
+    key = torch.randn(1, key_len, kv_heads, head_size, device=device, dtype=dtype)
+    value = torch.randn_like(key)
+    causal = mode in ("causal", "causal_padding")
+    key_mask = None
+    if mode == "causal_padding":
+        key_mask = torch.ones(1, key_len, device=device, dtype=torch.bool)
+        key_mask[:, -8:] = False
+    metadata = FlashInferMetadata(current_timestep=0, attn_mask=key_mask)
+    impl = FlashInferImpl(num_heads=query_heads,
+                          num_kv_heads=kv_heads,
+                          head_size=head_size,
+                          causal=causal,
+                          softmax_scale=scale)
+
+    actual = impl.forward(query, key, value, metadata)
+    if causal and key_mask is not None:
+        full_mask = key_mask[:, None, :].expand(-1, query_len, -1)
+        causal_mask = torch.ones((query_len, key_len), dtype=torch.bool,
+                                 device=device).tril(key_len - query_len)
+        full_mask = full_mask & causal_mask
+        expected = F.scaled_dot_product_attention(query.transpose(1, 2),
+                                                  key.transpose(1, 2),
+                                                  value.transpose(1, 2),
+                                                  attn_mask=full_mask[:, None, :, :],
+                                                  scale=scale).transpose(1, 2)
+    else:
+        expected = _sdpa_reference(query, key, value, scale=scale, causal=causal, key_mask=key_mask)
+
+    torch.cuda.synchronize(device)
+    torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2)