diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 98e3920d11..f44aa4d888 100644 --- a/docs/inference/optimizations.md +++ b/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 diff --git a/fastvideo/attention/backends/flashinfer.py b/fastvideo/attention/backends/flashinfer.py new file mode 100644 index 0000000000..0dceaf94d5 --- /dev/null +++ b/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 diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 86abc1bcde..dad899f9e6 100644 --- a/fastvideo/envs.py +++ b/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 diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index b6f8a9cb3f..74cbf6df23 100644 --- a/fastvideo/platforms/cuda.py +++ b/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 diff --git a/fastvideo/platforms/interface.py b/fastvideo/platforms/interface.py index ea55f77803..8c0d60caf5 100644 --- a/fastvideo/platforms/interface.py +++ b/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() diff --git a/fastvideo/tests/attention/test_flashinfer_backend.py b/fastvideo/tests/attention/test_flashinfer_backend.py new file mode 100644 index 0000000000..8d5551d87b --- /dev/null +++ b/fastvideo/tests/attention/test_flashinfer_backend.py @@ -0,0 +1,207 @@ +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: + monkeypatch.setenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "single") + 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 test_cudnn_forward_flattens_batch_and_passes_token_offsets(monkeypatch) -> None: + monkeypatch.setenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "cudnn") + FlashInferImpl._cudnn_workspaces.clear() + calls = [] + + def fake_cudnn(q, k, v, scale, workspace, **kwargs): + calls.append((q.shape, k.shape, v.shape, scale, workspace.shape, kwargs)) + return q + 1, None + + prefill = types.ModuleType("flashinfer.prefill") + prefill.cudnn_batch_prefill_with_kv_cache = fake_cudnn + flashinfer = types.ModuleType("flashinfer") + flashinfer.prefill = prefill + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.prefill", prefill) + monkeypatch.setattr(torch, "empty", lambda size, **kwargs: torch.zeros(size, dtype=kwargs["dtype"])) + + impl = FlashInferImpl(num_heads=4, head_size=128, causal=False, softmax_scale=0.125) + query = torch.randn(2, 3, 4, 128, dtype=torch.bfloat16) + key = torch.randn(2, 5, 2, 128, dtype=torch.bfloat16) + value = torch.randn_like(key) + output = impl.forward(query, key, value, FlashInferMetadata(current_timestep=0)) + + assert output.shape == query.shape + assert calls[0][0] == (6, 4, 128) + assert calls[0][1] == (10, 2, 128) + assert calls[0][5]["batch_offsets_units"] == "tokens" + torch.testing.assert_close(calls[0][5]["batch_offsets_q"], torch.tensor([0, 3, 6], dtype=torch.int32)) + torch.testing.assert_close(calls[0][5]["batch_offsets_k"], torch.tensor([0, 5, 10], dtype=torch.int32)) + + +def test_cudnn_rejects_custom_mask_and_unsupported_head_size(monkeypatch) -> None: + monkeypatch.setenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "cudnn") + with pytest.raises(ValueError, match="head size 128"): + FlashInferImpl(num_heads=4, head_size=64, causal=False, softmax_scale=0.125) + + impl = FlashInferImpl(num_heads=4, head_size=128, causal=False, softmax_scale=0.125) + query = torch.randn(1, 3, 4, 128, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="arbitrary attention masks"): + impl.forward(query, query, query, + FlashInferMetadata(current_timestep=0, attn_mask=torch.ones(1, 3, dtype=torch.bool))) + + +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(monkeypatch, dtype: torch.dtype, head_size: int) -> None: + """Launch the real single-GPU FlashInfer kernel for every supported head size.""" + monkeypatch.setenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "single") + _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(monkeypatch, mode: str) -> None: + """Exercise native causal, GQA/cross-attention, and combined custom masks.""" + monkeypatch.setenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "single") + _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) + + +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_flashinfer_real_cudnn_kernel_matches_sdpa(monkeypatch, batch_size: int) -> None: + """Launch FlashInfer's dedicated batched cuDNN SDPA entry point.""" + _require_flashinfer_cuda() + pytest.importorskip("flashinfer.prefill").cudnn_batch_prefill_with_kv_cache + monkeypatch.setenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "cudnn") + torch.manual_seed(2) + device = torch.device("cuda", torch.cuda.current_device()) + query = torch.randn(batch_size, 128, 4, 128, device=device, dtype=torch.bfloat16) + key = torch.randn(batch_size, 192, 2, 128, device=device, dtype=torch.bfloat16) + value = torch.randn_like(key) + scale = 128**-0.5 + impl = FlashInferImpl(num_heads=4, num_kv_heads=2, head_size=128, 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) + torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2) diff --git a/fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py b/fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py new file mode 100644 index 0000000000..1c1452451c --- /dev/null +++ b/fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for CudaPlatformBase.get_attn_backend_cls's FLASHINFER branch. + +Covers the ways FLASHINFER resolution can fail before ever touching a GPU +kernel: missing sm80 capability, a head size outside FlashInfer's safe list, +a head size the cuDNN prefill arm specifically can't handle, and +flashinfer-python not being importable. All device/import probes are +monkeypatched, so this runs without a real GPU or the flashinfer wheel. +""" +from __future__ import annotations + +import sys +import types + +import pytest +import torch + +from fastvideo.platforms.cuda import CudaPlatformBase +from fastvideo.platforms.interface import AttentionBackendEnum + + +def _patch_capability(monkeypatch, *, supported: bool) -> None: + monkeypatch.setattr(CudaPlatformBase, "has_device_capability", + classmethod(lambda cls, capability, device_id=0: supported)) + + +def _install_fake_flashinfer(monkeypatch) -> None: + prefill = types.ModuleType("flashinfer.prefill") + prefill.single_prefill_with_kv_cache = lambda *a, **k: None + flashinfer = types.ModuleType("flashinfer") + flashinfer.prefill = prefill + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.prefill", prefill) + + +def _install_fake_flashinfer_with_cudnn(monkeypatch) -> None: + prefill = types.ModuleType("flashinfer.prefill") + prefill.single_prefill_with_kv_cache = lambda *a, **k: None + prefill.cudnn_batch_prefill_with_kv_cache = lambda *a, **k: None + flashinfer = types.ModuleType("flashinfer") + flashinfer.prefill = prefill + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.prefill", prefill) + + +def _block_flashinfer_import(monkeypatch) -> None: + # A None entry in sys.modules makes `import flashinfer.prefill` raise + # ImportError, regardless of whether the real package is installed. + monkeypatch.setitem(sys.modules, "flashinfer.prefill", None) + monkeypatch.setitem(sys.modules, "flashinfer", None) + + +def test_flashinfer_requires_sm80(monkeypatch) -> None: + _patch_capability(monkeypatch, supported=False) + with pytest.raises(RuntimeError, match="sm80"): + CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER, 128, torch.bfloat16) + + +def test_flashinfer_rejects_unsupported_head_size(monkeypatch) -> None: + _patch_capability(monkeypatch, supported=True) + _install_fake_flashinfer(monkeypatch) + with pytest.raises(ValueError, match="does not safely support head size"): + CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER, 96, torch.bfloat16) + + +def test_flashinfer_missing_install_raises_actionable_error(monkeypatch) -> None: + _patch_capability(monkeypatch, supported=True) + _block_flashinfer_import(monkeypatch) + with pytest.raises(ImportError, match="flashinfer-python is not importable"): + CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER, 128, torch.bfloat16) + + +@pytest.mark.parametrize("head_size", [64, 128, 256]) +def test_flashinfer_resolves_for_supported_head_sizes(monkeypatch, head_size: int) -> None: + _patch_capability(monkeypatch, supported=True) + _install_fake_flashinfer(monkeypatch) + backend_cls = CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER, head_size, torch.bfloat16) + assert backend_cls == "fastvideo.attention.backends.flashinfer.FlashInferBackend" + + +def test_flashinfer_warns_on_dtype_cast(monkeypatch, caplog) -> None: + _patch_capability(monkeypatch, supported=True) + _install_fake_flashinfer(monkeypatch) + with caplog.at_level("WARNING"): + CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER, 128, torch.float32) + assert any("cast" in record.message for record in caplog.records) + + +@pytest.mark.parametrize("head_size", [64, 256]) +def test_flashinfer_cudnn_rejects_non_128_head_size_at_dispatch(monkeypatch, head_size: int) -> None: + # FlashInferBackend.get_supported_head_sizes() allows 64/128/256 generally, + # but the cudnn prefill arm narrows that to 128 only (FlashInferImpl.__init__). + # This must fail here, at backend selection, not later per-layer during + # model construction. + monkeypatch.setenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "cudnn") + _patch_capability(monkeypatch, supported=True) + _install_fake_flashinfer_with_cudnn(monkeypatch) + with pytest.raises(ValueError, match="cuDNN prefill requires head size 128"): + CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER, head_size, torch.bfloat16) + + +def test_flashinfer_cudnn_resolves_for_head_size_128(monkeypatch) -> None: + monkeypatch.setenv("FASTVIDEO_FLASHINFER_PREFILL_BACKEND", "cudnn") + _patch_capability(monkeypatch, supported=True) + _install_fake_flashinfer_with_cudnn(monkeypatch) + backend_cls = CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER, 128, torch.bfloat16) + assert backend_cls == "fastvideo.attention.backends.flashinfer.FlashInferBackend" diff --git a/fastvideo/tests/attention/test_selector_role_override.py b/fastvideo/tests/attention/test_selector_role_override.py index d894f37196..cb992c4d19 100644 --- a/fastvideo/tests/attention/test_selector_role_override.py +++ b/fastvideo/tests/attention/test_selector_role_override.py @@ -69,6 +69,37 @@ def test_unsupported_role_override_honors_layer_default(monkeypatch) -> None: selector._cached_get_attn_backend.cache_clear() +def test_flashinfer_requires_explicit_layer_support(monkeypatch) -> None: + """A layer must list FLASHINFER itself; declaring FLASH_ATTN is not enough. + + FlashInfer's dense kernel has not been vetted per-model (head sizes, + masking conventions, SP contract), so unlike the old bridge, a layer that + only declares FLASH_ATTN support falls back rather than silently + running FlashInfer. + """ + monkeypatch.setattr(platforms, "_current_platform", _FakePlatform()) + monkeypatch.setattr(selector, "resolve_obj_by_qualname", lambda name: name) + + try: + assert selector.get_attn_backend( + head_size=128, + dtype=torch.bfloat16, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + default_backend=AttentionBackendEnum.FLASH_ATTN, + requested=AttentionBackendEnum.FLASHINFER, + ) == "FLASH_ATTN" + + assert selector.get_attn_backend( + head_size=128, + dtype=torch.bfloat16, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.FLASHINFER, + AttentionBackendEnum.TORCH_SDPA), + requested=AttentionBackendEnum.FLASHINFER, + ) == "FLASHINFER" + finally: + selector._cached_get_attn_backend.cache_clear() + + def test_explicit_backend_config_rejects_typos() -> None: try: selector.coerce_attn_backend("attn_qat_typo") diff --git a/pyproject.toml b/pyproject.toml index da3eadaea3..44178fc203 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,7 @@ dependencies = [ # CUDA-only: no wheels for non-Linux platforms, and the sdist build shells # out to nvcc unconditionally. Gate on sys_platform so installs on macOS # (MPS) and other non-Linux platforms don't try to build it from source. - "flashinfer-python; sys_platform == 'linux'", + "flashinfer-python>=0.6.18; sys_platform == 'linux'", # Acceleration & Optimization "accelerate==1.0.1",