Compare commits
14
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
42a4bd7a18 | ||
|
|
d2dd78bc42 | ||
|
|
e095388095 | ||
|
|
138f057ed7 | ||
|
|
ff0444a959 | ||
|
|
20be8c3dfa | ||
|
|
dc178be814 | ||
|
|
d72db9193b | ||
|
|
a9095abdc1 | ||
|
|
3c0792de21 | ||
|
|
cae8fa18dc | ||
|
|
821e5a0832 | ||
|
|
c1abc42782 | ||
|
|
ef15ea2391 |
@@ -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
|
||||
|
||||
@@ -15,8 +15,14 @@ class SDPABackend(AttentionBackend):
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [32, 64, 96, 128, 160, 192, 224, 256]
|
||||
def get_supported_head_sizes() -> list[int] | None:
|
||||
# torch.nn.functional.scaled_dot_product_attention is head-size
|
||||
# agnostic: the math backend handles any size and the fused kernels
|
||||
# fall back internally when they cannot. The previous list
|
||||
# ([32, 64, ..., 256]) was copied from FlashAttentionBackend in the
|
||||
# v1 refactor (#270) and wrongly excluded sizes such as 80 (e.g.
|
||||
# CLIP vision encoders). None means "no restriction".
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -72,8 +72,13 @@ def _rms_norm_dispatch(
|
||||
"""Use QuACK RMSNorm only for LTX-2 refine stage, else Torch RMSNorm."""
|
||||
if _is_ltx2_refine_stage():
|
||||
return _quack_rmsnorm(x, weight=weight, eps=eps)
|
||||
# Official LTX-2 runs the DiT natively in bf16 with no autocast, so its
|
||||
# rms_norm outputs stay bf16 on every torch version. torch 2.12 added
|
||||
# rms_norm to autocast's float32 cast policy, which under our
|
||||
# autocast(bfloat16) forward returns float32 instead. Restore the input
|
||||
# dtype to match the official numerics (and torch <= 2.11 behavior).
|
||||
return torch.nn.functional.rms_norm(
|
||||
x, (x.shape[-1], ), weight=weight, eps=eps)
|
||||
x, (x.shape[-1], ), weight=weight, eps=eps).to(x.dtype)
|
||||
|
||||
|
||||
class StageAwareRMSNorm(nn.RMSNorm):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
+11
-1
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
@@ -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:
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""CPU tests: SDPA must not restrict attention head sizes.
|
||||
|
||||
``torch.nn.functional.scaled_dot_product_attention`` is head-size
|
||||
agnostic, but ``SDPABackend.get_supported_head_sizes()`` used to declare
|
||||
a list copied from FlashAttentionBackend that omitted e.g. 80 (the head
|
||||
size of CLIP-L vision encoders such as matrixgame2's). That made
|
||||
capability validation flag SDPA as incompatible with head-80 layers even
|
||||
though it handles them fine.
|
||||
|
||||
CPU-only; no GPU or flash-attn dependency.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.attention.backends.sdpa import SDPABackend, SDPAImpl, SDPAMetadata
|
||||
from fastvideo.platforms.cuda import CudaPlatformBase
|
||||
from fastvideo.platforms.interface import AttentionBackendEnum
|
||||
|
||||
|
||||
def test_sdpa_head_sizes_unrestricted() -> None:
|
||||
"""SDPA declares no head-size restriction (None), so 80 is allowed."""
|
||||
assert SDPABackend.get_supported_head_sizes() is None
|
||||
# Once the capability-validation API (#1494) lands, head-80 must pass.
|
||||
validate = getattr(SDPABackend, "validate_compatibility", None)
|
||||
if validate is not None:
|
||||
assert validate(head_size=80, dtype=torch.bfloat16) is None
|
||||
|
||||
|
||||
def test_explicit_torch_sdpa_pin_head_80_resolves() -> None:
|
||||
"""An explicit TORCH_SDPA pin on a head-80 layer resolves to SDPABackend."""
|
||||
cls_str = CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.TORCH_SDPA, head_size=80, dtype=torch.bfloat16)
|
||||
assert cls_str == "fastvideo.attention.backends.sdpa.SDPABackend"
|
||||
|
||||
|
||||
def test_sdpa_impl_computes_head_80_attention() -> None:
|
||||
"""SDPAImpl produces correct output for head_size=80 (not in the old list)."""
|
||||
bs, seq, heads, head_dim = 2, 8, 2, 80
|
||||
softmax_scale = head_dim**-0.5
|
||||
impl = SDPAImpl(num_heads=heads, head_size=head_dim, causal=False, softmax_scale=softmax_scale)
|
||||
torch.manual_seed(0)
|
||||
q, k, v = (torch.randn(bs, seq, heads, head_dim) for _ in range(3))
|
||||
|
||||
out = impl.forward(q, k, v, SDPAMetadata(current_timestep=0))
|
||||
|
||||
assert out.shape == (bs, seq, heads, head_dim)
|
||||
qt, kt, vt = (t.transpose(1, 2) for t in (q, k, v))
|
||||
ref = torch.softmax(qt @ kt.transpose(-2, -1) * softmax_scale, dim=-1) @ vt
|
||||
torch.testing.assert_close(out, ref.transpose(1, 2), rtol=1e-4, atol=1e-4)
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user