From c29c7163fc56c4f293a0c501e0d080efc12e80e5 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 28 Jun 2026 17:45:44 -0700 Subject: [PATCH] [bugfix]: fix attention backend routing and coverage (#1494) --- docs/inference/support_matrix.md | 2 +- .../gradio/serving/ray_serve_backend.py | 11 ++- ..._vsa.yaml => fast_ti2v_fullattn_lora.yaml} | 12 +-- fastvideo/attention/selector.py | 12 ++- fastvideo/configs/models/dits/base.py | 47 +++++++-- fastvideo/configs/pipelines/base.py | 8 +- fastvideo/configs/pipelines/wan.py | 11 ++- fastvideo/models/dits/wanvideo.py | 27 +++--- fastvideo/registry.py | 19 ++-- .../test_attention_backend_selector_cache.py | 90 +++++++++++++++++ .../attention/test_cuda_backend_missing.py | 97 +++++++++++++++++++ .../attention/test_fastwan_vsa_scoping.py | 82 ++++++++++++++-- fastvideo/tests/modal/pr_test.py | 2 +- ...> inference_wan_FullAttn_DMD_5B_720P.yaml} | 2 +- 14 files changed, 360 insertions(+), 62 deletions(-) rename examples/train/configs/fine_tuning/wan/{fast_ti2v_fullattn_lora_vsa.yaml => fast_ti2v_fullattn_lora.yaml} (86%) create mode 100644 fastvideo/tests/attention/test_attention_backend_selector_cache.py create mode 100644 fastvideo/tests/attention/test_cuda_backend_missing.py rename scripts/inference/{inference_wan_VSA_DMD_5B_720P.yaml => inference_wan_FullAttn_DMD_5B_720P.yaml} (93%) diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index d0eec1564..f46e8dfc7 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -56,7 +56,7 @@ pipeline initialization and sampling. | Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sliding Tile Attn (Legacy Branch) | Sage Attn | VSA | BSA | |------------|---------------------|-------------|----------|-------------------|-----------|-----|-----| | FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | -| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | +| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ❌ | ⭕ | | Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | | DreamX-World 5B Cam | `FastVideo/DreamX-World-5B-Cam-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | DreamX-World 5B AR | `FastVideo/DreamX-World-5B-Diffusers` | 704px1280p | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | diff --git a/examples/inference/gradio/serving/ray_serve_backend.py b/examples/inference/gradio/serving/ray_serve_backend.py index 2a783207b..d0a64e351 100644 --- a/examples/inference/gradio/serving/ray_serve_backend.py +++ b/examples/inference/gradio/serving/ray_serve_backend.py @@ -127,11 +127,12 @@ def save_image_from_base64(image_data: str, output_dir: str) -> Optional[str]: def setup_model_environment(model_path: str) -> None: - # if "fullattn" in model_path.lower(): - # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" - # else: - # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" + from fastvideo.platforms import AttentionBackendEnum + from fastvideo.registry import get_pipeline_config_cls_from_name + + pipeline_config = get_pipeline_config_cls_from_name(model_path)() + attention_backend = pipeline_config.dit_config.required_attention_backend or AttentionBackendEnum.FLASH_ATTN + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = attention_backend.name os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" diff --git a/examples/train/configs/fine_tuning/wan/fast_ti2v_fullattn_lora_vsa.yaml b/examples/train/configs/fine_tuning/wan/fast_ti2v_fullattn_lora.yaml similarity index 86% rename from examples/train/configs/fine_tuning/wan/fast_ti2v_fullattn_lora_vsa.yaml rename to examples/train/configs/fine_tuning/wan/fast_ti2v_fullattn_lora.yaml index 1d57661a8..deb979708 100644 --- a/examples/train/configs/fine_tuning/wan/fast_ti2v_fullattn_lora_vsa.yaml +++ b/examples/train/configs/fine_tuning/wan/fast_ti2v_fullattn_lora.yaml @@ -1,6 +1,8 @@ -# FastWan 2.2 TI2V 5B FullAttn LoRA finetune with VSA. +# FastWan 2.2 TI2V 5B FullAttn LoRA finetune with dense attention. # # Notes: +# - The FullAttn checkpoint has no VSA gate-compression weights. Train its +# dense attention projections without enabling VIDEO_SPARSE_ATTN. # - This model still uses T2V-style preprocessed parquet data in the new # training stack. # - Validation/inference should use WanPipeline rather than the I2V pipeline; @@ -21,7 +23,6 @@ models: - to_k - to_v - to_out - - to_gate_compress method: _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod @@ -59,21 +60,18 @@ training: gradient_accumulation_steps: 5 checkpoint: - output_dir: outputs/fast_wan_ti2v_fullattn_lora_vsa + output_dir: outputs/fast_wan_ti2v_fullattn_lora training_state_checkpointing_steps: 5 checkpoints_total_limit: 2 resume_from_checkpoint: null tracker: project_name: fastvideo - run_name: fast_wan_ti2v_fullattn_lora_vsa + run_name: fast_wan_ti2v_fullattn_lora model: enable_gradient_checkpointing_type: full - vsa: - sparsity: 0.8 - callbacks: grad_clip: _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback diff --git a/fastvideo/attention/selector.py b/fastvideo/attention/selector.py index ffbee6ee1..baeacb354 100644 --- a/fastvideo/attention/selector.py +++ b/fastvideo/attention/selector.py @@ -144,21 +144,23 @@ def get_attn_backend( | None = None, default_backend: AttentionBackendEnum | None = None, ) -> type[AttentionBackend]: - return _cached_get_attn_backend(head_size, dtype, supported_attention_backends, default_backend) + # Resolve the global force / environment override before entering the + # cached function so backend changes participate in the cache key. This + # also validates the environment on every lookup, including cache hits. + selected_backend = get_selected_attn_backend() + return _cached_get_attn_backend(head_size, dtype, supported_attention_backends, selected_backend, default_backend) @cache def _cached_get_attn_backend( head_size: int, dtype: torch.dtype, - supported_attention_backends: tuple[AttentionBackendEnum, ...] - | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None, + selected_backend: AttentionBackendEnum | None, default_backend: AttentionBackendEnum | None = None, ) -> type[AttentionBackend]: if not supported_attention_backends: raise ValueError("supported_attention_backends is empty") - # A global force overrides the FASTVIDEO_ATTENTION_BACKEND env var. - selected_backend = get_selected_attn_backend() # Layer-level default (e.g. a checkpoint that requires a specific sparse # backend). Lower precedence than the global force and the env var, so diff --git a/fastvideo/configs/models/dits/base.py b/fastvideo/configs/models/dits/base.py index d799033e2..6d746a7dd 100644 --- a/fastvideo/configs/models/dits/base.py +++ b/fastvideo/configs/models/dits/base.py @@ -53,21 +53,54 @@ class DiTConfig(ModelConfig): # fails loudly if the selected FASTVIDEO_ATTENTION_BACKEND does not match it. # None = no requirement (the common case). required_attention_backend: AttentionBackendEnum | None = None + # A checkpoint can support the model's ordinary attention implementation + # while still being incompatible with a particular backend-specific block + # layout. Keep those checkpoint constraints separate from the architecture's + # general _supported_attention_backends list. + incompatible_attention_backends: tuple[AttentionBackendEnum, ...] = () + + @staticmethod + def _parse_attention_backend(backend: AttentionBackendEnum | str, field_name: str) -> AttentionBackendEnum: + if isinstance(backend, AttentionBackendEnum): + return backend + if isinstance(backend, str): + try: + return AttentionBackendEnum[backend] + except KeyError as exc: + valid_backends = ", ".join(AttentionBackendEnum.__members__) + raise ValueError( + f"Unknown attention backend {backend!r} in {field_name}; expected one of: {valid_backends}" + ) from exc + raise TypeError(f"{field_name} must contain AttentionBackendEnum values or names, got {type(backend).__name__}") def update_model_config(self, source_model_dict: dict[str, Any]) -> None: source_model_dict = source_model_dict.copy() required_backend = source_model_dict.get("required_attention_backend") if isinstance(required_backend, str): - try: - source_model_dict["required_attention_backend"] = AttentionBackendEnum[required_backend] - except KeyError as exc: - valid_backends = ", ".join(AttentionBackendEnum.__members__) - raise ValueError( - f"Unknown required attention backend {required_backend!r}; expected one of: {valid_backends}" - ) from exc + source_model_dict["required_attention_backend"] = self._parse_attention_backend( + required_backend, "required_attention_backend") + + incompatible_backends = source_model_dict.get("incompatible_attention_backends") + if incompatible_backends is not None: + if not isinstance(incompatible_backends, list | tuple): + raise TypeError("incompatible_attention_backends must be a list or tuple") + source_model_dict["incompatible_attention_backends"] = tuple( + self._parse_attention_backend(backend, "incompatible_attention_backends") + for backend in incompatible_backends) super().update_model_config(source_model_dict) + def validate_attention_backend_compatibility( + self, + selected_backend: AttentionBackendEnum | None, + *, + model_name: str = "This model", + ) -> None: + if selected_backend in self.incompatible_attention_backends: + assert selected_backend is not None + raise ValueError(f"{model_name} is incompatible with the {selected_backend.name} attention backend. " + "Select a compatible attention backend before loading the model.") + @staticmethod def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any: """Add CLI arguments for DiTConfig fields""" diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 4384d2999..a059a8f3c 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -284,8 +284,12 @@ class PipelineConfig: model_dict = asdict(value) # Model Arch Config should be hidden away from the users model_dict.pop("arch_config") - if isinstance(value, DiTConfig) and value.required_attention_backend is not None: - model_dict["required_attention_backend"] = value.required_attention_backend.name + if isinstance(value, DiTConfig): + if value.required_attention_backend is not None: + model_dict["required_attention_backend"] = value.required_attention_backend.name + model_dict["incompatible_attention_backends"] = [ + backend.name for backend in value.incompatible_attention_backends + ] output_dict[key] = model_dict elif isinstance(value, tuple) and all(isinstance(v, ModelConfig) for v in value): model_dicts = [] diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 2773e72d6..9dba53233 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -290,17 +290,20 @@ class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config): @dataclass -class FastWan2_2_TI2V_5B_FullAttn_Config(FastWan2_2_TI2V_5B_Config): +class FastWan2_2_TI2V_5B_FullAttn_Config(Wan2_2_TI2V_5B_Config): """FullAttn (dense-attention) DMD variant of FastWan 2.2 TI2V 5B. - Shares FastWan's DMD denoising schedule but runs dense attention, so -- unlike - the sparse-distilled FastWan2.2-TI2V-5B-Diffusers -- it does NOT require VSA. + Shares FastWan's DMD denoising schedule but uses the published dense + checkpoint layout, which has no VSA gate-compression weights. """ + flow_shift: float | None = 5.0 + dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522]) + def __post_init__(self) -> None: super().__post_init__() - # This checkpoint is dense; clear the VSA requirement inherited from FastWan. self.dit_config.required_attention_backend = None + self.dit_config.incompatible_attention_backends = (FASTWAN_REQUIRED_ATTENTION_BACKEND, ) @dataclass diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 6ad3119c5..027117e03 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -7,7 +7,6 @@ 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.attention.selector import check_attn_backend_requirement @@ -608,6 +607,21 @@ class WanTransformerBlock_VSA(nn.Module): return hidden_states +def _select_wan_transformer_block( + config: WanVideoConfig, + *, + model_name: str, +) -> type[WanTransformerBlock] | type[WanTransformerBlock_VSA]: + # Resolve checkpoint-level requirements before choosing a block layout. + # Sparse FastWan checkpoints require VSA's extra gate-compression weights; + # dense checkpoints explicitly reject that layout. + attn_backend = check_attn_backend_requirement(config.required_attention_backend, model_name=model_name) + config.validate_attention_backend_compatibility(attn_backend, model_name=model_name) + if attn_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: + return WanTransformerBlock_VSA + return WanTransformerBlock + + class WanTransformer3DModel(BaseDiT): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions @@ -620,6 +634,7 @@ class WanTransformer3DModel(BaseDiT): def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: super().__init__(config=config, hf_config=hf_config) + transformer_block = _select_wan_transformer_block(config, model_name=type(self).__name__) self.quant_config = config.quant_config inner_dim = config.num_attention_heads * config.attention_head_dim @@ -652,16 +667,6 @@ class WanTransformer3DModel(BaseDiT): ) # 3. Transformer blocks - # Some model families (e.g. FastWan) are only correct with a specific - # attention backend. Fail loudly if the required backend is not the one - # the user selected, rather than silently forcing it -- a force would - # only reach block construction here and leave the denoising stage - # resolving a different backend. Returns the selected backend - # (global force > FASTVIDEO_ATTENTION_BACKEND). - attn_backend = check_attn_backend_requirement(config.required_attention_backend, - model_name=type(self).__name__) - transformer_block = (WanTransformerBlock_VSA - if attn_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN else WanTransformerBlock) self.blocks = nn.ModuleList([ transformer_block(inner_dim, config.ffn_dim, diff --git a/fastvideo/registry.py b/fastvideo/registry.py index cb4fc778a..aa85c02f4 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -45,7 +45,6 @@ from fastvideo.configs.pipelines.turbodiffusion import ( ) from fastvideo.configs.pipelines.wan import ( FastWan2_1_T2V_480P_Config, - FastWan2_2_TI2V_5B_Config, FastWan2_2_TI2V_5B_FullAttn_Config, LucyEditDevConfig, SelfForcingWan2_2_T2V480PConfig, @@ -934,8 +933,14 @@ def _register_configs() -> None: workload_types=(WorkloadType.T2V, WorkloadType.I2V), hf_model_paths=[ "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", + "FastVideo/FastWan2.2-TI2V-5B-Diffusers", + ], + model_detectors=[ + lambda path: any(slug in path.lower() for slug in ( + "fastwan2.2-ti2v-5b-fullattn-diffusers", + "fastwan2.2-ti2v-5b-diffusers", + )), ], - model_detectors=[lambda path: "fastwan2.2-ti2v-5b-fullattn" in path.lower()], model_family="wan", default_preset="fast_wan_2_2_ti2v_5b", ) @@ -995,16 +1000,6 @@ def _register_configs() -> None: model_family="dreamx_world", default_preset="dreamx_world_5b_ar", ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=FastWan2_2_TI2V_5B_Config, - workload_types=(WorkloadType.T2V, WorkloadType.I2V), - hf_model_paths=[ - "FastVideo/FastWan2.2-TI2V-5B-Diffusers", - ], - model_family="wan", - default_preset="fast_wan_2_2_ti2v_5b", - ) register_configs( sampling_param_cls=None, pipeline_config_cls=LucyEditDevConfig, diff --git a/fastvideo/tests/attention/test_attention_backend_selector_cache.py b/fastvideo/tests/attention/test_attention_backend_selector_cache.py new file mode 100644 index 000000000..24235c76b --- /dev/null +++ b/fastvideo/tests/attention/test_attention_backend_selector_cache.py @@ -0,0 +1,90 @@ +from types import SimpleNamespace + +import pytest +import torch + +import fastvideo.platforms as platforms +from fastvideo.attention import selector +from fastvideo.platforms.interface import AttentionBackendEnum + + +class FlashBackend: + pass + + +class VideoSparseBackend: + pass + + +@pytest.fixture(autouse=True) +def reset_attention_backend_selector(): + selector.global_force_attn_backend(None) + yield + selector.global_force_attn_backend(None) + + +@pytest.fixture +def stub_backend_resolution(monkeypatch): + requested_backends = [] + + def get_attn_backend_cls(selected_backend, head_size, dtype): + requested_backends.append(selected_backend) + return selected_backend.name + + platform = SimpleNamespace( + device_name="test device", + get_attn_backend_cls=get_attn_backend_cls, + ) + backend_classes = { + AttentionBackendEnum.FLASH_ATTN.name: FlashBackend, + AttentionBackendEnum.VIDEO_SPARSE_ATTN.name: VideoSparseBackend, + } + monkeypatch.setattr(platforms, "_current_platform", platform) + monkeypatch.setattr(selector, "resolve_obj_by_qualname", backend_classes.__getitem__) + return requested_backends + + +def test_changing_env_backend_cannot_reuse_cached_backend(monkeypatch, stub_backend_resolution): + supported_backends = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + ) + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", AttentionBackendEnum.FLASH_ATTN.name) + + first_backend = selector.get_attn_backend(64, torch.float16, supported_backends) + + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", AttentionBackendEnum.VIDEO_SPARSE_ATTN.name) + second_backend = selector.get_attn_backend(64, torch.float16, supported_backends) + + assert first_backend is FlashBackend + assert second_backend is VideoSparseBackend + assert stub_backend_resolution == [ + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + ] + + +def test_invalid_env_backend_cannot_be_hidden_by_cache_hit(monkeypatch, stub_backend_resolution): + supported_backends = (AttentionBackendEnum.FLASH_ATTN, ) + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", AttentionBackendEnum.FLASH_ATTN.name) + assert selector.get_attn_backend(64, torch.float16, supported_backends) is FlashBackend + + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATN") + + with pytest.raises(ValueError, match="FLASH_ATN"): + selector.get_attn_backend(64, torch.float16, supported_backends) + assert stub_backend_resolution == [AttentionBackendEnum.FLASH_ATTN] + + +def test_global_force_still_takes_precedence_over_env(monkeypatch, stub_backend_resolution): + supported_backends = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + ) + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", AttentionBackendEnum.VIDEO_SPARSE_ATTN.name) + selector.global_force_attn_backend(AttentionBackendEnum.FLASH_ATTN) + + backend = selector.get_attn_backend(64, torch.float16, supported_backends) + + assert backend is FlashBackend + assert stub_backend_resolution == [AttentionBackendEnum.FLASH_ATTN] diff --git a/fastvideo/tests/attention/test_cuda_backend_missing.py b/fastvideo/tests/attention/test_cuda_backend_missing.py new file mode 100644 index 000000000..c43578b2f --- /dev/null +++ b/fastvideo/tests/attention/test_cuda_backend_missing.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import builtins +import sys +from types import ModuleType + +import pytest +import torch + +from fastvideo.platforms.cuda import CudaPlatformBase +from fastvideo.platforms.interface import AttentionBackendEnum + + +class _GpuIndependentCudaPlatform(CudaPlatformBase): + + @classmethod + def has_device_capability( + cls, + capability: tuple[int, int] | int, + device_id: int = 0, + ) -> bool: + return False + + +@pytest.mark.parametrize( + ("backend", "missing_module", "package_name", "install_hint"), + [ + ( + AttentionBackendEnum.SAGE_ATTN, + "sageattention", + "sageattention", + "uv pip install sageattention", + ), + ( + AttentionBackendEnum.SAGE_ATTN_THREE, + "sageattn3", + "sageattn3", + "uv pip install sageattn3", + ), + ( + AttentionBackendEnum.ATTN_QAT_INFER, + "fastvideo.attention.backends.attn_qat_infer", + "attn_qat_infer", + "https://hao-ai-lab.github.io/FastVideo/", + ), + ], +) +def test_explicit_backend_import_failure_does_not_fallback( + monkeypatch, + backend: AttentionBackendEnum, + missing_module: str, + package_name: str, + install_hint: str, +) -> None: + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", backend.name) + real_import = builtins.__import__ + + def import_with_missing_backend(name, globals=None, locals=None, fromlist=(), level=0): + if name == missing_module: + raise ImportError(f"No module named {name!r}") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", import_with_missing_backend) + + with pytest.raises(ImportError) as excinfo: + _GpuIndependentCudaPlatform.get_attn_backend_cls(backend, head_size=64, dtype=torch.float16) + + message = str(excinfo.value) + assert "explicitly requested" in message + assert backend.name in message + assert package_name in message + assert install_hint in message + assert excinfo.value.__cause__ is not None + + +def test_explicit_attn_qat_infer_unavailable_does_not_fallback(monkeypatch) -> None: + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", AttentionBackendEnum.ATTN_QAT_INFER.name) + module_name = "fastvideo.attention.backends.attn_qat_infer" + unavailable_backend = ModuleType(module_name) + unavailable_backend.AttnQatInferBackend = object + unavailable_backend.is_attn_qat_infer_available = lambda: False + monkeypatch.setitem(sys.modules, module_name, unavailable_backend) + + with pytest.raises(ImportError) as excinfo: + _GpuIndependentCudaPlatform.get_attn_backend_cls( + AttentionBackendEnum.ATTN_QAT_INFER, + head_size=64, + dtype=torch.float16, + ) + + message = str(excinfo.value) + assert "explicitly requested" in message + assert AttentionBackendEnum.ATTN_QAT_INFER.name in message + assert "attn_qat_infer" in message + assert "not built" in message + assert "Install the FastVideo Attn-QAT inference kernel" in message + assert "https://hao-ai-lab.github.io/FastVideo/" in message diff --git a/fastvideo/tests/attention/test_fastwan_vsa_scoping.py b/fastvideo/tests/attention/test_fastwan_vsa_scoping.py index 8436fa87d..2232673ab 100644 --- a/fastvideo/tests/attention/test_fastwan_vsa_scoping.py +++ b/fastvideo/tests/attention/test_fastwan_vsa_scoping.py @@ -16,6 +16,11 @@ from fastvideo.configs.pipelines.wan import ( FastWan2_2_TI2V_5B_FullAttn_Config, WanT2V480PConfig, ) +from fastvideo.models.dits.wanvideo import ( + WanTransformerBlock, + WanTransformerBlock_VSA, + _select_wan_transformer_block, +) from fastvideo.platforms.interface import AttentionBackendEnum from fastvideo.registry import get_pipeline_config_cls_from_name @@ -62,9 +67,21 @@ def test_fastwan_2_2_required_vsa_does_not_mutate_global_backend(): def test_fastwan_2_2_fullattn_config_does_not_require_vsa(): - # The dense FullAttn checkpoint shares FastWan's DMD schedule but must NOT be - # forced onto VSA -- it runs dense attention. - assert FastWan2_2_TI2V_5B_FullAttn_Config().dit_config.required_attention_backend is None + config = FastWan2_2_TI2V_5B_FullAttn_Config() + + assert config.dit_config.required_attention_backend is None + assert config.dit_config.incompatible_attention_backends == (VSA, ) + + +@pytest.mark.parametrize( + "model_id", + [ + "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", + "FastVideo/FastWan2.2-TI2V-5B-Diffusers", + ], +) +def test_fastwan_2_2_fullattn_hf_ids_resolve_dense_config(model_id: str) -> None: + assert get_pipeline_config_cls_from_name(model_id) is FastWan2_2_TI2V_5B_FullAttn_Config def _write_minimal_wan_dmd_repo(model_dir: Path) -> None: @@ -84,7 +101,7 @@ def _write_minimal_wan_dmd_repo(model_dir: Path) -> None: "relative_model_path", [ Path("models--FastVideo--FastWan2.2-TI2V-5B-FullAttn-Diffusers") / "snapshots" / "deadbeef", - Path("renamed-FastWan2.2-TI2V-5B-FullAttn-Diffusers-copy"), + Path("models--FastVideo--FastWan2.2-TI2V-5B-Diffusers") / "snapshots" / "deadbeef", ], ) def test_fastwan_2_2_fullattn_local_path_resolves_dense_config(tmp_path: Path, @@ -97,6 +114,29 @@ def test_fastwan_2_2_fullattn_local_path_resolves_dense_config(tmp_path: Path, assert resolved_cls is FastWan2_2_TI2V_5B_FullAttn_Config +def test_fastwan_2_2_fullattn_rejects_vsa_before_block_construction(): + config = FastWan2_2_TI2V_5B_FullAttn_Config().dit_config + global_force_attn_backend(VSA) + + with pytest.raises(ValueError, match=r"incompatible with the VIDEO_SPARSE_ATTN attention backend"): + _select_wan_transformer_block(config, model_name="FullAttn") + + +@pytest.mark.parametrize("backend", [AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA]) +def test_fastwan_2_2_fullattn_accepts_dense_backends(backend): + config = FastWan2_2_TI2V_5B_FullAttn_Config().dit_config + global_force_attn_backend(backend) + + assert _select_wan_transformer_block(config, model_name="FullAttn") is WanTransformerBlock + + +def test_sparse_fastwan_selects_vsa_block(): + config = FastWan2_2_TI2V_5B_Config().dit_config + global_force_attn_backend(VSA) + + assert _select_wan_transformer_block(config, model_name="FastWan") is WanTransformerBlock_VSA + + @pytest.mark.parametrize("config_cls", [FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config]) def test_sparse_fastwan_pipeline_config_json_roundtrip(tmp_path, config_cls): config_path = tmp_path / "pipeline_config.json" @@ -112,15 +152,45 @@ def test_sparse_fastwan_pipeline_config_json_roundtrip(tmp_path, config_cls): assert restored.dit_config.required_attention_backend is VSA -def test_pipeline_config_load_restores_required_backend_enum(tmp_path): +def test_fullattn_pipeline_config_json_roundtrip(tmp_path): config_path = tmp_path / "pipeline_config.json" - config = PipelineConfig(dit_config=DiTConfig(required_attention_backend=VSA)) + config = FastWan2_2_TI2V_5B_FullAttn_Config() + + config.dump_to_json(str(config_path)) + + payload = json.loads(config_path.read_text(encoding="utf-8")) + assert payload["dit_config"]["required_attention_backend"] is None + assert payload["dit_config"]["incompatible_attention_backends"] == [VSA.name] + + restored = FastWan2_2_TI2V_5B_FullAttn_Config() + restored.load_from_json(str(config_path)) + assert restored.dit_config.required_attention_backend is None + assert restored.dit_config.incompatible_attention_backends == (VSA, ) + + +def test_fullattn_config_cannot_be_overridden_to_require_vsa(): + config = FastWan2_2_TI2V_5B_FullAttn_Config() + + config.update_pipeline_config({"dit_config": {"required_attention_backend": VSA.name}}) + + assert config.dit_config.required_attention_backend is None + assert config.dit_config.incompatible_attention_backends == (VSA, ) + + +def test_pipeline_config_load_restores_attention_backend_constraints(tmp_path): + config_path = tmp_path / "pipeline_config.json" + config = PipelineConfig( + dit_config=DiTConfig( + required_attention_backend=VSA, + incompatible_attention_backends=(AttentionBackendEnum.FLASH_ATTN, ), + )) config.dump_to_json(str(config_path)) restored = PipelineConfig() restored.load_from_json(str(config_path)) assert restored.dit_config.required_attention_backend is VSA + assert restored.dit_config.incompatible_attention_backends == (AttentionBackendEnum.FLASH_ATTN, ) def test_check_requirement_returns_none_when_unrequired_and_unset(monkeypatch): diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index a35325c57..44170ae98 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -277,7 +277,7 @@ def run_self_forcing_tests(): @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret]) def run_unit_test(): run_test( - "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs" + "pytest ./fastvideo/tests/api/ ./fastvideo/tests/attention/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs" ) diff --git a/scripts/inference/inference_wan_VSA_DMD_5B_720P.yaml b/scripts/inference/inference_wan_FullAttn_DMD_5B_720P.yaml similarity index 93% rename from scripts/inference/inference_wan_VSA_DMD_5B_720P.yaml rename to scripts/inference/inference_wan_FullAttn_DMD_5B_720P.yaml index d0f6dc0c5..603c4a11e 100644 --- a/scripts/inference/inference_wan_VSA_DMD_5B_720P.yaml +++ b/scripts/inference/inference_wan_FullAttn_DMD_5B_720P.yaml @@ -1,5 +1,5 @@ # Run with: -# FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN fastvideo generate --config scripts/inference/inference_wan_VSA_DMD_5B_720P.yaml +# FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN fastvideo generate --config scripts/inference/inference_wan_FullAttn_DMD_5B_720P.yaml generator: model_path: FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers engine: