diff --git a/fastvideo/attention/backends/flashinfer.py b/fastvideo/attention/backends/flashinfer.py new file mode 100644 index 0000000000..96b4f75de5 --- /dev/null +++ b/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 diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 4c971ca9e4..a116eb126f 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/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): diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index b6f8a9cb3f..728537160e 100644 --- a/fastvideo/platforms/cuda.py +++ b/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 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..74da860008 --- /dev/null +++ b/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) 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..1b782f308e --- /dev/null +++ b/fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for CudaPlatformBase.get_attn_backend_cls's FLASHINFER branch. + +Covers the three ways FLASHINFER resolution can fail before ever touching a +GPU kernel: missing sm80 capability, a head size outside FlashInfer's safe +list, 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 _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) 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")