Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
757e024dc7 | ||
|
|
d3ff912ac0 |
@@ -47,11 +47,7 @@ COPY . /opt/FastVideo
|
||||
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
source /opt/venv/bin/activate \
|
||||
&& if [[ "${UV_TORCH_BACKEND}" == "cu130" ]]; then \
|
||||
uv pip install "/opt/FastVideo[dreamverse]" "nvidia-cutlass-dsl[cu13]"; \
|
||||
else \
|
||||
uv pip install "/opt/FastVideo[dreamverse]"; \
|
||||
fi
|
||||
&& uv pip install "/opt/FastVideo[dreamverse]"
|
||||
|
||||
# Standard docker build does not expose GPUs, while fastvideo-kernel/build.sh
|
||||
# detects the CUDA architecture with torch at build time. The FastVideo package
|
||||
|
||||
@@ -14,7 +14,7 @@ dependencies = [
|
||||
[project.optional-dependencies]
|
||||
server = [
|
||||
"cerebras-cloud-sdk",
|
||||
"flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@82d6441eec5d4dfec120153db2c0145ae855a083#subdirectory=flash_attn/cute",
|
||||
"flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute",
|
||||
"flashinfer-python",
|
||||
"openai>=1.40",
|
||||
]
|
||||
|
||||
@@ -39,16 +39,15 @@ export FASTVIDEO_GENERATION_SEGMENT_CAP="${FASTVIDEO_GENERATION_SEGMENT_CAP:-6}"
|
||||
export FASTVIDEO_PROMPT_AUTO_SLEEP_MS="${FASTVIDEO_PROMPT_AUTO_SLEEP_MS:-120}"
|
||||
export FASTVIDEO_PROMPT_AUTO_TIMEOUT_MS="${FASTVIDEO_PROMPT_AUTO_TIMEOUT_MS:-1800}"
|
||||
|
||||
if [[ "${ENABLE_TORCH_COMPILE}" == "1" ]]; then
|
||||
# Persist Inductor, AOTAutograd, and Triton artifacts across launches.
|
||||
export DREAMVERSE_TORCH_COMPILE_CACHE_ROOT="${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT:-${HOME}/.cache/dreamverse/torch_compile}"
|
||||
export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}/inductor}"
|
||||
export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}/triton}"
|
||||
export TORCHINDUCTOR_FX_GRAPH_CACHE="${TORCHINDUCTOR_FX_GRAPH_CACHE:-1}"
|
||||
export TORCHINDUCTOR_AUTOGRAD_CACHE="${TORCHINDUCTOR_AUTOGRAD_CACHE:-1}"
|
||||
mkdir -p "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}"
|
||||
echo "[launch-demo] torch.compile cache: ${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}"
|
||||
fi
|
||||
# Persistent torch.compile cache so warmup pays the full autotune/codegen cost only on the first launch.
|
||||
# Later launches on the same box reload the Inductor FX-graph + AOTAutograd + Triton artifacts from disk.
|
||||
export DREAMVERSE_TORCH_COMPILE_CACHE_ROOT="${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT:-${HOME}/.cache/dreamverse/torch_compile}"
|
||||
export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}/inductor}"
|
||||
export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}/triton}"
|
||||
export TORCHINDUCTOR_FX_GRAPH_CACHE="${TORCHINDUCTOR_FX_GRAPH_CACHE:-1}"
|
||||
export TORCHINDUCTOR_AUTOGRAD_CACHE="${TORCHINDUCTOR_AUTOGRAD_CACHE:-1}"
|
||||
mkdir -p "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}"
|
||||
echo "[launch-demo] torch.compile cache: ${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT} (first launch builds it, later launches reload it)"
|
||||
|
||||
cd "${DREAMVERSE_ROOT}"
|
||||
|
||||
|
||||
+9
-15
@@ -65,8 +65,8 @@ ARG FLASH_ATTN_WHEEL_RELEASE=https://github.com/mjun0812/flash-attention-prebuil
|
||||
ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22
|
||||
|
||||
# The prebuilt flash-attn wheel ships an FA4 `flash_attn.cute` built against the
|
||||
# cutlass-4.4 `cute.core.ThrMma` API, which crashes on the CuTe DSL 4.6 that
|
||||
# flashinfer/quack pull in. After the wheel install we overlay this 4.6.0.dev0-compatible
|
||||
# cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
|
||||
# flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
|
||||
# upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
|
||||
ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
|
||||
|
||||
@@ -169,24 +169,18 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
# flash_attn/__init__.py), so FA2/varlen/bert_padding stay from the install above;
|
||||
# rmtree clears the wheel's stale cute files first to avoid an install conflict.
|
||||
# Then verify both survive so a broken overlay fails the build instead of shipping
|
||||
# an FA2-less image. CUDA 13 selects the matching runtime extra on both amd64 and
|
||||
# arm64 (validated on GB200). CUDA 12.6 amd64 keeps the bare/default runtime;
|
||||
# CUDA 12.6 arm64 remains unvalidated and skips the overlay. FA4 stays opt-in, so
|
||||
# installing it does not auto-enable the backend on GB10 (sm_121).
|
||||
# an FA2-less image. x86 only: the FA4 stack (quack-kernels etc.) is unvalidated on
|
||||
# arm64 / GB10 (sm_121), so there we skip the overlay; FA4 is opt-in
|
||||
# (FASTVIDEO_FA4=1) and errors if set without the overlay, so leave it unset on
|
||||
# arm64 and the image runs FA3/FA2 as usual.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
source $HOME/.local/bin/env && \
|
||||
source /opt/venv/bin/activate && \
|
||||
if [ "${UV_TORCH_BACKEND}" = "cu130" ]; then \
|
||||
FA4_PACKAGE="flash-attn-4[cu13]"; \
|
||||
elif [ "${TARGETARCH}" = "arm64" ]; then \
|
||||
echo "Skipping FA4 cute overlay on CUDA 12.6 arm64 (do not set FASTVIDEO_FA4)"; \
|
||||
FA4_PACKAGE=""; \
|
||||
if [ "${TARGETARCH}" = "arm64" ]; then \
|
||||
echo "Skipping FA4 cute overlay on arm64 (FA4 stack unvalidated there; do not set FASTVIDEO_FA4)"; \
|
||||
else \
|
||||
FA4_PACKAGE="flash-attn-4"; \
|
||||
fi && \
|
||||
if [ -n "${FA4_PACKAGE}" ]; then \
|
||||
python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
|
||||
uv pip install "${FA4_PACKAGE} @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
|
||||
uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
|
||||
python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'), 'FA2 was clobbered by the cute overlay'; import flash_attn.cute; print('FA2 + FA4 cute OK')"; \
|
||||
fi
|
||||
|
||||
|
||||
@@ -78,8 +78,8 @@ python setup.py install
|
||||
|
||||
FastVideo never auto-selects FlashAttention-4 (`flash_attn.cute`) just because it
|
||||
is installed: its CuTeDSL kernels JIT-compile per shape family and can fail at
|
||||
runtime on some GPU/shape combinations. To use FA4 on CUDA 13, install the pinned
|
||||
`flash-attn-4[cu13]` build (see the `flash-attn-4` source in `pyproject.toml`) and set:
|
||||
runtime on some GPU/shape combinations. To use FA4, install the pinned
|
||||
`flash-attn-4` build (see the `flash-attn-4` source in `pyproject.toml`) and set:
|
||||
|
||||
```bash
|
||||
export FASTVIDEO_FA4=1
|
||||
@@ -110,18 +110,14 @@ See the [Attn-QAT paper](https://arxiv.org/abs/2603.00040) and [flash-attention-
|
||||
Install the FP4 flash attention kernel (without upgrading your existing torch):
|
||||
|
||||
```bash
|
||||
# Commit 940bf7e5 carries the CuTe DSL 4.5 fix (cute.core.ThrMma -> cute.ThrMma).
|
||||
pip install --no-deps "git+https://github.com/hao-ai-lab/flash-attention-fp4.git@940bf7e511375ec160bc2d7188bef35915ded1e3#subdirectory=flash_attn/cute"
|
||||
pip install "nvidia-cutlass-dsl[cu13]==4.5.2" "quack-kernels==0.5.0" apache-tvm-ffi flashinfer-python torch-c-dlpack-ext
|
||||
# branch fix/cutlass-dsl-4.5 carries the cutlass-dsl 4.5 fix (cute.core.ThrMma
|
||||
# -> cute.ThrMma); switch back to @fp4 once hao-ai-lab/flash-attention-fp4#2 merges.
|
||||
pip install --no-deps "git+ssh://git@github.com/hao-ai-lab/flash-attention-fp4.git@fix/cutlass-dsl-4.5#subdirectory=flash_attn/cute"
|
||||
pip install "nvidia-cutlass-dsl>=4.5.2" apache-tvm-ffi flashinfer-python
|
||||
```
|
||||
|
||||
The `--no-deps` flag prevents upgrading torch/torchvision. Use the supported
|
||||
PyTorch 2.12.0 and CUDA 13 environment for this kernel. Keep this private FP4
|
||||
overlay in a separate inference environment from FastVideo's upstream dense FA4
|
||||
dependency: the two distributions provide the same `flash_attn.cute` package but
|
||||
require different CuTe DSL versions. Keep the CUTLASS and Quack pins together;
|
||||
Quack 0.5.1 and newer use CuTe DSL 4.6, whose API is incompatible with this FP4
|
||||
kernel revision.
|
||||
PyTorch 2.12.0 and CUDA 13 environment for this kernel.
|
||||
|
||||
#### Usage
|
||||
|
||||
|
||||
@@ -75,13 +75,13 @@ fully usable without it).
|
||||
The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`,
|
||||
`flash_attn.cute.interface._flash_attn_fwd`) are provided upstream by
|
||||
[Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to
|
||||
commit `82d6441eec5d4dfec120153db2c0145ae855a083`, the revision FastVideo pins as
|
||||
commit `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as
|
||||
the `flash-attn-4` source in the repo-root `pyproject.toml`; other revisions may
|
||||
have an incompatible `_flash_attn_fwd` signature.
|
||||
|
||||
```bash
|
||||
pip install torchvision
|
||||
pip install "flash-attn-4[cu13] @ git+https://github.com/Dao-AILab/flash-attention.git@82d6441eec5d4dfec120153db2c0145ae855a083#subdirectory=flash_attn/cute"
|
||||
pip install "nvidia-cutlass-dsl>=4.5.0" torchvision
|
||||
pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"
|
||||
```
|
||||
|
||||
The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against
|
||||
|
||||
@@ -42,8 +42,7 @@ except ImportError:
|
||||
_flash_attn_2_func = None
|
||||
_flash_attn_2_varlen_func = None
|
||||
|
||||
# Dynamo unwraps functools caches, so keep this cache inside the opaque helper.
|
||||
_SM90_OR_NEWER_BY_DEVICE: dict[int, bool] = {}
|
||||
_IS_SM90_OR_NEWER = current_platform.has_device_capability(90)
|
||||
|
||||
|
||||
def _check_dropout(dropout_p: float) -> None:
|
||||
@@ -51,17 +50,13 @@ def _check_dropout(dropout_p: float) -> None:
|
||||
raise NotImplementedError(f"flash_attn.cute does not support dropout (got dropout_p={dropout_p})")
|
||||
|
||||
|
||||
@torch.compiler.assume_constant_result
|
||||
def _sm90_or_newer(device_id: int) -> bool:
|
||||
if device_id not in _SM90_OR_NEWER_BY_DEVICE:
|
||||
_SM90_OR_NEWER_BY_DEVICE[device_id] = (current_platform.has_device_capability(90, device_id))
|
||||
return _SM90_OR_NEWER_BY_DEVICE[device_id]
|
||||
def _sm90_or_newer() -> bool:
|
||||
# Compiled callers only read a boolean instead of tracing through a memoized platform query.
|
||||
return _IS_SM90_OR_NEWER
|
||||
|
||||
|
||||
def _use_fa2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool:
|
||||
device_id = q.device.index
|
||||
assert device_id is not None
|
||||
if _sm90_or_newer(device_id):
|
||||
if _sm90_or_newer():
|
||||
return False
|
||||
# Pre-sm90 FA4 cute limitations, both served by FA2 (deterministic
|
||||
# capability gate, not a runtime fallback):
|
||||
|
||||
@@ -5,7 +5,6 @@ LTX-2 transformer implementation
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
from enum import Enum
|
||||
import functools
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -768,9 +767,14 @@ def _apply_ltx_split_rotary_emb(
|
||||
return output
|
||||
|
||||
|
||||
def _generate_ltx_freq_grid_float64(
|
||||
def generate_ltx_freq_grid_float64(
|
||||
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
|
||||
) -> torch.Tensor:
|
||||
"""Generate LTX-2 rotary frequencies with float64 precision.
|
||||
|
||||
Implemented in pure torch and free of Python memoization so it is safe
|
||||
for callers that may be captured and reused by torch.compile.
|
||||
"""
|
||||
theta = positional_embedding_theta
|
||||
start = 1
|
||||
end = theta
|
||||
@@ -787,24 +791,10 @@ def _generate_ltx_freq_grid_float64(
|
||||
return (pow_indices * math.pi / 2).to(dtype=torch.float32)
|
||||
|
||||
|
||||
_cached_generate_ltx_freq_grid_float64 = functools.lru_cache(maxsize=5)(
|
||||
_generate_ltx_freq_grid_float64)
|
||||
|
||||
|
||||
def generate_ltx_freq_grid_float64(
|
||||
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
|
||||
) -> torch.Tensor:
|
||||
"""Generate an eager-cached, compile-safe float64 LTX-2 rotary grid."""
|
||||
if torch.compiler.is_compiling():
|
||||
return _generate_ltx_freq_grid_float64(
|
||||
positional_embedding_theta, positional_embedding_max_pos_count, inner_dim)
|
||||
return _cached_generate_ltx_freq_grid_float64(
|
||||
positional_embedding_theta, positional_embedding_max_pos_count, inner_dim)
|
||||
|
||||
|
||||
def _generate_ltx_freq_grid_pytorch(
|
||||
def generate_ltx_freq_grid_pytorch(
|
||||
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
|
||||
) -> torch.Tensor:
|
||||
"""Generate LTX-2 rotary frequencies in torch for speed."""
|
||||
theta = positional_embedding_theta
|
||||
start = 1
|
||||
end = theta
|
||||
@@ -821,21 +811,6 @@ def _generate_ltx_freq_grid_pytorch(
|
||||
return indices * math.pi / 2
|
||||
|
||||
|
||||
_cached_generate_ltx_freq_grid_pytorch = functools.lru_cache(maxsize=5)(
|
||||
_generate_ltx_freq_grid_pytorch)
|
||||
|
||||
|
||||
def generate_ltx_freq_grid_pytorch(
|
||||
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
|
||||
) -> torch.Tensor:
|
||||
"""Generate an eager-cached, compile-safe float32 LTX-2 rotary grid."""
|
||||
if torch.compiler.is_compiling():
|
||||
return _generate_ltx_freq_grid_pytorch(
|
||||
positional_embedding_theta, positional_embedding_max_pos_count, inner_dim)
|
||||
return _cached_generate_ltx_freq_grid_pytorch(
|
||||
positional_embedding_theta, positional_embedding_max_pos_count, inner_dim)
|
||||
|
||||
|
||||
def _ltx_get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor:
|
||||
n_pos_dims = indices_grid.shape[1]
|
||||
if n_pos_dims != len(max_pos):
|
||||
|
||||
@@ -5,7 +5,6 @@ from __future__ import annotations
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -56,44 +55,6 @@ def _extract_out(x):
|
||||
return x
|
||||
|
||||
|
||||
def test_capability_gate_is_lazy_and_device_keyed(monkeypatch) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA is required to import flash_attn.cute")
|
||||
try:
|
||||
mod = _load_local_flash_attn_cute_module()
|
||||
except ImportError as exc:
|
||||
pytest.skip(f"flash_attn.cute is not available: {exc}")
|
||||
calls: list[tuple[int, int]] = []
|
||||
|
||||
def has_device_capability(capability: int, device_id: int) -> bool:
|
||||
calls.append((capability, device_id))
|
||||
return device_id == 1
|
||||
|
||||
def tensor_stub(device_id: int, heads: int):
|
||||
return SimpleNamespace(
|
||||
device=torch.device("cuda", device_id),
|
||||
shape=(1, 1, heads, 64),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
mod._SM90_OR_NEWER_BY_DEVICE.clear()
|
||||
monkeypatch.setattr(
|
||||
mod.current_platform,
|
||||
"has_device_capability",
|
||||
has_device_capability,
|
||||
)
|
||||
try:
|
||||
q0, k0, v0 = tensor_stub(0, 4), tensor_stub(0, 2), tensor_stub(0, 2)
|
||||
assert mod._use_fa2(q0, k0, v0)
|
||||
assert mod._use_fa2(q0, k0, v0)
|
||||
|
||||
q1, k1, v1 = tensor_stub(1, 4), tensor_stub(1, 2), tensor_stub(1, 2)
|
||||
assert not mod._use_fa2(q1, k1, v1)
|
||||
assert calls == [(90, 0), (90, 1)]
|
||||
finally:
|
||||
mod._SM90_OR_NEWER_BY_DEVICE.clear()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def flash_attn_impls():
|
||||
if not torch.cuda.is_available():
|
||||
|
||||
@@ -1,50 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Keep every FA4 installation surface on the same tested revision."""
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
|
||||
|
||||
def _match(path: str, pattern: str) -> str:
|
||||
source = (REPO_ROOT / path).read_text(encoding="utf-8")
|
||||
match = re.search(pattern, source)
|
||||
assert match is not None, f"FA4 pin not found in {path}"
|
||||
return match.group(1)
|
||||
|
||||
|
||||
def test_fa4_dependency_pins_match() -> None:
|
||||
pins = {
|
||||
"pyproject.toml": _match(
|
||||
"pyproject.toml",
|
||||
r'flash-attn-4 = .*rev = "([0-9a-f]{40})"'),
|
||||
"docker/Dockerfile": _match(
|
||||
"docker/Dockerfile", r"ARG FA4_CUTE_REF=([0-9a-f]{40})"),
|
||||
"apps/dreamverse/pyproject.toml": _match(
|
||||
"apps/dreamverse/pyproject.toml",
|
||||
r"flash-attention\.git@([0-9a-f]{40})"),
|
||||
"fastvideo-kernel/README.md": _match(
|
||||
"fastvideo-kernel/README.md",
|
||||
r"flash-attention\.git@([0-9a-f]{40})"),
|
||||
}
|
||||
|
||||
assert len(set(pins.values())) == 1, f"FA4 dependency pins disagree: {pins}"
|
||||
|
||||
|
||||
def test_fa4_blackwell_installs_select_cuda13_runtime() -> None:
|
||||
root_pyproject = (REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8")
|
||||
dreamverse_pyproject = (REPO_ROOT / "apps/dreamverse/pyproject.toml").read_text(encoding="utf-8")
|
||||
dockerfile = (REPO_ROOT / "docker/Dockerfile").read_text(encoding="utf-8")
|
||||
dreamverse_dockerfile = (REPO_ROOT / "apps/dreamverse/docker/Dockerfile").read_text(encoding="utf-8")
|
||||
kernel_readme = (REPO_ROOT / "fastvideo-kernel/README.md").read_text(encoding="utf-8")
|
||||
|
||||
assert '"flash-attn-4",' in root_pyproject
|
||||
assert '"flash-attn-4 @ git+' in dreamverse_pyproject
|
||||
assert 'if [ "${UV_TORCH_BACKEND}" = "cu130" ]' in dockerfile
|
||||
assert 'FA4_PACKAGE="flash-attn-4[cu13]"' in dockerfile
|
||||
assert 'FA4_PACKAGE="flash-attn-4"' in dockerfile
|
||||
assert 'elif [ "${TARGETARCH}" = "arm64" ]' in dockerfile
|
||||
assert '"${UV_TORCH_BACKEND}" == "cu130"' in dreamverse_dockerfile
|
||||
assert '"nvidia-cutlass-dsl[cu13]"' in dreamverse_dockerfile
|
||||
assert "flash-attn-4[cu13]" in kernel_readme
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Contract tests for LTX2 torch.compile integration.
|
||||
"""Contract test: LTX2's per-block torch.compile wiring stays intact.
|
||||
|
||||
``ComposedPipelineBase._compile_with_conditions`` compiles the submodules
|
||||
matched by ``model._compile_conditions`` and falls back to compiling the
|
||||
@@ -11,14 +11,13 @@ wiring and the block matcher so a refactor can't reintroduce the fallback
|
||||
unnoticed. CPU-only; no model instantiation.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig, is_ltx2_blocks
|
||||
from fastvideo.models.dits import ltx2
|
||||
|
||||
|
||||
def test_ltx2_model_exposes_compile_conditions():
|
||||
conditions = ltx2.LTX2Transformer3DModel._compile_conditions
|
||||
from fastvideo.models.dits.ltx2 import LTX2Transformer3DModel
|
||||
|
||||
conditions = LTX2Transformer3DModel._compile_conditions
|
||||
assert conditions, (
|
||||
"LTX2Transformer3DModel._compile_conditions is empty: "
|
||||
"torch.compile falls back to a whole-model fullgraph")
|
||||
@@ -34,29 +33,3 @@ def test_is_ltx2_blocks_matches_only_top_level_blocks():
|
||||
assert not is_ltx2_blocks("transformer_blocks", None)
|
||||
assert not is_ltx2_blocks("transformer_blocks.0.attn1", None)
|
||||
assert not is_ltx2_blocks("audio_transformer_blocks.0", None)
|
||||
|
||||
|
||||
def test_freq_grid_cache_is_bypassed_while_compiling():
|
||||
args = (10000.0, 3, 96)
|
||||
for generator, cached_generator in (
|
||||
(
|
||||
ltx2.generate_ltx_freq_grid_float64,
|
||||
ltx2._cached_generate_ltx_freq_grid_float64,
|
||||
),
|
||||
(
|
||||
ltx2.generate_ltx_freq_grid_pytorch,
|
||||
ltx2._cached_generate_ltx_freq_grid_pytorch,
|
||||
),
|
||||
):
|
||||
cached_generator.cache_clear()
|
||||
eager = generator(*args)
|
||||
assert generator(*args) is eager
|
||||
|
||||
cached_generator.cache_clear()
|
||||
compiled_generator = torch.compile(
|
||||
lambda zero: generator(*args) + zero,
|
||||
backend="eager",
|
||||
fullgraph=True,
|
||||
)
|
||||
torch.testing.assert_close(compiled_generator(torch.tensor(0.0)), eager)
|
||||
assert cached_generator.cache_info().currsize == 0
|
||||
|
||||
+1
-1
@@ -114,7 +114,7 @@ fastvideo-kernel = [
|
||||
# uv reads UV_TORCH_BACKEND / --torch-backend (not a pyproject setting). The
|
||||
# supported backends (cu126, cu130) both provide torch 2.12.0.
|
||||
imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
|
||||
# FA4 cute, pinned to the CuTe DSL 4.6.0.dev0-compatible revision. torch.compile support
|
||||
# FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support
|
||||
# comes from FastVideo's own custom_op wrappers.
|
||||
flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" }
|
||||
|
||||
|
||||
Reference in New Issue
Block a user