From dd53cdd27628972a605a7e245eddf05256d53473 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 12:48:32 -0700 Subject: [PATCH 01/19] [feat]support flashinfer attention kernel of prefill --- fastvideo/attention/backends/flashinfer.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 fastvideo/attention/backends/flashinfer.py diff --git a/fastvideo/attention/backends/flashinfer.py b/fastvideo/attention/backends/flashinfer.py new file mode 100644 index 0000000000..e69de29bb2 From fb179b1ac83a75e081a3b64f74371ec728d60929 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 12:48:49 -0700 Subject: [PATCH 02/19] [feat]support flashinfer attention kernel of prefill --- fastvideo/attention/backends/flashinfer.py | 131 +++++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/fastvideo/attention/backends/flashinfer.py b/fastvideo/attention/backends/flashinfer.py index e69de29bb2..f625f9deb9 100644 --- a/fastvideo/attention/backends/flashinfer.py +++ b/fastvideo/attention/backends/flashinfer.py @@ -0,0 +1,131 @@ +# 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. +""" + +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 From 5dcef8321f23460166a02324d30c85b5f5a9b44e Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 12:50:52 -0700 Subject: [PATCH 03/19] add flashinfer enum --- fastvideo/platforms/interface.py | 1 + 1 file changed, 1 insertion(+) 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() From 37bfe598109d6b721f3a145139fe257742bf5934 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 12:53:28 -0700 Subject: [PATCH 04/19] check flashinfer env --- fastvideo/platforms/cuda.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index b6f8a9cb3f..fa7437acc5 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.info("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 From b95b1602a7bbb774b8c1bbfc463bb91d961d3a10 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 16:08:49 -0700 Subject: [PATCH 05/19] add flashinfer backend in selector --- fastvideo/attention/selector.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/fastvideo/attention/selector.py b/fastvideo/attention/selector.py index 93924cf84b..6297cebdec 100644 --- a/fastvideo/attention/selector.py +++ b/fastvideo/attention/selector.py @@ -276,7 +276,14 @@ def _cached_get_attn_backend( # get device-specific attn_backend from fastvideo.platforms import current_platform - if (selected_backend is not None and selected_backend not in supported_attention_backends): + # FLASHINFER implements the same dense BSHD semantic contract as + # FLASH_ATTN. Treat every layer that already declares FLASH_ATTN support + # as FlashInfer-capable, avoiding model-by-model tuple churn while keeping + # sparse-only layers protected by their existing declarations. + selected_is_supported = (selected_backend in supported_attention_backends or + (selected_backend == AttentionBackendEnum.FLASHINFER + and AttentionBackendEnum.FLASH_ATTN in supported_attention_backends)) + if selected_backend is not None and not selected_is_supported: fallback_backend = (default_backend if default_backend in supported_attention_backends else None) logger.warning( "Requested attention backend %s is not supported by this " From ca396157344ba51b507f09527918132620d98bcc Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 16:13:21 -0700 Subject: [PATCH 06/19] [test]add tests for flashinfer --- .../attention/test_flashinfer_backend.py | 140 ++++++++++++++++++ .../attention/test_selector_role_override.py | 15 ++ 2 files changed, 155 insertions(+) create mode 100644 fastvideo/tests/attention/test_flashinfer_backend.py 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_selector_role_override.py b/fastvideo/tests/attention/test_selector_role_override.py index d894f37196..058249b93f 100644 --- a/fastvideo/tests/attention/test_selector_role_override.py +++ b/fastvideo/tests/attention/test_selector_role_override.py @@ -69,6 +69,21 @@ def test_unsupported_role_override_honors_layer_default(monkeypatch) -> None: selector._cached_get_attn_backend.cache_clear() +def test_flashinfer_is_compatible_with_dense_flash_attn_layers(monkeypatch) -> None: + 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), + 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") From 5868eefe9a41443a0df72e69d3621578d48dbb1e Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 16:34:57 -0700 Subject: [PATCH 07/19] beanchmark for test --- .../inference/benchmark_attention_backends.py | 213 ++++++++++++++++++ 1 file changed, 213 insertions(+) create mode 100644 examples/inference/benchmark_attention_backends.py diff --git a/examples/inference/benchmark_attention_backends.py b/examples/inference/benchmark_attention_backends.py new file mode 100644 index 0000000000..573f782f1e --- /dev/null +++ b/examples/inference/benchmark_attention_backends.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Compare complete FastVideo inference with FlashAttention and FlashInfer. + +Each backend runs in a fresh subprocess because attention selection is resolved +when model components are constructed and CUDA state/model memory must not leak +between benchmark arms. + +Example: + CUDA_VISIBLE_DEVICES=0 python examples/inference/benchmark_attention_backends.py \ + --config scripts/inference/inference_wan.yaml \ + --override request.prompt="A fox running through snow" \ + --override request.inputs.prompt_path=null \ + --warmups 1 --repeats 3 --save-outputs +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import sys +import time +from copy import deepcopy +from dataclasses import asdict +from pathlib import Path +from typing import Any + +_BACKENDS = ("FLASH_ATTN", "FLASHINFER") +_RESULT_PREFIX = "FASTVIDEO_BACKEND_BENCHMARK_RESULT=" + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", required=True, help="Nested FastVideo inference YAML/JSON config") + parser.add_argument("--override", + action="append", + default=[], + help="Dotted generator/request override; repeat as needed") + parser.add_argument("--output-dir", default="outputs/attention_backend_comparison") + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--save-outputs", + action=argparse.BooleanOptionalAction, + default=True, + help="Save measured videos for visual comparison (included in wall time)") + parser.add_argument("--backend-order", nargs=2, choices=_BACKENDS, default=list(_BACKENDS)) + parser.add_argument("--worker-backend", choices=_BACKENDS, help=argparse.SUPPRESS) + return parser.parse_args() + + +def _synchronize_cuda(torch_module: Any) -> None: + if torch_module.cuda.is_available(): + torch_module.cuda.synchronize() + + +def _run_worker(args: argparse.Namespace) -> None: + # This must happen before importing FastVideo: backend selection is folded + # into component construction and must stay fixed for the worker lifetime. + backend = args.worker_backend + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = backend + + import torch + + from fastvideo import VideoGenerator + from fastvideo.entrypoints.cli.inference_config import build_generate_run_config + + if not torch.cuda.is_available(): + raise RuntimeError("This benchmark requires an NVIDIA CUDA GPU") + if backend == "FLASHINFER" and torch.cuda.get_device_capability() < (8, 0): + raise RuntimeError("FLASHINFER requires an NVIDIA GPU with compute capability sm80 or newer") + + config_args = argparse.Namespace(config=args.config) + config_overrides = [item if item.startswith("--") else f"--{item}" for item in args.override] + run_config = build_generate_run_config(config_args, overrides=config_overrides) + if run_config.generator.engine.num_gpus != 1: + raise ValueError("This comparison script currently requires generator.engine.num_gpus=1") + + backend_dir = Path(args.output_dir).resolve() / backend.lower() + backend_dir.mkdir(parents=True, exist_ok=True) + + load_started = time.perf_counter() + generator = VideoGenerator.from_config(run_config.generator) + load_seconds = time.perf_counter() - load_started + + def run_once(phase: str, index: int, *, measured: bool) -> float: + request = deepcopy(run_config.request) + request.output.output_path = str(backend_dir) + request.output.output_video_name = f"{backend.lower()}_{phase}_{index:02d}" + request.output.save_video = args.save_outputs if measured else False + request.output.return_frames = False + + _synchronize_cuda(torch) + started = time.perf_counter() + generator.generate(request) + _synchronize_cuda(torch) + return time.perf_counter() - started + + warmup_seconds = [run_once("warmup", index, measured=False) for index in range(args.warmups)] + measured_seconds = [run_once("run", index, measured=True) for index in range(args.repeats)] + result = { + "backend": backend, + "device": torch.cuda.get_device_name(torch.cuda.current_device()), + "device_capability": list(torch.cuda.get_device_capability()), + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "config": str(Path(args.config).resolve()), + "generator": asdict(run_config.generator), + "request": asdict(run_config.request), + "load_seconds": load_seconds, + "warmup_seconds": warmup_seconds, + "measured_seconds": measured_seconds, + "median_seconds": statistics.median(measured_seconds), + "mean_seconds": statistics.mean(measured_seconds), + "save_outputs": args.save_outputs, + "output_dir": str(backend_dir), + } + print(f"{_RESULT_PREFIX}{json.dumps(result, default=str)}", flush=True) + + +def _run_backend_subprocess(args: argparse.Namespace, backend: str) -> dict[str, Any]: + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--config", + args.config, + "--output-dir", + args.output_dir, + "--warmups", + str(args.warmups), + "--repeats", + str(args.repeats), + "--worker-backend", + backend, + "--save-outputs" if args.save_outputs else "--no-save-outputs", + ] + for override in args.override: + command.extend(("--override", override)) + + print(f"\n===== {backend} =====", flush=True) + process = subprocess.Popen(command, + cwd=Path.cwd(), + env=os.environ.copy(), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1) + result: dict[str, Any] | None = None + assert process.stdout is not None + for line in process.stdout: + print(line, end="", flush=True) + if line.startswith(_RESULT_PREFIX): + result = json.loads(line[len(_RESULT_PREFIX):]) + return_code = process.wait() + if return_code != 0: + raise RuntimeError(f"{backend} worker failed with exit code {return_code}") + if result is None: + raise RuntimeError(f"{backend} worker exited without a benchmark result") + return result + + +def _write_summary(args: argparse.Namespace, results: list[dict[str, Any]]) -> Path: + by_backend = {result["backend"]: result for result in results} + flash_seconds = by_backend["FLASH_ATTN"]["median_seconds"] + flashinfer_seconds = by_backend["FLASHINFER"]["median_seconds"] + summary = { + "results": by_backend, + "comparison": { + "flash_attn_median_seconds": flash_seconds, + "flashinfer_median_seconds": flashinfer_seconds, + "flashinfer_speedup": flash_seconds / flashinfer_seconds, + "flashinfer_time_change_percent": (flashinfer_seconds / flash_seconds - 1.0) * 100.0, + }, + "timing_scope": "VideoGenerator.generate wall time with CUDA synchronization", + "notes": [ + "Warmup runs are excluded from statistics.", + "Saved-video encoding is included when --save-outputs is enabled.", + "A speedup greater than 1.0 means FLASHINFER was faster.", + ], + } + output_dir = Path(args.output_dir).resolve() + output_dir.mkdir(parents=True, exist_ok=True) + summary_path = output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2, default=str) + "\n", encoding="utf-8") + return summary_path + + +def main() -> None: + args = _parse_args() + if args.warmups < 1: + raise ValueError("--warmups must be at least 1 so JIT/caches are excluded") + if args.repeats < 1: + raise ValueError("--repeats must be at least 1") + if set(args.backend_order) != set(_BACKENDS): + raise ValueError("--backend-order must contain FLASH_ATTN and FLASHINFER exactly once") + if args.worker_backend is not None: + _run_worker(args) + return + + results = [_run_backend_subprocess(args, backend) for backend in args.backend_order] + summary_path = _write_summary(args, results) + comparison = json.loads(summary_path.read_text(encoding="utf-8"))["comparison"] + print("\n===== Comparison =====") + print(f"FLASH_ATTN median: {comparison['flash_attn_median_seconds']:.3f} s") + print(f"FLASHINFER median: {comparison['flashinfer_median_seconds']:.3f} s") + print(f"FLASHINFER speedup: {comparison['flashinfer_speedup']:.3f}x") + print(f"Summary: {summary_path}") + + +if __name__ == "__main__": + main() From 86c7c61a1381644e4fc0a0e9aec84c00cab52179 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 16:35:21 -0700 Subject: [PATCH 08/19] delete benchmark --- .../inference/benchmark_attention_backends.py | 213 ------------------ 1 file changed, 213 deletions(-) delete mode 100644 examples/inference/benchmark_attention_backends.py diff --git a/examples/inference/benchmark_attention_backends.py b/examples/inference/benchmark_attention_backends.py deleted file mode 100644 index 573f782f1e..0000000000 --- a/examples/inference/benchmark_attention_backends.py +++ /dev/null @@ -1,213 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -"""Compare complete FastVideo inference with FlashAttention and FlashInfer. - -Each backend runs in a fresh subprocess because attention selection is resolved -when model components are constructed and CUDA state/model memory must not leak -between benchmark arms. - -Example: - CUDA_VISIBLE_DEVICES=0 python examples/inference/benchmark_attention_backends.py \ - --config scripts/inference/inference_wan.yaml \ - --override request.prompt="A fox running through snow" \ - --override request.inputs.prompt_path=null \ - --warmups 1 --repeats 3 --save-outputs -""" - -from __future__ import annotations - -import argparse -import json -import os -import statistics -import subprocess -import sys -import time -from copy import deepcopy -from dataclasses import asdict -from pathlib import Path -from typing import Any - -_BACKENDS = ("FLASH_ATTN", "FLASHINFER") -_RESULT_PREFIX = "FASTVIDEO_BACKEND_BENCHMARK_RESULT=" - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--config", required=True, help="Nested FastVideo inference YAML/JSON config") - parser.add_argument("--override", - action="append", - default=[], - help="Dotted generator/request override; repeat as needed") - parser.add_argument("--output-dir", default="outputs/attention_backend_comparison") - parser.add_argument("--warmups", type=int, default=1) - parser.add_argument("--repeats", type=int, default=3) - parser.add_argument("--save-outputs", - action=argparse.BooleanOptionalAction, - default=True, - help="Save measured videos for visual comparison (included in wall time)") - parser.add_argument("--backend-order", nargs=2, choices=_BACKENDS, default=list(_BACKENDS)) - parser.add_argument("--worker-backend", choices=_BACKENDS, help=argparse.SUPPRESS) - return parser.parse_args() - - -def _synchronize_cuda(torch_module: Any) -> None: - if torch_module.cuda.is_available(): - torch_module.cuda.synchronize() - - -def _run_worker(args: argparse.Namespace) -> None: - # This must happen before importing FastVideo: backend selection is folded - # into component construction and must stay fixed for the worker lifetime. - backend = args.worker_backend - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = backend - - import torch - - from fastvideo import VideoGenerator - from fastvideo.entrypoints.cli.inference_config import build_generate_run_config - - if not torch.cuda.is_available(): - raise RuntimeError("This benchmark requires an NVIDIA CUDA GPU") - if backend == "FLASHINFER" and torch.cuda.get_device_capability() < (8, 0): - raise RuntimeError("FLASHINFER requires an NVIDIA GPU with compute capability sm80 or newer") - - config_args = argparse.Namespace(config=args.config) - config_overrides = [item if item.startswith("--") else f"--{item}" for item in args.override] - run_config = build_generate_run_config(config_args, overrides=config_overrides) - if run_config.generator.engine.num_gpus != 1: - raise ValueError("This comparison script currently requires generator.engine.num_gpus=1") - - backend_dir = Path(args.output_dir).resolve() / backend.lower() - backend_dir.mkdir(parents=True, exist_ok=True) - - load_started = time.perf_counter() - generator = VideoGenerator.from_config(run_config.generator) - load_seconds = time.perf_counter() - load_started - - def run_once(phase: str, index: int, *, measured: bool) -> float: - request = deepcopy(run_config.request) - request.output.output_path = str(backend_dir) - request.output.output_video_name = f"{backend.lower()}_{phase}_{index:02d}" - request.output.save_video = args.save_outputs if measured else False - request.output.return_frames = False - - _synchronize_cuda(torch) - started = time.perf_counter() - generator.generate(request) - _synchronize_cuda(torch) - return time.perf_counter() - started - - warmup_seconds = [run_once("warmup", index, measured=False) for index in range(args.warmups)] - measured_seconds = [run_once("run", index, measured=True) for index in range(args.repeats)] - result = { - "backend": backend, - "device": torch.cuda.get_device_name(torch.cuda.current_device()), - "device_capability": list(torch.cuda.get_device_capability()), - "torch_version": torch.__version__, - "cuda_version": torch.version.cuda, - "config": str(Path(args.config).resolve()), - "generator": asdict(run_config.generator), - "request": asdict(run_config.request), - "load_seconds": load_seconds, - "warmup_seconds": warmup_seconds, - "measured_seconds": measured_seconds, - "median_seconds": statistics.median(measured_seconds), - "mean_seconds": statistics.mean(measured_seconds), - "save_outputs": args.save_outputs, - "output_dir": str(backend_dir), - } - print(f"{_RESULT_PREFIX}{json.dumps(result, default=str)}", flush=True) - - -def _run_backend_subprocess(args: argparse.Namespace, backend: str) -> dict[str, Any]: - command = [ - sys.executable, - str(Path(__file__).resolve()), - "--config", - args.config, - "--output-dir", - args.output_dir, - "--warmups", - str(args.warmups), - "--repeats", - str(args.repeats), - "--worker-backend", - backend, - "--save-outputs" if args.save_outputs else "--no-save-outputs", - ] - for override in args.override: - command.extend(("--override", override)) - - print(f"\n===== {backend} =====", flush=True) - process = subprocess.Popen(command, - cwd=Path.cwd(), - env=os.environ.copy(), - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - bufsize=1) - result: dict[str, Any] | None = None - assert process.stdout is not None - for line in process.stdout: - print(line, end="", flush=True) - if line.startswith(_RESULT_PREFIX): - result = json.loads(line[len(_RESULT_PREFIX):]) - return_code = process.wait() - if return_code != 0: - raise RuntimeError(f"{backend} worker failed with exit code {return_code}") - if result is None: - raise RuntimeError(f"{backend} worker exited without a benchmark result") - return result - - -def _write_summary(args: argparse.Namespace, results: list[dict[str, Any]]) -> Path: - by_backend = {result["backend"]: result for result in results} - flash_seconds = by_backend["FLASH_ATTN"]["median_seconds"] - flashinfer_seconds = by_backend["FLASHINFER"]["median_seconds"] - summary = { - "results": by_backend, - "comparison": { - "flash_attn_median_seconds": flash_seconds, - "flashinfer_median_seconds": flashinfer_seconds, - "flashinfer_speedup": flash_seconds / flashinfer_seconds, - "flashinfer_time_change_percent": (flashinfer_seconds / flash_seconds - 1.0) * 100.0, - }, - "timing_scope": "VideoGenerator.generate wall time with CUDA synchronization", - "notes": [ - "Warmup runs are excluded from statistics.", - "Saved-video encoding is included when --save-outputs is enabled.", - "A speedup greater than 1.0 means FLASHINFER was faster.", - ], - } - output_dir = Path(args.output_dir).resolve() - output_dir.mkdir(parents=True, exist_ok=True) - summary_path = output_dir / "summary.json" - summary_path.write_text(json.dumps(summary, indent=2, default=str) + "\n", encoding="utf-8") - return summary_path - - -def main() -> None: - args = _parse_args() - if args.warmups < 1: - raise ValueError("--warmups must be at least 1 so JIT/caches are excluded") - if args.repeats < 1: - raise ValueError("--repeats must be at least 1") - if set(args.backend_order) != set(_BACKENDS): - raise ValueError("--backend-order must contain FLASH_ATTN and FLASHINFER exactly once") - if args.worker_backend is not None: - _run_worker(args) - return - - results = [_run_backend_subprocess(args, backend) for backend in args.backend_order] - summary_path = _write_summary(args, results) - comparison = json.loads(summary_path.read_text(encoding="utf-8"))["comparison"] - print("\n===== Comparison =====") - print(f"FLASH_ATTN median: {comparison['flash_attn_median_seconds']:.3f} s") - print(f"FLASHINFER median: {comparison['flashinfer_median_seconds']:.3f} s") - print(f"FLASHINFER speedup: {comparison['flashinfer_speedup']:.3f}x") - print(f"Summary: {summary_path}") - - -if __name__ == "__main__": - main() From a52292f5d30967fb355455834ebf0ebeef5daaf1 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 31 Aug 2026 17:25:23 -0700 Subject: [PATCH 09/19] fix some bugs --- fastvideo/attention/backends/flashinfer.py | 11 + fastvideo/attention/selector.py | 9 +- fastvideo/models/dits/wanvideo.py | 773 ++++++++++++++++++ fastvideo/platforms/cuda.py | 2 +- .../test_flashinfer_cuda_dispatch.py | 75 ++ .../attention/test_selector_role_override.py | 18 +- 6 files changed, 878 insertions(+), 10 deletions(-) create mode 100644 fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py diff --git a/fastvideo/attention/backends/flashinfer.py b/fastvideo/attention/backends/flashinfer.py index f625f9deb9..96b4f75de5 100644 --- a/fastvideo/attention/backends/flashinfer.py +++ b/fastvideo/attention/backends/flashinfer.py @@ -6,6 +6,17 @@ 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 diff --git a/fastvideo/attention/selector.py b/fastvideo/attention/selector.py index 6297cebdec..93924cf84b 100644 --- a/fastvideo/attention/selector.py +++ b/fastvideo/attention/selector.py @@ -276,14 +276,7 @@ def _cached_get_attn_backend( # get device-specific attn_backend from fastvideo.platforms import current_platform - # FLASHINFER implements the same dense BSHD semantic contract as - # FLASH_ATTN. Treat every layer that already declares FLASH_ATTN support - # as FlashInfer-capable, avoiding model-by-model tuple churn while keeping - # sparse-only layers protected by their existing declarations. - selected_is_supported = (selected_backend in supported_attention_backends or - (selected_backend == AttentionBackendEnum.FLASHINFER - and AttentionBackendEnum.FLASH_ATTN in supported_attention_backends)) - if selected_backend is not None and not selected_is_supported: + if (selected_backend is not None and selected_backend not in supported_attention_backends): fallback_backend = (default_backend if default_backend in supported_attention_backends else None) logger.warning( "Requested attention backend %s is not supported by this " diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 50310c451a..dc236694ba 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +<<<<<<< HEAD """Compatibility imports for the canonical :mod:`fastvideo.models.wan.transformer`.""" from fastvideo.models.wan.transformer import ( @@ -28,3 +29,775 @@ "WanTransformerBlock", "WanTransformerBlock_VSA", ] +======= + +import copy +import math +from typing import Any + +import torch +import torch.nn as nn + +import fastvideo.envs as envs +from fastvideo.attention import (DistributedAttention, DistributedAttention_VSA, LocalAttention) +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard) +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, + ScaleResidualLayerNormScaleShift) +from fastvideo.layers.linear import ReplicatedLinear +# from torch.nn import RMSNorm +# TODO: RMSNorm .... +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import get_rotary_pos_embed +from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder) +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum, current_platform +from fastvideo.layers.quantization import QuantizationConfig + +from fastvideo.distributed.parallel_state import get_sp_world_size + +logger = init_logger(__name__) + + +class WanImageEmbedding(torch.nn.Module): + + def __init__(self, in_features: int, out_features: int): + super().__init__() + + self.norm1 = FP32LayerNorm(in_features) + self.ff = MLP(in_features, in_features, out_features, act_type="gelu") + self.norm2 = FP32LayerNorm(out_features) + + def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: + dtype = encoder_hidden_states_image.dtype + hidden_states = self.norm1(encoder_hidden_states_image) + hidden_states = self.ff(hidden_states) + hidden_states = self.norm2(hidden_states).to(dtype) + return hidden_states + + +class WanTimeTextImageEmbedding(nn.Module): + + def __init__( + self, + dim: int, + time_freq_dim: int, + text_embed_dim: int, + image_embed_dim: int | None = None, + *, + r_embedder: bool = False, + r_embedder_fusion: str = "additive", + r_embedder_gate_value: float = 0.25, + r_embedder_deltatime_type: str = "r", + ): + super().__init__() + + self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") + self.text_embedder = MLP(text_embed_dim, dim, dim, bias=True, + act_type="gelu_pytorch_tanh") if text_embed_dim > 0 else None + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = WanImageEmbedding(image_embed_dim, dim) + + # AnyFlow dual-timestep support. When r_embedder is False the forward + # path bypasses delta_embedder entirely and the output is byte-identical + # to the legacy single-timestep implementation. + self._r_embedder_enabled = bool(r_embedder) + self._r_embedder_fusion = r_embedder_fusion + self._r_embedder_deltatime_type = r_embedder_deltatime_type + if self._r_embedder_enabled: + if r_embedder_fusion not in ("additive", "gated"): + raise ValueError("r_embedder_fusion must be one of {additive, gated}, " + f"got {r_embedder_fusion!r}") + if r_embedder_deltatime_type not in ("r", "t-r"): + raise ValueError("r_embedder_deltatime_type must be one of {r, t-r}, " + f"got {r_embedder_deltatime_type!r}") + # Deep-copy preserves identical initialization with time_embedder, + # matching AnyFlow reference setup_flowmap_model() behavior. + self.delta_embedder = copy.deepcopy(self.time_embedder) + # Non-persistent buffer — gate is a hyperparameter, not learned. + self.register_buffer( + "_r_embedder_gate", + torch.tensor(float(r_embedder_gate_value)), + persistent=False, + ) + else: + self.delta_embedder = None + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, + r_timestep: torch.Tensor | None = None, + ): + temb = self.time_embedder(timestep, timestep_seq_len) + + if self._r_embedder_enabled and r_timestep is not None: + assert self.delta_embedder is not None + if self._r_embedder_deltatime_type == "r": + delta_input = r_timestep + else: + delta_input = timestep - r_timestep + delta_emb = self.delta_embedder(delta_input, timestep_seq_len) + gate = self._r_embedder_gate + if self._r_embedder_fusion == "gated": + temb = (1.0 - gate) * temb + gate * delta_emb + else: + # Additive (additional channel, no convex blend). + temb = temb + gate * delta_emb + + timestep_proj = self.time_modulation(temb) + + if self.text_embedder is not None: + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + else: + encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), + device=temb.device, + dtype=temb.dtype) + if encoder_hidden_states_image is not None: + assert self.image_embedder is not None + encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image) + + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + + +class WanSelfAttention(nn.Module): + + def __init__(self, + dim: int, + num_heads: int, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6, + parallel_attention=False, + quant_config: QuantizationConfig | None = None, + prefix: str = "") -> None: + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + self.parallel_attention = parallel_attention + + # layers + self.to_q = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_q") + self.to_k = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_k") + self.to_v = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_v") + self.to_out = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_out") + 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 + 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.FLASHINFER, + AttentionBackendEnum.TORCH_SDPA)) + + def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + pass + + +class WanT2VCrossAttention(WanSelfAttention): + + def forward(self, x, context, context_lens, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + + # compute attention + x = self.attn(q, k, v) if k.size(1) > 0 else torch.zeros_like(q) + + # output + x = x.flatten(2) + x, _ = self.to_out(x) + return x + + +class WanI2VCrossAttention(WanSelfAttention): + + def __init__( + self, + dim: int, + num_heads: int, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__(dim, + num_heads, + window_size, + qk_norm, + eps, + supported_attention_backends, + quant_config=quant_config, + prefix=prefix) + + self.add_k_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_k_proj") + self.add_v_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_v_proj") + self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, context, context_lens): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + context_img = context[:, :257] + context = context[:, 257:] + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(b, -1, n, d) + v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + img_x = self.attn(q, k_img, v_img) + # compute attention + x = self.attn(q, k, v) if k.size(1) > 0 else torch.zeros_like(q) + + # output + x = x.flatten(2) + img_x = img_x.flatten(2) + x = x + img_x + x, _ = self.to_out(x) + return x + + +class WanTransformerBlock(nn.Module): + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q") + self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k") + self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v") + + self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out") + self.attn1 = DistributedAttention(num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1") + self.hidden_dim = dim + self.num_attention_heads = num_heads + dim_head = dim // num_heads + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + # LTX applies qk norm across all heads + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + print("QK Norm type not supported") + raise Exception + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32) + + # 2. Cross-attention + if added_kv_proj_dim is not None: + # I2V + self.attn2 = WanI2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.attn2") + else: + # T2V + self.attn2 = WanT2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.attn2") + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + original_seq_len: int, + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + # assert orig_dtype != torch.float32 + + if temb.dim() == 4: + # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) + # batch_size, seq_len, 1, inner_dim + shift_msa = shift_msa.squeeze(2) + scale_msa = scale_msa.squeeze(2) + gate_msa = gate_msa.squeeze(2) + c_shift_msa = c_shift_msa.squeeze(2) + c_scale_msa = c_scale_msa.squeeze(2) + c_gate_msa = c_gate_msa.squeeze(2) + else: + # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) + e = self.scale_shift_table + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) + assert shift_msa.dtype == torch.float32 + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + attn_output, _ = self.attn1( + query, + key, + value, + original_seq_len, + freqs_cis=freqs_cis, + ) + attn_output = attn_output.flatten(2) + attn_output, _ = self.to_out(attn_output) + attn_output = attn_output.squeeze(1) + + null_shift = null_scale = torch.tensor([0], device=hidden_states.device) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + + +class WanTransformerBlock_VSA(nn.Module): + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q") + self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k") + self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v") + self.to_gate_compress = ReplicatedLinear(dim, + dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_gate_compress") + self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out") + self.attn1 = DistributedAttention_VSA(num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1") + self.hidden_dim = dim + self.num_attention_heads = num_heads + dim_head = dim // num_heads + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + # LTX applies qk norm across all heads + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + print("QK Norm type not supported") + raise Exception + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32) + + # 2. Cross-attention + if added_kv_proj_dim is not None: + # I2V + self.attn2 = WanI2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.attn2") + else: + # T2V + self.attn2 = WanT2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.attn2") + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + original_seq_len: int, + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + # assert orig_dtype != torch.float32 + e = self.scale_shift_table + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) + assert shift_msa.dtype == torch.float32 + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + gate_compress, _ = self.to_gate_compress(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + gate_compress = gate_compress.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + attn_output, _ = self.attn1( + query, + key, + value, + original_seq_len, + freqs_cis=freqs_cis, + gate_compress=gate_compress, + ) + attn_output = attn_output.flatten(2) + attn_output, _ = self.to_out(attn_output) + attn_output = attn_output.squeeze(1) + + null_shift = null_scale = torch.tensor([0], device=hidden_states.device) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + + +class WanTransformer3DModel(BaseDiT): + _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions + _compile_conditions = WanVideoConfig()._compile_conditions + _supported_attention_backends = WanVideoConfig()._supported_attention_backends + 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 + + def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + self.quant_config = config.quant_config + + inner_dim = config.num_attention_heads * config.attention_head_dim + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + self.text_len = config.text_len + + assert config.num_attention_heads % get_sp_world_size( + ) == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})" + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed(in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False) + + # 2. Condition embeddings + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + text_embed_dim=config.text_dim, + image_embed_dim=config.image_dim, + r_embedder=config.r_embedder, + r_embedder_fusion=config.r_embedder_fusion, + r_embedder_gate_value=config.r_embedder_gate_value, + r_embedder_deltatime_type=config.r_embedder_deltatime_type, + ) + + # 3. Transformer blocks + attn_backend = envs.FASTVIDEO_ATTENTION_BACKEND + transformer_block = WanTransformerBlock_VSA if attn_backend == "VIDEO_SPARSE_ATTN" else WanTransformerBlock + self.blocks = nn.ModuleList([ + transformer_block(inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + quant_config=config.quant_config, + prefix=f"{config.prefix}.blocks.{i}") for i in range(config.num_layers) + ]) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift(inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + self.__post_init__() + + def forward(self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + guidance=None, + r_timestep: torch.Tensor | None = None, + **kwargs) -> torch.Tensor: + orig_dtype = hidden_states.dtype + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, list): + encoder_hidden_states_image = (encoder_hidden_states_image[0] + if len(encoder_hidden_states_image) > 0 else None) + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000) + freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float()) + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + # Shard with padding support - returns (sharded_tensor, original_seq_len) + hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) + + current_seq_len = hidden_states.shape[1] + sp_world_size = get_sp_world_size() + padded_seq_len = current_seq_len * sp_world_size + + # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) + if timestep.dim() == 2: + ts_seq_len = timestep.shape[1] + timestep = timestep.flatten() # batch_size * seq_len + else: + ts_seq_len = None + + # AnyFlow dual-timestep — match timestep's flattening so embedder + # sees aligned shapes. + if r_timestep is not None and r_timestep.dim() == 2: + r_timestep = r_timestep.flatten() + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, + encoder_hidden_states, + encoder_hidden_states_image, + timestep_seq_len=ts_seq_len, + r_timestep=r_timestep) + if ts_seq_len is not None: + # batch_size, seq_len, 6, inner_dim + timestep_proj = timestep_proj.unflatten(2, (6, -1)) + else: + # batch_size, 6, inner_dim + timestep_proj = timestep_proj.unflatten(1, (6, -1)) + + if encoder_hidden_states_image is not None: + if encoder_hidden_states is not None: + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) + else: + encoder_hidden_states = encoder_hidden_states_image + + if current_platform.is_mps() or current_platform.is_npu(): + encoder_hidden_states = encoder_hidden_states.to(orig_dtype) + else: + encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU + + assert encoder_hidden_states.dtype == orig_dtype + + # 4. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, original_seq_len) + else: + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, original_seq_len) + # 5. Output norm, projection & unpatchify + if temb.dim() == 3: + # batch_size, seq_len, inner_dim (wan 2.2 ti2v) + shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2) + shift = shift.squeeze(2) + scale = scale.squeeze(2) + else: + # batch_size, inner_dim + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) + + hidden_states = self.norm_out(hidden_states, shift, scale) + + # Gather and unpad in one operation + hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width, + p_t, p_h, p_w, -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output + + +# Entry point for model registry +EntryClass = WanTransformer3DModel +>>>>>>> 9b568284 (fix some bugs) diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index fa7437acc5..728537160e 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -177,7 +177,7 @@ def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, hea 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.info("FLASHINFER will cast %s inputs to bfloat16 for the kernel.", dtype) + 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 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 058249b93f..cb992c4d19 100644 --- a/fastvideo/tests/attention/test_selector_role_override.py +++ b/fastvideo/tests/attention/test_selector_role_override.py @@ -69,7 +69,14 @@ def test_unsupported_role_override_honors_layer_default(monkeypatch) -> None: selector._cached_get_attn_backend.cache_clear() -def test_flashinfer_is_compatible_with_dense_flash_attn_layers(monkeypatch) -> None: +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) @@ -78,6 +85,15 @@ def test_flashinfer_is_compatible_with_dense_flash_attn_layers(monkeypatch) -> N 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: From 0ac91518731787769e74837e62abfd9f58091012 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Sun, 6 Sep 2026 22:43:35 -0700 Subject: [PATCH 10/19] rebase --- fastvideo/models/dits/wanvideo.py | 773 ------------------------------ 1 file changed, 773 deletions(-) diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index dc236694ba..50310c451a 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 -<<<<<<< HEAD """Compatibility imports for the canonical :mod:`fastvideo.models.wan.transformer`.""" from fastvideo.models.wan.transformer import ( @@ -29,775 +28,3 @@ "WanTransformerBlock", "WanTransformerBlock_VSA", ] -======= - -import copy -import math -from typing import Any - -import torch -import torch.nn as nn - -import fastvideo.envs as envs -from fastvideo.attention import (DistributedAttention, DistributedAttention_VSA, LocalAttention) -from fastvideo.configs.models.dits import WanVideoConfig -from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, - sequence_model_parallel_shard) -from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, - ScaleResidualLayerNormScaleShift) -from fastvideo.layers.linear import ReplicatedLinear -# from torch.nn import RMSNorm -# TODO: RMSNorm .... -from fastvideo.layers.mlp import MLP -from fastvideo.layers.rotary_embedding import get_rotary_pos_embed -from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder) -from fastvideo.logger import init_logger -from fastvideo.models.dits.base import BaseDiT -from fastvideo.platforms import AttentionBackendEnum, current_platform -from fastvideo.layers.quantization import QuantizationConfig - -from fastvideo.distributed.parallel_state import get_sp_world_size - -logger = init_logger(__name__) - - -class WanImageEmbedding(torch.nn.Module): - - def __init__(self, in_features: int, out_features: int): - super().__init__() - - self.norm1 = FP32LayerNorm(in_features) - self.ff = MLP(in_features, in_features, out_features, act_type="gelu") - self.norm2 = FP32LayerNorm(out_features) - - def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: - dtype = encoder_hidden_states_image.dtype - hidden_states = self.norm1(encoder_hidden_states_image) - hidden_states = self.ff(hidden_states) - hidden_states = self.norm2(hidden_states).to(dtype) - return hidden_states - - -class WanTimeTextImageEmbedding(nn.Module): - - def __init__( - self, - dim: int, - time_freq_dim: int, - text_embed_dim: int, - image_embed_dim: int | None = None, - *, - r_embedder: bool = False, - r_embedder_fusion: str = "additive", - r_embedder_gate_value: float = 0.25, - r_embedder_deltatime_type: str = "r", - ): - super().__init__() - - self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu") - self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") - self.text_embedder = MLP(text_embed_dim, dim, dim, bias=True, - act_type="gelu_pytorch_tanh") if text_embed_dim > 0 else None - - self.image_embedder = None - if image_embed_dim is not None: - self.image_embedder = WanImageEmbedding(image_embed_dim, dim) - - # AnyFlow dual-timestep support. When r_embedder is False the forward - # path bypasses delta_embedder entirely and the output is byte-identical - # to the legacy single-timestep implementation. - self._r_embedder_enabled = bool(r_embedder) - self._r_embedder_fusion = r_embedder_fusion - self._r_embedder_deltatime_type = r_embedder_deltatime_type - if self._r_embedder_enabled: - if r_embedder_fusion not in ("additive", "gated"): - raise ValueError("r_embedder_fusion must be one of {additive, gated}, " - f"got {r_embedder_fusion!r}") - if r_embedder_deltatime_type not in ("r", "t-r"): - raise ValueError("r_embedder_deltatime_type must be one of {r, t-r}, " - f"got {r_embedder_deltatime_type!r}") - # Deep-copy preserves identical initialization with time_embedder, - # matching AnyFlow reference setup_flowmap_model() behavior. - self.delta_embedder = copy.deepcopy(self.time_embedder) - # Non-persistent buffer — gate is a hyperparameter, not learned. - self.register_buffer( - "_r_embedder_gate", - torch.tensor(float(r_embedder_gate_value)), - persistent=False, - ) - else: - self.delta_embedder = None - - def forward( - self, - timestep: torch.Tensor, - encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: torch.Tensor | None = None, - timestep_seq_len: int | None = None, - r_timestep: torch.Tensor | None = None, - ): - temb = self.time_embedder(timestep, timestep_seq_len) - - if self._r_embedder_enabled and r_timestep is not None: - assert self.delta_embedder is not None - if self._r_embedder_deltatime_type == "r": - delta_input = r_timestep - else: - delta_input = timestep - r_timestep - delta_emb = self.delta_embedder(delta_input, timestep_seq_len) - gate = self._r_embedder_gate - if self._r_embedder_fusion == "gated": - temb = (1.0 - gate) * temb + gate * delta_emb - else: - # Additive (additional channel, no convex blend). - temb = temb + gate * delta_emb - - timestep_proj = self.time_modulation(temb) - - if self.text_embedder is not None: - encoder_hidden_states = self.text_embedder(encoder_hidden_states) - else: - encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), - device=temb.device, - dtype=temb.dtype) - if encoder_hidden_states_image is not None: - assert self.image_embedder is not None - encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image) - - return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image - - -class WanSelfAttention(nn.Module): - - def __init__(self, - dim: int, - num_heads: int, - window_size=(-1, -1), - qk_norm=True, - eps=1e-6, - parallel_attention=False, - quant_config: QuantizationConfig | None = None, - prefix: str = "") -> None: - assert dim % num_heads == 0 - super().__init__() - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - self.window_size = window_size - self.qk_norm = qk_norm - self.eps = eps - self.parallel_attention = parallel_attention - - # layers - self.to_q = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_q") - self.to_k = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_k") - self.to_v = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_v") - self.to_out = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_out") - 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 - 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.FLASHINFER, - AttentionBackendEnum.TORCH_SDPA)) - - def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int): - r""" - Args: - x(Tensor): Shape [B, L, num_heads, C / num_heads] - seq_lens(Tensor): Shape [B] - grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) - freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] - """ - pass - - -class WanT2VCrossAttention(WanSelfAttention): - - def forward(self, x, context, context_lens, crossattn_cache=None): - r""" - Args: - x(Tensor): Shape [B, L1, C] - context(Tensor): Shape [B, L2, C] - context_lens(Tensor): Shape [B] - """ - b, n, d = x.size(0), self.num_heads, self.head_dim - - # compute query, key, value - q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) - - if crossattn_cache is not None: - if not crossattn_cache["is_init"]: - crossattn_cache["is_init"] = True - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) - v = self.to_v(context)[0].view(b, -1, n, d) - crossattn_cache["k"] = k - crossattn_cache["v"] = v - else: - k = crossattn_cache["k"] - v = crossattn_cache["v"] - else: - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) - v = self.to_v(context)[0].view(b, -1, n, d) - - # compute attention - x = self.attn(q, k, v) if k.size(1) > 0 else torch.zeros_like(q) - - # output - x = x.flatten(2) - x, _ = self.to_out(x) - return x - - -class WanI2VCrossAttention(WanSelfAttention): - - def __init__( - self, - dim: int, - num_heads: int, - window_size=(-1, -1), - qk_norm=True, - eps=1e-6, - supported_attention_backends: tuple[AttentionBackendEnum, ...] - | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__(dim, - num_heads, - window_size, - qk_norm, - eps, - supported_attention_backends, - quant_config=quant_config, - prefix=prefix) - - self.add_k_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_k_proj") - self.add_v_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_v_proj") - self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - - def forward(self, x, context, context_lens): - r""" - Args: - x(Tensor): Shape [B, L1, C] - context(Tensor): Shape [B, L2, C] - context_lens(Tensor): Shape [B] - """ - context_img = context[:, :257] - context = context[:, 257:] - b, n, d = x.size(0), self.num_heads, self.head_dim - - # compute query, key, value - q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) - v = self.to_v(context)[0].view(b, -1, n, d) - k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(b, -1, n, d) - v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) - img_x = self.attn(q, k_img, v_img) - # compute attention - x = self.attn(q, k, v) if k.size(1) > 0 else torch.zeros_like(q) - - # output - x = x.flatten(2) - img_x = img_x.flatten(2) - x = x + img_x - x, _ = self.to_out(x) - return x - - -class WanTransformerBlock(nn.Module): - - def __init__(self, - dim: int, - ffn_dim: int, - num_heads: int, - qk_norm: str = "rms_norm_across_heads", - cross_attn_norm: bool = False, - eps: float = 1e-6, - added_kv_proj_dim: int | None = None, - supported_attention_backends: tuple[AttentionBackendEnum, ...] - | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = ""): - super().__init__() - - # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q") - self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k") - self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v") - - self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out") - self.attn1 = DistributedAttention(num_heads=num_heads, - head_size=dim // num_heads, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn1") - self.hidden_dim = dim - self.num_attention_heads = num_heads - dim_head = dim // num_heads - if qk_norm == "rms_norm": - self.norm_q = RMSNorm(dim_head, eps=eps) - self.norm_k = RMSNorm(dim_head, eps=eps) - elif qk_norm == "rms_norm_across_heads": - # LTX applies qk norm across all heads - self.norm_q = RMSNorm(dim, eps=eps) - self.norm_k = RMSNorm(dim, eps=eps) - else: - print("QK Norm type not supported") - raise Exception - assert cross_attn_norm is True - self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) - - # 2. Cross-attention - if added_kv_proj_dim is not None: - # I2V - self.attn2 = WanI2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps, - quant_config=quant_config, - prefix=f"{prefix}.attn2") - else: - # T2V - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps, - quant_config=quant_config, - prefix=f"{prefix}.attn2") - self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - - # 3. Feed-forward - self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn") - self.mlp_residual = ScaleResidual() - - self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - temb: torch.Tensor, - freqs_cis: tuple[torch.Tensor, torch.Tensor], - original_seq_len: int, - ) -> torch.Tensor: - if hidden_states.dim() == 4: - hidden_states = hidden_states.squeeze(1) - bs, seq_length, _ = hidden_states.shape - orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 - - if temb.dim() == 4: - # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) - # batch_size, seq_len, 1, inner_dim - shift_msa = shift_msa.squeeze(2) - scale_msa = scale_msa.squeeze(2) - gate_msa = gate_msa.squeeze(2) - c_shift_msa = c_shift_msa.squeeze(2) - c_scale_msa = c_scale_msa.squeeze(2) - c_gate_msa = c_gate_msa.squeeze(2) - else: - # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) - e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) - assert shift_msa.dtype == torch.float32 - - # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) - query, _ = self.to_q(norm_hidden_states) - key, _ = self.to_k(norm_hidden_states) - value, _ = self.to_v(norm_hidden_states) - - if self.norm_q is not None: - query = self.norm_q(query) - if self.norm_k is not None: - key = self.norm_k(key) - - query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - - attn_output, _ = self.attn1( - query, - key, - value, - original_seq_len, - freqs_cis=freqs_cis, - ) - attn_output = attn_output.flatten(2) - attn_output, _ = self.to_out(attn_output) - attn_output = attn_output.squeeze(1) - - null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, - null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) - - # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, - c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) - - # 3. Feed-forward - ff_output = self.ffn(norm_hidden_states) - hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) - hidden_states = hidden_states.to(orig_dtype) - - return hidden_states - - -class WanTransformerBlock_VSA(nn.Module): - - def __init__(self, - dim: int, - ffn_dim: int, - num_heads: int, - qk_norm: str = "rms_norm_across_heads", - cross_attn_norm: bool = False, - eps: float = 1e-6, - added_kv_proj_dim: int | None = None, - supported_attention_backends: tuple[AttentionBackendEnum, ...] - | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = ""): - super().__init__() - - # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q") - self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k") - self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v") - self.to_gate_compress = ReplicatedLinear(dim, - dim, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.to_gate_compress") - self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out") - self.attn1 = DistributedAttention_VSA(num_heads=num_heads, - head_size=dim // num_heads, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn1") - self.hidden_dim = dim - self.num_attention_heads = num_heads - dim_head = dim // num_heads - if qk_norm == "rms_norm": - self.norm_q = RMSNorm(dim_head, eps=eps) - self.norm_k = RMSNorm(dim_head, eps=eps) - elif qk_norm == "rms_norm_across_heads": - # LTX applies qk norm across all heads - self.norm_q = RMSNorm(dim, eps=eps) - self.norm_k = RMSNorm(dim, eps=eps) - else: - print("QK Norm type not supported") - raise Exception - assert cross_attn_norm is True - self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) - - # 2. Cross-attention - if added_kv_proj_dim is not None: - # I2V - self.attn2 = WanI2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps, - quant_config=quant_config, - prefix=f"{prefix}.attn2") - else: - # T2V - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps, - quant_config=quant_config, - prefix=f"{prefix}.attn2") - self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - - # 3. Feed-forward - self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn") - self.mlp_residual = ScaleResidual() - - self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - temb: torch.Tensor, - freqs_cis: tuple[torch.Tensor, torch.Tensor], - original_seq_len: int, - ) -> torch.Tensor: - if hidden_states.dim() == 4: - hidden_states = hidden_states.squeeze(1) - bs, seq_length, _ = hidden_states.shape - orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 - e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) - assert shift_msa.dtype == torch.float32 - - # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) - query, _ = self.to_q(norm_hidden_states) - key, _ = self.to_k(norm_hidden_states) - value, _ = self.to_v(norm_hidden_states) - gate_compress, _ = self.to_gate_compress(norm_hidden_states) - - if self.norm_q is not None: - query = self.norm_q(query) - if self.norm_k is not None: - key = self.norm_k(key) - - query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - gate_compress = gate_compress.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - - attn_output, _ = self.attn1( - query, - key, - value, - original_seq_len, - freqs_cis=freqs_cis, - gate_compress=gate_compress, - ) - attn_output = attn_output.flatten(2) - attn_output, _ = self.to_out(attn_output) - attn_output = attn_output.squeeze(1) - - null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, - null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) - - # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, - c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) - - # 3. Feed-forward - ff_output = self.ffn(norm_hidden_states) - hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) - hidden_states = hidden_states.to(orig_dtype) - - return hidden_states - - -class WanTransformer3DModel(BaseDiT): - _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions - _compile_conditions = WanVideoConfig()._compile_conditions - _supported_attention_backends = WanVideoConfig()._supported_attention_backends - 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 - - def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: - super().__init__(config=config, hf_config=hf_config) - self.quant_config = config.quant_config - - inner_dim = config.num_attention_heads * config.attention_head_dim - self.hidden_size = config.hidden_size - self.num_attention_heads = config.num_attention_heads - self.in_channels = config.in_channels - self.out_channels = config.out_channels - self.num_channels_latents = config.num_channels_latents - self.patch_size = config.patch_size - self.text_len = config.text_len - - assert config.num_attention_heads % get_sp_world_size( - ) == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})" - - # 1. Patch & position embedding - self.patch_embedding = PatchEmbed(in_chans=config.in_channels, - embed_dim=inner_dim, - patch_size=config.patch_size, - flatten=False) - - # 2. Condition embeddings - self.condition_embedder = WanTimeTextImageEmbedding( - dim=inner_dim, - time_freq_dim=config.freq_dim, - text_embed_dim=config.text_dim, - image_embed_dim=config.image_dim, - r_embedder=config.r_embedder, - r_embedder_fusion=config.r_embedder_fusion, - r_embedder_gate_value=config.r_embedder_gate_value, - r_embedder_deltatime_type=config.r_embedder_deltatime_type, - ) - - # 3. Transformer blocks - attn_backend = envs.FASTVIDEO_ATTENTION_BACKEND - transformer_block = WanTransformerBlock_VSA if attn_backend == "VIDEO_SPARSE_ATTN" else WanTransformerBlock - self.blocks = nn.ModuleList([ - transformer_block(inner_dim, - config.ffn_dim, - config.num_attention_heads, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - quant_config=config.quant_config, - prefix=f"{config.prefix}.blocks.{i}") for i in range(config.num_layers) - ]) - - # 4. Output norm & projection - self.norm_out = LayerNormScaleShift(inner_dim, - norm_type="layer", - eps=config.eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) - self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) - - self.gradient_checkpointing = False - self.__post_init__() - - def forward(self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] - | None = None, - guidance=None, - r_timestep: torch.Tensor | None = None, - **kwargs) -> torch.Tensor: - orig_dtype = hidden_states.dtype - if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): - encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, list): - encoder_hidden_states_image = (encoder_hidden_states_image[0] - if len(encoder_hidden_states_image) > 0 else None) - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t, p_h, p_w = self.patch_size - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p_h - post_patch_width = width // p_w - - # Get rotary embeddings - d = self.hidden_size // self.num_attention_heads - rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000) - freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float()) - - hidden_states = self.patch_embedding(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) - - # Shard with padding support - returns (sharded_tensor, original_seq_len) - hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) - - current_seq_len = hidden_states.shape[1] - sp_world_size = get_sp_world_size() - padded_seq_len = current_seq_len * sp_world_size - - # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) - if timestep.dim() == 2: - ts_seq_len = timestep.shape[1] - timestep = timestep.flatten() # batch_size * seq_len - else: - ts_seq_len = None - - # AnyFlow dual-timestep — match timestep's flattening so embedder - # sees aligned shapes. - if r_timestep is not None and r_timestep.dim() == 2: - r_timestep = r_timestep.flatten() - - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, - encoder_hidden_states, - encoder_hidden_states_image, - timestep_seq_len=ts_seq_len, - r_timestep=r_timestep) - if ts_seq_len is not None: - # batch_size, seq_len, 6, inner_dim - timestep_proj = timestep_proj.unflatten(2, (6, -1)) - else: - # batch_size, 6, inner_dim - timestep_proj = timestep_proj.unflatten(1, (6, -1)) - - if encoder_hidden_states_image is not None: - if encoder_hidden_states is not None: - encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) - else: - encoder_hidden_states = encoder_hidden_states_image - - if current_platform.is_mps() or current_platform.is_npu(): - encoder_hidden_states = encoder_hidden_states.to(orig_dtype) - else: - encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU - - assert encoder_hidden_states.dtype == orig_dtype - - # 4. Transformer blocks - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.blocks: - hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, original_seq_len) - else: - for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, original_seq_len) - # 5. Output norm, projection & unpatchify - if temb.dim() == 3: - # batch_size, seq_len, inner_dim (wan 2.2 ti2v) - shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2) - shift = shift.squeeze(2) - scale = scale.squeeze(2) - else: - # batch_size, inner_dim - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) - - hidden_states = self.norm_out(hidden_states, shift, scale) - - # Gather and unpad in one operation - hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width, - p_t, p_h, p_w, -1) - hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) - output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - - return output - - -# Entry point for model registry -EntryClass = WanTransformer3DModel ->>>>>>> 9b568284 (fix some bugs) From e8db45c3bf66b6d60fa6659f85df9855a0339058 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 00:19:39 -0700 Subject: [PATCH 11/19] change version compatibility for cudnn kernel in flashinfer --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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", From 05b2379b259bbcbe9cb2925fac4517d97b19b5c4 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 00:20:51 -0700 Subject: [PATCH 12/19] add flashinfer implementation --- fastvideo/attention/backends/flashinfer.py | 66 ++++++++++++++++++++-- 1 file changed, 60 insertions(+), 6 deletions(-) diff --git a/fastvideo/attention/backends/flashinfer.py b/fastvideo/attention/backends/flashinfer.py index 96b4f75de5..d411ba302e 100644 --- a/fastvideo/attention/backends/flashinfer.py +++ b/fastvideo/attention/backends/flashinfer.py @@ -1,9 +1,13 @@ # 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. +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. @@ -23,6 +27,7 @@ import torch +import fastvideo.envs as envs from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder) @@ -94,6 +99,11 @@ def _mask_for_sample(attn_mask: torch.Tensor, sample: int, query_len: int, key_l 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, @@ -102,17 +112,52 @@ def __init__(self, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args) -> None: - del num_heads, head_size, num_kv_heads, prefix, extra_impl_args + 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.") - 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) @@ -121,6 +166,15 @@ def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, 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 From ba0812e6029a9f1f79416ca2925dcca33b21b77e Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 00:22:15 -0700 Subject: [PATCH 13/19] import cudnn kernels --- fastvideo/platforms/cuda.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index 728537160e..b6d8b1408f 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -181,6 +181,9 @@ def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, hea 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: From e8b1669b8bfc8e2fee8074fadc17206407ff7bbf Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 00:24:18 -0700 Subject: [PATCH 14/19] add prefill backend for flashinfer --- fastvideo/envs.py | 7 +++++++ 1 file changed, 7 insertions(+) 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 From 78a737113d60e96bd0c85976f82b70f6e6769949 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 00:25:07 -0700 Subject: [PATCH 15/19] add cudnn test --- .../attention/test_flashinfer_backend.py | 71 ++++++++++++++++++- 1 file changed, 69 insertions(+), 2 deletions(-) diff --git a/fastvideo/tests/attention/test_flashinfer_backend.py b/fastvideo/tests/attention/test_flashinfer_backend.py index 74da860008..8d5551d87b 100644 --- a/fastvideo/tests/attention/test_flashinfer_backend.py +++ b/fastvideo/tests/attention/test_flashinfer_backend.py @@ -18,6 +18,7 @@ def test_padding_mask_is_front_padded_and_expanded() -> None: 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): @@ -49,6 +50,49 @@ def fake_kernel(q, k, v, **kwargs): 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") @@ -77,8 +121,9 @@ def _sdpa_reference(query: torch.Tensor, @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: +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()) @@ -97,8 +142,9 @@ def test_flashinfer_real_cuda_kernel_matches_sdpa(dtype: torch.dtype, head_size: @pytest.mark.parametrize("mode", ["causal", "cross_gqa", "causal_padding"]) -def test_flashinfer_real_cuda_kernel_attention_modes(mode: str) -> None: +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()) @@ -138,3 +184,24 @@ def test_flashinfer_real_cuda_kernel_attention_modes(mode: str) -> None: 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) From 880a7734dd2805415b7db23b733b65047e902ad4 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 01:06:57 -0700 Subject: [PATCH 16/19] pre-commit --- fastvideo/attention/backends/flashinfer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/fastvideo/attention/backends/flashinfer.py b/fastvideo/attention/backends/flashinfer.py index d411ba302e..0dceaf94d5 100644 --- a/fastvideo/attention/backends/flashinfer.py +++ b/fastvideo/attention/backends/flashinfer.py @@ -123,6 +123,7 @@ def __init__(self, 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 From 01fcc8c65ddc539073dccd8d19eebfe3862f5e00 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 01:08:38 -0700 Subject: [PATCH 17/19] update docs for flashinfer kernel --- docs/inference/optimizations.md | 63 +++++++++++++++++++++++++-------- 1 file changed, 48 insertions(+), 15 deletions(-) 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 From ed03e31f80d519dfc7559e6f5220fbf3a76651b4 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 17:35:06 -0700 Subject: [PATCH 18/19] update test --- .../test_flashinfer_cuda_dispatch.py | 40 +++++++++++++++++-- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py b/fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py index 1b782f308e..1c1452451c 100644 --- a/fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py +++ b/fastvideo/tests/attention/test_flashinfer_cuda_dispatch.py @@ -1,10 +1,11 @@ # 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. +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 @@ -32,6 +33,16 @@ def _install_fake_flashinfer(monkeypatch) -> None: 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. @@ -73,3 +84,24 @@ def test_flashinfer_warns_on_dtype_cast(monkeypatch, caplog) -> None: 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" From 981a04bb9eb87f8f8380397c827f4b71a37f63d2 Mon Sep 17 00:00:00 2001 From: klhhhhh <1412841649@qq.com> Date: Mon, 7 Sep 2026 18:11:13 -0700 Subject: [PATCH 19/19] add cudnn head size check --- fastvideo/platforms/cuda.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index b6d8b1408f..74cbf6df23 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -193,6 +193,13 @@ def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, hea 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: