[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/wan/transformer.py

@@ -260,6 +260,12 @@ def forward(
 
 
 class WanSelfAttention(nn.Module):
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+        AttentionBackendEnum.ATTN_QAT_INFER,
+        AttentionBackendEnum.FLASHINFER,
+    )
 
     def __init__(self,
                  dim: int,
@@ -269,7 +275,8 @@ def __init__(self,
                  eps=1e-6,
                  parallel_attention=False,
                  quant_config: QuantizationConfig | None = None,
-                 prefix: str = "") -> None:
+                 prefix: str = "",
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None) -> None:
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
@@ -288,15 +295,18 @@ def __init__(self,
         self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
-        # Scaled dot product attention
+        # A kernel must be supported by both this layer and the calling model.
+        # Unscoped callers retain the shared Wan family's capabilities.
+        if supported_attention_backends is None:
+            supported_attention_backends = WanVideoConfig()._supported_attention_backends
+        cross_backends = tuple(backend for backend in self._supported_attention_backends
+                               if backend in supported_attention_backends)
         self.attn = LocalAttention(num_heads=num_heads,
                                    head_size=self.head_dim,
                                    dropout_rate=0,
                                    softmax_scale=None,
                                    causal=False,
-                                   supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
-                                                                 AttentionBackendEnum.TORCH_SDPA,
-                                                                 AttentionBackendEnum.ATTN_QAT_INFER))
+                                   supported_attention_backends=cross_backends)
 
     def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int):
         r"""
@@ -365,7 +375,7 @@ def __init__(
                          window_size,
                          qk_norm,
                          eps,
-                         supported_attention_backends,
+                         supported_attention_backends=supported_attention_backends,
                          quant_config=quant_config,
                          prefix=prefix)
 
@@ -461,15 +471,17 @@ def __init__(self,
                                               qk_norm=qk_norm,
                                               eps=eps,
                                               quant_config=quant_config,
-                                              prefix=f"{prefix}.attn2")
+                                              prefix=f"{prefix}.attn2",
+                                              supported_attention_backends=supported_attention_backends)
         else:
             # T2V
             self.attn2 = WanT2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
                                               eps=eps,
                                               quant_config=quant_config,
-                                              prefix=f"{prefix}.attn2")
+                                              prefix=f"{prefix}.attn2",
+                                              supported_attention_backends=supported_attention_backends)
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
                                                                          norm_type="layer",
                                                                          eps=eps,
@@ -633,15 +645,17 @@ def __init__(self,
                                               qk_norm=qk_norm,
                                               eps=eps,
                                               quant_config=quant_config,
-                                              prefix=f"{prefix}.attn2")
+                                              prefix=f"{prefix}.attn2",
+                                              supported_attention_backends=supported_attention_backends)
         else:
             # T2V
             self.attn2 = WanT2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
                                               eps=eps,
                                               quant_config=quant_config,
-                                              prefix=f"{prefix}.attn2")
+                                              prefix=f"{prefix}.attn2",
+                                              supported_attention_backends=supported_attention_backends)
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
                                                                          norm_type="layer",
                                                                          eps=eps,
@@ -729,7 +743,20 @@ def forward(
 class WanTransformer3DModel(BaseDiT):
     _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
     _compile_conditions = WanVideoConfig()._compile_conditions
-    _supported_attention_backends = WanVideoConfig()._supported_attention_backends
+    # Model-level opt-in: shared arch configs retain the family defaults.
+    _supported_attention_backends = (
+        AttentionBackendEnum.SAGE_ATTN,
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+        AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+        AttentionBackendEnum.VMOBA_ATTN,
+        AttentionBackendEnum.SAGE_ATTN_THREE,
+        AttentionBackendEnum.ATTN_QAT_INFER,
+        AttentionBackendEnum.ATTN_QAT_TRAIN,
+        AttentionBackendEnum.SLA_ATTN,
+        AttentionBackendEnum.SAGE_SLA_ATTN,
+        AttentionBackendEnum.FLASHINFER,
+    )
     param_names_mapping = WanVideoConfig().param_names_mapping
     reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
     lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping

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.get())
         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, device_id=torch.cuda.current_device()):
+                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

@@ -25,6 +25,7 @@ class AttentionBackendEnum(enum.Enum):
     SAGE_SLA_ATTN = enum.auto()
     NABLA_ATTN = enum.auto()
     NO_ATTENTION = enum.auto()
+    FLASHINFER = enum.auto()
 
 
 class PlatformEnum(enum.Enum):