[bugfix]: address FastWan FullAttn review comments

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 5B Full Attn T2V-only* | `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 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
@@ -79,6 +79,9 @@ pipeline initialization and sampling.
 | Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ |
 
+*FastWan2.2 5B Full Attn is T2V-only and uses dense attention; do not run it with
+`FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`.
+
 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
 
 ***Lucy Edit Dev uses a non-commercial model license. FastVideo support is

fastvideo/configs/pipelines/wan.py

@@ -12,6 +12,7 @@ from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConf
 from fastvideo.configs.models.vaes import WanVAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEArchConfig
 from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.platforms import AttentionBackendEnum
 
 
 def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
@@ -64,7 +65,7 @@ class WanT2V480PConfig(PipelineConfig):
 
     # WanConfig-specific added parameters
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         self.vae_config.load_encoder = False
         self.vae_config.load_decoder = True
 
@@ -278,6 +279,33 @@ class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config):
     dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522])
 
 
+@dataclass
+class FastWan2_2_TI2V_5B_FullAttn_Config(FastWan2_2_TI2V_5B_Config):
+    dit_config: DiTConfig = field(
+        default_factory=lambda: WanVideoConfig(arch_config=WanVideoArchConfig(_supported_attention_backends=(
+            AttentionBackendEnum.FLASH_ATTN,
+            AttentionBackendEnum.TORCH_SDPA,
+            AttentionBackendEnum.SAGE_ATTN,
+            AttentionBackendEnum.SAGE_ATTN_THREE,
+            AttentionBackendEnum.ATTN_QAT_INFER,
+        ))))
+    ti2v_task: bool = False
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.vae_config.load_encoder = False
+
+    def check_pipeline_config(self) -> None:
+        super().check_pipeline_config()
+
+        from fastvideo import envs
+        if envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
+            raise ValueError("FastWan2.2-TI2V-5B-FullAttn-Diffusers is a dense FullAttn T2V-only checkpoint "
+                             "and does not support FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN. "
+                             "Unset FASTVIDEO_ATTENTION_BACKEND or use a dense backend such as FLASH_ATTN or "
+                             "TORCH_SDPA.")
+
+
 @dataclass
 class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
     flow_shift: float | None = 12.0

fastvideo/models/dits/wanvideo.py

@@ -607,6 +607,19 @@ class WanTransformerBlock_VSA(nn.Module):
         return hidden_states
 
 
+def _select_wan_transformer_block(config: WanVideoConfig) -> type[nn.Module]:
+    attn_backend = envs.FASTVIDEO_ATTENTION_BACKEND
+    if attn_backend == "VIDEO_SPARSE_ATTN":
+        supported_backends = config._supported_attention_backends or ()
+        if AttentionBackendEnum.VIDEO_SPARSE_ATTN not in supported_backends:
+            raise ValueError("FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN requires a Wan VSA checkpoint/config, "
+                             "but this Wan config does not support VIDEO_SPARSE_ATTN. "
+                             "For FastWan2.2-TI2V-5B-FullAttn-Diffusers, unset FASTVIDEO_ATTENTION_BACKEND or use "
+                             "a dense backend such as FLASH_ATTN or TORCH_SDPA.")
+        return WanTransformerBlock_VSA
+    return WanTransformerBlock
+
+
 class WanTransformer3DModel(BaseDiT):
     _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
     _compile_conditions = WanVideoConfig()._compile_conditions
@@ -651,8 +664,7 @@ class WanTransformer3DModel(BaseDiT):
         )
 
         # 3. Transformer blocks
-        attn_backend = envs.FASTVIDEO_ATTENTION_BACKEND
-        transformer_block = WanTransformerBlock_VSA if attn_backend == "VIDEO_SPARSE_ATTN" else WanTransformerBlock
+        transformer_block = _select_wan_transformer_block(config)
         self.blocks = nn.ModuleList([
             transformer_block(inner_dim,
                               config.ffn_dim,
@@ -661,7 +673,7 @@ class WanTransformer3DModel(BaseDiT):
                               config.cross_attn_norm,
                               config.eps,
                               config.added_kv_proj_dim,
-                              self._supported_attention_backends,
+                              config._supported_attention_backends,
                               quant_config=config.quant_config,
                               prefix=f"{config.prefix}.blocks.{i}")
             for i in range(config.num_layers)

fastvideo/models/encoders/qwen2_5_vl_custom.py

@@ -75,6 +75,25 @@ if not is_flash_attn_2_available():
 
 logger = logging.get_logger(__name__)
 
+
+def _resolve_torch_dtype(dtype, default: torch.dtype = torch.float32) -> torch.dtype:
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    if isinstance(dtype, str):
+        normalized = dtype.strip().removeprefix("torch.").lower()
+        return {
+            "bfloat16": torch.bfloat16,
+            "bf16": torch.bfloat16,
+            "float16": torch.float16,
+            "fp16": torch.float16,
+            "half": torch.float16,
+            "float32": torch.float32,
+            "fp32": torch.float32,
+            "float": torch.float32,
+        }.get(normalized, default)
+    return default
+
+
 class Qwen2_5_VLMLP(nn.Module):
     def __init__(self, config, bias: bool = False):
         super().__init__()
@@ -304,10 +323,13 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module):
     config_class = Qwen2_5_VLVisionConfig
     _no_split_modules = ["Qwen2_5_VLVisionBlock"]
 
-    def __init__(self, config) -> None:
+    def __init__(self, config, parent_torch_dtype=None) -> None:
         super().__init__()
 
-        self.dtype = torch.bfloat16 if config.torch_dtype == "bfloat16" else torch.float32
+        config_torch_dtype = getattr(config, "torch_dtype", None)
+        self.dtype = _resolve_torch_dtype(
+            config_torch_dtype if config_torch_dtype is not None else parent_torch_dtype
+        )
 
         self.spatial_merge_size = config.spatial_merge_size
         self.patch_size = config.patch_size
@@ -1458,7 +1480,10 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module):
         super().__init__()
         config = _flatten_text_config(config)
         self.config = config
-        self.visual = Qwen2_5_VisionTransformerPretrainedModel(config.vision_config)
+        self.visual = Qwen2_5_VisionTransformerPretrainedModel(
+            config.vision_config,
+            parent_torch_dtype=getattr(config, "torch_dtype", None),
+        )
 
         self.model = Qwen2_5_VLModel(config)
         self.vocab_size = config.vocab_size

fastvideo/registry.py

@@ -46,6 +46,7 @@ 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,
     SelfForcingWanT2V480PConfig,
@@ -985,12 +986,21 @@ def _register_configs() -> None:
         pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
         workload_types=(WorkloadType.T2V, WorkloadType.I2V),
         hf_model_paths=[
-            "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
             "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=FastWan2_2_TI2V_5B_FullAttn_Config,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+        ],
+        model_family="wan",
+        default_preset="fast_wan_2_2_ti2v_5b",
+    )
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=LucyEditDevConfig,

fastvideo/tests/api/test_fastwan_fullattn.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch.nn as nn
+
+import pytest
+
+from fastvideo.configs.pipelines.wan import (
+    FastWan2_2_TI2V_5B_Config,
+    FastWan2_2_TI2V_5B_FullAttn_Config,
+)
+from fastvideo.models.dits import wanvideo as wanvideo_module
+from fastvideo.models.dits.wanvideo import (
+    WanTransformerBlock,
+    WanTransformerBlock_VSA,
+    WanTransformer3DModel,
+    _select_wan_transformer_block,
+)
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.registry import (
+    get_pipeline_config_cls_from_name,
+    get_registered_models_with_workloads,
+)
+
+FULLATTN_MODEL_ID = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
+FULLATTN_SHORT_NAME = "FastWan2.2-TI2V-5B-FullAttn-Diffusers"
+FASTWAN_TI2V_MODEL_ID = "FastVideo/FastWan2.2-TI2V-5B-Diffusers"
+
+
+def _workloads_for(model_id: str) -> list[str]:
+    for model in get_registered_models_with_workloads():
+        if model["id"] == model_id:
+            return model["workload_types"]
+    raise AssertionError(f"{model_id} was not registered")
+
+
+def test_fastwan_fullattn_registry_is_t2v_only() -> None:
+    assert get_pipeline_config_cls_from_name(FULLATTN_MODEL_ID) is FastWan2_2_TI2V_5B_FullAttn_Config
+    assert get_pipeline_config_cls_from_name(FULLATTN_SHORT_NAME) is FastWan2_2_TI2V_5B_FullAttn_Config
+    assert get_pipeline_config_cls_from_name(FASTWAN_TI2V_MODEL_ID) is FastWan2_2_TI2V_5B_Config
+
+    assert _workloads_for(FULLATTN_MODEL_ID) == ["t2v"]
+    i2v_model_ids = {model["id"] for model in get_registered_models_with_workloads("i2v")}
+    assert FULLATTN_MODEL_ID not in i2v_model_ids
+    assert FASTWAN_TI2V_MODEL_ID in i2v_model_ids
+
+
+def test_fastwan_fullattn_config_is_dense_t2v_only(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_ATTENTION_BACKEND", raising=False)
+
+    config = FastWan2_2_TI2V_5B_FullAttn_Config()
+
+    assert config.ti2v_task is False
+    assert config.vae_config.load_encoder is False
+    assert config.vae_config.load_decoder is True
+    assert config.dit_config.expand_timesteps is True
+
+    supported_backends = config.dit_config._supported_attention_backends
+    assert AttentionBackendEnum.VIDEO_SPARSE_ATTN not in supported_backends
+    assert AttentionBackendEnum.FLASH_ATTN in supported_backends
+    assert AttentionBackendEnum.TORCH_SDPA in supported_backends
+
+    config.check_pipeline_config()
+
+
+def test_fastwan_fullattn_config_rejects_vsa_backend(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "VIDEO_SPARSE_ATTN")
+
+    with pytest.raises(ValueError, match="FullAttn.*VIDEO_SPARSE_ATTN"):
+        FastWan2_2_TI2V_5B_FullAttn_Config().check_pipeline_config()
+
+
+def test_wan_block_selection_rejects_vsa_for_fullattn(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "VIDEO_SPARSE_ATTN")
+
+    with pytest.raises(ValueError, match="VIDEO_SPARSE_ATTN.*FullAttn"):
+        _select_wan_transformer_block(FastWan2_2_TI2V_5B_FullAttn_Config().dit_config)
+
+
+def test_wan_block_selection_rejects_vsa_with_missing_backend_list(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "VIDEO_SPARSE_ATTN")
+    config = FastWan2_2_TI2V_5B_FullAttn_Config().dit_config
+    config._supported_attention_backends = None  # type: ignore[assignment]
+
+    with pytest.raises(ValueError, match="VIDEO_SPARSE_ATTN.*FullAttn"):
+        _select_wan_transformer_block(config)
+
+
+def test_wan_block_selection_preserves_existing_vsa_config(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "VIDEO_SPARSE_ATTN")
+    assert _select_wan_transformer_block(FastWan2_2_TI2V_5B_Config().dit_config) is WanTransformerBlock_VSA
+
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+    assert _select_wan_transformer_block(FastWan2_2_TI2V_5B_Config().dit_config) is WanTransformerBlock
+
+
+def test_wan_transformer_uses_config_supported_backends(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_ATTENTION_BACKEND", raising=False)
+    captured_backends = []
+
+    class CapturingWanTransformerBlock(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,
+                     **kwargs) -> None:
+            super().__init__()
+            captured_backends.append(supported_attention_backends)
+
+    monkeypatch.setattr(wanvideo_module, "WanTransformerBlock", CapturingWanTransformerBlock)
+    monkeypatch.setattr(wanvideo_module, "get_sp_world_size", lambda: 1)
+
+    config = FastWan2_2_TI2V_5B_FullAttn_Config().dit_config
+    config.arch_config.num_attention_heads = 1
+    config.arch_config.attention_head_dim = 4
+    config.arch_config.in_channels = 4
+    config.arch_config.out_channels = 4
+    config.arch_config.num_layers = 2
+    config.arch_config.ffn_dim = 8
+    config.arch_config.text_dim = 4
+    config.arch_config.freq_dim = 4
+    config.arch_config.patch_size = (1, 1, 1)
+    config.arch_config.__post_init__()
+
+    WanTransformer3DModel(config=config, hf_config={})
+
+    assert captured_backends == [
+        config._supported_attention_backends,
+        config._supported_attention_backends,
+    ]
+    assert AttentionBackendEnum.VIDEO_SPARSE_ATTN not in captured_backends[0]

fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py

@@ -50,9 +50,10 @@ def test_resolver_skips_cute_without_opt_in(monkeypatch) -> None:
     monkeypatch.setattr(builtins, "__import__", spying_import)
 
     mod = _reload_resolver_module()
-    resolved = mod._resolve_flash_attn_varlen_func()
+    resolved, version = mod._resolve_flash_attn_varlen_func()
     assert resolved is not None
     assert resolved.__name__ == "flash_attn_varlen_func"
+    assert version in ("2", "3")
     assert "fastvideo.attention.utils.flash_attn_cute" not in attempted
 
 
@@ -92,9 +93,10 @@ def test_resolver_returns_flash_attn_when_interface_unavailable(monkeypatch) ->
     monkeypatch.setattr(builtins, "__import__", patched_import)
 
     mod = _reload_resolver_module()
-    resolved = mod._resolve_flash_attn_varlen_func()
+    resolved, version = mod._resolve_flash_attn_varlen_func()
     assert resolved is not None
     assert resolved.__module__.startswith("flash_attn")
+    assert version == "2"
 
 
 def test_module_level_impl_is_resolver_output() -> None:

fastvideo/tests/encoders/test_qwen2_5_vl_vision_dtype.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import torch
+
+from fastvideo.models.encoders.qwen2_5_vl_custom import (
+    Qwen2_5_VisionTransformerPretrainedModel,
+    Qwen2_5_VLForConditionalGenerationSimple,
+)
+
+
+def _vision_config(torch_dtype=None):
+    return SimpleNamespace(
+        torch_dtype=torch_dtype,
+        spatial_merge_size=2,
+        patch_size=14,
+        fullatt_block_indexes=[],
+        window_size=112,
+        temporal_patch_size=2,
+        in_channels=3,
+        hidden_size=8,
+        num_heads=2,
+        depth=0,
+        _attn_implementation="sdpa",
+        out_hidden_size=8,
+    )
+
+
+def _full_config(torch_dtype):
+    return SimpleNamespace(
+        vision_config=_vision_config(),
+        torch_dtype=torch_dtype,
+        vocab_size=16,
+        hidden_size=8,
+        num_hidden_layers=0,
+        num_attention_heads=2,
+        pad_token_id=0,
+        _attn_implementation="sdpa",
+        rms_norm_eps=1e-6,
+        max_position_embeddings=32,
+        rope_theta=1_000_000.0,
+        rope_scaling=None,
+    )
+
+
+def test_vision_dtype_uses_parent_bfloat16_string_when_vision_dtype_missing():
+    model = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(),
+        parent_torch_dtype="bfloat16",
+    )
+
+    assert model.dtype == torch.bfloat16
+
+
+def test_vision_dtype_accepts_parent_torch_dtype_object():
+    model = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(),
+        parent_torch_dtype=torch.bfloat16,
+    )
+
+    assert model.dtype == torch.bfloat16
+
+
+def test_vision_dtype_strips_parent_dtype_string_whitespace():
+    model = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(),
+        parent_torch_dtype=" torch.bfloat16 ",
+    )
+
+    assert model.dtype == torch.bfloat16
+
+
+def test_vision_dtype_prefers_explicit_vision_dtype_over_parent_dtype():
+    model = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(torch_dtype="float16"),
+        parent_torch_dtype="bfloat16",
+    )
+
+    assert model.dtype == torch.float16
+
+
+def test_vision_dtype_falls_back_to_float32_for_missing_or_unknown_dtype():
+    missing = Qwen2_5_VisionTransformerPretrainedModel(_vision_config())
+    unknown = Qwen2_5_VisionTransformerPretrainedModel(
+        _vision_config(),
+        parent_torch_dtype="not-a-real-dtype",
+    )
+
+    assert missing.dtype == torch.float32
+    assert unknown.dtype == torch.float32
+
+
+def test_conditional_generation_passes_parent_dtype_to_visual_tower():
+    model = Qwen2_5_VLForConditionalGenerationSimple(_full_config("torch.bfloat16"))
+
+    assert model.visual.dtype == torch.bfloat16

scripts/inference/inference_wan_VSA_DMD_5B_720P.yaml

@@ -1,5 +1,6 @@
 # Run with:
 #   FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN fastvideo generate --config scripts/inference/inference_wan_VSA_DMD_5B_720P.yaml
+# FullAttn checkpoints use dense attention; do not run this config with VIDEO_SPARSE_ATTN.
 generator:
   model_path: FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers
   engine: