[bugfix]: install the FA4 CUDA 13 runtime

Select flash-attn-4's cu13 extra on Blackwell install surfaces so revision 82d6441 receives its exact CuTe DSL 4.6.0.dev0 runtime. Enable the opt-in overlay for arm64 CUDA 13 images, including GB200, while CUDA 12.6 continues to use FA3/FA2.

.buildkite/pipeline.yml

@@ -404,7 +404,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile"
               config:
-                command: "timeout 25m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                 label: ":test_tube: Training Tests"
                 env:
                   - TEST_TYPE=training
@@ -415,7 +415,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile"
               config:
-                command: "timeout 25m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                 label: ":test_tube: Distillation DMD Tests"
                 env:
                   - TEST_TYPE=distillation_dmd

.github/workflows/infra-build-image.yml

@@ -13,23 +13,16 @@ on:
         required: false
         default: false
         type: boolean
-  # Auto-rebuild the CUDA images when a repository-controlled image input
-  # changes on main. This includes the trusted SM89 kernel artifact's source,
-  # metadata/key helper, ABI dependency metadata, and build orchestration.
-  # Dreamverse (apps/dreamverse/docker/Dockerfile) and the ROCm Dockerfile stay
-  # manual-dispatch only.
+  # Auto-rebuild the CUDA images when their Dockerfile changes on main. The CUDA
+  # matrix is the only lane that builds from docker/Dockerfile, so a path-scoped
+  # push trigger is a sufficient change detector on its own -- no separate
+  # detect-changes/paths-filter job is needed now that there is a single
+  # in-scope Dockerfile. Dreamverse (apps/dreamverse/docker/Dockerfile) and the
+  # rocm Dockerfile stay manual-dispatch only.
   push:
     branches: [main]
     paths:
-      - '.dockerignore'
-      - '.github/workflows/_template-build-image.yml'
-      - '.github/workflows/infra-build-image.yml'
-      - '.gitmodules'
       - 'docker/Dockerfile'
-      - 'docker/uv-excludes'
-      - 'fastvideo-kernel/**'
-      - 'fastvideo/tests/modal/kernel_build_cache.py'
-      - 'pyproject.toml'
 
 
 permissions:
@@ -57,7 +50,7 @@ jobs:
   # 2.8.3 comes from the architecture-specific prebuilt releases.
   build-cuda-images:
     # Runs on a manual dispatch when build_cuda_matrix is set, or automatically
-    # on an in-scope main push (inputs are null on push). The
+    # on a push that changed docker/Dockerfile (inputs are null on push). The
     # repository guard keeps fork syncs from auto-building; manual dispatch
     # still works in forks.
     if: ${{ (github.event_name == 'push' && github.repository == 'hao-ai-lab/FastVideo') || github.event.inputs.build_cuda_matrix == 'true' }}

.gitignore

@@ -75,7 +75,6 @@ docs/distillation/examples/
 *.pkl
 
 # Reference videos (negations must come after the catch-all on line below)
-!fastvideo/tests/nightly/reference_video_*.mp4
 
 # Static images
 !docs/assets/images/**/*.png

apps/dreamverse/docker/Dockerfile

@@ -47,7 +47,11 @@ COPY . /opt/FastVideo
 
 RUN --mount=type=cache,target=/opt/uv/cache \
     source /opt/venv/bin/activate \
-    && uv pip install "/opt/FastVideo[dreamverse]"
+    && if [[ "${UV_TORCH_BACKEND}" == "cu130" ]]; then \
+        uv pip install "/opt/FastVideo[dreamverse]" "nvidia-cutlass-dsl[cu13]"; \
+    else \
+        uv pip install "/opt/FastVideo[dreamverse]"; \
+    fi
 
 # Standard docker build does not expose GPUs, while fastvideo-kernel/build.sh
 # detects the CUDA architecture with torch at build time. The FastVideo package

apps/dreamverse/pyproject.toml

@@ -14,7 +14,7 @@ dependencies = [
 [project.optional-dependencies]
 server = [
     "cerebras-cloud-sdk",
-    "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute",
+    "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@82d6441eec5d4dfec120153db2c0145ae855a083#subdirectory=flash_attn/cute",
     "flashinfer-python",
     "openai>=1.40",
 ]

docker/Dockerfile

@@ -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 cutlass-dsl 4.5 that
-# flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
+# 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
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
 ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
 
@@ -169,57 +169,38 @@ 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. 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.
+# 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).
 RUN --mount=type=cache,target=/opt/uv/cache \
     source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    if [ "${TARGETARCH}" = "arm64" ]; then \
-        echo "Skipping FA4 cute overlay on arm64 (FA4 stack unvalidated there; do not set FASTVIDEO_FA4)"; \
+    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=""; \
     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 "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
+        uv pip install "${FA4_PACKAGE} @ 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
 
 COPY . .
 
-# Build immutable FastVideo kernel wheels for the published image. The requested
-# architecture remains installed for normal image users; amd64 images also carry
-# an SM89 artifact so the predominant L40S Modal lanes can reuse it exactly.
-ARG FASTVIDEO_KERNEL_PREBUILT_DIR=/opt/fastvideo-kernel-prebuilt
+# Install FastVideo Unified Kernel exactly once. build.sh initializes only its
+# CUTLASS/ThunderKittens submodules, then compiles for TORCH_CUDA_ARCH_LIST
+# (default Hopper sm_90a) without requiring a live GPU.
 RUN --mount=type=cache,target=/opt/uv/cache \
     source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    default_arch="${TORCH_CUDA_ARCH_LIST}" && \
-    default_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/${default_arch}" && \
-    export TORCH_CUDA_ARCH_LIST="${default_arch}" && \
     cd fastvideo-kernel && \
-    CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
-    ./build.sh --wheel-dir "${default_wheel_dir}" && \
-    cd /FastVideo && \
-    CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
-        --wheel-dir "${default_wheel_dir}" \
-        --output "${default_wheel_dir}/metadata.json" && \
-    if [[ "${TARGETARCH:-amd64}" == "amd64" && "${default_arch}" != "8.9" ]]; then \
-        export TORCH_CUDA_ARCH_LIST=8.9 && \
-        l40s_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/8.9" && \
-        cd /FastVideo/fastvideo-kernel && \
-        CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
-        ./build.sh --wheel-dir "${l40s_wheel_dir}" && \
-        cd /FastVideo && \
-        CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
-            --wheel-dir "${l40s_wheel_dir}" \
-            --output "${l40s_wheel_dir}/metadata.json"; \
-    fi && \
-    export TORCH_CUDA_ARCH_LIST="${default_arch}" && \
-    default_wheel="$(find "${default_wheel_dir}" -maxdepth 1 -type f \
-        \( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) \
-        | sort | tail -n 1)" && \
-    uv pip install "${default_wheel}" \
-        --reinstall-package fastvideo-kernel --no-deps
+    CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
+    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} ./build.sh
 
 # Install FastVideo itself (editable) now that the source is present, and set up
 # shell configuration. Dependencies and the local kernel are already installed,

docs/contributing/ci_architecture.md

@@ -269,26 +269,16 @@ The docs job:
 ### Docker Images
 
 `.github/workflows/infra-build-image.yml` supports manual `workflow_dispatch`
-runs and automatically rebuilds the CUDA matrix when a repository-controlled
-image input changes on `main` in the canonical repository. Those inputs include
-the CUDA Dockerfile and reusable workflow, dependency metadata, Docker context
-policy, `fastvideo-kernel/**`, and the kernel artifact metadata/key helper.
-Manual runs let maintainers choose which image families to build. The
-`fastvideo-dev` matrix builds Python 3.12 images for CUDA 12.6 and CUDA 13 on
-native `amd64` and `arm64` runners, then publishes one multi-platform manifest
-per CUDA version. CUDA 12.6 owns the `py3.12-latest` and global `latest` tags, as
-well as the explicit `py3.12-cuda12.6.3-latest` alias. CUDA 13 is published under
-the explicit `py3.12-cuda13.0.0-latest` tag. This publication policy does not
-change the unparameterized `docker/Dockerfile` build defaults, which remain CUDA
-13 and `cu130`.
-
-Published amd64 development images keep their configured Hopper kernel wheel
-installed and also carry an immutable SM89 wheel under
-`/opt/fastvideo-kernel-prebuilt`. Modal PR and SSIM jobs select the exact
-source, ABI, and GPU-architecture match from that directory, so L40S jobs reuse
-the trusted image artifact while kernel-changing PRs still build locally. Once
-a kernel or artifact-key change reaches `main`, the image workflow republishes
-the matching trusted artifact before later jobs consume the updated image tag.
+runs and automatically rebuilds the CUDA matrix when `docker/Dockerfile`
+changes on `main` in the canonical repository. Manual runs let maintainers
+choose which image families to build. The `fastvideo-dev` matrix builds Python
+3.12 images for CUDA 12.6 and CUDA 13 on native `amd64` and `arm64` runners,
+then publishes one multi-platform manifest per CUDA version. CUDA 12.6 owns the
+`py3.12-latest` and global `latest` tags, as well as the explicit
+`py3.12-cuda12.6.3-latest` alias. CUDA 13 is published under the explicit
+`py3.12-cuda13.0.0-latest` tag. This publication policy does not change the
+unparameterized `docker/Dockerfile` build defaults, which remain CUDA 13 and
+`cu130`.
 
 The optional Dreamverse matrix builds backend and UI images for CUDA 12.6 and
 CUDA 13 on `amd64`. Dreamverse remains `amd64`-only because its FA4 dependency

docs/getting_started/installation/gpu.md

@@ -134,4 +134,4 @@ If you're planning to contribute to FastVideo please see the following page:
 
 If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
 
-You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ) for additional support.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.

docs/getting_started/installation/mps.md

@@ -100,4 +100,4 @@ If you're planning to contribute to FastVideo please see the following page:
 
 If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
 
-You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ) for additional support.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.

docs/getting_started/installation/spark.md

@@ -134,7 +134,7 @@ uv pip install "https://github.com/mjun0812/flash-attention-prebuild-wheels/rele
 
 If you hit other issues, please open an issue on our
 [GitHub repository](https://github.com/hao-ai-lab/FastVideo). You can also join
-our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ)
+our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ)
 for additional support.
 
 ## Development Environment Setup

docs/inference/add_pipeline.md

@@ -11,7 +11,7 @@ This guide explains how to implement a custom diffusion pipeline in FastVideo, l
 4. **Register Your Pipeline** - Make it discoverable by the framework
 5. **Configure Your Pipeline** - (Coming soon)
 
-Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ).
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
 
 ## Step 1: Pipeline Modules
 

docs/inference/inference_quick_start.md

@@ -134,4 +134,5 @@ If the generated video doesn't match your prompt:
 - Learn about [Advanced Inference Configurations](configuration.md)
 - Learn about using [Optimizations](optimizations.md)
 - See [Examples](examples/examples_inference_index.md) for more usage scenarios
-- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ).
+- Join our [Community Discord](https://discord.gg/JA7cksDz86).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).

docs/inference/optimizations.md

@@ -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, install the pinned
-`flash-attn-4` build (see the `flash-attn-4` source in `pyproject.toml`) and set:
+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:
 
 ```bash
 export FASTVIDEO_FA4=1
@@ -103,44 +103,25 @@ See the [Attn-QAT paper](https://arxiv.org/abs/2603.00040) and [flash-attention-
 
 - **GPU**: NVIDIA Blackwell (sm100a or sm103a) — B200, B300, GB200, GB300
 - **CUDA**: 13.0+
-- **Python**: 3.10 – 3.12
+- **Python**: 3.10 or 3.11
 
 #### Installation
 
 Install the FP4 flash attention kernel (without upgrading your existing torch):
 
 ```bash
-# 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+https://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
+# 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
 ```
 
 The `--no-deps` flag prevents upgrading torch/torchvision. Use the supported
-PyTorch 2.12.0 and CUDA 13 environment for this kernel.
-
-The `fix/cutlass-dsl-4.5` branch above is the recommended install because it
-works with the pip-available `nvidia-cutlass-dsl>=4.5.2` (it carries the
-`cute.ThrMma` rename fix). The hardware-validated receipt set — the runs
-behind the published GB200/GB300 numbers — used the `fp4` branch with the
-pinned `4.4.2` toolchain listed in the compatibility table below; both
-configurations run the same kernels.
-
-Branch-to-`nvidia-cutlass-dsl` compatibility (the fork tracks the CuTe DSL API
-surface closely):
-
-| fork branch | cutlass-dsl | notes |
-|---|---|---|
-| `fp4` | `==4.4.2` (+ `nvidia-cutlass-dsl-libs-base==4.4.2`) | validated set on GB200: `quack-kernels==0.4.1`, `flashinfer-python==0.6.8`, `CUTE_DSL_ENABLE_TVM_FFI=1`, `FASTVIDEO_FA4=1` |
-| `fix/cutlass-dsl-4.5` | `>=4.5.2` | carries the `cute.core.ThrMma` -> `cute.ThrMma` fix |
-| any | 4.6-era | unsupported: `cute.make_fragment` was removed at module level; fails at CuTe JIT trace |
-
-`FASTVIDEO_FA4=1` is required alongside the fork: it ships no compiled
-FlashAttention-2, so dense attention paths raise ImportError without the FA4
-opt-in. The same kernel also serves `ATTN_QAT_INFER` on sm_100a/sm_103a
-(datacenter Blackwell) — the selection log's receipt line
-(`ATTN_QAT_INFER resolved: ...`) records the arch, kernel, and quantization
-mode that actually bound.
+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.
 
 #### Usage
 

docs/inference/support_matrix.md

@@ -13,92 +13,6 @@ For the canonical, code-level list of model IDs recognized by
     We do this because we believe VSA is strictly better than STA for the
     actively maintained `main` inference path.
 
-## Registered Model IDs
-
-Every Hugging Face model ID registered in `fastvideo/registry.py` on `main`
-(commit `8d89f30d`), grouped by family. Any ID below can
-be passed to `VideoGenerator.from_pretrained(...)`; FastVideo resolves the
-matching pipeline and sampling defaults. The **Family** column is a
-documentation grouping: it follows each registration's declared `model_family`,
-except `black-forest-labs/FLUX.1-dev`, which declares none and is listed under
-`flux` for readability. The **Workloads** column shows each
-registration's declared `workload_types`; `—` means the entry is registered
-without a UI workload option but is still loadable by ID. The **Example**
-column links a runnable script in `examples/inference/basic/` where one exists.
-
-| Family | HuggingFace Model ID | Workloads | Example |
-|--------|----------------------|-----------|---------|
-| cosmos | `nvidia/Cosmos-Predict2-2B-Video2World` | T2V | — |
-| cosmos25 | `KyleShao/Cosmos-Predict2.5-2B-Diffusers` | T2V | [basic_cosmos2_5_t2w.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_cosmos2_5_t2w.py) |
-| cosmos25 | `nvidia/Cosmos-Predict2.5-14B` | T2V | [basic_cosmos2_5_t2w.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_cosmos2_5_t2w.py) |
-| dreamx_world | `FastVideo/DreamX-World-5B-Cam-Diffusers` | I2V | [basic_dreamx_world.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_dreamx_world.py) |
-| dreamx_world | `FastVideo/DreamX-World-5B-Diffusers` | I2V | [basic_dreamx_world.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_dreamx_world.py) |
-| flux | `black-forest-labs/FLUX.1-dev` | T2I | [basic_flux_dev.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_flux_dev.py) |
-| flux2 | `black-forest-labs/FLUX.2-klein-4B`<br>`black-forest-labs/FLUX.2-klein-9B` | T2I | [basic_flux2_klein.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_flux2_klein.py) |
-| flux2 | `black-forest-labs/FLUX.2-dev` | T2I | [basic_flux2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_flux2.py) |
-| gamecraft | `FastVideo/HunyuanGameCraft-Diffusers` | I2V | [basic_gamecraft.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_gamecraft.py) |
-| gen3c | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | T2V | [basic_gen3c.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_gen3c.py) |
-| glm_image | `zai-org/GLM-Image` | T2I | [basic_glm_image.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_glm_image.py) |
-| hunyuan | `hunyuanvideo-community/HunyuanVideo` | T2V | — |
-| hunyuan | `FastVideo/FastHunyuan-diffusers` | T2V | — |
-| hunyuan15 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | T2V | [basic_hy15.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_hy15.py) |
-| hunyuan15 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled` | I2V | [basic_hy15.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_hy15.py) |
-| hunyuan15 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | T2V | [basic_hy15.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_hy15.py) |
-| hunyuan15 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v_distilled` | I2V | [basic_hy15.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_hy15.py) |
-| hunyuan15 | `weizhou03/HunyuanVideo-1.5-Diffusers-1080p`<br>`weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR` | — | [basic_hy15_1080p.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_hy15_1080p.py) |
-| hyworld | `FastVideo/HY-WorldPlay-Bidirectional-Diffusers` | — | [basic_hyworld.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_hyworld.py) |
-| kandinsky5 | `kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers` | T2V | [basic_kandinsky5_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_t2v.py) |
-| kandinsky5 | `kandinskylab/Kandinsky-5.0-T2V-Pro-sft-5s-Diffusers` | T2V | [basic_kandinsky5_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_t2v.py) |
-| kandinsky5 | `kandinskylab/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers` | T2V | [basic_kandinsky5_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_t2v.py) |
-| kandinsky5 | `kandinskylab/Kandinsky-5.0-T2V-Pro-distilled-5s-Diffusers` | T2V | [basic_kandinsky5_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_t2v.py) |
-| kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) |
-| kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) |
-| kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) |
-| lingbot_video | `FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers` | T2V | [basic_lingbot_video.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbot_video.py) |
-| lingbot_video | `FastVideo/LingBot-Video-Dense-1.3B-Diffusers` | T2V | [basic_lingbot_video.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbot_video.py) |
-| lingbotworld | `FastVideo/LingBot-World-Base-Cam-Diffusers` | I2V | [basic_lingbotworld_base_cam.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbotworld_base_cam.py) |
-| lingbotworld2 | `robbyant/lingbot-world-v2-14b-causal-fast` | I2V | [basic_lingbotworld2_causal_fast.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbotworld2_causal_fast.py) |
-| longcat | `FastVideo/LongCat-Video-T2V-Diffusers` | T2V | [basic_longcat_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_t2v.py) |
-| longcat | `FastVideo/LongCat-Video-I2V-Diffusers` | I2V | [basic_longcat_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_i2v.py) |
-| longcat | `FastVideo/LongCat-Video-VC-Diffusers` | — | [basic_longcat_vc.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_vc.py) |
-| ltx2 | `FastVideo/LTX2-Distilled-Diffusers`<br>`FastVideo/LTX2.3-Distilled-Diffusers`<br>`FastVideo/LTX-2.3-Distilled-Diffusers` | T2V | [basic_ltx2_distilled.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_distilled.py) |
-| ltx2 | `Lightricks/LTX-2.3`<br>`FastVideo/LTX2.3-base`<br>`FastVideo/LTX2.3-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) |
-| ltx2 | `Lightricks/LTX-2`<br>`FastVideo/LTX2-base`<br>`FastVideo/LTX2-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) |
-| matrixgame | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-Base-Diffusers`<br>`FastVideo/Matrix-Game-2.0-GTA-Diffusers`<br>`FastVideo/Matrix-Game-2.0-TempleRun-Diffusers`<br>`mignonjia/mg_longtuning_distilled_zelda`<br>`mignonjia/mg_sf_distilled_zelda_1k_steps`<br>`mignonjia/mg_sf_distilled_zelda`<br>`mignonjia/mg_causal_zelda`<br>`mignonjia/mg_bidirectional_zelda` | I2V | [basic_matrixgame2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame2.py) |
-| matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) |
-| sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) |
-| stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) |
-| stable_audio | `FastVideo/stable-audio-open-small-Diffusers` | T2V | [basic_stable_audio_small.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio_small.py) |
-| turbodiffusion | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | T2V | [basic_turbodiffusion.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_turbodiffusion.py) |
-| turbodiffusion | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | T2V | [basic_turbodiffusion_14b.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_turbodiffusion_14b.py) |
-| turbodiffusion | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | I2V | [basic_turbodiffusion_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_turbodiffusion_i2v.py) |
-| wan | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | T2V | [basic.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic.py) |
-| wan | `Wan-AI/Wan2.1-T2V-14B-Diffusers`<br>`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` | T2V | — |
-| wan | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | I2V | — |
-| wan | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | I2V | — |
-| wan | `weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers` | I2V | — |
-| wan | `IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers` | — | [basic_wan2_2_Fun.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_wan2_2_Fun.py) |
-| wan | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers`<br>`FastVideo/FastWan2.1-T2V-14B-480P-Diffusers` | T2V | [basic_dmd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_dmd.py) |
-| wan | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | T2V, I2V | [basic_wan2_2_ti2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_wan2_2_ti2v.py) |
-| wan | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers`<br>`FastVideo/FastWan2.2-TI2V-5B-Diffusers` | T2V, I2V | [basic_dmd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_dmd.py) |
-| wan | `decart-ai/Lucy-Edit-Dev`<br>`decart-ai/Lucy-Edit-1.1-Dev` | — | [basic_lucy_edit.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lucy_edit.py) |
-| wan | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | T2V | [basic_wan2_2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_wan2_2.py) |
-| wan | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | I2V | [basic_wan2_2_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_wan2_2_i2v.py) |
-| wan | `wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers` | T2V | [basic_self_forcing_causal.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal.py) |
-| wan | `rand0nmr/SFWan2.2-T2V-A14B-Diffusers` | T2V | [basic_self_forcing_causal_wan2_2_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py) |
-| wan | `FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers` | I2V | [basic_self_forcing_causal_wan2_2_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py) |
-| zimage | `Tongyi-MAI/Z-Image-Turbo` | T2I | [basic_zimage.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_zimage.py) |
-
-**Note (stable_audio)**: the Stable Audio Open pipelines generate audio
-(`StableAudioT2AConfig` / `StableAudioOpenSmallConfig`); they are registered
-under the generic T2V workload option in the registry.
-
-**Note (Wan-VACE)**: not currently supported — no VACE pipeline or registered
-model ID exists on `main`
-([#1435](https://github.com/hao-ai-lab/FastVideo/issues/1435)). The closest
-supported path is the Wan2.1-Fun control pipeline
-(`IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers`).
-
 The symbols used have the following meanings:
 
 - ✅ = Full compatibility
@@ -111,9 +25,6 @@ The `HuggingFace Model ID` can be passed directly to
 `from_pretrained()`. FastVideo then uses model-specific default settings for
 pipeline initialization and sampling.
 
-Registered models absent from this table have not been validated against these
-optimizations: absence means **untested**, not incompatible.
-
 <style>
   /* Target tables in this section */
   #models-x-optimization + p + table {
@@ -145,7 +56,7 @@ optimizations: absence means **untested**, not incompatible.
 | 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 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
@@ -154,14 +65,14 @@ optimizations: absence means **untested**, not incompatible.
 | Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
 | FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
-| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
-| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
-| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
+| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ | ⭕ |
+| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅* | ✅ | ⭕ | ⭕ |
+| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ | ⭕ |
 | Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
 | TurboWan2.1 T2V 1.3B | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | TurboWan2.1 T2V 14B | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | 480P, 720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | TurboWan2.2 I2V A14B | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
-| LongCat T2V 13.6B | `FastVideo/LongCat-Video-T2V-Diffusers` | 480P<br>720P | ❌ | ❌ | ❌ | ⭕ | ✅ |
+| LongCat T2V 13.6B | See note** | 480P<br>720P | ❌ | ❌ | ❌ | ⭕ | ✅ |
 | Matrix Game 2.0 Base Distilled | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Matrix Game 2.0 GTA Distilled | `FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Matrix Game 2.0 TempleRun Distilled | `FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
@@ -187,21 +98,6 @@ resolve default pipeline and sampling configuration for it.
 `FastVideo/GEN3C-Cosmos-7B-Diffusers`) or convert locally with
 `scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py`.
 
-## Hardware and OS
-
-Per the installation guides:
-
-- **NVIDIA GPU (x86_64)** — CUDA 12.6 or 13.0; see the
-  [GPU install guide](../getting_started/installation/gpu.md).
-- **NVIDIA DGX Spark (GB10, aarch64)** — CUDA 13, from-source kernel build; see
-  the [DGX Spark install guide](../getting_started/installation/spark.md).
-- **Apple silicon (MPS)** — macOS 14 or newer; see the
-  [MPS install guide](../getting_started/installation/mps.md) and
-  [`basic_mps.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mps.py).
-
-Optimization-specific hardware constraints (e.g. STA requiring Hopper) are
-listed under [Special requirements](#special-requirements).
-
 ## Special requirements
 
 ### Sliding Tile Attention

docs/training/train_infra.md

@@ -212,29 +212,6 @@ pipeline:
   flow_shift: 8
 ```
 
-Registered transformer linear-quantization configs can also be selected by
-name. For example, the LTX-2 NVFP4-QAT recipe applies real FP4 forward GEMMs
-with a straight-through-estimator backward to its deployment-targeted
-attention/FFN projections:
-
-```yaml
-pipeline:
-  dit_config:
-    quant_config: nvfp4_qat_train
-```
-
-The LTX-2 recipe in
-`examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml` combines that linear
-configuration with `models.student.attention_backend: ATTN_QAT_TRAIN` for
-video-attention forward/backward. On sm120, its validation callback temporarily
-switches those layers to `ATTN_QAT_INFER`.
-On GB200, set `callbacks.validation.attn_qat_infer: false` to keep validation on
-the train-time QAT backend; the inference kernel is sm120-only.
-
-User-adaptable LTX-2 fine-tuning recipes (full, LoRA, and NVFP4 QAT) live in
-`examples/train/configs/fine_tuning/ltx2/`, alongside the other model
-families under `examples/train/configs/fine_tuning/`.
-
 ---
 
 ## Training Methods

docs/utilities/debugging.md

@@ -77,9 +77,7 @@ If forcing a backend fails, verify optional dependencies are installed:
 - `SAGE_ATTN`: SageAttention package
 - `SAGE_ATTN_THREE`: upstream `sageattn3` package
 - `ATTN_QAT_INFER`: `fastvideo-kernel` checkout/source install that exposes
-  `attn_qat_infer`, AND a consumer-Blackwell (sm_120/sm_121) GPU -- on any
-  other device the backend reports unavailable (even if a CUDA 13 wheel
-  bundles the extension) and selection falls back to FlashAttention
+  `attn_qat_infer`
 - `ATTN_QAT_TRAIN`: `fastvideo-kernel`; its runtime-JIT Triton implementation
   selects an optimized route on SM100, joins the quantized and STE P@V paths on
   SM120, and retains the previous route for unsupported configurations. See

examples/inference/basic/basic_ltx2_distilled_fast_profile.py

@@ -191,22 +191,6 @@ def main() -> None:
     print(f"Using refine upsampler: {refine_upsampler_path}")
 
     pipeline_config = PipelineConfig.from_pretrained(model_root)
-    # LTX-2 NVFP4 deploy contract (train==deploy surface):
-    #   * Linears: NVFP4 block-scaled GEMMs (per-16 E2M1 + E4M3 SFs) on every
-    #     arch, via flashinfer.
-    #   * ATTN_QAT_INFER attention differs per arch: sm_120a/sm_121a use the
-    #     fastvideo-kernel CUTLASS (SageAttention3-FP4) scheme that
-    #     ATTN_QAT_TRAIN simulates; sm_100a (GB200) / sm_103a (GB300) use the
-    #     FP4 FA4 kernel (flash-attention-fp4) with per-16 block-scaled NVFP4
-    #     Q/K and BF16 P/V -- a train-sim mismatch that is gated by MS-SSIM
-    #     measurement, not assumed equal. The selection receipt is logged at
-    #     backend resolution ("ATTN_QAT_INFER resolved: ...").
-    # Original-weight retention: the default purges the always-FP4 layers'
-    # bf16 originals after conversion. Refine-only layers (the cross-modal
-    # AV projections) always keep theirs: the base stage profile runs them
-    # dense by deployment contract -- in the two-stage fast profile AND the
-    # distilled single-stage deploy. retain_original_weights=True keeps
-    # everything (debugging).
     pipeline_config.dit_config.quant_config = NVFP4Config()
     OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
     torch_compile_kwargs = {

examples/inference/gradio/local/gradio_local_demo_ltx2_3/README.md

@@ -6,6 +6,10 @@ original single-file version to make each concern independently reviewable.
 The folder name matches the sibling `gradio_local_demo*.py` demos in this
 directory so both flat and packaged demos read consistently.
 
+> **Status: draft.** This package is structurally in place but will not run
+> against the current upstream `fastvideo` package. See *Blocking prereqs*
+> below.
+
 ## Layout
 
 | File | Purpose |
@@ -34,21 +38,27 @@ python -m gradio_local_demo_ltx2_3 --port 7860
 GPU requirement: a single FP4-capable GPU (B200 or comparable) for the
 "real-time 1080p" speed claim. Lower tiers will still run but slower.
 
-## Upstream surfaces the demo relies on
+## Blocking prereqs (why this draft PR cannot be merged yet)
 
-Everything the demo needs ships in the upstream `fastvideo` package:
+The upstream `fastvideo` package is missing three pieces that the demo
+currently depends on verbatim. Each needs its own upstreaming PR before this
+demo can actually boot:
 
-- **`fastvideo.layers.quantization.nvfp4_config.NVFP4Config`** — the demo sets
-  `pipeline_config.dit_config.quant_config = NVFP4Config()` in `app.py` (same
-  pattern as `examples/inference/basic/basic_ltx2_distilled_fast_profile.py`).
-- **LTX-2.3 refine kwargs on `VideoGenerator`** — `ltx2_refine_enabled`,
-  `ltx2_refine_upsampler_path`, `ltx2_refine_lora_path`,
-  `ltx2_refine_num_inference_steps`, `ltx2_refine_guidance_scale`,
-  `ltx2_refine_add_noise`, and `ltx2_vae_tiling` are wired in
-  `fastvideo/fastvideo_args.py`, backed by
-  `fastvideo/pipelines/basic/ltx2/stages/ltx2_refine.py`.
-- **`SamplingParam`** — imported from the public `fastvideo` package root
-  (`fastvideo.api.sampling_param` under the hood).
+1. **`fastvideo.layers.quantization.fp4_config.FP4Config`** — the demo sets
+   `pipeline_config.dit_config.quant_config = FP4Config()` in `app.py`.
+   Upstream only ships `absmax_fp8.py` and `base_config.py` under
+   `fastvideo/layers/quantization/`.
+2. **LTX-2.3 refine / image-conditioning kwargs on `VideoGenerator`** —
+   `ltx2_refine_enabled`, `ltx2_refine_upsampler_path`, `ltx2_refine_lora_path`,
+   `ltx2_refine_num_inference_steps`, `ltx2_refine_guidance_scale`,
+   `ltx2_refine_add_noise`, `ltx2_images`, `ltx2_image_crf`. Upstream
+   `fastvideo/fastvideo_args.py` currently wires only `ltx2_vae_tiling`.
+   The backing stages (`ltx2_refine.py`, `ltx2_i2v_conditioning.py`) are
+   also missing from `fastvideo/pipelines/stages/`.
+3. **`fastvideo.configs.sample.base.SamplingParam`** — the import path used
+   by this demo. Upstream moved sampling params to
+   `fastvideo.api.sampling_param`. A re-export shim at the old path, or an
+   import update here once the other two prereqs land, will resolve it.
 
 ## Environment variables
 

examples/inference/gradio/local/gradio_local_demo_ltx2_3/app.py

@@ -4,10 +4,10 @@ from pathlib import Path
 
 import gradio as gr
 
-from fastvideo import SamplingParam
 from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.sample.base import SamplingParam
 from fastvideo.entrypoints.video_generator import VideoGenerator
-from fastvideo.layers.quantization.nvfp4_config import NVFP4Config
+from fastvideo.layers.quantization.fp4_config import FP4Config
 from fastvideo.utils import maybe_download_model
 
 from .config import (
@@ -44,7 +44,7 @@ def main():
         resolved_model_path = Path(model_root)
 
         pipeline_config = PipelineConfig.from_pretrained(str(resolved_model_path))
-        pipeline_config.dit_config.quant_config = NVFP4Config()
+        pipeline_config.dit_config.quant_config = FP4Config()
         refine_upsampler_path = resolve_refine_upsampler_path(resolved_model_path)
         print(f"Using refine upsampler: {refine_upsampler_path}")
 

examples/inference/gradio/local/gradio_local_demo_ltx2_3/config.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import torch
 import torch._inductor.config
 
-from fastvideo import SamplingParam
+from fastvideo.configs.sample.base import SamplingParam
 
 LOCAL_DEMO_DIR = Path(__file__).resolve().parent
 CLASSIFIER_DIR = Path(

examples/inference/gradio/local/gradio_local_demo_ltx2_3/ui.py

@@ -5,7 +5,7 @@ from copy import deepcopy
 
 import gradio as gr
 
-from fastvideo import SamplingParam
+from fastvideo.configs.sample.base import SamplingParam
 from fastvideo.entrypoints.video_generator import VideoGenerator
 
 from .config import (

examples/inference/ltx2_3/optimized_nvfp4_t2v.py

@@ -1,176 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""LTX-2.3 distilled text-to-video with the optimized NVFP4 inference stack.
-
-Runs `FastVideo/LTX-2.3-Distilled-Diffusers` on a single GPU with the full
-validated optimization stack:
-
-* NVFP4 block-scaled linear layers (per-16 E2M1 weights + E4M3 scale factors),
-* ATTN_QAT_INFER FP4 attention (arch-resolved kernel, receipt logged),
-* torch.compile (fullgraph) over the DiT, text encoder, and VAE,
-* single-stage 8-step distilled sampling at guidance 1.0.
-
-Quick start
------------
-    # On GB200-class ARM hosts, unset LD_LIBRARY_PATH (see Hardware notes):
-    env -u LD_LIBRARY_PATH python examples/inference/ltx2_3/optimized_nvfp4_t2v.py
-
-    # Optional overrides:
-    #   export LTX23_MODEL_PATH=/path/to/local/snapshot
-    #   export LTX23_T2V_PROMPT="a red fox running through fresh snow"
-    #   export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_nvfp4_t2v
-
-Hardware notes
---------------
-- On GB200 / Blackwell, run with `env -u LD_LIBRARY_PATH ...` to avoid a
-  system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM (some ARM
-  container images ship an LD_LIBRARY_PATH that breaks torch.compile's
-  toolchain discovery). The `_inductor.shape_padding = False` line below
-  also avoids a pad_mm landmine on the same generation of cards.
-"""
-from __future__ import annotations
-
-import os
-import time
-from pathlib import Path
-
-import torch._inductor.config as _inductor
-
-from fastvideo import VideoGenerator
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.layers.quantization.nvfp4_config import NVFP4Config
-from fastvideo.utils import maybe_download_model
-
-# ATTN_QAT_INFER is the FP4 attention half of the NVFP4 deploy contract. It
-# resolves per arch (CUTLASS SageAttention3-FP4 on sm_120a/sm_121a, FP4 FA4
-# on sm_100a/sm_103a) and logs a one-line receipt of what actually bound.
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER")
-os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1")
-
-# Inductor knobs. The first one (shape_padding=False) is mandatory on
-# Blackwell to avoid a cuBLAS INVALID_VALUE crash inside pad_mm. The rest
-# are the same matmul-friendliness flags the sibling LTX-2 examples use.
-_inductor.shape_padding = False
-_inductor.conv_1x1_as_mm = True  # treat 1x1 convolutions as matrix muls
-_inductor.coordinate_descent_tuning = True
-_inductor.coordinate_descent_check_all_directions = True
-_inductor.epilogue_fusion = False  # do not fuse pointwise ops into matmuls
-
-MODEL_ID = os.path.expandvars(
-    os.path.expanduser(
-        os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers")
-    )
-)
-OUTPUT_DIR = Path(os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_nvfp4_t2v"))
-DEFAULT_PROMPT = (
-    "A fashion model takes a slow step forward and shifts her weight, "
-    "the soft fabric of her clothing swaying and rippling with the "
-    "motion, her hair shifting gently, soft even studio lighting on a "
-    "clean light background, elegant slow-motion runway feel."
-)
-PROMPT = os.getenv("LTX23_T2V_PROMPT", DEFAULT_PROMPT)
-
-
-def main() -> None:
-    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-    model_root = maybe_download_model(MODEL_ID)
-    print(f"Model:      {model_root}")
-    print(f"Output dir: {OUTPUT_DIR.resolve()}")
-
-    # Loading the pipeline config *with model_path* binds model-specific
-    # tuning (notably VAE precision/decoder defaults) into the config.
-    pipeline_config = PipelineConfig.from_pretrained(model_root)
-
-    # NVFP4 linear layers on the DiT. The default purges the original BF16
-    # weights of the always-FP4 linears right after conversion — a large
-    # peak-memory reduction. Refine-only layers (the cross-modal AV
-    # projections) always keep theirs: the base stage profile runs them
-    # dense by deployment contract. retain_original_weights=True keeps
-    # everything (debugging).
-    pipeline_config.dit_config.quant_config = NVFP4Config()
-
-    # fullgraph=True is supported on the NVFP4 path: the FP4 quantize step
-    # is a registered custom op (fastvideo::nvfp4_quantize_fa4), so dynamo
-    # traces through it without graph breaks. mode="default" — the two
-    # CUDAGraph modes ("reduce-overhead" / "max-autotune") are a separate
-    # opt-in, not part of this validated preset.
-    torch_compile_kwargs = {
-        "backend": "inductor",
-        "fullgraph": True,
-        "mode": "default",
-        "dynamic": False,
-    }
-
-    generator = VideoGenerator.from_pretrained(
-        model_root,
-        num_gpus=1,
-        pipeline_config=pipeline_config,
-        # Compile the DiT, text encoder, and VAE — all three stages benefit,
-        # and the VAE's codec submodules compile cleanly under fullgraph.
-        enable_torch_compile=True,
-        enable_torch_compile_text_encoder=True,
-        enable_torch_compile_vae=True,
-        torch_compile_kwargs=torch_compile_kwargs,
-        torch_compile_kwargs_vae=torch_compile_kwargs,
-        # Keep everything resident — no CPU offload for serving-style runs.
-        dit_cpu_offload=False,
-        text_encoder_cpu_offload=False,
-        vae_cpu_offload=False,
-        ltx2_vae_tiling=False,
-    )
-
-    common_kwargs = dict(
-        prompt=PROMPT,
-        negative_prompt="",      # distilled is CFG-free; no negative needed
-        guidance_scale=1.0,      # CFG=1 for distilled
-        height=1280, width=832,  # portrait runway aspect
-        num_frames=121, fps=24,  # ~5s clip
-        # Single-stage 8-step distilled sampling — the validated preset for
-        # this checkpoint (no two-stage refine; the NVFP4 deploy contract
-        # runs the distilled single-stage recipe).
-        num_inference_steps=8,
-        save_video=True,
-    )
-
-    try:
-        # Warmup: pays cold compile + first-shape guard work, untimed.
-        print("\n[warmup] compiling + generating…")
-        generator.generate_video(
-            output_path=str(OUTPUT_DIR / "_warmup.mp4"),
-            seed=7,
-            **common_kwargs,
-        )
-        (OUTPUT_DIR / "_warmup.mp4").unlink(missing_ok=True)
-
-        # Measured run.
-        out_path = OUTPUT_DIR / "output_ltx2_3_nvfp4_t2v.mp4"
-        print(f"\n[measured] generating: {out_path}")
-        t0 = time.perf_counter()
-        result = generator.generate_video(
-            output_path=str(out_path),
-            seed=2002,
-            **common_kwargs,
-        )
-        wall = time.perf_counter() - t0
-        e2e = (result.get("e2e_latency") if isinstance(result, dict) else None) or wall
-        print(f"[measured] e2e={e2e:.2f}s wall={wall:.2f}s")
-    finally:
-        generator.shutdown()
-
-    # Expected receipts — verify these two lines in your own run's log:
-    #
-    # 1. ATTN_QAT_INFER routing receipt (logged at backend resolution; on a
-    #    GB200-class part it reads):
-    #
-    #      ATTN_QAT_INFER resolved: arch=sm_100 kernel=flash-attention-fp4 \
-    #        qk_mode=nvfp4(per-16-e4m3-sf) pv_mode=bf16 train_sim_mismatch=measured
-    #
-    # 2. NVFP4 weight purge receipt (logged after model conversion; N/M/X
-    #    depend on the checkpoint):
-    #
-    #      NVFP4 weight purge receipt: purged N original bf16 weight tensors \
-    #        (X.XX GiB freed); retained M (refine-only dense fallback or \
-    #        retain_original_weights).
-
-
-if __name__ == "__main__":
-    main()

examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh

@@ -1,70 +0,0 @@
-#!/bin/bash
-# Stage-2 QAT-aware DMD2 distillation for Kandinsky5 T2V (480p).
-#
-# Usage:
-#   bash distill_dmd_qat.sh <stage1-checkpoint> [--dotted.key value ...]
-#
-# <stage1-checkpoint> is either:
-#   - a raw stage-1 DCP checkpoint written by
-#     examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh
-#     (a checkpoint-<N>/ dir with a dcp/ subdir, or the stage-1 output_dir
-#     itself -- dcp_to_diffusers auto-picks the latest checkpoint). It is
-#     converted to a diffusers export (<checkpoint>-diffusers) with
-#     `dcp_to_diffusers --role student --verify` before launching --
-#     models.*.init_from need a diffusers model directory (model_index.json
-#     + component subfolders), not the raw DCP layout; or
-#   - an already-exported diffusers model dir (model_index.json present),
-#     used as-is with no conversion.
-#
-# The export is passed to all three models.*.init_from overrides: student,
-# teacher, and critic all load the same weights. Teacher/critic are
-# automatically masked back to dense attention and full-precision weights
-# by the _loading_teacher_critic_model gate in
-# fastvideo/models/loader/component_loader.py (family-agnostic, no
-# Kandinsky5-specific handling needed) -- NOT by pointing them at a
-# different export.
-#
-# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN keeps the student's dense/local
-# attention fake-quantized during distillation.
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)"
-
-CHECKPOINT="${1:?Usage: $0 <stage1-checkpoint (DCP dir or diffusers export)> [extra --dotted.key overrides...]}"
-shift
-
-if [[ "${CHECKPOINT}" != /* ]]; then
-    CHECKPOINT="$(pwd)/${CHECKPOINT}"
-fi
-CHECKPOINT="${CHECKPOINT%/}"
-
-if [[ -f "${CHECKPOINT}/model_index.json" ]]; then
-    # Already a diffusers export -- use as-is.
-    INIT_FROM="${CHECKPOINT}"
-else
-    # Convert the raw DCP checkpoint. Runs BEFORE ATTN_QAT_TRAIN is
-    # exported below: backend selection for that env var refuses to start
-    # (ImportError) when the training kernel isn't built, and the
-    # conversion itself must not depend on it. --overwrite keeps relaunches
-    # from silently reusing a stale export for a newer checkpoint;
-    # --verify strictly reloads the exported transformer so a key-mapping
-    # bug fails here, not deep inside the training launch below.
-    INIT_FROM="${CHECKPOINT}-diffusers"
-    echo "Converting stage-1 DCP checkpoint ${CHECKPOINT} -> ${INIT_FROM}"
-    python -m fastvideo.train.entrypoint.dcp_to_diffusers \
-        --checkpoint "${CHECKPOINT}" \
-        --output-dir "${INIT_FROM}" \
-        --role student \
-        --overwrite \
-        --verify
-fi
-
-export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
-
-bash "${REPO_ROOT}/examples/train/run.sh" \
-    "${SCRIPT_DIR}/dmd2_t2v_480p_qat.yaml" \
-    --models.student.init_from "${INIT_FROM}" \
-    --models.teacher.init_from "${INIT_FROM}" \
-    --models.critic.init_from "${INIT_FROM}" \
-    "$@"

examples/train/configs/distribution_matching/kandinsky5/dmd2_t2v_480p_qat.yaml

@@ -1,113 +0,0 @@
-# DMD2 QAT-aware distillation: Kandinsky5 T2V Lite, 480p (teacher many-step
-# -> student 4-step), continuing from a stage-1 Attn-QAT finetune checkpoint.
-#
-# init_from below must be a diffusers model directory (model_index.json +
-# component subfolders), not the raw checkpoint-N/ stage 1 writes (dcp/ +
-# metadata/RNG state only). Launching via distill_dmd_qat.sh handles this:
-# it runs fastvideo.train.entrypoint.dcp_to_diffusers (--verify) on the
-# stage-1 checkpoint you pass it and overrides all three init_from values
-# with the export -- the placeholder paths below only apply when launching
-# the YAML directly.
-#
-# - Teacher: frozen, full-precision Kandinsky5 (loaded from the stage-1 QAT
-#   checkpoint's base weights, or the original pretrained checkpoint --
-#   the _loading_teacher_critic_model gate in component_loader.py masks
-#   quant_config and FASTVIDEO_ATTENTION_BACKEND for teacher/critic
-#   regardless, so they always run full precision / dense attention).
-# - Student: trainable, initialized from the stage-1 QAT finetune checkpoint.
-# - Critic:  trainable, initialized from the same checkpoint as the student.
-# - Validation: 4-step sampling via Kandinsky5DMDPipeline.
-#
-# Launch with FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN set (student only;
-# teacher/critic are masked to dense automatically) so the student continues
-# learning to absorb fake-quantized attention error during distillation.
-#
-# Scope: T2V, 480p only, dense/local attention only.
-
-models:
-  student:
-    _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
-    init_from: outputs/kandinsky5_t2v_qat_finetune/checkpoint-1000-diffusers
-    trainable: true
-  teacher:
-    _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
-    init_from: outputs/kandinsky5_t2v_qat_finetune/checkpoint-1000-diffusers
-    trainable: false
-    disable_custom_init_weights: true
-  critic:
-    _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
-    init_from: outputs/kandinsky5_t2v_qat_finetune/checkpoint-1000-diffusers
-    trainable: true
-    disable_custom_init_weights: true
-
-method:
-  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
-  rollout_mode: simulate
-  generator_update_interval: 5
-  real_score_guidance_scale: 4.5
-  dmd_denoising_steps: [1000, 750, 500, 250]
-
-  # Critic optimizer (required -- no fallback to training.optimizer)
-  fake_score_learning_rate: 8.0e-6
-  fake_score_betas: [0.0, 0.999]
-  fake_score_lr_scheduler: constant
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/kandinsky5_overfit_preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 31
-    num_height: 512
-    num_width: 768
-    num_frames: 121
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.0, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 4000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/kandinsky5_t2v_dmd2_4steps_qat
-    training_state_checkpointing_steps: 20
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_kandinsky5
-    run_name: kandinsky5_t2v_dmd2_4steps_qat
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-  validation:
-    pipeline_target: fastvideo.pipelines.basic.kandinsky5.kandinsky5_dmd_pipeline.Kandinsky5DMDPipeline
-    dataset_file: data/kandinsky5_overfit_preprocessed/validation_prompts.json
-    every_steps: 50
-    sampling_steps: [4]
-    sampling_timesteps: [1000, 750, 500, 250]
-    guidance_scale: 6.0
-  ema:
-    _target_: fastvideo.train.callbacks.ema.EMACallback
-    decay: 0.98
-    start_iter: 0
-
-pipeline:
-  flow_shift: 5

examples/train/configs/example.yaml

@@ -195,10 +195,6 @@ callbacks:
 # ------------------------------------------------------------------------------
 pipeline:
   flow_shift: 3                       # default: null (model-specific)
-  # LTX-2 linear QAT (SM100+ with FlashInfer); distinct from the role-local
-  # attention_backend option above, which selects attention math for Wan.
-  # dit_config:
-  #   quant_config: nvfp4_qat_train
   # flow_shift_sr: null               # default: null (super-resolution shift)
   # embedded_cfg_scale: 6.0           # default: 6.0
   # is_causal: false                  # default: false

examples/train/configs/fine_tuning/kandinsky5/README.md

@@ -1,198 +0,0 @@
-# Kandinsky5 QAD recipe (T2V, 480p)
-
-Quantization-aware distillation for Kandinsky-5.0 Lite T2V at
-**512×768, 121 frames, 24 fps** (the checkpoint's native `KANDINSKY5_T2V_LITE_5S`
-preset -- 121/24 ~= 5.04s, matching the "5s" in the checkpoint name),
-mirroring the Wan2.1 QAD recipe
-(`examples/training/finetune/wan_t2v_1.3B/mixkit/`) on the new
-`fastvideo/train/` stack. Scope is deliberately narrow: T2V only, 480p only,
-dense/local attention only. Kandinsky5's NABLA sparse attention path is
-never engaged at this resolution and is unsupported by this recipe --
-attempting a larger resolution will raise loudly from `Kandinsky5Model`
-rather than silently mis-scale the visual RoPE.
-
-## Data
-
-Kandinsky5 uses two text encoders (Qwen/Reason1 + CLIP), unlike Wan's single
-encoder. The CLIP pooled projection is zero-padded into a row and prepended
-to the Qwen sequence embeddings, then stored as a single `[seq+1, dim]`
-tensor in the existing `text_embedding` Parquet field -- the same trick
-`preprocess_hunyuan_overfit.py` uses for LLaMA+CLIP. No new Parquet schema
-is needed.
-
-Build a small Parquet dataset from raw videos + captions
-(`videos2caption.json` + `videos/*.mp4`, matching the Hunyuan overfit script's
-layout) with:
-
-```bash
-# from the repo root; reads data/kandinsky5_overfit and writes
-# data/kandinsky5_overfit_preprocessed by default -- override with
-# KANDINSKY5_OVERFIT_DATA_DIR / KANDINSKY5_OVERFIT_OUTPUT_DIR env vars
-python -m fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit
-```
-
-This loads Kandinsky5's VAE (shared with HunyuanVideo) and both text
-encoders via FastVideo's own loaders (`VAELoader`/`TextEncoderLoader`), and
-writes `data_00000.parquet` + `validation_prompts.json` into the configured
-output directory. Point `training.data.data_path` at that output directory
-in the YAML configs below.
-
-There is currently no scalable/production preprocessing pipeline (the kind
-built on `BasePreprocessPipeline` that Wan/LTX2 use) for Kandinsky5 -- like
-Hunyuan, its dual text encoders don't fit that base class's single-encoder
-assumption, and `BasePreprocessPipeline.preprocess_video_and_text` isn't
-cleanly overridable in pieces. If you need to scale past a handful of
-overfit samples, extend `preprocess_kandinsky5_overfit.py`'s loop rather
-than forking the production pipeline base class.
-
-## Train stage 1 (Attn-QAT finetune)
-
-```bash
-bash examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh \
-    --training.data.data_path data/kandinsky5_overfit_preprocessed
-```
-
-`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` (set by the script) routes
-Kandinsky5's dense/local attention through the fake-quantized
-straight-through-estimator Triton kernel, so the DiT learns to absorb
-quantization error instead of fighting it. This is purely env-var driven,
-model-agnostic, and does not touch weights during training -- weight-level
-FP4/FP8 quantization is applied post-hoc at inference time (see below), not
-during either training stage.
-
-## Train stage 2 (QAT-aware DMD distillation)
-
-Stage 1 writes a raw DCP checkpoint (`checkpoint-<N>/dcp` + metadata/RNG
-state) under `training.checkpoint.output_dir`
-(`outputs/kandinsky5_t2v_qat_finetune/` by default) -- `models.*.init_from`
-needs a diffusers model directory instead (`model_index.json` + component
-subfolders). Pass the stage-1 checkpoint to the launcher, which performs
-the conversion and injects the export into all three `init_from` overrides:
-
-```bash
-bash examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh \
-    outputs/kandinsky5_t2v_qat_finetune/checkpoint-<N>
-```
-
-Under the hood it runs exactly:
-
-```bash
-python -m fastvideo.train.entrypoint.dcp_to_diffusers \
-    --checkpoint outputs/kandinsky5_t2v_qat_finetune/checkpoint-<N> \
-    --output-dir outputs/kandinsky5_t2v_qat_finetune/checkpoint-<N>-diffusers \
-    --role student --overwrite --verify
-```
-
-then launches `dmd2_t2v_480p_qat.yaml` with
-`--models.student/teacher/critic.init_from` all pointing at the export.
-(Passing an already-exported diffusers directory to the launcher skips the
-conversion and uses it directly.)
-
-`--role student` exports the trained transformer weights (the only role
-that matters here: student/teacher/critic in stage 2 all load from the same
-checkpoint, and `_loading_teacher_critic_model` handles the full-precision
-masking for teacher/critic at load time, not at export time). `--verify`
-strictly reloads the exported transformer immediately, so a key-mapping bug
-fails here instead of deep inside the stage-2 launch.
-
-Only the student is quantized (Attn-QAT); the teacher and critic stay full
-precision / dense attention. This is enforced in the loader
-(`fastvideo/models/loader/component_loader.py`, via the
-`_loading_teacher_critic_model` flag), which masks `quant_config` and clears
-`FASTVIDEO_ATTENTION_BACKEND` for teacher/critic -- the same global env var
-reaches only the student, with no per-model flags or Kandinsky5-specific
-handling. Validation runs the distilled student through
-`Kandinsky5DMDPipeline` at the step counts configured in
-`callbacks.validation.sampling_steps`.
-
-## Export the stage-2 student for inference
-
-Stage 2 writes raw DCP checkpoints too, so the distilled student must be
-exported the same way stage 1 was before anything can load it -- the
-`path/to/kandinsky5_dmd_checkpoint` used below is exactly this export:
-
-```bash
-python -m fastvideo.train.entrypoint.dcp_to_diffusers \
-    --checkpoint outputs/kandinsky5_t2v_dmd2_4steps_qat/checkpoint-<N> \
-    --output-dir outputs/kandinsky5_t2v_dmd2_4steps_qat/checkpoint-<N>-diffusers \
-    --role student --verify
-```
-
-## Inference (NVFP4 / FP8 weight quantization)
-
-Kandinsky5's attention out-projection (`self_attention.out_layer` /
-`cross_attention.out_layer`) and FFN (`feed_forward.mlp.fc_in` /
-`feed_forward.mlp.fc_out`) names differ from Wan's (`to_out`, `ffn.fc_in` /
-`ffn.fc_out`); `to_query`/`to_key`/`to_value` already substring-match the
-`to_q`/`to_k`/`to_v` entries the quant configs use for Wan. Both
-`nvfp4_qat` and `FP8` now include Kandinsky5's out-projection/FFN names in
-their default layer lists (`fastvideo/layers/quantization/nvfp4_qat_config.py`,
-`fastvideo/layers/quantization/fp8_config.py`), so no `target_layers`
-override is needed:
-
-`dcp_to_diffusers` copies the base T2V checkpoint's `model_index.json`
-unchanged into every export (see "Train stage 2" above), so `_class_name`
-still says the base T2V pipeline and the registry
-(`fastvideo/registry.py`) has no way to auto-detect that a given directory
-is actually a DMD (four-step re-noise sampler) export -- it will resolve to
-`Kandinsky5T2VPipeline` and run the full-length sampler on DMD-distilled
-weights. Pass `override_pipeline_cls_name` and a `Kandinsky5DMDConfig`
-explicitly to select the right pipeline/sampler and get
-`dmd_denoising_steps` set (`Kandinsky5T2VConfig`'s default of `None` makes
-`Kandinsky5DmdDenoisingStage` raise immediately):
-
-```python
-from fastvideo import VideoGenerator
-from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig
-from fastvideo.layers.quantization import get_quantization_config
-
-gen = VideoGenerator.from_pretrained(
-    "path/to/kandinsky5_dmd_checkpoint", num_gpus=1,
-    override_pipeline_cls_name="Kandinsky5DMDPipeline",
-    pipeline_config=Kandinsky5DMDConfig(),
-    transformer_quant=get_quantization_config("nvfp4_qat")(),
-    use_fsdp_inference=False,
-)
-# Kandinsky5DmdDenoisingStage ignores request.sampling.num_inference_steps
-# and guidance_scale -- it's a fixed 4-step, no-CFG sampler driven entirely
-# by pipeline_config.dmd_denoising_steps (set above). Adjust
-# Kandinsky5DMDConfig(dmd_denoising_steps=[...]) instead if the checkpoint
-# was distilled/validated with a different schedule than the default
-# [1000, 750, 500, 250].
-gen.generate(request={"prompt": "...", "output": {"save_video": True}})
-```
-
-As with Wan, combine with `FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER` on an
-sm_120 GPU for the full 4-bit path; on other GPUs the attention falls back
-to a supported dense backend while the FP4 linear layers still run.
-`flashinfer` is required for the FP4 path.
-
-## Tests
-
-- `fastvideo/tests/train/models/test_load_kandinsky5.py` -- loads the real
-  checkpoint and runs one transformer forward pass.
-- `fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py` --
-  confirms `ATTN_QAT_TRAIN` actually engages (not a silent SDPA fallback)
-  and that gradients flow through a forward+backward pass.
-- `fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py` --
-  confirms an unmodified export resolves to `Kandinsky5T2VPipeline` (the bug)
-  and that `override_pipeline_cls_name="Kandinsky5DMDPipeline"` fixes it (the
-  documented workaround above), plus `Kandinsky5DMDConfig`'s
-  `dmd_denoising_steps` default.
-- `fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py` -- a few
-  steps of both training stages on a single synthetic sample, exercising the
-  whole recipe end to end: the `dcp_to_diffusers --verify` conversion
-  between the stages AND of the final stage-2 student, then reloading that
-  export through the documented `Kandinsky5DMDPipeline` /
-  `Kandinsky5DMDConfig` override, generating deterministically (fixed
-  seed), and comparing MS-SSIM against a committed reference video. The
-  test fails (does not skip) if the reference is missing; record it once on
-  a sanctioned GPU box with `KANDINSKY5_E2E_WRITE_REFERENCE=1`, review the
-  written video, and commit it (see the test's module docstring). Nightly
-  tests are not collected per-PR (`fastvideo/tests/contract/
-  test_ci_test_collection.py` allowlists the directory), so run it
-  explicitly:
-
-  ```bash
-  pytest fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py -vs
-  ```

examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh

@@ -1,33 +0,0 @@
-#!/bin/bash
-# QAD recipe -- quantization-aware finetune of Kandinsky5 T2V (480p) with
-# fake-quant (Attn-QAT) attention.
-#
-# The fake-quantized attention path is selected purely by env var
-# (config-driven, no monkey-patching): FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
-# routes Kandinsky5's dense/local attention through the fake-quantized
-# straight-through-estimator kernel, so the DiT learns to absorb the
-# quantization error instead of fighting it. No weight quantization happens
-# during this stage; that's applied post-hoc at inference time (see the
-# stage-2 config and README).
-#
-# Scope: T2V, 480p only, dense/local attention only -- NABLA sparse attention
-# is never engaged by Kandinsky5Model at this resolution.
-#
-# Data: preprocess with fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py
-# (or an equivalent parquet dataset matching pyarrow_schema_t2v) first.
-#
-# Next: this writes a raw DCP checkpoint (checkpoint-N/dcp + metadata/RNG
-# state), not something stage 2 can load directly. Pass that checkpoint to
-# ../../distribution_matching/kandinsky5/distill_dmd_qat.sh, which converts
-# it with fastvideo.train.entrypoint.dcp_to_diffusers (--verify) and feeds
-# the export to all three models.*.init_from overrides automatically.
-set -euo pipefail
-
-export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN   # <-- enables Attn-QAT training
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)"
-
-bash "${REPO_ROOT}/examples/train/run.sh" \
-    "${SCRIPT_DIR}/t2v_480p_qat.yaml" \
-    "$@"

examples/train/configs/fine_tuning/kandinsky5/t2v_480p_qat.yaml

@@ -1,85 +0,0 @@
-# Kandinsky5 T2V 480p stage-1 QAT finetune config.
-#
-# QAT here means attention-only fake-quant during training: launch with
-# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN set (see finetune_qat.sh), which
-# routes Kandinsky5's dense/local attention through the fake-quantized
-# straight-through-estimator kernel. No weight quantization happens during
-# training; that's applied post-hoc at inference time via the nvfp4_qat/FP8
-# quant configs (see the stage-2 config and README for that step).
-#
-# Scope: T2V, 480p only, dense/local attention only -- NABLA sparse attention
-# is never engaged by Kandinsky5Model at this resolution.
-#
-# Data must be preprocessed with Kandinsky5's VAE (shared with HunyuanVideo)
-# + dual text encoders (Qwen/Reason1 + CLIP) into parquet format before
-# training. See fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py.
-
-models:
-  student:
-    _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
-    init_from: kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers
-    trainable: true
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/kandinsky5_overfit_preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.1
-    seed: 1000
-    num_latent_t: 31
-    num_height: 512
-    num_width: 768
-    num_frames: 121
-
-  optimizer:
-    learning_rate: 5.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 1.0e-4
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 2000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/kandinsky5_t2v_qat_finetune
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 3
-    resume_from_checkpoint: latest
-
-  tracker:
-    project_name: fastvideo_kandinsky5
-    run_name: kandinsky5_t2v_qat_finetune
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-  validation:
-    _target_: fastvideo.train.callbacks.validation.ValidationCallback
-    pipeline_target: fastvideo.pipelines.basic.kandinsky5.kandinsky5_pipeline.Kandinsky5T2VPipeline
-    dataset_file: data/kandinsky5_overfit_preprocessed/validation_prompts.json
-    every_steps: 250
-    # Cut way down from 50 while shaking out validation-path bugs -- at
-    # 512x768x121 each step is ~30s, so 50 steps/sample x 6 samples/rank
-    # was ~2.5h per validation round. Bump back to 50 for real runs.
-    sampling_steps: [25]
-    guidance_scale: 6.0
-
-pipeline:
-  flow_shift: 5

examples/train/configs/fine_tuning/ltx2/nvfp4_qat_t2v.yaml

@@ -1,140 +0,0 @@
-# LTX-2 T2V NVFP4 quantization-aware fine-tune (QAT).
-#
-# Same recipe as fine_tuning/ltx2/t2v.yaml, with the deployment-matched
-# NVFP4 quantization from the validated
-# examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml run: the curated
-# attention/FFN linear set used by LTX-2 NVFP4 deployment runs its
-# forward through real NVFP4 GEMMs (flashinfer cutlass) with a
-# straight-through-estimator backward into the dense fp32 master weights.
-# No extra parameters or buffers are added, so FSDP sharding,
-# checkpointing, and export are identical to dense training. Video
-# attention uses ATTN_QAT_TRAIN for its quantized forward and STE
-# backward, then ATTN_QAT_INFER during validation. Head-dim-64 audio
-# attention and masked text attention remain dense.
-#
-# Validation-time ATTN_QAT_INFER is arch-aware:
-#   * sm_120 / sm_121: the bundled fastvideo-kernel CUTLASS extension.
-#   * sm_100 (GB200-class) / sm_103 (GB300-class): the FP4 FA4 kernel —
-#     requires the flash-attention-fp4 fork and its dependency pins; see
-#     docs/inference/optimizations.md.
-#   * any other GPU: falls back to flash attention with a logged receipt
-#     (validation still runs, just not with the quantized kernel).
-#
-# Prerequisites: preprocess your dataset into the trainer's latent format
-# (see docs/training/data_preprocess.md;
-# fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py is a minimal
-# LTX-2 reference to adapt). Data paths, step counts, and learning rate
-# below are placeholders you must adapt to your workload.
-#
-# Run (multi-GPU; see the validation note above for kernel availability):
-#   NUM_GPUS=4 bash examples/train/run.sh \
-#     examples/train/configs/fine_tuning/ltx2/nvfp4_qat_t2v.yaml
-#
-# `attn_qat_infer: true` below is the deployment-representative choice on
-# every arch with a quantized kernel available (see the validation note
-# above). Flip it off only when your GPU has no ATTN_QAT_INFER kernel and
-# you prefer validating with the train-time backend instead of the flash
-# fallback:
-#   NUM_GPUS=4 bash examples/train/run.sh \
-#     examples/train/configs/fine_tuning/ltx2/nvfp4_qat_t2v.yaml \
-#     --callbacks.validation.attn_qat_infer false
-
-models:
-  student:
-    _target_: fastvideo.train.models.ltx2.LTX2Model
-    # Validation settings below (8 steps, guidance 1.0) assume this
-    # distilled checkpoint.
-    init_from: FastVideo/LTX2-Distilled-Diffusers
-    trainable: true
-    enable_gradient_checkpointing_type: full
-    attention_backend: ATTN_QAT_TRAIN
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  distributed:
-    # The 13B video branch needs sharding; 4-way HSDP shard matches the
-    # validated LTX-2 runs. Scale num_gpus/hsdp_shard_dim together.
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    # REQUIRED: path to your preprocessed dataset.
-    data_path: data/your_dataset_preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    # LTX2Model requires 0.0: CFG dropout would zero post-connector
-    # embeddings, which is not the model's unconditional input.
-    training_cfg_rate: 0.0
-    seed: 42
-    # Must match your preprocessed resolution/length.
-    # num_latent_t = (num_frames - 1) / 8 + 1.
-    num_latent_t: 11
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    # Carried from the validated LTX-2 QAT overfit run as a starting
-    # point; tune for your dataset size and batch configuration.
-    learning_rate: 5.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    # Workload-dependent: set from your dataset size and target epochs.
-    max_train_steps: 2000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/ltx2_t2v_nvfp4_qat_finetune
-    # A full training-state checkpoint is ~150GB for the 13B trainable
-    # video branch — size the interval and total limit to your storage.
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 2
-    resume_from_checkpoint: latest
-
-  tracker:
-    trackers: [wandb]
-    project_name: fastvideo_ltx2
-    run_name: ltx2_t2v_nvfp4_qat_finetune
-
-  model:
-    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
-    # so the default noise-minus-clean target reproduces the official
-    # unweighted masked-MSE (mask is all-ones for plain T2V).
-    precondition_outputs: false
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-  validation:
-    _target_: fastvideo.train.callbacks.validation.ValidationCallback
-    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
-    # REQUIRED: prompts to sample during training.
-    dataset_file: data/your_dataset_preprocessed/validation_prompts.json
-    every_steps: 250
-    # 8-step single-pass sampling matches the distilled checkpoint
-    # (validated in the LTX-2 overfit runs).
-    sampling_steps: [8]
-    guidance_scale: 1.0
-    num_frames: 81
-    # Validation reuses the live transformer and temporarily replaces its
-    # ATTN_QAT_TRAIN implementations with the sm120 inference kernel.
-    # Set false on GB200 (see header).
-    attn_qat_infer: true
-
-# quant_config selects NVFP4 QAT for the same curated linear prefixes as
-# LTX-2 NVFP4 deployment. They fake-quantize through real FP4 GEMMs with
-# an STE backward. The string resolves to NVFP4QATTrainConfig at parse.
-pipeline:
-  dit_config:
-    quant_config: nvfp4_qat_train

examples/train/configs/fine_tuning/ltx2/t2v.yaml

@@ -1,110 +0,0 @@
-# LTX-2 T2V full fine-tune.
-#
-# Starting point for fine-tuning LTX-2 on your own dataset with the
-# modular trainer. Model wiring (checkpoint, CFG rate, output
-# preconditioning, sharding) follows the validated
-# examples/train/configs/overfit_ltx2_t2v.yaml run; data paths, step
-# counts, and learning rate are placeholders you must adapt — the right
-# values depend on your dataset and are marked below.
-#
-# Prerequisites: preprocess your dataset into the trainer's latent format
-# (see docs/training/data_preprocess.md;
-# fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py is a minimal
-# LTX-2 reference to adapt).
-#
-# Run:
-#   NUM_GPUS=4 bash examples/train/run.sh \
-#     examples/train/configs/fine_tuning/ltx2/t2v.yaml
-
-models:
-  student:
-    _target_: fastvideo.train.models.ltx2.LTX2Model
-    # Validation settings below (8 steps, guidance 1.0) assume this
-    # distilled checkpoint.
-    init_from: FastVideo/LTX2-Distilled-Diffusers
-    trainable: true
-    enable_gradient_checkpointing_type: full
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  distributed:
-    # The 13B video branch needs sharding; 4-way HSDP shard matches the
-    # validated LTX-2 runs. Scale num_gpus/hsdp_shard_dim together.
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    # REQUIRED: path to your preprocessed dataset.
-    data_path: data/your_dataset_preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    # LTX2Model requires 0.0: CFG dropout would zero post-connector
-    # embeddings, which is not the model's unconditional input.
-    training_cfg_rate: 0.0
-    seed: 42
-    # Must match your preprocessed resolution/length.
-    # num_latent_t = (num_frames - 1) / 8 + 1.
-    num_latent_t: 11
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    # Carried from the validated LTX-2 overfit run as a starting point;
-    # tune for your dataset size and batch configuration.
-    learning_rate: 5.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    # Workload-dependent: set from your dataset size and target epochs.
-    max_train_steps: 2000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/ltx2_t2v_finetune
-    # A full training-state checkpoint is ~150GB for the 13B trainable
-    # video branch — size the interval and total limit to your storage.
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 2
-    resume_from_checkpoint: latest
-
-  tracker:
-    trackers: [wandb]
-    project_name: fastvideo_ltx2
-    run_name: ltx2_t2v_finetune
-
-  model:
-    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
-    # so the default noise-minus-clean target reproduces the official
-    # unweighted masked-MSE (mask is all-ones for plain T2V).
-    precondition_outputs: false
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-  validation:
-    _target_: fastvideo.train.callbacks.validation.ValidationCallback
-    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
-    # REQUIRED: prompts to sample during training.
-    dataset_file: data/your_dataset_preprocessed/validation_prompts.json
-    every_steps: 250
-    # 8-step single-pass sampling matches the distilled checkpoint
-    # (validated in the LTX-2 overfit runs).
-    sampling_steps: [8]
-    guidance_scale: 1.0
-    num_frames: 81
-
-# Required so the LTX2T2VConfig pipeline config is resolved from
-# init_from (without a `pipeline:` key the loader falls back to a
-# generic PipelineConfig and the LTX-2 DiT cannot be constructed).
-pipeline: {}

examples/train/configs/fine_tuning/ltx2/t2v_lora.yaml

@@ -1,114 +0,0 @@
-# LTX-2 T2V LoRA fine-tune.
-#
-# Same recipe as fine_tuning/ltx2/t2v.yaml, but only LoRA adapters on the
-# attention projections train — the base weights stay frozen, so
-# training-state checkpoints are adapter-sized instead of ~150GB.
-# See that file's header for prerequisites and placeholder notes.
-#
-# Run:
-#   NUM_GPUS=4 bash examples/train/run.sh \
-#     examples/train/configs/fine_tuning/ltx2/t2v_lora.yaml
-
-models:
-  student:
-    _target_: fastvideo.train.models.ltx2.LTX2Model
-    # Validation settings below (8 steps, guidance 1.0) assume this
-    # distilled checkpoint.
-    init_from: FastVideo/LTX2-Distilled-Diffusers
-    trainable: true
-    enable_gradient_checkpointing_type: full
-    lora:
-      enable: true
-      # rank/alpha and the attention-projection target set follow the
-      # Wan/Hunyuan LoRA recipes; targets match by substring against
-      # module names (LTX-2 blocks name theirs to_q/to_k/to_v/to_out).
-      rank: 16
-      alpha: 32
-      target_modules:
-        - to_q
-        - to_k
-        - to_v
-        - to_out
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  distributed:
-    # The 13B video branch needs sharding even with frozen base weights;
-    # 4-way HSDP shard matches the validated LTX-2 runs.
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    # REQUIRED: path to your preprocessed dataset.
-    data_path: data/your_dataset_preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    # LTX2Model requires 0.0: CFG dropout would zero post-connector
-    # embeddings, which is not the model's unconditional input.
-    training_cfg_rate: 0.0
-    seed: 42
-    # Must match your preprocessed resolution/length.
-    # num_latent_t = (num_frames - 1) / 8 + 1.
-    num_latent_t: 11
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    # LoRA-only training usually tolerates a higher LR than full
-    # fine-tuning (the Wan/Hunyuan LoRA recipes use 1.0e-4); tune for
-    # your dataset.
-    learning_rate: 1.0e-4
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    # Workload-dependent: set from your dataset size and target epochs.
-    max_train_steps: 2000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/ltx2_t2v_lora_finetune
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 2
-    resume_from_checkpoint: latest
-
-  tracker:
-    trackers: [wandb]
-    project_name: fastvideo_ltx2
-    run_name: ltx2_t2v_lora_finetune
-
-  model:
-    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
-    # so the default noise-minus-clean target reproduces the official
-    # unweighted masked-MSE (mask is all-ones for plain T2V).
-    precondition_outputs: false
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-  validation:
-    _target_: fastvideo.train.callbacks.validation.ValidationCallback
-    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
-    # REQUIRED: prompts to sample during training.
-    dataset_file: data/your_dataset_preprocessed/validation_prompts.json
-    every_steps: 250
-    # 8-step single-pass sampling matches the distilled checkpoint
-    # (validated in the LTX-2 overfit runs).
-    sampling_steps: [8]
-    guidance_scale: 1.0
-    num_frames: 81
-
-# Required so the LTX2T2VConfig pipeline config is resolved from
-# init_from (without a `pipeline:` key the loader falls back to a
-# generic PipelineConfig and the LTX-2 DiT cannot be constructed).
-pipeline: {}

examples/train/configs/fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml

@@ -1,153 +0,0 @@
-# LTX-2.3 T2V NVFP4 quantization-aware fine-tune (QAT).
-#
-# Same recipe as fine_tuning/ltx2_3/t2v.yaml, with the deployment-matched
-# NVFP4 quantization wiring carried from the validated LTX-2.0 QAT run
-# (examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml): the curated
-# attention/FFN linear set runs its forward through real NVFP4 GEMMs
-# (flashinfer cutlass) with a straight-through-estimator backward into
-# the dense fp32 master weights. No extra parameters or buffers are
-# added, so FSDP sharding, checkpointing, and export are identical to
-# dense training. Video attention uses ATTN_QAT_TRAIN for its quantized
-# forward and STE backward, then ATTN_QAT_INFER during validation.
-# Head-dim-64 audio attention and masked text attention remain dense.
-#
-# LTX-2.3 applicability: the curated linear prefix set
-# (fastvideo/layers/quantization/nvfp4_config.py) targets
-# ltx2.blocks.{0..47}.* — LTX-2.3 uses the same LTX-2 transformer class,
-# module naming, and 48-block depth, so the same 576 linears attach.
-# The set was curated against LTX-2.0 NVFP4 deployment; a validated
-# LTX-2.3 QAT overfit run exists (300 steps, MS-SSIM within the bf16
-# reference population — an overfit convergence check, not a
-# generalization claim). The 2.3
-# gated-attention to_gate_logits projections are not in the set and
-# remain dense.
-#
-# Validation-time ATTN_QAT_INFER is arch-aware:
-#   * sm_120 / sm_121: the bundled fastvideo-kernel CUTLASS extension.
-#   * sm_100 (GB200-class) / sm_103 (GB300-class): the FP4 FA4 kernel —
-#     requires the flash-attention-fp4 fork and its dependency pins; see
-#     docs/inference/optimizations.md.
-#   * any other GPU: falls back to flash attention with a logged receipt
-#     (validation still runs, just not with the quantized kernel).
-#
-# Prerequisites: preprocess your dataset into the trainer's latent format
-# (see docs/training/data_preprocess.md;
-# fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py is a minimal
-# LTX-2 reference to adapt — set
-# LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers). Data paths,
-# step counts, and learning rate below are placeholders you must adapt
-# to your workload.
-#
-# Run (multi-GPU; see the validation note above for kernel availability):
-#   NUM_GPUS=4 bash examples/train/run.sh \
-#     examples/train/configs/fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml
-#
-# `attn_qat_infer: true` below is the deployment-representative choice on
-# every arch with a quantized kernel available (see the validation note
-# above). Flip it off only when your GPU has no ATTN_QAT_INFER kernel and
-# you prefer validating with the train-time backend instead of the flash
-# fallback:
-#   NUM_GPUS=4 bash examples/train/run.sh \
-#     examples/train/configs/fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml \
-#     --callbacks.validation.attn_qat_infer false
-
-models:
-  student:
-    _target_: fastvideo.train.models.ltx2.LTX2Model
-    # Validation settings below (8 steps, guidance 1.0) assume this
-    # distilled checkpoint.
-    init_from: FastVideo/LTX-2.3-Distilled-Diffusers
-    trainable: true
-    enable_gradient_checkpointing_type: full
-    attention_backend: ATTN_QAT_TRAIN
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  distributed:
-    # The 13B video branch needs sharding; 4-way HSDP shard matches the
-    # validated LTX-2.3 runs. Scale num_gpus/hsdp_shard_dim together.
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    # REQUIRED: path to your preprocessed dataset.
-    data_path: data/your_dataset_preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    # LTX2Model requires 0.0: CFG dropout would zero post-connector
-    # embeddings, which is not the model's unconditional input.
-    training_cfg_rate: 0.0
-    seed: 42
-    # Must match your preprocessed resolution/length.
-    # num_latent_t = (num_frames - 1) / 8 + 1.
-    num_latent_t: 11
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    # Carried from the validated LTX-2.0 QAT overfit run as a starting
-    # point; tune for your dataset size and batch configuration.
-    learning_rate: 5.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    # Workload-dependent: set from your dataset size and target epochs.
-    max_train_steps: 2000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/ltx2_3_t2v_nvfp4_qat_finetune
-    # A full training-state checkpoint is ~150GB for the 13B trainable
-    # video branch — size the interval and total limit to your storage.
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 2
-    resume_from_checkpoint: latest
-
-  tracker:
-    trackers: [wandb]
-    project_name: fastvideo_ltx2
-    run_name: ltx2_3_t2v_nvfp4_qat_finetune
-
-  model:
-    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
-    # so the default noise-minus-clean target reproduces the official
-    # unweighted masked-MSE (mask is all-ones for plain T2V).
-    precondition_outputs: false
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-  validation:
-    _target_: fastvideo.train.callbacks.validation.ValidationCallback
-    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
-    # REQUIRED: prompts to sample during training.
-    dataset_file: data/your_dataset_preprocessed/validation_prompts.json
-    every_steps: 250
-    # 8-step single-pass sampling matches the distilled checkpoint
-    # (validated in the LTX-2.3 overfit runs).
-    sampling_steps: [8]
-    guidance_scale: 1.0
-    num_frames: 81
-    # Validation reuses the live transformer and temporarily replaces its
-    # ATTN_QAT_TRAIN implementations with the sm120 inference kernel.
-    # Set false on GB200 (see header).
-    attn_qat_infer: true
-
-# quant_config selects NVFP4 QAT for the same curated linear prefixes as
-# LTX-2 NVFP4 deployment (see the LTX-2.3 applicability note in the
-# header). They fake-quantize through real FP4 GEMMs with an STE
-# backward. The string resolves to NVFP4QATTrainConfig at parse.
-pipeline:
-  dit_config:
-    quant_config: nvfp4_qat_train

examples/train/configs/fine_tuning/ltx2_3/t2v.yaml

@@ -1,112 +0,0 @@
-# LTX-2.3 T2V full fine-tune.
-#
-# Starting point for fine-tuning LTX-2.3 (gated attention, cross-attn
-# AdaLN, 4096-d post-connector text embeddings, no in-DiT caption
-# projection) on your own dataset with the modular trainer. Model wiring
-# (checkpoint, CFG rate, output preconditioning, sharding) follows the
-# validated examples/train/configs/overfit_ltx2_3_t2v.yaml run; data
-# paths, step counts, and learning rate are placeholders you must
-# adapt — the right values depend on your dataset and are marked below.
-#
-# Prerequisites: preprocess your dataset into the trainer's latent format
-# (see docs/training/data_preprocess.md;
-# fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py is a minimal
-# LTX-2 reference to adapt — set
-# LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers).
-#
-# Run:
-#   NUM_GPUS=4 bash examples/train/run.sh \
-#     examples/train/configs/fine_tuning/ltx2_3/t2v.yaml
-
-models:
-  student:
-    _target_: fastvideo.train.models.ltx2.LTX2Model
-    # Validation settings below (8 steps, guidance 1.0) assume this
-    # distilled checkpoint.
-    init_from: FastVideo/LTX-2.3-Distilled-Diffusers
-    trainable: true
-    enable_gradient_checkpointing_type: full
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  distributed:
-    # The 13B video branch needs sharding; 4-way HSDP shard matches the
-    # validated LTX-2.3 runs. Scale num_gpus/hsdp_shard_dim together.
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    # REQUIRED: path to your preprocessed dataset.
-    data_path: data/your_dataset_preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    # LTX2Model requires 0.0: CFG dropout would zero post-connector
-    # embeddings, which is not the model's unconditional input.
-    training_cfg_rate: 0.0
-    seed: 42
-    # Must match your preprocessed resolution/length.
-    # num_latent_t = (num_frames - 1) / 8 + 1.
-    num_latent_t: 11
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    # Carried from the validated LTX-2.3 overfit run as a starting point;
-    # tune for your dataset size and batch configuration.
-    learning_rate: 5.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    # Workload-dependent: set from your dataset size and target epochs.
-    max_train_steps: 2000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/ltx2_3_t2v_finetune
-    # A full training-state checkpoint is ~150GB for the 13B trainable
-    # video branch — size the interval and total limit to your storage.
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 2
-    resume_from_checkpoint: latest
-
-  tracker:
-    trackers: [wandb]
-    project_name: fastvideo_ltx2
-    run_name: ltx2_3_t2v_finetune
-
-  model:
-    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
-    # so the default noise-minus-clean target reproduces the official
-    # unweighted masked-MSE (mask is all-ones for plain T2V).
-    precondition_outputs: false
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-  validation:
-    _target_: fastvideo.train.callbacks.validation.ValidationCallback
-    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
-    # REQUIRED: prompts to sample during training.
-    dataset_file: data/your_dataset_preprocessed/validation_prompts.json
-    every_steps: 250
-    # 8-step single-pass sampling matches the distilled checkpoint
-    # (validated in the LTX-2.3 overfit runs).
-    sampling_steps: [8]
-    guidance_scale: 1.0
-    num_frames: 81
-
-# Required so the LTX2T2VConfig pipeline config is resolved from
-# init_from (without a `pipeline:` key the loader falls back to a
-# generic PipelineConfig and the LTX-2 DiT cannot be constructed).
-pipeline: {}

examples/train/configs/fine_tuning/ltx2_3/t2v_lora.yaml

@@ -1,116 +0,0 @@
-# LTX-2.3 T2V LoRA fine-tune.
-#
-# Same recipe as fine_tuning/ltx2_3/t2v.yaml, but only LoRA adapters on
-# the attention projections train — the base weights stay frozen, so
-# training-state checkpoints are adapter-sized instead of ~150GB.
-# See that file's header for prerequisites and placeholder notes.
-#
-# Run:
-#   NUM_GPUS=4 bash examples/train/run.sh \
-#     examples/train/configs/fine_tuning/ltx2_3/t2v_lora.yaml
-
-models:
-  student:
-    _target_: fastvideo.train.models.ltx2.LTX2Model
-    # Validation settings below (8 steps, guidance 1.0) assume this
-    # distilled checkpoint.
-    init_from: FastVideo/LTX-2.3-Distilled-Diffusers
-    trainable: true
-    enable_gradient_checkpointing_type: full
-    lora:
-      enable: true
-      # rank/alpha and the attention-projection target set follow the
-      # Wan/Hunyuan LoRA recipes; targets match by substring against
-      # module names (LTX-2.3 blocks name theirs to_q/to_k/to_v/to_out;
-      # the 2.3 gated-attention to_gate_logits projection is not
-      # targeted).
-      rank: 16
-      alpha: 32
-      target_modules:
-        - to_q
-        - to_k
-        - to_v
-        - to_out
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  distributed:
-    # The 13B video branch needs sharding even with frozen base weights;
-    # 4-way HSDP shard matches the validated LTX-2.3 runs.
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    # REQUIRED: path to your preprocessed dataset.
-    data_path: data/your_dataset_preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    # LTX2Model requires 0.0: CFG dropout would zero post-connector
-    # embeddings, which is not the model's unconditional input.
-    training_cfg_rate: 0.0
-    seed: 42
-    # Must match your preprocessed resolution/length.
-    # num_latent_t = (num_frames - 1) / 8 + 1.
-    num_latent_t: 11
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    # LoRA-only training usually tolerates a higher LR than full
-    # fine-tuning (the Wan/Hunyuan LoRA recipes use 1.0e-4); tune for
-    # your dataset.
-    learning_rate: 1.0e-4
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    # Workload-dependent: set from your dataset size and target epochs.
-    max_train_steps: 2000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/ltx2_3_t2v_lora_finetune
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 2
-    resume_from_checkpoint: latest
-
-  tracker:
-    trackers: [wandb]
-    project_name: fastvideo_ltx2
-    run_name: ltx2_3_t2v_lora_finetune
-
-  model:
-    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
-    # so the default noise-minus-clean target reproduces the official
-    # unweighted masked-MSE (mask is all-ones for plain T2V).
-    precondition_outputs: false
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-  validation:
-    _target_: fastvideo.train.callbacks.validation.ValidationCallback
-    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
-    # REQUIRED: prompts to sample during training.
-    dataset_file: data/your_dataset_preprocessed/validation_prompts.json
-    every_steps: 250
-    # 8-step single-pass sampling matches the distilled checkpoint
-    # (validated in the LTX-2.3 overfit runs).
-    sampling_steps: [8]
-    guidance_scale: 1.0
-    num_frames: 81
-
-# Required so the LTX2T2VConfig pipeline config is resolved from
-# init_from (without a `pipeline:` key the loader falls back to a
-# generic PipelineConfig and the LTX-2 DiT cannot be constructed).
-pipeline: {}

examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml

@@ -1,123 +0,0 @@
-# LTX-2 T2V overfitting test config with NVFP4 quantization-aware
-# training (QAT).
-#
-# Same recipe as overfit_ltx2_t2v.yaml, but the curated attention/FFN
-# linear set used by LTX-2 NVFP4 deployment runs its forward through
-# real NVFP4 GEMMs (flashinfer cutlass) with a straight-through-estimator
-# backward into the dense fp32 master weights. No extra parameters or
-# buffers are added, so FSDP sharding, checkpointing, and export are
-# identical to dense training. Video attention uses ATTN_QAT_TRAIN for its
-# quantized forward and STE backward, then ATTN_QAT_INFER during validation.
-# Head-dim-64 audio attention and masked text attention remain dense.
-#
-# Validation-time ATTN_QAT_INFER is arch-aware:
-#   * sm_120a/sm_121a: fastvideo-kernel CUTLASS extension -- the exact
-#     quantization scheme ATTN_QAT_TRAIN simulates.
-#   * sm_100a (GB200) / sm_103a (GB300): FP4 FA4 kernel
-#     (github.com/hao-ai-lab/flash-attention-fp4, branch fp4; per-16
-#     block-scaled NVFP4 Q/K, BF16 P/V; validated install set:
-#     nvidia-cutlass-dsl==4.4.2, quack-kernels==0.4.1,
-#     flashinfer-python==0.6.8, FASTVIDEO_FA4=1 — see
-#     docs/inference/optimizations.md). This scheme DIFFERS from the
-#     CUTLASS one the training simulation matches, so sm_100/sm_103
-#     validation and deployment carry a train-sim mismatch -- gate quality
-#     by MS-SSIM measurement rather than assuming parity. The resolution
-#     receipt ("ATTN_QAT_INFER resolved: ...") records arch + scheme.
-#
-# Preprocess data first (same data as the bf16 overfit):
-#   CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
-#
-# Run on multi-GPU sm_120 hardware:
-#   NUM_GPUS=4 \
-#     bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml
-#
-# On GB200 without flash-attention-fp4 installed (or any other arch with no
-# ATTN_QAT_INFER kernel), disable only the validation-time swap:
-#   NUM_GPUS=4 \
-#     bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml \
-#       --callbacks.validation.attn_qat_infer false
-
-models:
-  student:
-    _target_: fastvideo.train.models.ltx2.LTX2Model
-    init_from: FastVideo/LTX2-Distilled-Diffusers
-    trainable: true
-    enable_gradient_checkpointing_type: full
-    attention_backend: ATTN_QAT_TRAIN
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  distributed:
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    data_path: data/ltx2_overfit_preprocessed
-    dataloader_num_workers: 0
-    train_batch_size: 1
-    # LTX2Model requires 0.0: CFG dropout would zero post-connector
-    # embeddings, which is not the model's unconditional input.
-    training_cfg_rate: 0.0
-    seed: 42
-    num_latent_t: 11  # (81 - 1) / 8 + 1
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    learning_rate: 5.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 300
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/ltx2_overfit_nvfp4_attn_qat
-    # A full training-state checkpoint is ~150GB for the 13B trainable
-    # video branch; disable saves for the overfit smoke run.
-    training_state_checkpointing_steps: 0
-    checkpoints_total_limit: 1
-
-  tracker:
-    trackers: [wandb]
-    project_name: fastvideo_ltx2
-    run_name: ltx2_overfit_nvfp4_attn_qat
-
-  model:
-    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
-    # so the default noise-minus-clean target reproduces the official
-    # unweighted masked-MSE (mask is all-ones for plain T2V).
-    precondition_outputs: false
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-  validation:
-    _target_: fastvideo.train.callbacks.validation.ValidationCallback
-    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
-    dataset_file: data/ltx2_overfit_preprocessed/validation_prompts.json
-    every_steps: 50
-    sampling_steps: [8]
-    guidance_scale: 1.0
-    num_frames: 81
-    # Validation reuses the live transformer and temporarily replaces its
-    # ATTN_QAT_TRAIN implementations with the sm120 inference kernel.
-    attn_qat_infer: true
-
-# quant_config selects NVFP4 QAT for the same curated linear prefixes as
-# LTX-2 NVFP4 deployment. They fake-quantize through real FP4 GEMMs with
-# an STE backward. The string resolves to NVFP4QATTrainConfig at parse.
-pipeline:
-  dit_config:
-    quant_config: nvfp4_qat_train

fastvideo-kernel/README.md

@@ -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 `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as
+commit `82d6441eec5d4dfec120153db2c0145ae855a083`, 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 "nvidia-cutlass-dsl>=4.5.0" torchvision
-pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"
+pip install torchvision
+pip install "flash-attn-4[cu13] @ git+https://github.com/Dao-AILab/flash-attention.git@82d6441eec5d4dfec120153db2c0145ae855a083#subdirectory=flash_attn/cute"
 ```
 
 The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against

fastvideo-kernel/build.sh

@@ -4,7 +4,6 @@ set -ex
 # Simple build script wrapping uv/pip
 # Usage:
 #   ./build.sh  # local build (torch-based arch detection, TK only on SM90)
-#   ./build.sh --wheel-dir /tmp/wheels  # build + install a reusable wheel
 # Environment overrides (if set, they win over auto-detection):
 #   TORCH_CUDA_ARCH_LIST
 #   CMAKE_ARGS (for FASTVIDEO_KERNEL_BUILD_TK / CMAKE_CUDA_ARCHITECTURES / GPU_BACKEND)
@@ -68,25 +67,12 @@ fi
 # Install build dependencies
 uv pip install scikit-build-core cmake ninja
 
+RELEASE=0
 GPU_BACKEND=CUDA
-WHEEL_DIR=""
-while [[ $# -gt 0 ]]; do
-    case "$1" in
+for arg in "$@"; do
+    case "$arg" in
         --rocm)
             GPU_BACKEND=ROCM
-            shift
-            ;;
-        --wheel-dir)
-            if [[ $# -lt 2 ]]; then
-                echo "ERROR: --wheel-dir requires a directory path" >&2
-                exit 2
-            fi
-            WHEEL_DIR="$2"
-            shift 2
-            ;;
-        *)
-            echo "ERROR: unknown argument: $1" >&2
-            exit 2
             ;;
     esac
 done
@@ -179,15 +165,4 @@ echo "CMAKE_ARGS: ${CMAKE_ARGS:-<unset>}"
 echo "GPU_BACKEND: ${GPU_BACKEND:-<unset>}"
 # Build and install
 # Use -v for verbose output
-if [ -n "${WHEEL_DIR}" ]; then
-    mkdir -p "${WHEEL_DIR}"
-    uv build --wheel -v --no-build-isolation --out-dir "${WHEEL_DIR}" .
-    wheel_path="$(find "${WHEEL_DIR}" -maxdepth 1 -type f \( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) | sort | tail -n 1)"
-    if [ -z "${wheel_path}" ]; then
-        echo "ERROR: no fastvideo-kernel wheel was built in ${WHEEL_DIR}" >&2
-        exit 1
-    fi
-    uv pip install "${wheel_path}" --reinstall-package fastvideo-kernel --no-deps
-else
-    uv pip install . -v --no-build-isolation
-fi
+uv pip install . -v --no-build-isolation

fastvideo/attention/backends/attn_qat_infer.py

@@ -50,143 +50,8 @@ def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None:
     return _attn_qat_infer
 
 
-# Consumer-Blackwell compute capabilities the modified SageAttention3 FP4
-# kernel is compiled for (sm_120a / sm_121a -- see fastvideo-kernel/README.md).
-_SUPPORTED_DEVICE_CAPABILITIES = frozenset({(12, 0), (12, 1)})
-
-# Datacenter-Blackwell capabilities served by the FP4 FA4 kernel
-# (flash-attention-fp4 @ fp4, sm_100a/sm_103a) through #1221's plumbing:
-# per-16 block-scaled NVFP4 Q/K (E4M3 scale factors), BF16 P/V. This is a
-# DIFFERENT quantization scheme from the sm_12x CUTLASS extension above --
-# ATTN_QAT_TRAIN simulates the CUTLASS scheme, so sm_100/sm_103 deployment
-# carries a train-sim mismatch that is measured (MS-SSIM gate), not assumed.
-_FA4_FP4_CAPABILITIES = frozenset({(10, 0), (10, 3)})
-
-# The fork is written against the cutlass-dsl 4.4 API surface; the validated
-# install set (GB200-proven) is nvidia-cutlass-dsl==4.4.2 +
-# nvidia-cutlass-dsl-libs-base==4.4.2 + quack-kernels==0.4.1 +
-# flashinfer-python==0.6.8, with the fork on PYTHONPATH,
-# CUTE_DSL_ENABLE_TVM_FFI=1, and FASTVIDEO_FA4=1 (the fork ships no compiled
-# FA2, so dense attention paths need the FA4 opt-in). dsl 4.6-era installs
-# fail at CuTe JIT trace (cute.make_fragment was removed at module level).
-_FA4_INSTALL_HINT = ("install flash-attention-fp4 (branch fp4) from "
-                     "https://github.com/hao-ai-lab/flash-attention-fp4 with "
-                     "nvidia-cutlass-dsl==4.4.2, quack-kernels==0.4.1, "
-                     "flashinfer-python==0.6.8 and FASTVIDEO_FA4=1; "
-                     "see docs/inference/optimizations.md")
-
-_fa4_fp4_import_ok: bool | None = None
-
-
-def _fa4_fp4_available() -> bool:
-    """flash_attn.cute (FA4) import probe, cached. Reuses #1221's guarded
-    import chain in fastvideo.attention.utils.flash_attn_cute (which maps
-    cutlass-dsl version skew to ImportError with a loud warning)."""
-    global _fa4_fp4_import_ok
-    if _fa4_fp4_import_ok is None:
-        try:
-            from fastvideo.attention.utils.flash_attn_cute import (  # noqa: F401
-                flash_attn_fp4_func, )
-            _fa4_fp4_import_ok = True
-        except ImportError:
-            _fa4_fp4_import_ok = False
-    return _fa4_fp4_import_ok
-
-
-def _active_capability() -> tuple[int, int] | None:
-    if not torch.cuda.is_available():
-        return None
-    try:
-        return tuple(torch.cuda.get_device_capability())
-    except Exception:  # pragma: no cover - defensive: never break backend selection
-        return None
-
-
-def _resolved_kernel() -> str | None:
-    """Which ATTN_QAT_INFER kernel serves the active device, or None.
-
-    Per-arch resolution (single source of truth -- extend the capability
-    sets above, do not add equality checks elsewhere):
-      * sm_12x consumer Blackwell -> fastvideo-kernel CUTLASS extension
-        (modified SageAttention3 FP4).
-      * sm_100a/sm_103a datacenter Blackwell -> FP4 FA4 (flash-attention-fp4)
-        via the merged #1221 plumbing.
-    """
-    cap = _active_capability()
-    if cap in _SUPPORTED_DEVICE_CAPABILITIES and _get_attn_qat_infer() is not None:
-        return "cutlass_sm12x"
-    if cap in _FA4_FP4_CAPABILITIES and _fa4_fp4_available():
-        return "fa4_fp4"
-    return None
-
-
-def attn_qat_infer_receipt() -> str:
-    """One-line receipt of the resolution decision (arch + kernel + quant
-    knobs), for the selection log and for tooling. The FA4 knobs are the
-    repo's tuned defaults passed through verbatim: qk_mode=nvfp4
-    (per-16 E4M3 SFs), pv_mode=bf16 -- see flash_attn/cute/README.md in
-    the kernel repo."""
-    cap = _active_capability()
-    arch = f"sm_{cap[0]}{cap[1]}" if cap is not None else "no-cuda"
-    kernel = _resolved_kernel()
-    if kernel == "cutlass_sm12x":
-        return f"arch={arch} kernel=fastvideo-kernel-cutlass scheme=sage3-fp4-sm120"
-    if kernel == "fa4_fp4":
-        return (f"arch={arch} kernel=flash-attention-fp4 qk_mode=nvfp4(per-16-e4m3-sf) "
-                f"pv_mode=bf16 train_sim_mismatch=measured")
-    supported = "sm_120a/sm_121a via fastvideo-kernel build.sh; sm_100a/sm_103a via flash-attention-fp4"
-    if cap is not None and cap in _FA4_FP4_CAPABILITIES:
-        return f"arch={arch} kernel=none (flash_attn.cute not importable -- {_FA4_INSTALL_HINT})"
-    return f"arch={arch} kernel=none (supported: {supported})"
-
-
-_FA4_ROUTE_OPS: tuple | None = None
-
-
-def _import_fa4_route_ops() -> tuple:
-    """Slow path (own function so tests pin it runs once per process):
-    resolves the FA4 quantize helper and kernel entry point."""
-    from fastvideo.attention.backends.flash_attn import (
-        _nvfp4_quantize_for_fa4, )
-    from fastvideo.attention.utils.flash_attn_cute import (
-        flash_attn_fp4_func, )
-    return (_nvfp4_quantize_for_fa4, flash_attn_fp4_func)
-
-
-def _resolve_fa4_route_ops() -> tuple:
-    # Lazy but memoized: per-forward resolution graph-breaks dynamo every
-    # step and blocks fullgraph compilation of the NVFP4 path.
-    global _FA4_ROUTE_OPS
-    if _FA4_ROUTE_OPS is None:
-        _FA4_ROUTE_OPS = _import_fa4_route_ops()
-    return _FA4_ROUTE_OPS
-
-
-_receipt_logged = False
-
-
-def _log_receipt_once() -> None:
-    # One line per process, not per layer (the validation swap constructs
-    # one impl per attention layer).
-    global _receipt_logged
-    if not _receipt_logged:
-        _receipt_logged = True
-        logger.info("ATTN_QAT_INFER resolved: %s", attn_qat_infer_receipt())
-
-
 def is_attn_qat_infer_available() -> bool:
-    """True only when the active device has a built ATTN_QAT_INFER kernel.
-
-    The import check alone is not sufficient: CUDA 13 wheel builds can
-    carry the sm_12x extension on any host (e.g. H100, GB200), where the
-    import succeeds, backend selection picks this backend, and the first
-    kernel call then fails with an unsupported-capability error instead of
-    ever reaching the documented FlashAttention fallback in
-    fastvideo.platforms.cuda. Gating on the active device's capability
-    keeps that fallback working on every unsupported GPU, while
-    sm_100a/sm_103a now resolve to the FP4 FA4 kernel (#1221).
-    """
-    return _resolved_kernel() is not None
+    return _get_attn_qat_infer() is not None
 
 
 class AttnQatInferBackend(AttentionBackend):
@@ -232,12 +97,6 @@ class AttnQatInferImpl(AttentionImpl[AttentionMetadata]):
         if dropout_p > 0:
             raise NotImplementedError(f"attn_qat_infer does not support dropout (got dropout_p={dropout_p}). "
                                       "The QAT inference kernel applies no stochastic dropout.")
-        # Kernel resolution is per-forward, not per-construction: callers
-        # (the validation swap, backend selection) gate on
-        # is_attn_qat_infer_available() first, and constructing an impl on a
-        # host without the kernel must stay legal (pre-existing contract the
-        # validation-swap test pins).
-        _log_receipt_once()
 
     def forward(
         self,
@@ -246,18 +105,10 @@ class AttnQatInferImpl(AttentionImpl[AttentionMetadata]):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
-        # Dispatch on the single per-arch resolution: importability of the
-        # bundled sm_12x extension is NOT sufficient (CUDA 13 wheels carry it
-        # on unsupported hosts, where calling it is the wrong binary).
-        kernel = _resolved_kernel()
-        if kernel == "fa4_fp4":
-            return self._forward_fa4_fp4(query, key, value)
-        if kernel is None:
-            raise ImportError(f"attn_qat_infer is not available ({attn_qat_infer_receipt()}). "
-                              "Please ensure an ATTN_QAT_INFER kernel is installed for this device.")
-
         attn_qat_infer = _get_attn_qat_infer()
-        assert attn_qat_infer is not None  # kernel == "cutlass_sm12x" implies the import succeeded
+        if attn_qat_infer is None:
+            raise ImportError("attn_qat_infer is not available. Please ensure the "
+                              "attn_qat_infer kernel package is installed.")
 
         query = query.transpose(1, 2).contiguous()
         key = key.transpose(1, 2).contiguous()
@@ -272,39 +123,3 @@ class AttnQatInferImpl(AttentionImpl[AttentionMetadata]):
             sm_scale=self.softmax_scale,
         )
         return output.transpose(1, 2).contiguous()
-
-    def _forward_fa4_fp4(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-    ) -> torch.Tensor:
-        """sm_100a/sm_103a path: FP4 FA4 with the repo's tuned defaults
-        (NVFP4 per-16 block-scaled Q/K, BF16 V) -- mirrors
-        FlashAttentionImpl._forward_nvfp4 (#1221). Inputs/outputs are
-        (batch, seqlen, nheads, headdim); no transpose."""
-        _nvfp4_quantize_for_fa4, flash_attn_fp4_func = _resolve_fa4_route_ops()
-
-        orig_seqlen_q = query.shape[1]
-        orig_seqlen_k = key.shape[1]
-
-        q_fp4, q_sf = _nvfp4_quantize_for_fa4(query)
-        k_fp4, k_sf = _nvfp4_quantize_for_fa4(key)
-
-        # FP4/SF buffers are padded to a 128 multiple; FA4 masks to the
-        # original lengths so padding never biases the softmax.
-        q_fp4 = q_fp4[:, :orig_seqlen_q]
-        k_fp4 = k_fp4[:, :orig_seqlen_k]
-
-        output = flash_attn_fp4_func(
-            q_fp4,
-            k_fp4,
-            value,
-            q_sf,
-            k_sf,
-            softmax_scale=self.softmax_scale,
-            causal=self.causal,
-        )
-        if isinstance(output, tuple):
-            output = output[0]
-        return output

fastvideo/attention/backends/flash_attn.py

@@ -32,30 +32,8 @@ except ImportError:
     flash_attn_fp4_func = None
     _FA4_FP4_AVAILABLE = False
 
-_FA4_QUANT_OPS: tuple | None = None
 
-
-def _import_fa4_quant_ops() -> tuple:
-    """The slow path: resolves flashinfer's FP4 quantization entry points
-    (imports + JIT-module lookup, which probes the CUDA toolchain via a
-    subprocess on first use). Kept as its own function so tests can pin
-    that it runs at most once per process."""
-    from flashinfer.quantization import SfLayout, nvfp4_quantize
-    return (nvfp4_quantize, SfLayout)
-
-
-def _resolve_fa4_quant_ops() -> tuple:
-    # Lazy (construct-anywhere/fail-at-forward stays intact) but memoized:
-    # re-resolving per forward graph-breaks dynamo every step (making
-    # fullgraph compilation impossible for the NVFP4 path) and keeps eager
-    # dispatch overhead on the hot path.
-    global _FA4_QUANT_OPS
-    if _FA4_QUANT_OPS is None:
-        _FA4_QUANT_OPS = _import_fa4_quant_ops()
-    return _FA4_QUANT_OPS
-
-
-def _nvfp4_quantize_for_fa4_impl(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]:
     """Quantize a (batch, seqlen, nheads, headdim) BF16 tensor to FP4.
 
     Returns:
@@ -64,7 +42,7 @@ def _nvfp4_quantize_for_fa4_impl(tensor_4d: torch.Tensor) -> tuple[torch.Tensor,
             Caller should slice [:, :orig_seqlen] before passing to FA4.
         sf_tensor:  torch.uint8, shape (32, 4, rest_m, 4, rest_k, nheads, batch) with stride[3]=1
     """
-    nvfp4_quantize, SfLayout = _resolve_fa4_quant_ops()
+    from flashinfer.quantization import nvfp4_quantize, SfLayout
 
     batch, seqlen, nheads, headdim = tensor_4d.shape
     sf_vec_size = 16
@@ -104,43 +82,6 @@ def _nvfp4_quantize_for_fa4_impl(tensor_4d: torch.Tensor) -> tuple[torch.Tensor,
     return fp4_tensor, sf_mma
 
 
-# Dynamo boundary for the FP4 quantize path (same pattern as the masked
-# flash-attention entry points in flash_attn_no_pad.py): tracing a python
-# body that resolves flashinfer's JIT module descends into its toolchain
-# probe (a subprocess) regardless of any runtime memoization -- a cache hit
-# is invisible at trace time. Registering the whole quantize step as a
-# custom op makes it one opaque graph node, unlocking
-# torch.compile(fullgraph=True) for the NVFP4 attention path.
-@torch.library.custom_op(
-    "fastvideo::nvfp4_quantize_fa4",
-    mutates_args=(),
-    device_types="cuda",
-)
-def _nvfp4_quantize_fa4_op(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    return _nvfp4_quantize_for_fa4_impl(tensor_4d)
-
-
-@torch.library.register_fake("fastvideo::nvfp4_quantize_fa4")
-def _nvfp4_quantize_fa4_fake(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    batch, seqlen, nheads, headdim = tensor_4d.shape
-    seqlen_padded = (seqlen + 127) // 128 * 128
-    fp4 = tensor_4d.new_empty((batch, seqlen_padded, nheads, headdim // 2), dtype=torch.float4_e2m1fn_x2)
-    rest_m = seqlen_padded // 128
-    rest_k = (headdim // 16) // 4
-    # Must reproduce the impl's output STRIDES, not just its shape: the real
-    # sf is a permuted view of a contiguous (batch, nheads, rest_m, rest_k,
-    # 32, 4, 4) buffer, and torch.compile bakes the fake's strides into the
-    # generated code (a contiguous fake here asserts at runtime).
-    sf = tensor_4d.new_empty((batch, nheads, rest_m, rest_k, 32, 4, 4), dtype=torch.uint8).permute(4, 5, 2, 6, 3, 1, 0)
-    return fp4, sf
-
-
-def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize (batch, seqlen, nheads, headdim) BF16 to FP4 via the custom
-    op boundary; see _nvfp4_quantize_for_fa4_impl for the layout contract."""
-    return torch.ops.fastvideo.nvfp4_quantize_fa4(tensor_4d)
-
-
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 

fastvideo/attention/selector.py

@@ -150,15 +150,14 @@ def _cached_get_attn_backend(
     from fastvideo.platforms import current_platform
 
     if (selected_backend is not None and selected_backend not in supported_attention_backends):
-        fallback_backend = (default_backend if default_backend in supported_attention_backends else None)
         logger.warning(
             "Requested attention backend %s is not supported by this "
-            "layer; supported backends are %s. Falling back to %s.",
+            "layer; supported backends are %s. Falling back to automatic "
+            "selection.",
             selected_backend.name,
             [b.name for b in supported_attention_backends],
-            fallback_backend.name if fallback_backend is not None else "automatic selection",
         )
-        selected_backend = fallback_backend
+        selected_backend = None
     attention_cls = current_platform.get_attn_backend_cls(selected_backend, head_size, dtype)
     if not attention_cls:
         raise ValueError(f"Invalid attention backend for {current_platform.device_name}")

fastvideo/configs/models/dits/kandinsky5.py

@@ -13,15 +13,12 @@ def _is_kandinsky5_transformer_block(n: str, m) -> bool:
 class Kandinsky5ArchConfig(DiTArchConfig):
     _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_kandinsky5_transformer_block])
 
-    # NABLA block-sparse attention for attention_type="nabla" checkpoints, the
-    # dense backends every DiT supports, and the Attn-QAT backends needed for
-    # quantization-aware finetuning/distillation (fastvideo/train/models/kandinsky5/).
+    # NABLA block-sparse attention for attention_type="nabla" checkpoints, plus
+    # the dense backends every DiT supports.
     _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
         AttentionBackendEnum.NABLA_ATTN,
         AttentionBackendEnum.FLASH_ATTN,
         AttentionBackendEnum.TORCH_SDPA,
-        AttentionBackendEnum.ATTN_QAT_INFER,
-        AttentionBackendEnum.ATTN_QAT_TRAIN,
     )
 
     # Native FastVideo implementation uses the same parameter names as diffusers

fastvideo/configs/models/encoders/reason1.py

@@ -7,18 +7,6 @@ from typing import Any
 from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig
 
 
-def _is_transformer_layer(n: str, m) -> bool:
-    return "layers" in n and n.split(".")[-1].isdigit()
-
-
-def _is_embeddings(n: str, m) -> bool:
-    return n.endswith("embed_tokens")
-
-
-def _is_final_norm(n: str, m) -> bool:
-    return n.endswith("norm")
-
-
 @dataclass
 class Reason1ArchConfig(TextEncoderArchConfig):
     """Architecture settings (defaults match Qwen2.5-VL-7B-Instruct)."""
@@ -73,8 +61,6 @@ class Reason1ArchConfig(TextEncoderArchConfig):
 
     torch_dtype: str = "bfloat16"
     _attn_implementation: str = "flash_attention_2"
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm])
 
 
 @dataclass

fastvideo/configs/pipelines/__init__.py

@@ -6,7 +6,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
-from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig, Kandinsky5I2VConfig, Kandinsky5T2VConfig
+from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
 from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
 from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
@@ -21,6 +21,6 @@ __all__ = [
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
     "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
-    "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig",
-    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "Kandinsky5I2VConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig",
+    "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/kandinsky5.py

@@ -120,21 +120,3 @@ class Kandinsky5I2VConfig(Kandinsky5T2VConfig):
         super().__post_init__()
         # I2V needs the VAE encoder to encode the conditioning image.
         self.vae_config.load_encoder = True
-
-
-@dataclass
-class Kandinsky5DMDConfig(Kandinsky5T2VConfig):
-    """Kandinsky-5.0 DMD (few-step distilled) text-to-video pipeline configuration.
-
-    Checkpoints exported by ``fastvideo.train.entrypoint.dcp_to_diffusers``
-    copy their base T2V checkpoint's ``model_index.json`` unchanged, so
-    ``_class_name`` still says the base T2V pipeline and the registry
-    (``fastvideo/registry.py``) cannot auto-detect a DMD export -- pass this
-    config together with ``override_pipeline_cls_name="Kandinsky5DMDPipeline"``
-    explicitly to ``VideoGenerator.from_pretrained``/``from_config`` (see
-    ``examples/train/configs/fine_tuning/kandinsky5/README.md``). Without it,
-    ``Kandinsky5T2VConfig``'s ``dmd_denoising_steps=None`` makes
-    ``Kandinsky5DmdDenoisingStage`` raise immediately.
-    """
-
-    dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250])

fastvideo/entrypoints/cli/main.py

@@ -8,7 +8,6 @@ from fastvideo.entrypoints.cli.router_serve import (
 from fastvideo.entrypoints.cli.serve import cmd_init as serve_cmd_init
 from fastvideo.entrypoints.cli.bench import cmd_init as bench_cmd_init
 from fastvideo.entrypoints.cli.eval import cmd_init as eval_cmd_init
-from fastvideo.version import __version__
 
 
 def cmd_init() -> list[CLISubcommand]:
@@ -24,7 +23,7 @@ def cmd_init() -> list[CLISubcommand]:
 
 def main() -> None:
     parser = FlexibleArgumentParser(description="FastVideo CLI")
-    parser.add_argument('-v', '--version', action='version', version=__version__)
+    parser.add_argument('-v', '--version', action='version', version='0.1.0')
 
     subparsers = parser.add_subparsers(required=False, dest="subparser")
 

fastvideo/layers/quantization/fp8_config.py

@@ -26,9 +26,6 @@ FP8_DTYPE = torch.float8_e4m3fn
 FP8_MAX = float(torch.finfo(FP8_DTYPE).max)  # 448.0
 FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0)
 
-# Wan's "to_q"/"to_k"/"to_v" also substring-match Kandinsky5's "to_query"/
-# "to_key"/"to_value", so only Kandinsky5's out-projection and FFN names need
-# to be listed explicitly below.
 _FP8_SUFFIXES = (
     "ffn.fc_in",
     "ffn.fc_out",
@@ -36,11 +33,6 @@ _FP8_SUFFIXES = (
     "to_k",
     "to_v",
     "to_out",
-    # Kandinsky5
-    "self_attention.out_layer",
-    "cross_attention.out_layer",
-    "feed_forward.mlp.fc_in",
-    "feed_forward.mlp.fc_out",
 )
 
 

fastvideo/layers/quantization/nvfp4_config.py

@@ -58,29 +58,6 @@ _LTX2_REFINE_ONLY_SUFFIXES = (
     ".video_to_audio_attn.to_v",
 )
 
-_LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES = (
-    "attn1.to_q",
-    "attn1.to_k",
-    "attn1.to_v",
-    "attn1.to_out",
-    "attn2.to_q",
-    "attn2.to_out",
-    "audio_to_video_attn.to_q",
-    "audio_to_video_attn.to_out",
-    "video_to_audio_attn.to_k",
-    "video_to_audio_attn.to_v",
-    "ffn.fc_in",
-    "ffn.fc_out",
-)
-_LTX2_NVFP4_LINEAR_PREFIXES = frozenset(f"ltx2.blocks.{block_idx}.{suffix}" for block_idx in range(48)
-                                        for suffix in _LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES) | frozenset(
-                                            ("ltx2.adaln_single.linear", ))
-
-
-def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool:
-    """Return whether *prefix* belongs to the LTX-2 NVFP4 deployment set."""
-    return prefix in _LTX2_NVFP4_LINEAR_PREFIXES
-
 
 def _is_ltx2_refine_only_prefix(prefix: str) -> bool:
     return any(prefix.endswith(suffix) for suffix in _LTX2_REFINE_ONLY_SUFFIXES)
@@ -309,11 +286,6 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
         self.x_global_sf = torch.tensor(1.0, device="cuda", dtype=torch.float32)
         self.layer_prefix = layer_prefix
         self._is_refine_only_layer = _is_ltx2_refine_only_prefix(layer_prefix)
-        # Set from NVFP4Config.retain_original_weights in get_quant_method:
-        # True = retain every original bf16 weight; None/False (default) =
-        # purge the purgeable set. Refine-only layers are always retained --
-        # the base stage profile runs them dense by deployment contract.
-        self._retain_original_weights: bool | None = None
 
     def create_weights(self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int],
                        input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs):
@@ -354,11 +326,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
         | None = None,
     ) -> torch.Tensor:
         SfLayout, _, _ = _require_flashinfer()
-        # The original bf16 weight may have been purged after FP4 conversion
-        # (see convert_model_to_nvfp4); the packed FP4 weight keeps the
-        # output dim as its first dimension (only K is packed 2-per-byte).
-        weight = getattr(layer, "weight", None)
-        out_dim = weight.shape[0] if weight is not None else layer._nvfp4_weight.shape[0]
+        out_dim = layer.weight.shape[0]
         original_shape = x.shape
 
         # Stage-aware profile: keep refine-only FP4 layers in dense mode
@@ -366,13 +334,8 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
         # quantize/dequantize tax for layers it never touches.
         stage_profile = _get_ltx2_fp4_stage_profile(default="refine")
         if self._is_refine_only_layer and stage_profile == "base":
-            if weight is None:
-                raise RuntimeError(f"NVFP4 layer {self.layer_prefix!r} hit the stage-profile dense path, "
-                                   "but its original weights were purged "
-                                   "(NVFP4Config(retain_original_weights=False)). Streaming/two-stage "
-                                   "deploys must load with retain_original_weights left unset (auto) or True.")
-            out = (F.linear(x, weight, bias) if torch.cuda.is_available() or bias is None else F.linear(
-                x, weight, bias.to(x.dtype)))
+            out = (F.linear(x, layer.weight, bias) if torch.cuda.is_available() or bias is None else F.linear(
+                x, layer.weight, bias.to(x.dtype)))
             return out.view(*original_shape[:-1], out_dim)
         if pre_quantized is not None:
             x_fp4, x_scale, x_global_sf = pre_quantized
@@ -430,19 +393,11 @@ class NVFP4Config(QuantizationConfig):
     instead of hardcoding it here.
     """
 
-    def __init__(self, layer_profile: str = "refine", retain_original_weights: bool | None = None):
+    def __init__(self, layer_profile: str = "refine"):
         super().__init__()
         # ``base``: stage-1 set (no attn2.to_out, no cross-modal AV
         # projections). ``refine``: full stage-2 set.
         self.layer_profile = layer_profile
-        # Original bf16 ``layer.weight`` retention after FP4 conversion.
-        # Default (None/False): purge the purgeable originals -- every
-        # always-FP4 layer. Refine-only layers (the cross-modal AV
-        # projections) are ALWAYS retained: the ``base`` stage profile runs
-        # them dense by deployment contract, including the distilled
-        # single-stage deploy. True: retain everything (debugging /
-        # pre-purge behavior).
-        self.retain_original_weights = retain_original_weights
 
     def get_name(self):
         return "nvfp4"
@@ -460,20 +415,32 @@ class NVFP4Config(QuantizationConfig):
 
     @classmethod
     def from_config(cls, config: dict[str, Any]) -> NVFP4Config:
-        return cls(
-            layer_profile=config.get("layer_profile", "refine"),
-            retain_original_weights=config.get("retain_original_weights"),
-        )
+        return cls(layer_profile=config.get("layer_profile", "refine"))
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
         from fastvideo.layers.linear import LinearBase
 
         # Use the superset at build/load time, then switch active subset
         # dynamically in NVFP4QuantizeMethod.apply based on stage profile.
-        if isinstance(layer, LinearBase) and is_ltx2_nvfp4_linear_prefix(prefix):
-            method = NVFP4QuantizeMethod(layer_prefix=prefix)
-            method._retain_original_weights = self.retain_original_weights
-            return method
+        fp4_layers = [[
+            f"ltx2.blocks.{i}.attn1.to_q",
+            f"ltx2.blocks.{i}.attn1.to_k",
+            f"ltx2.blocks.{i}.attn1.to_v",
+            f"ltx2.blocks.{i}.attn1.to_out",
+            f"ltx2.blocks.{i}.attn2.to_q",
+            f"ltx2.blocks.{i}.attn2.to_out",
+            f"ltx2.blocks.{i}.audio_to_video_attn.to_q",
+            f"ltx2.blocks.{i}.audio_to_video_attn.to_out",
+            f"ltx2.blocks.{i}.video_to_audio_attn.to_k",
+            f"ltx2.blocks.{i}.video_to_audio_attn.to_v",
+            f"ltx2.blocks.{i}.ffn.fc_in",
+            f"ltx2.blocks.{i}.ffn.fc_out",
+        ] for i in range(48)]
+        fp4_layers.append([
+            "ltx2.adaln_single.linear",
+        ])
+        if isinstance(layer, LinearBase) and any(prefix in layer_names for layer_names in fp4_layers):
+            return NVFP4QuantizeMethod(layer_prefix=prefix)
         return None
 
 
@@ -481,9 +448,6 @@ def convert_model_to_nvfp4(model: torch.nn.Module) -> None:
     SfLayout, _, _ = _require_flashinfer()
     from torch.distributed.tensor import DTensor  # type: ignore
 
-    purged = 0
-    retained = 0
-    purged_bytes = 0
     for mod in model.modules():
         qm = getattr(mod, "quant_method", None)
         if isinstance(qm, NVFP4QuantizeMethod):
@@ -516,40 +480,9 @@ def convert_model_to_nvfp4(model: torch.nn.Module) -> None:
                 persistent=False,
             )
 
-            retain_flag = getattr(qm, "_retain_original_weights", None)
-            # Refine-only layers are NEVER purgeable: the "base" stage profile
-            # runs them dense by deployment contract (the distilled
-            # single-stage deploy included — its forward context is the base
-            # profile, so e.g. audio_to_video_attn routes dense every step).
-            # retain_original_weights therefore only widens retention
-            # (True = keep everything); it cannot narrow it below the
-            # dense-capable set.
-            retain = qm._is_refine_only_layer or retain_flag is True
-            if retain:
-                retained += 1
-            elif isinstance(weight, DTensor):
-                # ponytail: purging FSDP-sharded originals needs per-shard
-                # resharding bookkeeping; skip until a sharded deploy needs it.
-                retained += 1
-            else:
-                purged_bytes += weight.numel() * weight.element_size()
-                purged += 1
-                mod.register_parameter("weight", None)
-
-    if purged or retained:
-        logger.info(
-            "NVFP4 weight purge receipt: purged %d original bf16 weight tensors "
-            "(%.2f GiB freed); retained %d (refine-only dense fallback or "
-            "retain_original_weights).",
-            purged,
-            purged_bytes / (1 << 30),
-            retained,
-        )
-
 
 __all__ = [
     "NVFP4Config",
     "NVFP4QuantizeMethod",
     "convert_model_to_nvfp4",
-    "is_ltx2_nvfp4_linear_prefix",
 ]

fastvideo/layers/quantization/nvfp4_qat_config.py

@@ -46,11 +46,7 @@ from fastvideo.models.utils import set_weight_attrs
 logger = logging.getLogger(__name__)
 
 # Wan-style attention + FFN projection layers. Matched as substrings of the
-# layer prefix (e.g. "blocks.0.attn1.to_q" contains "to_q"). Wan's "to_q"/
-# "to_k"/"to_v" also substring-match Kandinsky5's "to_query"/"to_key"/
-# "to_value", so only Kandinsky5's out-projection and FFN names (which don't
-# share a substring with Wan's "to_out"/"ffn.fc_in"/"ffn.fc_out") need to be
-# listed explicitly below.
+# layer prefix (e.g. "blocks.0.attn1.to_q" contains "to_q").
 DEFAULT_FP4_LAYERS = (
     "ffn.fc_in",
     "ffn.fc_out",
@@ -58,11 +54,6 @@ DEFAULT_FP4_LAYERS = (
     "to_k",
     "to_v",
     "to_out",
-    # Kandinsky5
-    "self_attention.out_layer",
-    "cross_attention.out_layer",
-    "feed_forward.mlp.fc_in",
-    "feed_forward.mlp.fc_out",
 )
 
 

fastvideo/layers/quantization/nvfp4_qat_train_config.py

@@ -9,11 +9,10 @@ straight-through estimator), so the model learns to absorb FP4 linear error.
 
 That STE already exists in ``fastvideo.layers.fp4linear._LinearFWD4BWD16Fn``
 (FP4 forward, full-precision backward) but is otherwise unwired. This method
-bridges it into the standard ``quant_config`` path. Wan-style prefixes retain
-the broad attention/FFN projection profile used by ``nvfp4_qat``; LTX-2
-prefixes use the exact curated target set from its ``NVFP4`` deployment config.
-No ``convert_model_to_fp4`` is needed: the weight is kept in full precision and
-quantized on the fly each step.
+bridges it into the standard ``quant_config`` path, so it activates via
+``transformer_quant="nvfp4_qat_train"`` on the same Wan-2.1 layers as nvfp4_qat
+(to_q/k/v/out + ffn). No ``convert_model_to_fp4`` is needed: the weight is kept
+in full precision and quantized on the fly each step.
 """
 import logging
 from typing import Any
@@ -22,8 +21,6 @@ import torch
 from torch.nn.parameter import Parameter
 
 from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase
-from fastvideo.layers.quantization.nvfp4_config import is_ltx2_nvfp4_linear_prefix
-from fastvideo.layers.quantization.nvfp4_qat_config import DEFAULT_FP4_LAYERS
 from fastvideo.models.utils import set_weight_attrs
 
 logger = logging.getLogger(__name__)
@@ -75,10 +72,7 @@ class NVFP4QATTrainConfig(QuantizationConfig):
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
         from fastvideo.layers.linear import LinearBase
-        if not isinstance(layer, LinearBase):
-            return None
-        matches = (is_ltx2_nvfp4_linear_prefix(prefix) if prefix.startswith("ltx2.") else any(
-            layer_name in prefix for layer_name in DEFAULT_FP4_LAYERS))
-        if matches:
+        fp4_layers = ["ffn.fc_in", "ffn.fc_out", "to_q", "to_k", "to_v", "to_out"]
+        if isinstance(layer, LinearBase) and any(layer_name in prefix for layer_name in fp4_layers):
             return NVFP4QATTrainQuantizeMethod()
         return None

fastvideo/models/dits/kandinsky5.py

@@ -15,7 +15,6 @@ from fastvideo.configs.models.dits import Kandinsky5VideoConfig
 from fastvideo.layers.layernorm import LayerNormScaleShift
 from fastvideo.layers.linear import ReplicatedLinear
 from fastvideo.layers.mlp import MLP
-from fastvideo.layers.quantization import QuantizationConfig
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum
@@ -286,7 +285,6 @@ class Kandinsky5Attention(nn.Module):
         supported_attention_backends: tuple[AttentionBackendEnum, ...] | None,
         prefix: str = "",
         use_nabla: bool = False,
-        quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
         assert num_channels % head_dim == 0
@@ -295,24 +293,20 @@ class Kandinsky5Attention(nn.Module):
         self.to_query = ReplicatedLinear(num_channels,
                                          num_channels,
                                          bias=True,
-                                         quant_config=quant_config,
                                          prefix=f"{prefix}.to_query")
         self.to_key = ReplicatedLinear(num_channels,
                                        num_channels,
                                        bias=True,
-                                       quant_config=quant_config,
                                        prefix=f"{prefix}.to_key")
         self.to_value = ReplicatedLinear(num_channels,
                                          num_channels,
                                          bias=True,
-                                         quant_config=quant_config,
                                          prefix=f"{prefix}.to_value")
         self.query_norm = nn.RMSNorm(head_dim)
         self.key_norm = nn.RMSNorm(head_dim)
         self.out_layer = ReplicatedLinear(num_channels,
                                           num_channels,
                                           bias=True,
-                                          quant_config=quant_config,
                                           prefix=f"{prefix}.out_layer")
         self.local_attention = LocalAttention(
             num_heads=self.num_heads,
@@ -419,11 +413,9 @@ class Kandinsky5Attention(nn.Module):
 
 class Kandinsky5FeedForward(nn.Module):
 
-    def __init__(self, dim: int, ff_dim: int, prefix: str = "",
-                 quant_config: QuantizationConfig | None = None):
+    def __init__(self, dim: int, ff_dim: int):
         super().__init__()
-        self.mlp = MLP(dim, ff_dim, bias=False, act_type="gelu",
-                      quant_config=quant_config, prefix=f"{prefix}.mlp")
+        self.mlp = MLP(dim, ff_dim, bias=False, act_type="gelu")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.mlp(x)
@@ -475,8 +467,7 @@ class Kandinsky5TransformerEncoderBlock(nn.Module):
                  head_dim: int,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
-                 prefix: str = "",
-                 quant_config: QuantizationConfig | None = None):
+                 prefix: str = ""):
         super().__init__()
         self.text_modulation = Kandinsky5Modulation(time_dim, model_dim, 6)
 
@@ -491,8 +482,7 @@ class Kandinsky5TransformerEncoderBlock(nn.Module):
             model_dim,
             head_dim,
             supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.self_attention",
-            quant_config=quant_config)
+            prefix=f"{prefix}.self_attention")
 
         self.feed_forward_norm = LayerNormScaleShift(
             model_dim,
@@ -501,9 +491,7 @@ class Kandinsky5TransformerEncoderBlock(nn.Module):
             elementwise_affine=False,
             dtype=torch.float32,
             compute_dtype=torch.float32)
-        self.feed_forward = Kandinsky5FeedForward(
-            model_dim, ff_dim, prefix=f"{prefix}.feed_forward",
-            quant_config=quant_config)
+        self.feed_forward = Kandinsky5FeedForward(model_dim, ff_dim)
 
     def forward(self, x: torch.Tensor, time_embed: torch.Tensor,
                 rope: torch.Tensor) -> torch.Tensor:
@@ -535,8 +523,7 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
                  prefix: str = "",
-                 use_nabla: bool = False,
-                 quant_config: QuantizationConfig | None = None):
+                 use_nabla: bool = False):
         super().__init__()
         self.visual_modulation = Kandinsky5Modulation(time_dim, model_dim, 9)
 
@@ -552,8 +539,7 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
             head_dim,
             supported_attention_backends=supported_attention_backends,
             prefix=f"{prefix}.self_attention",
-            use_nabla=use_nabla,
-            quant_config=quant_config)
+            use_nabla=use_nabla)
 
         self.cross_attention_norm = LayerNormScaleShift(
             model_dim,
@@ -566,8 +552,7 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
             model_dim,
             head_dim,
             supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.cross_attention",
-            quant_config=quant_config)
+            prefix=f"{prefix}.cross_attention")
 
         self.feed_forward_norm = LayerNormScaleShift(
             model_dim,
@@ -576,9 +561,7 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
             elementwise_affine=False,
             dtype=torch.float32,
             compute_dtype=torch.float32)
-        self.feed_forward = Kandinsky5FeedForward(
-            model_dim, ff_dim, prefix=f"{prefix}.feed_forward",
-            quant_config=quant_config)
+        self.feed_forward = Kandinsky5FeedForward(model_dim, ff_dim)
 
     def forward(self, visual_embed: torch.Tensor, text_embed: torch.Tensor,
                 time_embed: torch.Tensor, rope: torch.Tensor,
@@ -653,8 +636,6 @@ class Kandinsky5Transformer3DModel(BaseDiT):
                  hf_config: dict[str, Any]) -> None:
         super().__init__(config=config, hf_config=hf_config)
         arch = config.arch_config
-        quant_config = config.quant_config
-        self.quant_config = quant_config
 
         head_dim = sum(arch.axes_dims)
         self.in_visual_dim = arch.in_visual_dim
@@ -683,8 +664,7 @@ class Kandinsky5Transformer3DModel(BaseDiT):
                                               arch.ff_dim,
                                               head_dim,
                                               self._supported_attention_backends,
-                                              prefix=f"{config.prefix}.text_transformer_blocks.{i}",
-                                              quant_config=quant_config)
+                                              prefix=f"{config.prefix}.text_transformer_blocks.{i}")
             for i in range(arch.num_text_blocks)
         ])
         self.visual_transformer_blocks = nn.ModuleList([
@@ -693,8 +673,7 @@ class Kandinsky5Transformer3DModel(BaseDiT):
                                               head_dim,
                                               self._supported_attention_backends,
                                               prefix=f"{config.prefix}.visual_transformer_blocks.{i}",
-                                              use_nabla=arch.attention_type == "nabla",
-                                              quant_config=quant_config)
+                                              use_nabla=arch.attention_type == "nabla")
             for i in range(arch.num_visual_blocks)
         ])
 

fastvideo/models/dits/ltx2.py

@@ -62,34 +62,6 @@ def _is_ltx2_refine_stage() -> bool:
     return forward_batch.extra.get("ltx2_fp4_stage_profile") == "refine"
 
 
-def _functional_rms_norm(
-    x: torch.Tensor,
-    eps: float,
-    weight: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Out-of-place RMSNorm equivalent to ``torch.nn.functional.rms_norm``.
-
-    ``F.rms_norm`` / ``nn.RMSNorm`` decompose under AOTAutograd into an
-    in-place ``variance.add_(eps)`` (an ``aten.copy_`` into the computed
-    ``mean``). AOTAutograd's ``assert_functional_graph`` requires every
-    ``copy_`` destination to be a graph input, so this trips when the DiT is
-    compiled through the ``torch.compile(backend="torch_tensorrt")`` path
-    (Dynamo -> AOTAutograd functionalization -> torch-tensorrt), breaking
-    TensorRT + sequence-parallel (Ulysses) export. Computing ``variance + eps``
-    out-of-place keeps the graph functional. Numerically matches
-    ``F.rms_norm`` (fp32 accumulation, dtype restored to the input)."""
-    input_dtype = x.dtype
-    hidden = x.to(torch.float32)
-    variance = hidden.pow(2).mean(-1, keepdim=True)
-    hidden = hidden * torch.rsqrt(variance + eps)
-    out = hidden.to(input_dtype)
-    if weight is not None:
-        # A weight dtype differing from x (e.g. fp32 params under autocast)
-        # would promote the output; keep the input-dtype guarantee.
-        out = out * weight
-    return out.to(input_dtype)
-
-
 def _rms_norm_dispatch(
     x: torch.Tensor,
     eps: float,
@@ -103,7 +75,8 @@ def _rms_norm_dispatch(
     # 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 _functional_rms_norm(x, eps, weight)
+    return torch.nn.functional.rms_norm(
+        x, (x.shape[-1], ), weight=weight, eps=eps).to(x.dtype)
 
 
 class StageAwareRMSNorm(nn.RMSNorm):
@@ -128,7 +101,7 @@ class StageAwareRMSNorm(nn.RMSNorm):
         # norm now returns float32 where torch <= 2.11 returned bfloat16.
         # Restore the input dtype so autocast-blind consumers downstream
         # (e.g. the flash-attention custom op) don't receive upcast q/k.
-        return _functional_rms_norm(x, self.eps, self.weight)
+        return super().forward(x).to(x.dtype)
 
 
 def _supports_prequantized_input(linear: ReplicatedLinear) -> bool:
@@ -1571,7 +1544,6 @@ class LTXSelfAttention(nn.Module):
             softmax_scale=None,
             causal=False,
             supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA,),
-            default_backend=AttentionBackendEnum.TORCH_SDPA,
         )
 
     def forward(
@@ -1854,15 +1826,12 @@ class BasicAVTransformerBlock(torch.nn.Module):
         # Text cross-attention uses LocalAttention (text embeddings are replicated)
         SelfAttnCls = LTXDistributedSelfAttention if use_distributed_attention else LTXSelfAttention
         CrossAttnCls = LTXSelfAttention  # Text cross-attention is always local
-        video_attn_backends = (
+        video_self_attn_backends = (
             AttentionBackendEnum.FLASH_ATTN,
             AttentionBackendEnum.TORCH_SDPA,
-            AttentionBackendEnum.ATTN_QAT_INFER,
-            AttentionBackendEnum.ATTN_QAT_TRAIN,
+            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
         )
-        video_self_attn_backends = video_attn_backends + (
-            AttentionBackendEnum.VIDEO_SPARSE_ATTN, )
-        audio_attn_backends = (
+        dense_attn_backends = (
             AttentionBackendEnum.FLASH_ATTN,
             AttentionBackendEnum.TORCH_SDPA,
         )
@@ -1889,7 +1858,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=video.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=video_attn_backends,
+                supported_attention_backends=dense_attn_backends,
                 apply_gated_attention=video.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.attn2",
@@ -1915,7 +1884,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=audio_attn_backends,
+                supported_attention_backends=dense_attn_backends,
                 apply_gated_attention=audio.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.audio_attn1",
@@ -1928,7 +1897,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=audio_attn_backends,
+                supported_attention_backends=dense_attn_backends,
                 apply_gated_attention=audio.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.audio_attn2",
@@ -1953,7 +1922,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=audio_attn_backends,
+                supported_attention_backends=dense_attn_backends,
                 apply_gated_attention=video.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.audio_to_video_attn",
@@ -1967,7 +1936,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=audio_attn_backends,
+                supported_attention_backends=dense_attn_backends,
                 apply_gated_attention=audio.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.video_to_audio_attn",

fastvideo/models/loader/fsdp_load.py

@@ -38,14 +38,8 @@ def _maybe_quantize_model(model: nn.Module) -> None:
     The walk returns on the first quantized layer found so unquantized callers
     pay only an ``isinstance`` check per module. Both imports are deferred so
     this is a no-op on hosts without the relevant backends.
-
-    QAT-*train* linears (``nvfp4_qat_train``) need no weight conversion — the
-    master weight stays full precision — but their attachment is otherwise
-    silent, so the same walk emits a one-line receipt with the count of
-    linears that actually carry the QAT method.
     """
     # Defer imports: these modules pull in heavy symbols at module-load time.
-    from fastvideo.layers.linear import LinearBase
     from fastvideo.layers.quantization.nvfp4_config import (
         NVFP4QuantizeMethod,
         convert_model_to_nvfp4,
@@ -54,15 +48,11 @@ def _maybe_quantize_model(model: nn.Module) -> None:
         NVFP4QATQuantizeMethod,
         convert_model_to_fp4,
     )
-    from fastvideo.layers.quantization.nvfp4_qat_train_config import (
-        NVFP4QATTrainQuantizeMethod, )
     from fastvideo.layers.quantization.fp8_config import (
         FP8QuantizeMethod,
         convert_model_to_fp8,
     )
 
-    qat_train_attached = 0
-    qat_train_skipped = 0
     for mod in model.modules():
         qm = getattr(mod, "quant_method", None)
         if isinstance(qm, NVFP4QuantizeMethod):
@@ -77,16 +67,6 @@ def _maybe_quantize_model(model: nn.Module) -> None:
             logger.info("Converting loaded model weights for FP8 linear layers")
             convert_model_to_fp8(model)
             return
-        # QAT-train configs are mutually exclusive with the inference schemes
-        # above (one quant_config per model), so when they're active the loop
-        # always runs to completion and the counts below are model-wide.
-        if isinstance(qm, NVFP4QATTrainQuantizeMethod):
-            qat_train_attached += 1
-        elif isinstance(mod, LinearBase):
-            qat_train_skipped += 1
-    if qat_train_attached:
-        logger.info("NVFP4 QAT: attached %d linears (%d skipped by prefix filter)", qat_train_attached,
-                    qat_train_skipped)
 
 
 # TODO(PY): move this to utils elsewhere

fastvideo/pipelines/basic/kandinsky5/kandinsky5_dmd_pipeline.py

@@ -1,82 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Kandinsky5 DMD (few-step distilled) validation pipeline.
-
-Composes the same stages as ``Kandinsky5T2VPipeline`` (text encoding, latent
-prep, decoding) but swaps in ``Kandinsky5DmdDenoisingStage`` for the
-denoising loop, mirroring ``wan_dmd_pipeline.py``'s split from the base T2V
-pipeline. This exists solely to drive the stage-2 DMD training config's
-validation callback.
-"""
-
-from __future__ import annotations
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.pipelines.stages.input_validation import InputValidationStage
-from fastvideo.pipelines.stages.kandinsky5 import (
-    Kandinsky5DecodingStage,
-    Kandinsky5DmdDenoisingStage,
-    Kandinsky5LatentPreparationStage,
-)
-from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
-from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
-
-
-class Kandinsky5DMDPipeline(ComposedPipelineBase):
-    """Kandinsky-5.0 Lite DMD (few-step) text-to-video pipeline."""
-
-    _required_config_modules = [
-        "scheduler",
-        "text_encoder",
-        "text_encoder_2",
-        "tokenizer",
-        "tokenizer_2",
-        "transformer",
-        "vae",
-    ]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
-
-        self.add_stage(
-            stage_name="text_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[
-                    self.get_module("text_encoder"),
-                    self.get_module("text_encoder_2"),
-                ],
-                tokenizers=[
-                    self.get_module("tokenizer"),
-                    self.get_module("tokenizer_2"),
-                ],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=Kandinsky5LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=Kandinsky5DmdDenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage",
-            stage=Kandinsky5DecodingStage(vae=self.get_module("vae"), pipeline=self),
-        )
-
-
-EntryClass = Kandinsky5DMDPipeline

fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py

@@ -1,309 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Preprocess Kandinsky5 overfit data into parquet format.
-
-Encodes videos with Kandinsky5's VAE (shared with HunyuanVideo) and
-captions with dual text encoders (Qwen/Reason1 + CLIP) into the t2v
-parquet schema expected by the training framework.
-
-Unlike Hunyuan's version of this script, text encoders are loaded via
-FastVideo's own ``TextEncoderLoader``/``VAELoader`` (matching
-``Kandinsky5Model.ensure_negative_conditioning``) rather than raw
-``transformers`` classes -- Kandinsky5's Qwen/Reason1 encoder is a
-native FastVideo wrapper around Qwen2.5-VL's language model, not a
-plain ``AutoModel.from_pretrained`` load.
-
-The CLIP pooled projection (``[768]``) is zero-padded into a row and
-prepended to the Qwen sequence embeddings, then stored as a single
-``[seq+1, dim]`` tensor in the existing ``text_embedding`` parquet
-field -- the same trick ``preprocess_hunyuan_overfit.py`` uses for
-LLaMA+CLIP. No schema change and no separate attention-mask field are
-needed: the standard collator derives ``text_attention_mask`` from
-padding against the stored (unpadded) row count, so the prepended
-pooled row is automatically counted as valid.
-
-Usage:
-    CUDA_VISIBLE_DEVICES=0 python -m fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit
-
-Input/output roots default to ``data/kandinsky5_overfit`` /
-``data/kandinsky5_overfit_preprocessed`` and can be overridden with the
-``KANDINSKY5_OVERFIT_DATA_DIR`` / ``KANDINSKY5_OVERFIT_OUTPUT_DIR`` env
-vars (used by the nightly e2e test to keep its disposable roots separate
-from a user's real dataset at the defaults).
-"""
-
-import json
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29520")
-
-import cv2
-import numpy as np
-import pyarrow as pa
-import pyarrow.parquet as pq
-import torch
-from transformers import AutoTokenizer
-
-from fastvideo.configs.pipelines.base import preprocess_text
-from fastvideo.configs.pipelines.kandinsky5 import (
-    Kandinsky5T2VConfig,
-    kandinsky5_clip_postprocess_text,
-    kandinsky5_qwen_postprocess_text,
-    kandinsky5_qwen_preprocess_text,
-)
-from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
-from fastvideo.distributed import maybe_init_distributed_environment_and_model_parallel
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.models.loader.component_loader import TextEncoderLoader, VAELoader
-from fastvideo.utils import maybe_download_model
-
-# --- Config ---
-# Matches KANDINSKY5_T2V_LITE_5S's native preset (fastvideo/pipelines/basic/
-# kandinsky5/presets.py): 121 frames @ 24fps ~= 5.04s, matching the "5s" in
-# the checkpoint name. 512/768 fall in the 480p RoPE scale_factor band
-# Kandinsky5Model asserts on, and satisfy the real divisor requirement
-# (spatial_compression_ratio(8) * patch_size[1](2) = 16 -- not just 8).
-NUM_FRAMES = 121  # 4k+1 for temporal compression ratio 4
-MAX_HEIGHT = 512
-MAX_WIDTH = 768
-TRAIN_FPS = 24.0
-
-MODEL_PATH = "kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers"
-# Overridable so automation (e.g. the nightly e2e test) can point at its own
-# test-owned directories instead of clobbering the documented default paths
-# a user may have populated with their real dataset.
-DATA_DIR = os.environ.get("KANDINSKY5_OVERFIT_DATA_DIR", "data/kandinsky5_overfit")
-OUTPUT_DIR = os.environ.get("KANDINSKY5_OVERFIT_OUTPUT_DIR", "data/kandinsky5_overfit_preprocessed")
-
-
-def load_video(path: str, num_frames: int, height: int, width: int) -> torch.Tensor:
-    """Load video as [1, C, T, H, W] in [-1, 1], resized to (height, width).
-
-    Source clips are frequently at native resolution (e.g. 1920x1080);
-    encoding that directly through the VAE (instead of at the target
-    training resolution) uses far more memory than intended and can OOM.
-    """
-    cap = cv2.VideoCapture(path)
-    frames: list[np.ndarray] = []
-    while len(frames) < num_frames:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        if frame.shape[0] != height or frame.shape[1] != width:
-            frame = cv2.resize(frame, (width, height), interpolation=cv2.INTER_AREA)
-        frames.append(frame)
-    cap.release()
-
-    if not frames:
-        # Without this, the repeat-last-frame fill below raises an opaque
-        # IndexError on frames[-1]; cv2.VideoCapture doesn't raise on a
-        # missing/corrupt file, it just decodes nothing.
-        raise ValueError(f"Could not decode any frames from video {path!r} -- "
-                         "the file is missing, empty, or not readable by OpenCV.")
-
-    if len(frames) < num_frames:
-        # Repeat last frame to fill
-        while len(frames) < num_frames:
-            frames.append(frames[-1])
-
-    frames = frames[:num_frames]
-    video = np.stack(frames, axis=0)
-    video = torch.from_numpy(video).float()
-    video = video / 127.5 - 1.0  # [0,255] -> [-1,1]
-    video = video.permute(3, 0, 1, 2).unsqueeze(0)  # [1,C,T,H,W]
-    return video
-
-
-def get_caption(item: dict) -> str:
-    """Extract a caption from one ``videos2caption.json`` entry.
-
-    The documented schema (``docs/training/data_preprocess.md``) stores
-    ``cap`` as a plain string; some producers (e.g. this repo's own e2e
-    test fixtures, mirroring ``preprocess_hunyuan_overfit.py``) instead
-    store a non-empty list of caption variants and use the first one.
-    Indexing unconditionally with ``item["cap"][0]`` silently takes the
-    first *character* of a string caption instead of erroring, so accept
-    and validate both forms explicitly here.
-    """
-    cap = item.get("cap")
-    if isinstance(cap, str):
-        if not cap:
-            raise ValueError(f"Empty 'cap' string for entry: {item}")
-        return cap
-    if isinstance(cap, list):
-        if not cap or not isinstance(cap[0], str) or not cap[0]:
-            raise ValueError(f"'cap' list must be non-empty with a non-empty first string, got: {item}")
-        return cap[0]
-    raise ValueError(f"'cap' must be a string or a list of strings, got {type(cap).__name__}: {item}")
-
-
-def main() -> None:
-    # FastVideo's native text-encoder layers (e.g. CLIPAttention's
-    # QKVParallelLinear) are tensor-parallel-aware and assert the TP process
-    # group is initialized, even for a single-process/single-GPU run.
-    maybe_init_distributed_environment_and_model_parallel(1, 1)
-
-    device = torch.device("cuda:0")
-    model_path = maybe_download_model(MODEL_PATH)
-
-    os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-    # Load captions. Validate up front (before spending minutes loading the
-    # VAE + two text encoders): an empty/malformed manifest would otherwise
-    # only surface as an IndexError on records[0] at parquet-write time.
-    manifest_path = os.path.join(DATA_DIR, "videos2caption.json")
-    with open(manifest_path) as f:
-        caption_data = json.load(f)
-    if not isinstance(caption_data, list) or not caption_data:
-        raise ValueError(f"Manifest {manifest_path} must be a non-empty JSON list of "
-                         f"{{'path', 'cap'}} entries, got: {type(caption_data).__name__} "
-                         f"with {len(caption_data) if isinstance(caption_data, list) else 'n/a'} entries")
-
-    pipeline_config = Kandinsky5T2VConfig()
-    # Kandinsky5T2VConfig.__post_init__ sets load_encoder=False by default --
-    # T2V inference only ever decodes generated latents, never encodes real
-    # video. Preprocessing needs the encoder to turn real clips into latents.
-    pipeline_config.vae_config.load_encoder = True
-    fastvideo_args = FastVideoArgs(
-        model_path=model_path,
-        dit_cpu_offload=False,
-        dit_layerwise_offload=False,
-        use_fsdp_inference=False,
-        text_encoder_cpu_offload=False,
-        vae_cpu_offload=False,
-        pipeline_config=pipeline_config,
-    )
-    fastvideo_args.device = device
-
-    # --- Load VAE (shared with HunyuanVideo) ---
-    print("Loading Kandinsky5 VAE...")
-    vae = VAELoader().load(os.path.join(model_path, "vae"), fastvideo_args)
-    vae = vae.to(device=device, dtype=torch.float16).eval()
-    print(f"VAE loaded ({sum(p.numel() for p in vae.parameters()) / 1e6:.0f}M)")
-
-    # --- Load text encoders ---
-    print("Loading Qwen/Reason1 text encoder...")
-    qwen_enc = TextEncoderLoader().load(
-        os.path.join(model_path, "text_encoder"),
-        fastvideo_args,
-    ).to(device).eval()
-    qwen_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
-    qwen_tok_kwargs = dict(pipeline_config.text_encoder_configs[0].tokenizer_kwargs)
-    # TextEncodingStage overrides max_length from pipeline_config.text_encoder_max_lengths
-    # at runtime (fastvideo/pipelines/stages/text_encoding.py) -- Reason1Config's static
-    # tokenizer_kwargs default (text_len=512) doesn't account for the 129-token Kandinsky
-    # system template ENCODE_START_IDX strips off afterwards, so mirror the runtime value
-    # here or training conditions on fewer caption tokens than inference does.
-    qwen_tok_kwargs["max_length"] = pipeline_config.text_encoder_max_lengths[0]
-
-    print("Loading CLIP text encoder...")
-    clip_enc = TextEncoderLoader().load(
-        os.path.join(model_path, "text_encoder_2"),
-        fastvideo_args,
-    ).to(device).eval()
-    clip_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer_2"))
-    clip_tok_kwargs = dict(pipeline_config.text_encoder_configs[1].tokenizer_kwargs)
-    clip_tok_kwargs["max_length"] = pipeline_config.text_encoder_max_lengths[1]
-
-    # --- Process each video ---
-    records = []
-    for item in caption_data:
-        video_name = item["path"]
-        caption = get_caption(item)
-        video_path = os.path.join(DATA_DIR, "videos", video_name)
-
-        print(f"\nProcessing: {video_name}")
-        print(f"  Caption: {caption[:80]}...")
-
-        # Encode video. Kandinsky5's VAE encode returns channel-first
-        # [B, C, T, H, W], matching the storage convention Kandinsky5Model
-        # expects (it permutes to channel-last only right before calling
-        # the transformer).
-        video = load_video(video_path, NUM_FRAMES, MAX_HEIGHT, MAX_WIDTH).to(device=device, dtype=torch.float16)
-        print(f"  Video shape: {video.shape}")
-
-        with torch.no_grad():
-            latent_dist = vae.encode(video)
-            latent = latent_dist.mean.squeeze(0).float().cpu()  # [C, T, H, W]
-        print(f"  Latent shape: {latent.shape}")
-
-        # Encode text with Qwen/Reason1.
-        qwen_text = kandinsky5_qwen_preprocess_text(caption)
-        with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None):
-            qwen_inputs = qwen_tok(qwen_text, **qwen_tok_kwargs).to(device)
-            qwen_out = qwen_enc(
-                input_ids=qwen_inputs.input_ids,
-                attention_mask=qwen_inputs.attention_mask,
-                output_hidden_states=True,
-            )
-            qwen_embeds, _qwen_mask = kandinsky5_qwen_postprocess_text(qwen_out, qwen_inputs.attention_mask)
-            qwen_embeds = qwen_embeds.squeeze(0)  # [seq, dim]
-
-        # Encode text with CLIP.
-        clip_text = preprocess_text(caption)
-        with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None):
-            clip_inputs = clip_tok(clip_text, **clip_tok_kwargs).to(device)
-            clip_out = clip_enc(
-                input_ids=clip_inputs.input_ids,
-                attention_mask=clip_inputs.attention_mask,
-            )
-            clip_pooled = kandinsky5_clip_postprocess_text(clip_out).squeeze(0)  # [768]
-
-        # Combine: [pooled_clip_row, qwen_embeds]
-        qwen_dim = qwen_embeds.shape[-1]
-        pooled_row = torch.zeros(qwen_dim, device=device, dtype=torch.float16)
-        pooled_row[:clip_pooled.shape[-1]] = clip_pooled
-        text_embedding = torch.cat(
-            [pooled_row.unsqueeze(0), qwen_embeds],
-            dim=0,
-        ).float().cpu()  # [seq+1, dim]
-        print(f"  Text embedding shape: {text_embedding.shape}")
-
-        record = {
-            "id": video_name,
-            "vae_latent_bytes": latent.numpy().tobytes(),
-            "vae_latent_shape": list(latent.shape),
-            "vae_latent_dtype": str(latent.dtype).replace("torch.", ""),
-            "text_embedding_bytes": text_embedding.numpy().tobytes(),
-            "text_embedding_shape": list(text_embedding.shape),
-            "text_embedding_dtype": str(text_embedding.dtype).replace("torch.", ""),
-            "file_name": video_name,
-            "caption": caption,
-            "media_type": "video",
-            "width": MAX_WIDTH,
-            "height": MAX_HEIGHT,
-            "num_frames": NUM_FRAMES,
-            "duration_sec": NUM_FRAMES / TRAIN_FPS,
-            "fps": TRAIN_FPS,
-        }
-        records.append(record)
-
-    # Clean up encoders
-    del qwen_enc, qwen_tok, clip_enc, clip_tok, vae
-
-    # Write parquet
-    table = pa.table(
-        {k: [r[k] for r in records]
-         for k in records[0]},
-        schema=pyarrow_schema_t2v,
-    )
-    output_path = os.path.join(OUTPUT_DIR, "data_00000.parquet")
-    pq.write_table(table, output_path)
-    print(f"\nWrote {len(records)} records to {output_path}")
-
-    # Write validation prompts for callback.
-    # Wrap in "data" key -- ValidationDataset expects field="data".
-    # Use "caption" field -- ValidationDataset aliases it to "prompt".
-    val_prompts = {"data": [{"caption": get_caption(item)} for item in caption_data]}
-    val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json")
-    with open(val_path, "w") as f:
-        json.dump(val_prompts, f, indent=2)
-    print(f"Wrote validation prompts to {val_path}")
-
-    print("\nDone! Use data_path: " + OUTPUT_DIR + " in training config.")
-
-
-if __name__ == "__main__":
-    main()

fastvideo/pipelines/stages/kandinsky5.py

@@ -10,18 +10,12 @@ import torch
 from diffusers.utils.torch_utils import randn_tensor
 from tqdm.auto import tqdm
 
-import fastvideo.envs as envs
-from fastvideo.attention import LocalAttention
 from fastvideo.attention.backends.nabla import NablaAttentionMetadataBuilder
-from fastvideo.attention.selector import backend_name_to_enum
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.loader.component_loader import TransformerLoader, VAELoader
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler, )
-from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.models.vaes.common import ParallelTiledVAE
 from fastvideo.models.vision_utils import normalize, numpy_to_pt, pil_to_numpy, resize
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
@@ -30,8 +24,7 @@ from fastvideo.pipelines.stages.decoding import DecodingStage
 from fastvideo.pipelines.stages.encoding import EncodingStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
-from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.utils import PRECISION_TO_TYPE, get_mixed_precision_state
+from fastvideo.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
 
@@ -173,115 +166,6 @@ class Kandinsky5DenoisingStage(PipelineStage):
         seq_len = int(mask.sum(1).max().item())
         return torch.arange(seq_len, device=device)
 
-    # Backends whose selection is all-or-nothing: fastvideo.platforms.cuda
-    # raises ImportError immediately if the ATTN_QAT_TRAIN kernel isn't
-    # available (training must never silently de-quantize), rather than the
-    # generic "requested backend unsupported by this layer, falling back to
-    # automatic selection" path _cached_get_attn_backend takes for most
-    # other backend names (see fastvideo/attention/selector.py).
-    # ATTN_QAT_TRAIN is also always in Kandinsky5ArchConfig's
-    # _supported_attention_backends, so that generic fallback never applies
-    # to it here. That combination makes an exact backend match a safe,
-    # false-positive-free signal *only* for ATTN_QAT_TRAIN.
-    # ATTN_QAT_INFER must NOT be in this set: fastvideo.platforms.cuda
-    # deliberately resolves it to FlashAttention when the sm_120 kernel
-    # isn't built, and the QAD recipe README documents that fallback ("on
-    # other GPUs the attention falls back to a supported dense backend
-    # while the FP4 linear layers still run") -- asserting an exact match
-    # would abort inference on every non-sm_120 machine relying on it.
-    # Any other FASTVIDEO_ATTENTION_BACKEND value (e.g. a backend meant for
-    # a different model family sharing the same process/env var) can
-    # likewise legitimately resolve to something other than what was
-    # requested.
-    _STRICT_BACKENDS = frozenset({
-        AttentionBackendEnum.ATTN_QAT_TRAIN,
-    })
-
-    def _assert_local_attention_backend_engaged(self) -> None:
-        """Guard against ``LocalAttention`` silently falling back to SDPA.
-
-        ``Kandinsky5Attention.forward`` catches the ``AssertionError``
-        ``LocalAttention`` raises when no pipeline forward context is set and
-        falls back to plain ``F.scaled_dot_product_attention`` -- silently
-        skipping whichever kernel ``FASTVIDEO_ATTENTION_BACKEND`` requested
-        (the fake-quantized ``ATTN_QAT_TRAIN`` kernel).
-        ``LocalAttention.backend`` is resolved once at module construction,
-        independent of whether a forward context is later set, so this only
-        catches a backend that failed to resolve to what the env var
-        requested -- it does not prove forward context was present for a
-        given forward call. Pair it with always wrapping the actual
-        transformer call in ``set_forward_context`` (see the denoising loops
-        below), which is what prevents the runtime fallback.
-
-        Scoped to ``_STRICT_BACKENDS``: other backend names -- including
-        ``ATTN_QAT_INFER``, whose FlashAttention fallback on non-sm_120
-        GPUs is documented, intentional behavior -- can legitimately differ
-        from what ``LocalAttention.backend`` resolved to (see
-        ``_STRICT_BACKENDS``'s comment), so asserting on those would flag
-        expected behavior as a bug.
-        """
-        backend_env = envs.FASTVIDEO_ATTENTION_BACKEND
-        if not backend_env:
-            return
-        expected = backend_name_to_enum(backend_env)
-        if expected is None or expected not in self._STRICT_BACKENDS:
-            return
-        for module in self.transformer.modules():
-            if isinstance(module, LocalAttention):
-                assert module.backend == expected, (
-                    f"Kandinsky5 local attention resolved to backend {module.backend}, expected "
-                    f"{expected} from FASTVIDEO_ATTENTION_BACKEND={backend_env}. This likely means "
-                    "LocalAttention's missing-forward-context guard silently fell back to SDPA.")
-                return
-
-    def _resolve_target_dtype(self, fastvideo_args: FastVideoArgs) -> torch.dtype:
-        """Resolve the transformer's actual compute dtype.
-
-        Trust ``pipeline_config.dit_precision`` directly for a normal,
-        non-FSDP load -- ``TransformerLoader.load()`` asserts every
-        parameter matches it exactly (``fastvideo/models/loader/
-        component_loader.py``), so this holds for standalone T2V/I2V
-        inference regardless of which precision (including fp32) was
-        requested.
-
-        FSDP2-wrapped modules are the one case where that assertion doesn't
-        carry forward: ``maybe_load_fsdp_model`` hardcodes its
-        ``MixedPrecisionPolicy`` to ``param_dtype=torch.bfloat16`` (with
-        ``cast_forward_inputs=False``) for every FSDP-wrapped load,
-        independent of ``dit_precision`` -- this covers both the live
-        transformer ``ValidationCallback`` reuses from training (whose
-        ``pipeline_config.dit_precision`` reflects the fp32 master-weight
-        load dtype, not the actual bf16 compute dtype) and multi-GPU
-        ``use_fsdp_inference=True`` runs.
-
-        For that FSDP case, read the policy itself
-        (``set_mixed_precision_policy`` records it right before
-        ``maybe_load_fsdp_model`` shards -- every FSDP wrap in this repo
-        goes through that path), NOT the parameter storage dtypes:
-        FSDP2 computes in the policy's ``param_dtype`` regardless of what
-        dtype the parameters are stored/loaded in, so e.g. an fp16-loaded
-        FSDP model still runs its forward in bf16 -- scanning parameters
-        would resolve fp16 and, with ``cast_forward_inputs=False`` and
-        autocast keyed off the wrong dtype, hand bf16-computing modules
-        fp16 inputs.
-        """
-        declared_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
-        try:
-            from torch.distributed.fsdp import FSDPModule
-        except Exception:  # pragma: no cover - FSDP not always available
-            return declared_dtype
-
-        if not isinstance(self.transformer, FSDPModule):
-            return declared_dtype
-
-        try:
-            policy_dtype = get_mixed_precision_state().param_dtype
-        except ValueError:
-            policy_dtype = None
-        # bf16 mirrors maybe_load_fsdp_model's hardcoded param_dtype if the
-        # (thread-local) policy state is somehow unset here.
-        return policy_dtype if policy_dtype is not None else torch.bfloat16
-
     @staticmethod
     def fast_sta_nabla(
         T: int,
@@ -383,10 +267,9 @@ class Kandinsky5DenoisingStage(PipelineStage):
             loader = TransformerLoader()
             self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
             fastvideo_args.model_loaded["transformer"] = True
-        self._assert_local_attention_backend_engaged()
 
         device = get_local_torch_device()
-        target_dtype = self._resolve_target_dtype(fastvideo_args)
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
         autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
         latents = batch.latents
         num_channels = getattr(
@@ -508,156 +391,6 @@ class Kandinsky5DenoisingStage(PipelineStage):
         return result
 
 
-class Kandinsky5DmdDenoisingStage(Kandinsky5DenoisingStage):
-    """DMD (few fixed steps, no CFG) variant of Kandinsky5DenoisingStage.
-
-    Reuses the parent's RoPE/scale_factor/sparse-params helpers; only the
-    denoising loop differs: a fixed short timestep schedule
-    (``pipeline_config.dmd_denoising_steps``) with a single forward pass per
-    step and no classifier-free-guidance branch, matching DMD's distilled
-    few-step generator.
-
-    ``dmd_denoising_steps`` values (e.g. ``[1000, 750, 500, 250]``) are
-    literal *final* target timesteps -- that's how ``DMD2Method`` on the
-    training side resolves them: nearest-sigma lookup against a scheduler
-    that has never had ``set_timesteps`` called on it, so e.g. "750" maps to
-    sigma 0.75. This stage therefore keeps a private
-    ``FlowMatchEulerDiscreteScheduler`` (same ``shift`` as the pipeline
-    scheduler) and drives it directly via predict-x0 + re-noise, exactly
-    mirroring ``DMD2Method._student_rollout``'s "simulate" branch and Wan's
-    own ``DmdDenoisingStage``. It must NOT reuse the pipeline's shared
-    ``scheduler`` object through ``scheduler.step()``: by the time this
-    stage runs, ``TimestepPreparationStage`` has already called
-    ``scheduler.set_timesteps(timesteps=dmd_denoising_steps)`` on it, which
-    re-applies the flow-match ``shift`` warp on top of values that are
-    already final (e.g. sigma 0.75 -> 0.9375 at shift=5) -- a double shift
-    that leaves every step far noisier than the student was trained for, so
-    the sampled video is still mostly noise after the last step instead of
-    converged.
-    """
-
-    def __init__(self, transformer, scheduler) -> None:
-        super().__init__(transformer, scheduler)
-        self._sample_scheduler = FlowMatchEulerDiscreteScheduler(shift=scheduler.shift)
-
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        if batch.latents is None:
-            raise ValueError("latents must be prepared before Kandinsky5 DMD denoising.")
-        if not fastvideo_args.model_loaded["transformer"]:
-            loader = TransformerLoader()
-            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
-            fastvideo_args.model_loaded["transformer"] = True
-        self._assert_local_attention_backend_engaged()
-
-        device = get_local_torch_device()
-        target_dtype = self._resolve_target_dtype(fastvideo_args)
-        autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
-        latents = batch.latents
-        num_channels = getattr(
-            self.transformer,
-            "in_visual_dim",
-            fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
-        )
-
-        prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
-        pooled = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
-        if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
-            raise ValueError("Kandinsky5 DMD requires Qwen prompt attention masks.")
-        text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask[0].to(device), device)
-
-        height = int(batch.height)
-        width = int(batch.width)
-        temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
-        spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
-        num_latent_frames = (int(batch.num_frames) - 1) // temporal_ratio + 1
-        visual_rope_pos = [
-            torch.arange(num_latent_frames, device=device),
-            torch.arange(height // spatial_ratio // 2, device=device),
-            torch.arange(width // spatial_ratio // 2, device=device),
-        ]
-        scale_factor = self._scale_factor(height, width)
-        sparse_params = self.get_sparse_params(latents, device)
-
-        dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps
-        if not dmd_steps:
-            raise ValueError("Kandinsky5 DMD denoising requires "
-                             "pipeline_config.dmd_denoising_steps to be set.")
-        timesteps = torch.tensor(dmd_steps, dtype=torch.long, device=device)
-
-        # I2V keeps the first (conditioning) frame fixed during denoising.
-        # Key off the actual image conditioning, not transformer.visual_cond:
-        # official T2V checkpoints also ship visual_cond=True, and skipping
-        # frame 0 for them leaves it as undenoised noise (see the same fix
-        # in Kandinsky5DenoisingStage.forward() above).
-        cond_frames = 1 if batch.image_latent is not None else 0
-
-        with tqdm(total=len(timesteps), desc="Kandinsky5 DMD Denoising") as progress_bar:
-            for i, timestep in enumerate(timesteps):
-                if hasattr(self, "interrupt") and self.interrupt:
-                    continue
-
-                t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
-                attn_metadata = None
-                if sparse_params is not None:
-                    attn_metadata = NablaAttentionMetadataBuilder().build(
-                        current_timestep=i,
-                        sta_mask=sparse_params["sta_mask"],
-                        P=sparse_params["P"],
-                        visual_shape=sparse_params["visual_shape"],
-                    )
-                autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
-                                if device.type == "cuda" else contextlib.nullcontext())
-                with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
-                                         forward_batch=batch), autocast_ctx:
-                    pred_velocity = self.transformer(
-                        hidden_states=latents.to(dtype=target_dtype),
-                        encoder_hidden_states=prompt_embeds,
-                        pooled_projections=pooled,
-                        timestep=t_expand,
-                        visual_rope_pos=visual_rope_pos,
-                        text_rope_pos=text_rope_pos,
-                        scale_factor=scale_factor,
-                        sparse_params=sparse_params,
-                        return_dict=True,
-                    ).sample
-
-                # Channel-last [B, T', H, W, C] -> channel-first [B, T', C, H, W]
-                # to match pred_noise_to_pred_video/add_noise's expected layout
-                # (same convention as Kandinsky5Model.add_noise/predict_x0).
-                sample_cf = latents[:, cond_frames:, :, :, :num_channels].permute(0, 1, 4, 2, 3)
-                velocity_cf = pred_velocity[:, cond_frames:].permute(0, 1, 4, 2, 3)
-                b, t = sample_cf.shape[:2]
-                step_timestep = timestep.reshape(1).to(device=device)
-                pred_x0_cf = pred_noise_to_pred_video(
-                    pred_noise=velocity_cf.flatten(0, 1),
-                    noise_input_latent=sample_cf.flatten(0, 1),
-                    timestep=step_timestep,
-                    scheduler=self._sample_scheduler,
-                ).unflatten(0, (b, t))
-
-                if i < len(timesteps) - 1:
-                    next_timestep = timesteps[i + 1].reshape(1).to(device=device)
-                    noise = randn_tensor(
-                        sample_cf.shape,
-                        generator=batch.generator,
-                        device=device,
-                        dtype=pred_x0_cf.dtype,
-                    )
-                    next_cf = self._sample_scheduler.add_noise(
-                        pred_x0_cf.flatten(0, 1),
-                        noise.flatten(0, 1),
-                        next_timestep,
-                    ).unflatten(0, (b, t))
-                else:
-                    next_cf = pred_x0_cf
-
-                latents[:, cond_frames:, :, :, :num_channels] = next_cf.permute(0, 1, 3, 4, 2).to(latents.dtype)
-                progress_bar.update()
-
-        batch.latents = latents[:, :, :, :, :num_channels]
-        return batch
-
-
 class Kandinsky5DecodingStage(DecodingStage):
 
     def __init__(self, vae: ParallelTiledVAE, pipeline=None) -> None:

fastvideo/platforms/cuda.py

@@ -142,13 +142,11 @@ class CudaPlatformBase(Platform):
                 logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.")
         elif selected_backend == AttentionBackendEnum.ATTN_QAT_INFER:
             from fastvideo.attention.backends.attn_qat_infer import (  # noqa: F401
-                AttnQatInferBackend, attn_qat_infer_receipt, is_attn_qat_infer_available)
+                AttnQatInferBackend, is_attn_qat_infer_available)
             if is_attn_qat_infer_available():
-                logger.info("Using Attn-QAT inference backend (%s).", attn_qat_infer_receipt())
+                logger.info("Using Attn-QAT inference (modified SageAttention3 FP4) backend.")
                 return "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
-            # Keep the trailing sentence stable: downstream receipts grep for it.
-            logger.info("Attn-QAT inference kernel is not built (%s). Fall back to Flash Attention.",
-                        attn_qat_infer_receipt())
+            logger.info("Attn-QAT inference kernel is not built. Fall back to Flash Attention.")
         elif selected_backend == AttentionBackendEnum.ATTN_QAT_TRAIN:
             from fastvideo.attention.backends.attn_qat_train import (  # noqa: F401
                 AttnQatTrainBackend, is_attn_qat_train_available)

fastvideo/tests/api/test_attn_qat_infer_capability_gate.py

@@ -1,124 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""ATTN_QAT_INFER selection must be gated on device capability, through
-the real platform resolver.
-
-``is_attn_qat_infer_available()`` used to test only whether the kernel
-extension imports. CUDA 13 wheel builds can carry the sm_120/sm_121
-extension on any host (e.g. H100 sm_90, GB200 sm_100): the import
-succeeds, ``CudaPlatformBase.get_attn_backend_cls`` selects the
-consumer-Blackwell backend, and the first kernel call fails with an
-unsupported-capability error -- instead of the FlashAttention fallback
-the QAD README documents for non-sm_120 GPUs.
-
-These tests drive the REAL resolver (``fastvideo.platforms.cuda``) and
-the REAL availability function; only the two physical facts are faked --
-"does the extension import" (``_get_attn_qat_infer``) and "what GPU is
-active" (``torch.cuda``). The stage-guard test
-(fastvideo/tests/stages/test_kandinsky5_attention_backend_guard.py)
-injects an already-resolved backend and by design cannot see this bug.
-
-CPU-only: the ATTN_QAT_INFER branch and its fallback never require a
-physical GPU to *resolve* (only to run).
-"""
-from __future__ import annotations
-
-import pytest
-import torch
-
-import fastvideo.attention.backends.attn_qat_infer as attn_qat_infer_module
-from fastvideo.attention.backends.attn_qat_infer import is_attn_qat_infer_available
-# The concrete resolver whose device queries go straight through torch.cuda
-# (the NVML variant only differs in device-info plumbing, not selection
-# logic) -- the abstract CudaPlatformBase can't run the fallthrough's
-# has_device_capability() check.
-from fastvideo.platforms.cuda import NonNvmlCudaPlatform
-from fastvideo.platforms.interface import AttentionBackendEnum
-
-ATTN_QAT_INFER_CLS = "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
-# What the resolver's fallthrough legitimately returns when ATTN_QAT_INFER
-# is unavailable: FlashAttention, or SDPA when flash_attn isn't installed
-# in the running environment (e.g. CPU-only CI).
-FALLBACK_CLASSES = {
-    "fastvideo.attention.backends.flash_attn.FlashAttentionBackend",
-    "fastvideo.attention.backends.sdpa.SDPABackend",
-}
-
-
-def _fake_gpu(monkeypatch,
-              *,
-              capability: tuple[int, int],
-              extension_imports: bool,
-              fa4_imports: bool = False) -> None:
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: capability)
-    monkeypatch.setattr(
-        attn_qat_infer_module,
-        "_get_attn_qat_infer",
-        lambda: (lambda *a, **k: None) if extension_imports else None,
-    )
-    # The FP4 FA4 probe (flash-attention-fp4, the sm_100a/sm_103a route) is
-    # a physical fact of the running environment; fake it like the others so
-    # these tests are deterministic on hosts with/without flash_attn.cute.
-    monkeypatch.setattr(attn_qat_infer_module, "_fa4_fp4_available", lambda: fa4_imports)
-
-
-def _resolve() -> str:
-    return NonNvmlCudaPlatform.get_attn_backend_cls(
-        AttentionBackendEnum.ATTN_QAT_INFER,
-        head_size=128,
-        dtype=torch.bfloat16,
-    )
-
-
-def test_sm90_host_with_bundled_extension_falls_back(monkeypatch):
-    """The reviewed failure: H100 + CUDA 13 wheel that bundles the sm_120
-    extension. Import succeeds; selection must still fall back."""
-    _fake_gpu(monkeypatch, capability=(9, 0), extension_imports=True)
-
-    assert not is_attn_qat_infer_available()
-    assert _resolve() in FALLBACK_CLASSES
-
-
-def test_sm100_host_with_bundled_extension_falls_back(monkeypatch):
-    """sm_100 with only the (unrunnable) bundled sm_12x extension and no
-    FP4 FA4 kernel still falls back -- the original reviewed failure class."""
-    _fake_gpu(monkeypatch, capability=(10, 0), extension_imports=True, fa4_imports=False)
-
-    assert not is_attn_qat_infer_available()
-    assert _resolve() in FALLBACK_CLASSES
-
-
-@pytest.mark.parametrize("capability", [(10, 0), (10, 3)])
-def test_datacenter_blackwell_with_fa4_selects_backend(monkeypatch, capability):
-    """sm_100a/sm_103a resolve ATTN_QAT_INFER through the FP4 FA4 route
-    when flash-attention-fp4 is installed (even without the sm_12x ext)."""
-    _fake_gpu(monkeypatch, capability=capability, extension_imports=False, fa4_imports=True)
-
-    assert is_attn_qat_infer_available()
-    assert _resolve() == ATTN_QAT_INFER_CLS
-
-
-@pytest.mark.parametrize("capability", [(12, 0), (12, 1)])
-def test_consumer_blackwell_with_extension_selects_backend(monkeypatch, capability):
-    _fake_gpu(monkeypatch, capability=capability, extension_imports=True)
-
-    assert is_attn_qat_infer_available()
-    assert _resolve() == ATTN_QAT_INFER_CLS
-
-
-def test_consumer_blackwell_without_extension_falls_back(monkeypatch):
-    _fake_gpu(monkeypatch, capability=(12, 0), extension_imports=False)
-
-    assert not is_attn_qat_infer_available()
-    assert _resolve() in FALLBACK_CLASSES
-
-
-def test_no_cuda_reports_unavailable(monkeypatch):
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
-    monkeypatch.setattr(
-        attn_qat_infer_module,
-        "_get_attn_qat_infer",
-        lambda: (lambda *a, **k: None),
-    )
-
-    assert not is_attn_qat_infer_available()

fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py

@@ -1,63 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regression test for exported Kandinsky5 DMD checkpoints resolving to the
-wrong pipeline.
-
-``fastvideo.train.entrypoint.dcp_to_diffusers`` copies the base T2V
-checkpoint's ``model_index.json`` unchanged into every export (see
-``examples/train/configs/fine_tuning/kandinsky5/README.md``), so
-``_class_name`` still says the base T2V pipeline. Without an explicit
-override, ``fastvideo.registry.get_model_info`` resolves such a directory to
-``Kandinsky5T2VPipeline`` (via the path-based "t2v" fallback detector in
-``fastvideo/registry.py``) instead of ``Kandinsky5DMDPipeline`` -- silently
-running the full-length multi-step sampler on DMD-distilled weights instead
-of the four-step re-noise sampler. ``override_pipeline_cls_name`` bypasses
-that resolution entirely and must be paired with ``Kandinsky5DMDConfig`` (not
-the default ``Kandinsky5T2VConfig``, whose ``dmd_denoising_steps=None`` makes
-``Kandinsky5DmdDenoisingStage`` raise immediately).
-"""
-
-import json
-
-from fastvideo.configs.pipelines.kandinsky5 import (
-    Kandinsky5DMDConfig,
-    Kandinsky5T2VConfig,
-)
-from fastvideo.pipelines.basic.kandinsky5.kandinsky5_dmd_pipeline import Kandinsky5DMDPipeline
-from fastvideo.pipelines.basic.kandinsky5.kandinsky5_pipeline import Kandinsky5T2VPipeline
-from fastvideo.pipelines.pipeline_registry import PipelineType
-from fastvideo.registry import get_model_info
-
-
-def _write_fake_export(tmp_path, name: str):
-    """Minimal on-disk stand-in for a dcp_to_diffusers export: just enough
-    for verify_model_config_and_directory to accept it."""
-    model_dir = tmp_path / name
-    (model_dir / "transformer").mkdir(parents=True)
-    (model_dir / "model_index.json").write_text(
-        json.dumps({
-            "_class_name": "Kandinsky5T2VPipeline",
-            "_diffusers_version": "0.30.0",
-        }))
-    return model_dir
-
-
-def test_kandinsky5_dmd_config_sets_denoising_steps():
-    assert Kandinsky5T2VConfig().dmd_denoising_steps is None
-    assert Kandinsky5DMDConfig().dmd_denoising_steps == [1000, 750, 500, 250]
-
-
-def test_exported_dmd_checkpoint_resolves_to_t2v_pipeline_without_override(tmp_path):
-    """Documents the bug: an unmodified export resolves to the T2V pipeline."""
-    model_dir = _write_fake_export(tmp_path, "kandinsky5_t2v_dmd2_4steps_qat")
-    info = get_model_info(model_path=str(model_dir), pipeline_type=PipelineType.BASIC)
-    assert info.pipeline_cls is Kandinsky5T2VPipeline
-
-
-def test_override_pipeline_cls_name_resolves_dmd_pipeline(tmp_path):
-    model_dir = _write_fake_export(tmp_path, "kandinsky5_t2v_dmd2_4steps_qat_override")
-    info = get_model_info(
-        model_path=str(model_dir),
-        pipeline_type=PipelineType.BASIC,
-        override_pipeline_cls_name="Kandinsky5DMDPipeline",
-    )
-    assert info.pipeline_cls is Kandinsky5DMDPipeline

fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py

@@ -1,184 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""CPU unit tests for the per-arch ATTN_QAT_INFER kernel resolution.
-
-The capability sets route:
-  * sm_120a/sm_121a -> fastvideo-kernel CUTLASS extension,
-  * sm_100a/sm_103a -> FP4 FA4 (flash-attention-fp4, #1221 plumbing),
-  * anything else  -> unavailable (honest FlashAttention fallback upstream).
-
-All device/import probes are monkeypatched; no GPU or kernel install needed.
-"""
-from __future__ import annotations
-
-import pytest
-
-import fastvideo.attention.backends.attn_qat_infer as aqi
-
-
-def _patch(monkeypatch, *, cap, cutlass, fa4) -> None:
-    monkeypatch.setattr(aqi, "_active_capability", lambda: cap)
-    monkeypatch.setattr(aqi, "_get_attn_qat_infer", lambda: (lambda *a, **k: None) if cutlass else None)
-    monkeypatch.setattr(aqi, "_fa4_fp4_available", lambda: fa4)
-
-
-@pytest.mark.parametrize(
-    "cap,cutlass,fa4,expected_kernel",
-    [
-        ((12, 0), True, False, "cutlass_sm12x"),
-        ((12, 1), True, False, "cutlass_sm12x"),
-        ((12, 0), False, False, None),  # ext not built
-        ((10, 0), False, True, "fa4_fp4"),  # GB200
-        ((10, 3), False, True, "fa4_fp4"),  # GB300
-        ((10, 0), False, False, None),  # fork not installed
-        ((9, 0), False, True, None),  # H100: FA4-FP4 set excludes it
-        ((8, 9), True, True, None),  # Ada: neither set
-        (None, True, True, None),  # no CUDA
-    ],
-)
-def test_arch_resolution(monkeypatch, cap, cutlass, fa4, expected_kernel) -> None:
-    _patch(monkeypatch, cap=cap, cutlass=cutlass, fa4=fa4)
-    assert aqi._resolved_kernel() == expected_kernel
-    assert aqi.is_attn_qat_infer_available() == (expected_kernel is not None)
-
-
-def test_receipt_records_fa4_quant_knobs(monkeypatch) -> None:
-    _patch(monkeypatch, cap=(10, 0), cutlass=False, fa4=True)
-    receipt = aqi.attn_qat_infer_receipt()
-    assert "arch=sm_100" in receipt
-    assert "qk_mode=nvfp4(per-16-e4m3-sf)" in receipt
-    assert "pv_mode=bf16" in receipt
-    assert "train_sim_mismatch=measured" in receipt
-
-
-def test_receipt_records_cutlass_scheme(monkeypatch) -> None:
-    _patch(monkeypatch, cap=(12, 0), cutlass=True, fa4=False)
-    receipt = aqi.attn_qat_infer_receipt()
-    assert "arch=sm_120" in receipt
-    assert "cutlass" in receipt
-
-
-def test_receipt_names_install_hint_on_uninstalled_fa4_arch(monkeypatch) -> None:
-    _patch(monkeypatch, cap=(10, 3), cutlass=False, fa4=False)
-    receipt = aqi.attn_qat_infer_receipt()
-    assert "arch=sm_103" in receipt
-    assert "flash-attention-fp4" in receipt
-
-
-def test_unsupported_arch_receipt_lists_support_matrix(monkeypatch) -> None:
-    _patch(monkeypatch, cap=(9, 0), cutlass=True, fa4=True)
-    receipt = aqi.attn_qat_infer_receipt()
-    assert "kernel=none" in receipt
-    assert "sm_100a/sm_103a" in receipt
-
-
-def test_unsupported_arch_forward_never_calls_bundled_extension(monkeypatch) -> None:
-    """An unsupported GPU (e.g. sm_90) carrying an importable sm_12x wheel
-    must fail cleanly at forward — never dispatch into the wrong binary."""
-    calls = []
-    monkeypatch.setattr(aqi, "_active_capability", lambda: (9, 0))
-    monkeypatch.setattr(aqi, "_get_attn_qat_infer", lambda: (lambda *a, **k: calls.append(1)))
-    monkeypatch.setattr(aqi, "_fa4_fp4_available", lambda: False)
-
-    import torch
-
-    impl = aqi.AttnQatInferImpl(num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5)
-    q = torch.zeros(1, 4, 1, 128)
-    with pytest.raises(ImportError, match="not available"):
-        impl.forward(q, q, q, attn_metadata=None)
-    assert not calls, "unsupported arch dispatched into the bundled CUTLASS extension"
-
-
-def test_fa4_route_resolution_runs_once_across_forwards(monkeypatch) -> None:
-    """The FA4 kernel/quantizer resolution is lazy but memoized: across N
-    forward calls the slow resolution path (imports + toolchain probe)
-    executes exactly once — per-forward re-resolution graph-breaks dynamo
-    every step and blocks fullgraph compilation."""
-    import torch
-
-    _patch(monkeypatch, cap=(10, 0), cutlass=False, fa4=True)
-
-    resolves = []
-
-    def fake_quant(t):
-        return t, torch.zeros(1)
-
-    def fake_kernel(q, k, v, sfq, sfk, softmax_scale=None, causal=False):
-        return v
-
-    monkeypatch.setattr(aqi, "_import_fa4_route_ops", lambda: (resolves.append(1) or (fake_quant, fake_kernel)))
-    monkeypatch.setattr(aqi, "_FA4_ROUTE_OPS", None)
-
-    impl = aqi.AttnQatInferImpl(num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5)
-    q = torch.zeros(1, 4, 1, 128)
-    for _ in range(3):
-        out = impl.forward(q, q, q, attn_metadata=None)
-    assert out.shape == q.shape
-    assert len(resolves) == 1, f"resolution ran {len(resolves)}x across 3 forwards"
-
-
-def _register_fa4_quantize_cpu_kernel():
-    """Register a CPU kernel for fastvideo::nvfp4_quantize_fa4 that honors the
-    production output contract — including STRIDES: the real sf output is a
-    permuted view of a contiguous (batch, nheads, rest_m, rest_k, 32, 4, 4)
-    buffer, and torch.compile bakes output strides into generated code."""
-    import torch
-
-    def _cpu_kernel(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        batch, seqlen, nheads, headdim = tensor_4d.shape
-        seqlen_padded = (seqlen + 127) // 128 * 128
-        fp4 = torch.zeros(batch, seqlen_padded, nheads, headdim // 2,
-                          dtype=torch.int8).view(torch.float4_e2m1fn_x2)
-        rest_m, rest_k = seqlen_padded // 128, (headdim // 16) // 4
-        sf = torch.zeros(batch, nheads, rest_m, rest_k, 32, 4, 4,
-                         dtype=torch.uint8).permute(4, 5, 2, 6, 3, 1, 0)
-        return fp4, sf
-
-    try:
-        torch.library.register_kernel("fastvideo::nvfp4_quantize_fa4", "cpu")(_cpu_kernel)
-    except RuntimeError:
-        pass  # already registered by a previous test/run
-
-
-def test_fa4_quantize_path_is_fullgraph_traceable() -> None:
-    """The FP4 quantize step must be a single opaque graph node: tracing its
-    python body descends into flashinfer's toolchain probe (a subprocess),
-    which torch.compile(fullgraph=True) rejects — a runtime memo cannot fix
-    that because cache hits are invisible at trace time. fullgraph=True over
-    the op-backed path must compile and run without a graph break."""
-    import torch
-
-    from fastvideo.attention.backends import flash_attn as fa
-
-    _register_fa4_quantize_cpu_kernel()
-
-    def path(x: torch.Tensor) -> torch.Tensor:
-        fp4, sf = fa._nvfp4_quantize_for_fa4(x)
-        return fp4[:, :x.shape[1]].view(torch.int8).float() + float(sf.shape[0])
-
-    compiled = torch.compile(path, fullgraph=True, backend="eager")
-    out = compiled(torch.randn(1, 64, 2, 128, dtype=torch.bfloat16))
-    assert out.shape == (1, 64, 2, 64)
-
-
-def test_fa4_quantize_op_fake_matches_real() -> None:
-    """torch.library.opcheck cross-checks the registered fake against a real
-    kernel run — shapes, dtypes, AND strides — so a fake whose output layout
-    drifts from the impl (e.g. contiguous fake vs permuted-view real sf, which
-    torch.compile turns into a runtime stride assertion) fails here on CPU.
-
-    The op is forward-only (no autograd registration), so restrict opcheck to
-    the non-autograd suites."""
-    import torch
-
-    import fastvideo.attention.backends.flash_attn  # noqa: F401  registers the op + fake
-
-    _register_fa4_quantize_cpu_kernel()
-
-    # batch>1 and seqlen>128 so every sf dim (incl. batch, rest_m) is
-    # non-trivial and its stride actually participates in the check.
-    x = torch.randn(2, 200, 2, 128, dtype=torch.bfloat16)
-    torch.library.opcheck(
-        torch.ops.fastvideo.nvfp4_quantize_fa4,
-        (x,),
-        test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"),
-    )

fastvideo/tests/attention/test_selector_role_override.py

@@ -54,24 +54,6 @@ def test_role_override_invalidates_cached_backend(monkeypatch) -> None:
         selector.global_force_attn_backend(None)
 
 
-def test_unsupported_role_override_honors_layer_default(monkeypatch) -> None:
-    monkeypatch.setattr(platforms, "_current_platform", _FakePlatform())
-    monkeypatch.setattr(selector, "resolve_obj_by_qualname", lambda name: name)
-    selector.global_force_attn_backend(None)
-
-    try:
-        with selector.global_force_attn_backend_context_manager(
-                AttentionBackendEnum.ATTN_QAT_TRAIN):
-            assert selector.get_attn_backend(
-                head_size=128,
-                dtype=torch.bfloat16,
-                supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA, ),
-                default_backend=AttentionBackendEnum.TORCH_SDPA,
-            ) == "TORCH_SDPA"
-    finally:
-        selector.global_force_attn_backend(None)
-
-
 def test_explicit_backend_config_rejects_typos() -> None:
     try:
         selector.coerce_attn_backend("attn_qat_typo")

fastvideo/tests/contract/test_fa4_dependency_pins.py

@@ -0,0 +1,50 @@
+# 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

fastvideo/tests/modal/kernel_build_cache.py

@@ -1,571 +0,0 @@
-"""Build and reuse FastVideo kernel wheels in Modal CI.
-
-This module stays standalone because Modal launchers execute it from a freshly
-cloned checkout after dependency installation. Shared cache consumers are
-read-only; only the explicit ``produce`` command may mutate cache state.
-"""
-
-import argparse
-import datetime
-import hashlib
-import json
-import os
-import platform
-import re
-import shlex
-import shutil
-import subprocess
-import sys
-import sysconfig
-import tempfile
-import zipfile
-from email.parser import Parser
-from pathlib import Path
-
-
-CACHE_SCHEMA_VERSION = 3
-DEFAULT_PREBUILT_INFO_PATH = "/opt/fastvideo-kernel-prebuilt"
-KERNEL_RELATIVE_DIR = "fastvideo-kernel"
-DEFAULT_BUILD_INFO_OUTPUT = "/opt/fastvideo-kernel-prebuilt/default/metadata.json"
-METADATA_FILE = "metadata.json"
-EXPECTED_DISTRIBUTION = "fastvideo-kernel"
-
-
-def _utc_now_isoformat() -> str:
-    return datetime.datetime.now(datetime.timezone.utc).isoformat()
-
-
-def _log(message: str) -> None:
-    print(f"[fastvideo-kernel-cache] {message}", flush=True)
-
-
-def _run(args: list[str],
-         *,
-         cwd: Path | None = None,
-         env: dict[str, str] | None = None,
-         capture: bool = False) -> str:
-    _log("$ " + " ".join(str(arg) for arg in args))
-    result = subprocess.run(
-        args,
-        cwd=str(cwd) if cwd else None,
-        env=env,
-        check=True,
-        text=True,
-        stdout=subprocess.PIPE if capture else None,
-        stderr=subprocess.STDOUT if capture else None,
-    )
-    if capture:
-        return result.stdout.strip()
-    return ""
-
-
-def _run_optional(args: list[str], *, cwd: Path | None = None) -> str:
-    try:
-        return _run(args, cwd=cwd, capture=True)
-    except (FileNotFoundError, subprocess.CalledProcessError, OSError) as error:
-        return f"<unavailable: {error}>"
-
-
-def _sha256_text(text: str) -> str:
-    return hashlib.sha256(text.encode("utf-8", errors="replace")).hexdigest()
-
-
-def _sha256_file(path: Path) -> str:
-    hasher = hashlib.sha256()
-    with path.open("rb") as file:
-        for chunk in iter(lambda: file.read(1024 * 1024), b""):
-            hasher.update(chunk)
-    return hasher.hexdigest()
-
-
-def _hash_directory(root: Path) -> str:
-    hasher = hashlib.sha256()
-    skip_dirs = {".git", ".mypy_cache", ".pytest_cache", "__pycache__", "build", "dist"}
-    for current_root, dirnames, filenames in os.walk(root):
-        dirnames[:] = sorted(
-            dirname for dirname in dirnames if dirname not in skip_dirs and not dirname.endswith(".egg-info"))
-        for filename in sorted(filenames):
-            path = Path(current_root) / filename
-            hasher.update(path.relative_to(root).as_posix().encode("utf-8"))
-            hasher.update(b"\0")
-            with path.open("rb") as file:
-                for chunk in iter(lambda: file.read(1024 * 1024), b""):
-                    hasher.update(chunk)
-            hasher.update(b"\0")
-    return hasher.hexdigest()
-
-
-def _kernel_source_hash(repo_root: Path) -> dict[str, str]:
-    kernel_root = repo_root / KERNEL_RELATIVE_DIR
-    try:
-        tree_hash = _run(["git", "rev-parse", f"HEAD:{KERNEL_RELATIVE_DIR}"], cwd=repo_root, capture=True)
-        diff = _run_optional(["git", "diff", "--binary", "--", KERNEL_RELATIVE_DIR], cwd=repo_root)
-        cached_diff = _run_optional(["git", "diff", "--cached", "--binary", "--", KERNEL_RELATIVE_DIR],
-                                    cwd=repo_root)
-        submodules = _run_optional(
-            [
-                "git",
-                "submodule",
-                "status",
-                "--",
-                "fastvideo-kernel/include/cutlass",
-                "fastvideo-kernel/include/tk",
-            ],
-            cwd=repo_root,
-        )
-        descriptor = {
-            "mode": "git",
-            "tree_hash": tree_hash,
-            "diff_hash": _sha256_text(diff) if diff else "",
-            "cached_diff_hash": _sha256_text(cached_diff) if cached_diff else "",
-            "submodules_hash": _sha256_text(submodules),
-        }
-        descriptor["source_hash"] = _sha256_text(json.dumps(descriptor, sort_keys=True))
-        return descriptor
-    except (FileNotFoundError, subprocess.CalledProcessError):
-        descriptor = {
-            "mode": "directory",
-            "tree_hash": _hash_directory(kernel_root),
-            "diff_hash": "",
-            "cached_diff_hash": "",
-            "submodules_hash": "",
-        }
-        descriptor["source_hash"] = descriptor["tree_hash"]
-        return descriptor
-
-
-def _read_kernel_version(repo_root: Path) -> str:
-    pyproject = repo_root / KERNEL_RELATIVE_DIR / "pyproject.toml"
-    for line in pyproject.read_text(encoding="utf-8").splitlines():
-        stripped = line.strip()
-        if stripped.startswith("version"):
-            _, _, value = stripped.partition("=")
-            return value.strip().strip("\"'")
-    return "<unknown>"
-
-
-def _detect_arch_from_torch() -> str:
-    try:
-        import torch
-
-        if not torch.cuda.is_available():
-            return "<unavailable>"
-        major, minor = torch.cuda.get_device_capability(0)
-        if major == 9 and minor == 0:
-            return "9.0a"
-        if major == 12 and minor == 0:
-            return "12.0a"
-        return f"{major}.{minor}"
-    except Exception as error:  # noqa: BLE001 - metadata should explain fallback
-        return f"<unavailable: {error}>"
-
-
-def _torch_metadata() -> dict[str, object]:
-    try:
-        import torch
-
-        compiled_with_cxx11_abi = getattr(torch, "compiled_with_cxx11_abi", None)
-        if callable(compiled_with_cxx11_abi):
-            cxx11_abi: object = bool(compiled_with_cxx11_abi())
-        else:
-            cxx11_abi = getattr(getattr(torch, "_C", object()), "_GLIBCXX_USE_CXX11_ABI", "<unavailable>")
-        return {
-            "torch_version": str(torch.__version__),
-            "torch_cuda_version": str(torch.version.cuda),
-            "torch_file": str(getattr(torch, "__file__", "")),
-            "torch_config": str(torch.__config__.show()),
-            "cxx11_abi": cxx11_abi,
-        }
-    except Exception as error:  # noqa: BLE001 - report in metadata
-        return {
-            "torch_version": f"<unavailable: {error}>",
-            "torch_cuda_version": "<unavailable>",
-            "torch_file": "<unavailable>",
-            "torch_config": "<unavailable>",
-            "cxx11_abi": "<unavailable>",
-        }
-
-
-def _selected_command_metadata(environment_name: str, default_command: str) -> dict[str, str]:
-    raw = os.environ.get(environment_name, "").strip() or default_command
-    try:
-        command = shlex.split(raw)
-    except ValueError as error:
-        return {
-            "raw": raw,
-            "command": "<invalid>",
-            "resolved_executable": "<invalid>",
-            "version": f"<unavailable: {error}>",
-        }
-    if not command:
-        command = [default_command]
-    resolved_executable = shutil.which(command[0]) or command[0]
-    return {
-        "raw": raw,
-        "command": json.dumps(command),
-        "resolved_executable": resolved_executable,
-        "version": _run_optional([*command, "--version"]),
-    }
-
-
-def _compiler_libc_metadata() -> dict[str, object]:
-    return {
-        "compiler": {
-            "cc": _selected_command_metadata("CC", "cc"),
-            "cxx": _selected_command_metadata("CXX", "c++"),
-        },
-        "build_tools": {
-            "cmake_version": _run_optional(["cmake", "--version"]),
-            "ninja_version": _run_optional(["ninja", "--version"]),
-            "linker": _selected_command_metadata("LD", "ld"),
-        },
-        "libc": {
-            "platform_libc": list(platform.libc_ver()),
-            "ldd_version": _run_optional(["ldd", "--version"]),
-        },
-    }
-
-
-def _build_metadata(repo_root: Path) -> dict[str, object]:
-    explicit_arch = os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()
-    resolved_arch = explicit_arch or _detect_arch_from_torch()
-    cache_key_build = {
-        "gpu_backend": os.environ.get("GPU_BACKEND", "CUDA"),
-        "resolved_torch_cuda_arch_list": resolved_arch,
-        "cmake_args": os.environ.get("CMAKE_ARGS", ""),
-        "cflags": os.environ.get("CFLAGS", ""),
-        "cxxflags": os.environ.get("CXXFLAGS", ""),
-        "ldflags": os.environ.get("LDFLAGS", ""),
-    }
-    cache_key_metadata: dict[str, object] = {
-        "schema_version": CACHE_SCHEMA_VERSION,
-        "kernel_version": _read_kernel_version(repo_root),
-        "source": _kernel_source_hash(repo_root),
-        "python": {
-            "version": platform.python_version(),
-            "implementation": platform.python_implementation(),
-            "cache_tag": sys.implementation.cache_tag,
-            "platform": sysconfig.get_platform(),
-            "machine": platform.machine(),
-        },
-        "torch": _torch_metadata(),
-        "cuda": {
-            "cuda_home": os.environ.get("CUDA_HOME", ""),
-            "nvcc": _selected_command_metadata("CUDACXX", "nvcc"),
-        },
-        "abi": _compiler_libc_metadata(),
-        "build": cache_key_build,
-    }
-    metadata = {
-        **cache_key_metadata,
-        "build": {
-            **cache_key_build,
-            "torch_cuda_arch_list": explicit_arch,
-        },
-    }
-    metadata["cache_key"] = _sha256_text(json.dumps(cache_key_metadata, sort_keys=True, separators=(",", ":")))
-    return metadata
-
-
-def _find_wheel(directory: Path) -> Path:
-    candidates = sorted(
-        path
-        for pattern in ("fastvideo_kernel-*.whl", "fastvideo-kernel-*.whl")
-        for path in directory.glob(pattern)
-    )
-    if not candidates:
-        raise RuntimeError(f"No fastvideo-kernel wheel found in {directory}")
-    return candidates[-1]
-
-
-def _normalize_distribution(name: str) -> str:
-    return re.sub(r"[-_.]+", "-", name).lower()
-
-
-def _wheel_artifact(wheel: Path) -> dict[str, object]:
-    if not wheel.is_file() or wheel.is_symlink():
-        raise RuntimeError(f"wheel is missing or is not a regular file: {wheel}")
-    try:
-        with zipfile.ZipFile(wheel) as archive:
-            bad_member = archive.testzip()
-            if bad_member is not None:
-                raise RuntimeError(f"wheel archive has a corrupt member: {bad_member}")
-            names = archive.namelist()
-            wheel_metadata = [name for name in names if name.endswith(".dist-info/WHEEL")]
-            package_metadata = [name for name in names if name.endswith(".dist-info/METADATA")]
-            if len(wheel_metadata) != 1 or len(package_metadata) != 1:
-                raise RuntimeError("wheel archive must contain one WHEEL and one METADATA file")
-            parsed_metadata = Parser().parsestr(archive.read(package_metadata[0]).decode("utf-8"))
-    except (OSError, UnicodeDecodeError, zipfile.BadZipFile) as error:
-        raise RuntimeError(f"invalid wheel archive: {error}") from error
-
-    distribution = _normalize_distribution(parsed_metadata.get("Name", ""))
-    if distribution != EXPECTED_DISTRIBUTION:
-        raise RuntimeError(f"unexpected wheel distribution {distribution!r}")
-    return {
-        "wheel_name": wheel.name,
-        "size_bytes": wheel.stat().st_size,
-        "sha256": _sha256_file(wheel),
-        "distribution": distribution,
-    }
-
-
-def _validate_wheel(wheel: Path, expected_artifact: object) -> tuple[bool, str]:
-    if not isinstance(expected_artifact, dict):
-        return False, "artifact metadata is missing"
-    try:
-        actual_artifact = _wheel_artifact(wheel)
-    except RuntimeError as error:
-        return False, str(error)
-    for field in ("wheel_name", "size_bytes", "sha256", "distribution"):
-        if actual_artifact[field] != expected_artifact.get(field):
-            return False, f"artifact {field} mismatch"
-    return True, ""
-
-
-def _install_wheel(wheel: Path) -> None:
-    _log(f"installing wheel {wheel}")
-    _run([
-        "uv",
-        "pip",
-        "install",
-        str(wheel),
-        "--reinstall-package",
-        "fastvideo-kernel",
-        "--no-deps",
-    ])
-
-
-def _load_metadata(path: Path) -> tuple[dict[str, object] | None, str]:
-    if path.is_symlink():
-        return None, "metadata is a symlink"
-    try:
-        metadata = json.loads(path.read_text(encoding="utf-8"))
-    except (OSError, json.JSONDecodeError) as error:
-        return None, f"metadata is unreadable: {error}"
-    if not isinstance(metadata, dict):
-        return None, "metadata is not an object"
-    return metadata, ""
-
-
-def _validated_payload_wheel(payload: dict[str, object], wheel: Path, cache_key: str) -> tuple[Path | None, str]:
-    if payload.get("schema_version") != CACHE_SCHEMA_VERSION:
-        return None, "schema version mismatch"
-    if payload.get("cache_key") != cache_key:
-        return None, "cache key mismatch"
-    valid, reason = _validate_wheel(wheel, payload.get("artifact"))
-    if not valid:
-        return None, reason
-    return wheel, ""
-
-
-def _cache_entry_wheel(cache_entry: Path, cache_key: str) -> tuple[Path | None, str]:
-    if cache_entry.is_symlink():
-        return None, "cache entry is a symlink"
-    metadata, reason = _load_metadata(cache_entry / METADATA_FILE)
-    if metadata is None:
-        return None, reason
-    artifact = metadata.get("artifact")
-    wheel_name = artifact.get("wheel_name") if isinstance(artifact, dict) else None
-    if not isinstance(wheel_name, str):
-        return None, "artifact wheel name is missing"
-    wheel = cache_entry / wheel_name
-    return _validated_payload_wheel(metadata, wheel, cache_key)
-
-
-def _prebuilt_metadata_paths(prebuilt_info_path: Path) -> list[Path]:
-    if prebuilt_info_path.is_dir() and not prebuilt_info_path.is_symlink():
-        return sorted(prebuilt_info_path.glob(f"*/{METADATA_FILE}"))
-    if prebuilt_info_path.exists() or prebuilt_info_path.is_symlink():
-        return [prebuilt_info_path]
-    return []
-
-
-def _try_install_prebuilt(metadata: dict[str, object], prebuilt_info_path: Path) -> bool:
-    cache_key = str(metadata["cache_key"])
-    for metadata_path in _prebuilt_metadata_paths(prebuilt_info_path):
-        prebuilt, reason = _load_metadata(metadata_path)
-        if prebuilt is None:
-            _log(f"Docker-prebuilt kernel metadata is invalid at {metadata_path}: {reason}")
-            continue
-        wheel_path = Path(str(prebuilt.get("wheel_path", "")))
-        wheel, reason = _validated_payload_wheel(prebuilt, wheel_path, cache_key)
-        if wheel is None:
-            _log(f"Docker-prebuilt kernel artifact rejected at {metadata_path}: {reason}")
-            continue
-        try:
-            _log(f"Docker-prebuilt cache hit for key {cache_key}: {wheel}")
-            _install_wheel(wheel)
-        except (FileNotFoundError, OSError, subprocess.CalledProcessError) as error:
-            _log(f"Docker-prebuilt kernel installation failed: {error}; falling back")
-            return False
-        return True
-    _log(f"Docker-prebuilt cache miss for key {cache_key}")
-    return False
-
-
-def _store_cache_entry(cache_root: Path, metadata: dict[str, object], wheel: Path) -> Path:
-    cache_key = str(metadata["cache_key"])
-    cache_entry = cache_root / cache_key
-    if cache_entry.exists() or cache_entry.is_symlink():
-        cached_wheel, reason = _cache_entry_wheel(cache_entry, cache_key)
-        if cached_wheel is not None:
-            _log(f"cache entry {cache_entry} appeared after build; leaving existing entry in place")
-            return cache_entry
-        _log(f"evicting invalid cache entry {cache_entry}: {reason}")
-        if cache_entry.is_dir() and not cache_entry.is_symlink():
-            shutil.rmtree(cache_entry)
-        else:
-            cache_entry.unlink()
-
-    temp_entry = Path(tempfile.mkdtemp(prefix=f".{cache_key}.tmp-", dir=cache_root))
-    try:
-        stored_wheel = temp_entry / wheel.name
-        shutil.copy2(wheel, stored_wheel)
-        stored_metadata = {
-            **metadata,
-            "created_at_utc": _utc_now_isoformat(),
-            "artifact": _wheel_artifact(stored_wheel),
-        }
-        (temp_entry / METADATA_FILE).write_text(json.dumps(stored_metadata, indent=2, sort_keys=True),
-                                                encoding="utf-8")
-        temp_entry.rename(cache_entry)
-    except OSError as error:
-        shutil.rmtree(temp_entry, ignore_errors=True)
-        cached_wheel, _ = _cache_entry_wheel(cache_entry, cache_key)
-        if cached_wheel is not None:
-            _log(f"cache entry {cache_entry} was created concurrently; leaving existing entry in place")
-            return cache_entry
-        raise RuntimeError(f"Failed to store kernel cache entry at {cache_entry}") from error
-
-    cached_wheel, reason = _cache_entry_wheel(cache_entry, cache_key)
-    if cached_wheel is None:
-        shutil.rmtree(cache_entry, ignore_errors=True)
-        raise RuntimeError(f"Stored kernel cache entry {cache_entry} could not be validated: {reason}")
-    return cache_entry
-
-
-def _build_wheel(repo_root: Path, metadata: dict[str, object]) -> Path:
-    with tempfile.TemporaryDirectory(prefix="fastvideo-kernel-wheel-") as temp_dir:
-        wheel_dir = Path(temp_dir)
-        env = os.environ.copy()
-        build_metadata = metadata["build"]
-        if not isinstance(build_metadata, dict):
-            raise RuntimeError("build metadata is invalid")
-        resolved_arch = str(build_metadata["resolved_torch_cuda_arch_list"])
-        if not env.get("TORCH_CUDA_ARCH_LIST") and not resolved_arch.startswith("<unavailable"):
-            env["TORCH_CUDA_ARCH_LIST"] = resolved_arch
-        _run(["./build.sh", "--wheel-dir", str(wheel_dir)], cwd=repo_root / KERNEL_RELATIVE_DIR, env=env)
-        wheel = _find_wheel(wheel_dir)
-        persisted = Path(tempfile.mkdtemp(prefix="fastvideo-kernel-built-wheel-")) / wheel.name
-        shutil.copy2(wheel, persisted)
-        return persisted
-
-
-def _build_and_install_local(repo_root: Path, metadata: dict[str, object]) -> None:
-    wheel = _build_wheel(repo_root, metadata)
-    try:
-        _wheel_artifact(wheel)
-        _install_wheel(wheel)
-    finally:
-        shutil.rmtree(wheel.parent, ignore_errors=True)
-
-
-def install_cached_or_build(repo_root: Path, cache_root: Path | None, prebuilt_info_path: Path) -> None:
-    """Install from a validated artifact or build locally without writing shared state."""
-    metadata = _build_metadata(repo_root)
-    cache_key = str(metadata["cache_key"])
-    _log(f"resolved cache key {cache_key}")
-
-    if _try_install_prebuilt(metadata, prebuilt_info_path):
-        return
-
-    if cache_root is not None:
-        cache_entry = cache_root / cache_key
-        cached_wheel, reason = _cache_entry_wheel(cache_entry, cache_key)
-        if cached_wheel is not None:
-            try:
-                _log(f"cache hit: {cached_wheel}")
-                _install_wheel(cached_wheel)
-                return
-            except (FileNotFoundError, OSError, subprocess.CalledProcessError) as error:
-                _log(f"cached wheel installation failed: {error}; building locally")
-        else:
-            _log(f"cache miss: {cache_entry} ({reason})")
-    else:
-        _log("shared cache disabled; using Docker-prebuilt artifact or a job-local build")
-
-    _build_and_install_local(repo_root, metadata)
-
-
-def produce_cached_wheel(repo_root: Path, cache_root: Path) -> None:
-    """Build and store one validated cache entry from a trusted producer."""
-    metadata = _build_metadata(repo_root)
-    cache_key = str(metadata["cache_key"])
-    _log(f"resolved cache key {cache_key}")
-    cache_root.mkdir(parents=True, exist_ok=True)
-    cache_entry = cache_root / cache_key
-    cached_wheel, reason = _cache_entry_wheel(cache_entry, cache_key)
-    if cached_wheel is not None:
-        _log(f"cache hit: {cached_wheel}")
-        return
-    if cache_entry.exists() or cache_entry.is_symlink():
-        _log(f"producer will repair invalid cache entry {cache_entry}: {reason}")
-
-    _log(f"cache miss: {cache_entry}")
-    wheel = _build_wheel(repo_root, metadata)
-    try:
-        _wheel_artifact(wheel)
-        _install_wheel(wheel)
-        stored_entry = _store_cache_entry(cache_root, metadata, wheel)
-        _log(f"stored wheel cache entry: {stored_entry}")
-    finally:
-        shutil.rmtree(wheel.parent, ignore_errors=True)
-
-
-def write_build_info(repo_root: Path, wheel_dir: Path, output: Path) -> None:
-    metadata = _build_metadata(repo_root)
-    wheel = _find_wheel(wheel_dir)
-    output.parent.mkdir(parents=True, exist_ok=True)
-    payload = {
-        **metadata,
-        "created_at_utc": _utc_now_isoformat(),
-        "artifact": _wheel_artifact(wheel),
-        "wheel_path": str(wheel),
-    }
-    output.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-    _log(f"wrote Docker-prebuilt kernel metadata: {output}")
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("command", choices=("install", "produce", "write-build-info", "print-key"))
-    parser.add_argument("--repo-root", default=os.getcwd())
-    parser.add_argument("--cache-root", default=os.environ.get("FASTVIDEO_KERNEL_CACHE_ROOT", ""))
-    parser.add_argument("--prebuilt-info-path",
-                        default=os.environ.get("FASTVIDEO_KERNEL_PREBUILT_INFO", DEFAULT_PREBUILT_INFO_PATH))
-    parser.add_argument("--wheel-dir", default="")
-    parser.add_argument("--output", default=DEFAULT_BUILD_INFO_OUTPUT)
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    repo_root = Path(args.repo_root).resolve()
-    cache_root = Path(args.cache_root) if args.cache_root else None
-    if args.command == "install":
-        install_cached_or_build(repo_root, cache_root, Path(args.prebuilt_info_path))
-    elif args.command == "produce":
-        if cache_root is None:
-            raise RuntimeError("--cache-root is required for produce")
-        produce_cached_wheel(repo_root, cache_root)
-    elif args.command == "write-build-info":
-        if not args.wheel_dir:
-            raise RuntimeError("--wheel-dir is required for write-build-info")
-        write_build_info(repo_root, Path(args.wheel_dir), Path(args.output))
-    elif args.command == "print-key":
-        print(json.dumps(_build_metadata(repo_root), indent=2, sort_keys=True))
-
-
-if __name__ == "__main__":
-    main()

fastvideo/tests/modal/kernel_cache_smoke.py

@@ -1,253 +0,0 @@
-"""Cross-container smoke for the Modal fastvideo-kernel build cache.
-
-Example:
-    python -m modal run fastvideo/tests/modal/kernel_cache_smoke.py \
-        --git-repo https://github.com/macthecadillac/FastVideo.git \
-        --git-commit <commit>
-
-This is intentionally opt-in: it builds fastvideo-kernel on a real Modal GPU.
-"""
-
-import os
-import shlex
-import shutil
-import subprocess
-import sys
-import time
-from typing import Any
-
-import modal
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-try:
-    from modal_image_utils import (  # noqa: E402
-        resolve_image_ref, resolve_uv_torch_backend)
-except ModuleNotFoundError:
-    # Remote Modal containers re-import this module but mount only the
-    # entrypoint file; the digest resolution already happened at local launch.
-    def resolve_image_ref(image_ref: str) -> str:
-        return image_ref
-
-    def resolve_uv_torch_backend(image_tag: str) -> str | None:
-        return os.environ.get("UV_TORCH_BACKEND")
-
-
-app = modal.App("fastvideo-kernel-cache-smoke")
-
-REPO_DIR = "/FastVideo"
-KERNEL_CACHE_VOLUME_PATH = "/root/fastvideo-kernel-cache"
-KERNEL_CACHE_VOLUME_NAME = "fastvideo-kernel-build-cache-smoke-v3"
-IMAGE_VERSION = os.getenv("IMAGE_VERSION", "latest")
-IMAGE_TAG = os.environ.get(
-    "FASTVIDEO_MODAL_IMAGE",
-    f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{IMAGE_VERSION}",
-)
-IMAGE_REF = resolve_image_ref(IMAGE_TAG)
-
-print(f"Using image: {IMAGE_REF}")
-print(f"Using kernel cache volume: {KERNEL_CACHE_VOLUME_NAME}")
-
-kernel_cache_vol = modal.Volume.from_name(KERNEL_CACHE_VOLUME_NAME, create_if_missing=True)
-uv_torch_backend_override = resolve_uv_torch_backend(IMAGE_TAG)
-
-image = (
-    modal.Image.from_registry(IMAGE_REF, add_python="3.12")
-    .apt_install(
-        "cmake",
-        "pkg-config",
-        "build-essential",
-        "curl",
-        "git",
-        "libssl-dev",
-        "ffmpeg",
-    )
-    .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
-    .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
-    .env({
-        "PATH": "/root/.cargo/bin:$PATH",
-        **({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}),
-    })
-)
-
-PRODUCER_KWARGS = dict(
-    gpu="L40S:1",
-    image=image,
-    timeout=5400,
-    volumes={KERNEL_CACHE_VOLUME_PATH: kernel_cache_vol},
-)
-CONSUMER_KWARGS = dict(
-    gpu="L40S:1",
-    image=image,
-    timeout=5400,
-    volumes={KERNEL_CACHE_VOLUME_PATH: kernel_cache_vol.with_mount_options(read_only=True)},
-)
-
-
-def _run_local_git_command(args: list[str]) -> str:
-    result = subprocess.run(
-        ["git", *args],
-        check=True,
-        capture_output=True,
-        text=True,
-    )
-    return result.stdout.strip()
-
-
-def _normalize_git_repo_url(git_repo: str) -> str:
-    if git_repo.startswith("git@github.com:"):
-        return "https://github.com/" + git_repo[len("git@github.com:"):]
-    if git_repo.startswith("ssh://git@github.com/"):
-        return "https://github.com/" + git_repo[len("ssh://git@github.com/"):]
-    return git_repo
-
-
-def _resolve_git_repo(git_repo: str) -> str:
-    if git_repo.strip():
-        return _normalize_git_repo_url(git_repo.strip())
-    env_repo = os.environ.get("BUILDKITE_REPO", "").strip()
-    if env_repo:
-        return _normalize_git_repo_url(env_repo)
-    discovered_repo = _run_local_git_command(["config", "--get", "remote.origin.url"])
-    if discovered_repo:
-        return _normalize_git_repo_url(discovered_repo)
-    raise RuntimeError("Could not resolve git repo URL. Pass --git-repo or set BUILDKITE_REPO.")
-
-
-def _resolve_git_commit(git_commit: str) -> str:
-    if git_commit.strip():
-        return git_commit.strip()
-    env_commit = os.environ.get("BUILDKITE_COMMIT", "").strip()
-    if env_commit:
-        return env_commit
-    discovered_commit = _run_local_git_command(["rev-parse", "HEAD"])
-    if discovered_commit:
-        return discovered_commit
-    raise RuntimeError("Could not resolve git commit. Pass --git-commit or set BUILDKITE_COMMIT.")
-
-
-def _run(args: list[str], cwd: str | None = None) -> str:
-    print("$ " + " ".join(shlex.quote(arg) for arg in args), flush=True)
-    result = subprocess.run(
-        args,
-        cwd=cwd,
-        check=True,
-        stdout=subprocess.PIPE,
-        stderr=sys.stderr,
-        text=True,
-    )
-    if result.stdout:
-        print(result.stdout, end="", flush=True)
-    return result.stdout.strip()
-
-
-def _run_shell_capture(command: str, cwd: str) -> str:
-    print(f"$ {command}", flush=True)
-    result = subprocess.run(
-        ["/bin/bash", "-lc", command],
-        cwd=cwd,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        text=True,
-        check=False,
-    )
-    if result.stdout:
-        print(result.stdout, end="", flush=True)
-    if result.returncode != 0:
-        raise RuntimeError(f"Command failed with exit code {result.returncode}: {command}")
-    return result.stdout
-
-
-def _clone_checkout(git_repo: str, git_commit: str) -> str:
-    shutil.rmtree(REPO_DIR, ignore_errors=True)
-    _run(["git", "-c", "http.version=HTTP/1.1", "clone", git_repo, REPO_DIR], cwd="/")
-    _run(["git", "checkout", git_commit], cwd=REPO_DIR)
-    _run(["git", "submodule", "update", "--init", "--recursive"], cwd=REPO_DIR)
-    return _run(["git", "rev-parse", "HEAD"], cwd=REPO_DIR)
-
-
-def _cache_command_output(command: str, cache_root: str) -> str:
-    # This smoke exercises the Modal volume cache specifically, so do not let
-    # a matching Docker-prebuilt wheel satisfy the miss/hit assertions.
-    disabled_prebuilt_info_path = f"{cache_root}/.docker-prebuilt-disabled.json"
-    return _run_shell_capture(
-        "source $HOME/.local/bin/env && "
-        "source /opt/venv/bin/activate && "
-        f"python fastvideo/tests/modal/kernel_build_cache.py {shlex.quote(command)} "
-        f"--cache-root {shlex.quote(cache_root)} "
-        f"--prebuilt-info-path {shlex.quote(disabled_prebuilt_info_path)}",
-        cwd=REPO_DIR,
-    )
-
-
-def _assert_in_output(output: str, needle: str) -> None:
-    if needle not in output:
-        raise RuntimeError(f"Expected output to contain {needle!r}")
-
-
-def _assert_not_in_output(output: str, needle: str) -> None:
-    if needle in output:
-        raise RuntimeError(f"Expected output not to contain {needle!r}")
-
-
-@app.function(**PRODUCER_KWARGS)
-def build_and_store_cache(git_repo: str, git_commit: str, cache_root: str) -> dict[str, Any]:
-    checked_out_commit = _clone_checkout(git_repo, git_commit)
-    output = _cache_command_output("produce", cache_root)
-    _assert_in_output(output, "cache miss:")
-    _assert_in_output(output, "./build.sh --wheel-dir")
-    _assert_in_output(output, "stored wheel cache entry:")
-    print("Committing kernel cache volume after miss/store", flush=True)
-    kernel_cache_vol.commit()
-    return {
-        "checked_out_commit": checked_out_commit,
-        "cache_root": cache_root,
-        "cache_miss": True,
-        "stored": True,
-    }
-
-
-@app.function(**CONSUMER_KWARGS)
-def verify_cache_hit(git_repo: str, git_commit: str, cache_root: str) -> dict[str, Any]:
-    checked_out_commit = _clone_checkout(git_repo, git_commit)
-    output = _cache_command_output("install", cache_root)
-    _assert_in_output(output, "cache hit:")
-    _assert_not_in_output(output, "cache miss:")
-    _assert_not_in_output(output, "./build.sh --wheel-dir")
-    return {
-        "checked_out_commit": checked_out_commit,
-        "cache_root": cache_root,
-        "cache_hit": True,
-    }
-
-
-@app.function(image=image, timeout=300, volumes={KERNEL_CACHE_VOLUME_PATH: kernel_cache_vol})
-def cleanup_cache(cache_root: str) -> None:
-    shutil.rmtree(cache_root, ignore_errors=True)
-    print(f"Removed smoke cache root {cache_root}", flush=True)
-    kernel_cache_vol.commit()
-
-
-@app.local_entrypoint()
-def main(
-    git_repo: str = "",
-    git_commit: str = "",
-    cache_namespace: str = "",
-    cleanup: bool = True,
-):
-    resolved_git_repo = _resolve_git_repo(git_repo)
-    resolved_git_commit = _resolve_git_commit(git_commit)
-    namespace = cache_namespace.strip() or f"smoke-{resolved_git_commit[:12]}-{int(time.time())}"
-    cache_root = f"{KERNEL_CACHE_VOLUME_PATH}/{namespace}"
-
-    print("Launching cross-container kernel cache smoke")
-    print(f"Repo: {resolved_git_repo}")
-    print(f"Commit: {resolved_git_commit}")
-    print(f"Cache root: {cache_root}")
-
-    miss_result = build_and_store_cache.remote(resolved_git_repo, resolved_git_commit, cache_root)
-    hit_result = verify_cache_hit.remote(resolved_git_repo, resolved_git_commit, cache_root)
-    if cleanup:
-        cleanup_cache.remote(cache_root)
-
-    print("Kernel cache smoke passed")
-    print({"miss": miss_result, "hit": hit_result, "cleanup": cleanup})

fastvideo/tests/modal/pr_test.py

@@ -209,10 +209,6 @@ def run_test_command(test_command: str,
     lane would then test stale kernels. Pass install_command="" for commands
     that manage their own installs.
     """
-    import os
-    import subprocess
-    import sys
-
     git_repo = os.environ.get("BUILDKITE_REPO", "")
     git_commit = os.environ.get("BUILDKITE_COMMIT", "")
     pr_number = os.environ.get("BUILDKITE_PULL_REQUEST")
@@ -223,32 +219,23 @@ def run_test_command(test_command: str,
         print(f"PR number: {pr_number}")
     _checkout_repository(git_repo, git_commit, pr_number)
 
-    setup_steps = [
-        "source $HOME/.local/bin/env",
-        "source /opt/venv/bin/activate",
-        "cd /FastVideo",
-    ]
-    if install_command:
-        setup_steps.append(install_command)
-    if build_kernel:
-        setup_steps.append("python fastvideo/tests/modal/kernel_build_cache.py install")
-    setup_command = " &&\n    ".join(setup_steps)
+    build_kernel_command = """
+    cd fastvideo-kernel &&
+    ./build.sh &&
+    cd .. &&
+    """ if build_kernel else ""
 
-    setup_result = subprocess.run(["/bin/bash", "-c", setup_command],
-                                  stdout=sys.stdout,
-                                  stderr=sys.stderr,
-                                  check=False)
-    if setup_result.returncode != 0:
-        raise RuntimeError(
-            f"Setup command failed with exit code {setup_result.returncode}")
+    install_clause = f"{install_command} &&" if install_command else ""
 
+    command = f"""
+    source $HOME/.local/bin/env &&
+    source /opt/venv/bin/activate &&
+    cd /FastVideo &&
+    {install_clause}
+    {build_kernel_command}
+    {test_command}
+    """
 
-    command = " &&\n    ".join([
-        "source $HOME/.local/bin/env",
-        "source /opt/venv/bin/activate",
-        "cd /FastVideo",
-        test_command,
-    ])
     result = subprocess.run(["/bin/bash", "-c", command],
                             stdout=sys.stdout,
                             stderr=sys.stderr,
@@ -382,9 +369,7 @@ def run_unit_test():
         "./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ "
         "./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/worker/ "
         "./fastvideo/tests/training/test_trackers.py "
-        "./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py "
-        "./fastvideo/tests/modal/test_kernel_build_cache.py ./fastvideo/tests/modal/test_pr_test.py "
-        "./fastvideo/tests/modal/test_ssim_test.py "
+        "./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py ./fastvideo/tests/modal/test_pr_test.py "
         "--ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py "
         "--ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
     )

fastvideo/tests/modal/ssim_test.py

@@ -507,7 +507,9 @@ def _prepare_ssim_workspace(
     rm -rf fastvideo/tests/ssim/reference_videos
     git_retry git submodule update --init --recursive
     uv pip install -e ".[test]"
-    python fastvideo/tests/modal/kernel_build_cache.py install
+    cd fastvideo-kernel
+    ./build.sh
+    cd ..
     uv pip install git+https://github.com/microsoft/MoGe.git
     # Stable Audio Open 1.0 inference deps (optional in basic install,
     # required by `StableAudioDenoisingStage`; consumed by
@@ -525,11 +527,9 @@ def _prepare_ssim_workspace(
             "HF_API_KEY": hf_api_key,
         },
     )
-    if result.stdout:
-        print(result.stdout, end="" if result.stdout.endswith("\n") else "\n")
-    if result.stderr:
-        print(result.stderr, end="" if result.stderr.endswith("\n") else "\n", file=sys.stderr)
     if result.returncode != 0:
+        print(result.stdout)
+        print(result.stderr)
         raise RuntimeError(f"Workspace setup failed with exit code {result.returncode}")
 
     ssim_dir = os.path.join(repo_root, "fastvideo", "tests", "ssim")

fastvideo/tests/modal/test_kernel_build_cache.py

@@ -1,429 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import errno
-import importlib.util
-import json
-import subprocess
-import zipfile
-from pathlib import Path
-
-import pytest
-import yaml
-
-
-def _load_kernel_build_cache():
-    module_path = Path(__file__).with_name("kernel_build_cache.py")
-    spec = importlib.util.spec_from_file_location("kernel_build_cache_for_test", module_path)
-    assert spec is not None
-    assert spec.loader is not None
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
-
-
-kernel_build_cache = _load_kernel_build_cache()
-REPO_ROOT = Path(__file__).resolve().parents[3]
-WHEEL_NAME = "fastvideo_kernel-0.3.2-py3-none-any.whl"
-TRUSTED_KERNEL_IMAGE_INPUTS = {
-    ".dockerignore",
-    ".github/workflows/_template-build-image.yml",
-    ".github/workflows/infra-build-image.yml",
-    ".gitmodules",
-    "docker/Dockerfile",
-    "docker/uv-excludes",
-    "fastvideo-kernel/**",
-    "fastvideo/tests/modal/kernel_build_cache.py",
-    "pyproject.toml",
-}
-
-
-def _write_wheel(path: Path, *, distribution: str = "fastvideo-kernel", payload: str = "payload") -> Path:
-    dist_info = "fastvideo_kernel-0.3.2.dist-info"
-    with zipfile.ZipFile(path, "w") as archive:
-        archive.writestr("fastvideo_kernel/__init__.py", payload)
-        archive.writestr(
-            f"{dist_info}/WHEEL",
-            "Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: py3-none-any\n",
-        )
-        archive.writestr(
-            f"{dist_info}/METADATA",
-            f"Metadata-Version: 2.1\nName: {distribution}\nVersion: 0.3.2\n",
-        )
-    return path
-
-
-def _metadata(cache_key: str = "cache-key") -> dict[str, object]:
-    return {
-        "cache_key": cache_key,
-        "schema_version": kernel_build_cache.CACHE_SCHEMA_VERSION,
-    }
-
-
-def _write_cache_entry(cache_root: Path, cache_key: str, *, payload: str = "payload") -> Path:
-    cache_entry = cache_root / cache_key
-    cache_entry.mkdir()
-    wheel = _write_wheel(cache_entry / WHEEL_NAME, payload=payload)
-    stored_metadata = {
-        **_metadata(cache_key),
-        "artifact": kernel_build_cache._wheel_artifact(wheel),
-    }
-    (cache_entry / kernel_build_cache.METADATA_FILE).write_text(json.dumps(stored_metadata), encoding="utf-8")
-    return cache_entry
-
-
-def _write_prebuilt_artifact(prebuilt_root: Path, name: str, cache_key: str) -> Path:
-    artifact_dir = prebuilt_root / name
-    artifact_dir.mkdir(parents=True)
-    wheel = _write_wheel(artifact_dir / WHEEL_NAME)
-    payload = {
-        **_metadata(cache_key),
-        "artifact": kernel_build_cache._wheel_artifact(wheel),
-        "wheel_path": str(wheel),
-    }
-    (artifact_dir / kernel_build_cache.METADATA_FILE).write_text(json.dumps(payload), encoding="utf-8")
-    return wheel
-
-
-def _patch_stable_metadata(monkeypatch) -> None:
-    for name in (
-        "GPU_BACKEND",
-        "CMAKE_ARGS",
-        "CFLAGS",
-        "CXXFLAGS",
-        "LDFLAGS",
-        "TORCH_CUDA_ARCH_LIST",
-        "FASTVIDEO_CONTAINER_IMAGE_REF",
-        "CUDACXX",
-    ):
-        monkeypatch.delenv(name, raising=False)
-    monkeypatch.setattr(kernel_build_cache, "_kernel_source_hash", lambda repo_root: {"source_hash": "source"})
-    monkeypatch.setattr(kernel_build_cache, "_read_kernel_version", lambda repo_root: "0.3.2")
-    monkeypatch.setattr(
-        kernel_build_cache,
-        "_torch_metadata",
-        lambda: {
-            "torch_version": "2.9.0",
-            "torch_cuda_version": "12.8",
-            "torch_file": "/opt/venv/lib/python3.12/site-packages/torch/__init__.py",
-            "torch_config": "USE_CUDA=ON",
-            "cxx11_abi": True,
-        },
-    )
-    monkeypatch.setattr(
-        kernel_build_cache,
-        "_compiler_libc_metadata",
-        lambda: {
-            "compiler": {"cc": {"path": "/usr/bin/gcc", "version": "gcc 11.4.0"}},
-            "libc": {"ldd_version": "ldd 2.35"},
-        },
-    )
-    monkeypatch.setattr(
-        kernel_build_cache,
-        "_selected_command_metadata",
-        lambda environment_name, default_command: {
-            "raw": default_command,
-            "command": json.dumps([default_command]),
-            "resolved_executable": f"/usr/bin/{default_command}",
-            "version": f"{default_command} 12.8",
-        },
-    )
-
-
-def test_kernel_only_main_change_republishes_trusted_l40s_artifact() -> None:
-    workflow = yaml.load(
-        (REPO_ROOT / ".github/workflows/infra-build-image.yml").read_text(encoding="utf-8"),
-        Loader=yaml.BaseLoader,
-    )
-    push = workflow["on"]["push"]
-
-    assert push["branches"] == ["main"]
-    assert TRUSTED_KERNEL_IMAGE_INPUTS <= set(push["paths"])
-    assert "github.event_name == 'push'" in workflow["jobs"]["build-cuda-images"]["if"]
-
-    dockerfile = (REPO_ROOT / "docker/Dockerfile").read_text(encoding="utf-8")
-    assert 'l40s_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/8.9"' in dockerfile
-    assert 'export TORCH_CUDA_ARCH_LIST=8.9' in dockerfile
-    assert '--output "${l40s_wheel_dir}/metadata.json"' in dockerfile
-
-
-def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, tmp_path) -> None:
-    _patch_stable_metadata(monkeypatch)
-    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a")
-
-    monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", "9.0a")
-    explicit_hopper = kernel_build_cache._build_metadata(tmp_path)
-
-    monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False)
-    detected_hopper = kernel_build_cache._build_metadata(tmp_path)
-
-    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "8.9")
-    detected_l40s = kernel_build_cache._build_metadata(tmp_path)
-
-    assert explicit_hopper["cache_key"] == detected_hopper["cache_key"]
-    assert explicit_hopper["build"]["torch_cuda_arch_list"] == "9.0a"
-    assert detected_hopper["build"]["torch_cuda_arch_list"] == ""
-    assert detected_l40s["cache_key"] != detected_hopper["cache_key"]
-
-
-@pytest.mark.parametrize("environment_name", ["CFLAGS", "CXXFLAGS", "LDFLAGS"])
-def test_cache_key_changes_with_build_flags(monkeypatch, tmp_path, environment_name) -> None:
-    _patch_stable_metadata(monkeypatch)
-    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a")
-    baseline = kernel_build_cache._build_metadata(tmp_path)
-
-    monkeypatch.setenv(environment_name, "-DFASTVIDEO_ABI_VARIANT=1")
-    changed = kernel_build_cache._build_metadata(tmp_path)
-
-    assert baseline["cache_key"] != changed["cache_key"]
-
-
-def test_cache_key_changes_with_compiler_or_torch_abi(monkeypatch, tmp_path) -> None:
-    _patch_stable_metadata(monkeypatch)
-    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a")
-    baseline = kernel_build_cache._build_metadata(tmp_path)
-
-    monkeypatch.setattr(
-        kernel_build_cache,
-        "_compiler_libc_metadata",
-        lambda: {
-            "compiler": {"cc": {"path": "/opt/gcc-12/bin/gcc", "version": "gcc 12.3.0"}},
-            "libc": {"ldd_version": "ldd 2.35"},
-        },
-    )
-    compiler_changed = kernel_build_cache._build_metadata(tmp_path)
-
-    _patch_stable_metadata(monkeypatch)
-    torch_metadata = kernel_build_cache._torch_metadata()
-    monkeypatch.setattr(kernel_build_cache, "_torch_metadata", lambda: {**torch_metadata, "cxx11_abi": False})
-    torch_abi_changed = kernel_build_cache._build_metadata(tmp_path)
-
-    assert baseline["schema_version"] == 3
-    assert baseline["cache_key"] != compiler_changed["cache_key"]
-    assert baseline["cache_key"] != torch_abi_changed["cache_key"]
-
-
-def test_find_wheel_uses_fastvideo_kernel_wheel_patterns(tmp_path) -> None:
-    (tmp_path / "unrelated-0.1.0.whl").write_text("")
-    (tmp_path / "fastvideo_kernel-0.3.1-py3-none-any.whl").write_text("")
-    expected = tmp_path / WHEEL_NAME
-    expected.write_text("")
-
-    assert kernel_build_cache._find_wheel(tmp_path) == expected
-
-
-def test_wheel_artifact_validates_archive_digest_size_and_distribution(tmp_path) -> None:
-    wheel = _write_wheel(tmp_path / WHEEL_NAME)
-    artifact = kernel_build_cache._wheel_artifact(wheel)
-
-    valid, reason = kernel_build_cache._validate_wheel(wheel, artifact)
-
-    assert valid
-    assert reason == ""
-    assert artifact["distribution"] == "fastvideo-kernel"
-    assert artifact["size_bytes"] == wheel.stat().st_size
-    assert len(artifact["sha256"]) == 64
-
-
-@pytest.mark.parametrize("corruption", ["digest", "size", "archive", "distribution"])
-def test_wheel_validation_rejects_corrupt_or_unexpected_artifacts(tmp_path, corruption) -> None:
-    wheel = _write_wheel(tmp_path / WHEEL_NAME)
-    artifact = kernel_build_cache._wheel_artifact(wheel)
-    if corruption == "digest":
-        artifact["sha256"] = "0" * 64
-    elif corruption == "size":
-        artifact["size_bytes"] = int(artifact["size_bytes"]) + 1
-    elif corruption == "archive":
-        wheel.write_bytes(b"not a wheel")
-    else:
-        wheel = _write_wheel(tmp_path / "unexpected.whl", distribution="other-package")
-        artifact = {**artifact, "wheel_name": wheel.name}
-
-    valid, reason = kernel_build_cache._validate_wheel(wheel, artifact)
-
-    assert not valid
-    assert reason
-
-
-def test_store_cache_entry_handles_concurrent_valid_entry(monkeypatch, tmp_path) -> None:
-    cache_root = tmp_path / "cache"
-    cache_root.mkdir()
-    wheel = _write_wheel(tmp_path / WHEEL_NAME)
-    rename_calls = []
-
-    def fake_rename(self, target):
-        rename_calls.append((self, target))
-        _write_cache_entry(target.parent, target.name)
-        raise OSError(errno.ENOTEMPTY, "Directory not empty")
-
-    monkeypatch.setattr(Path, "rename", fake_rename)
-
-    cache_entry = kernel_build_cache._store_cache_entry(cache_root, _metadata(), wheel)
-
-    assert cache_entry == cache_root / "cache-key"
-    assert rename_calls
-    assert not list(cache_root.glob(".cache-key.tmp-*"))
-
-
-def test_store_cache_entry_evicts_existing_invalid_entry(tmp_path) -> None:
-    cache_root = tmp_path / "cache"
-    cache_root.mkdir()
-    invalid_entry = cache_root / "cache-key"
-    invalid_entry.mkdir()
-    sentinel = invalid_entry / "stale"
-    sentinel.write_text("invalid")
-    wheel = _write_wheel(tmp_path / WHEEL_NAME)
-
-    cache_entry = kernel_build_cache._store_cache_entry(cache_root, _metadata(), wheel)
-    cached_wheel, reason = kernel_build_cache._cache_entry_wheel(cache_entry, "cache-key")
-
-    assert cached_wheel == cache_entry / WHEEL_NAME
-    assert reason == ""
-    assert not sentinel.exists()
-
-
-def test_store_cache_entry_raises_when_renamed_entry_fails_validation(monkeypatch, tmp_path) -> None:
-    cache_root = tmp_path / "cache"
-    cache_root.mkdir()
-    wheel = _write_wheel(tmp_path / WHEEL_NAME)
-    original_rename = Path.rename
-
-    def fake_rename(self, target):
-        original_rename(self, target)
-        metadata_path = target / kernel_build_cache.METADATA_FILE
-        payload = json.loads(metadata_path.read_text(encoding="utf-8"))
-        payload["cache_key"] = "other-key"
-        metadata_path.write_text(json.dumps(payload), encoding="utf-8")
-
-    monkeypatch.setattr(Path, "rename", fake_rename)
-
-    with pytest.raises(RuntimeError, match="could not be validated"):
-        kernel_build_cache._store_cache_entry(cache_root, _metadata(), wheel)
-
-    assert not (cache_root / "cache-key").exists()
-
-
-def test_consumer_installs_valid_cache_hit_without_build(monkeypatch, tmp_path) -> None:
-    cache_root = tmp_path / "cache"
-    cache_root.mkdir()
-    cache_entry = _write_cache_entry(cache_root, "cache-key")
-    installed = []
-    monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata())
-    monkeypatch.setattr(kernel_build_cache, "_try_install_prebuilt", lambda metadata, path: False)
-    monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append)
-    monkeypatch.setattr(
-        kernel_build_cache,
-        "_build_and_install_local",
-        lambda repo_root, metadata: pytest.fail("cache hit unexpectedly built locally"),
-    )
-
-    kernel_build_cache.install_cached_or_build(tmp_path, cache_root, tmp_path / "missing-prebuilt.json")
-
-    assert installed == [cache_entry / WHEEL_NAME]
-
-
-def test_consumer_rejects_corrupt_cache_without_mutating_it(monkeypatch, tmp_path) -> None:
-    cache_root = tmp_path / "cache"
-    cache_root.mkdir()
-    cache_entry = _write_cache_entry(cache_root, "cache-key")
-    wheel = cache_entry / WHEEL_NAME
-    wheel.write_bytes(b"corrupt")
-    before_metadata = (cache_entry / kernel_build_cache.METADATA_FILE).read_bytes()
-    built = []
-    monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata())
-    monkeypatch.setattr(kernel_build_cache, "_try_install_prebuilt", lambda metadata, path: False)
-    monkeypatch.setattr(kernel_build_cache, "_build_and_install_local", lambda repo_root, metadata: built.append(True))
-
-    kernel_build_cache.install_cached_or_build(tmp_path, cache_root, tmp_path / "missing-prebuilt.json")
-
-    assert built == [True]
-    assert wheel.read_bytes() == b"corrupt"
-    assert (cache_entry / kernel_build_cache.METADATA_FILE).read_bytes() == before_metadata
-
-
-def test_consumer_falls_back_when_cached_install_fails(monkeypatch, tmp_path) -> None:
-    cache_root = tmp_path / "cache"
-    cache_root.mkdir()
-    _write_cache_entry(cache_root, "cache-key")
-    built = []
-    monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata())
-    monkeypatch.setattr(kernel_build_cache, "_try_install_prebuilt", lambda metadata, path: False)
-    monkeypatch.setattr(
-        kernel_build_cache,
-        "_install_wheel",
-        lambda wheel: (_ for _ in ()).throw(subprocess.CalledProcessError(1, ["uv"])),
-    )
-    monkeypatch.setattr(kernel_build_cache, "_build_and_install_local", lambda repo_root, metadata: built.append(True))
-
-    kernel_build_cache.install_cached_or_build(tmp_path, cache_root, tmp_path / "missing-prebuilt.json")
-
-    assert built == [True]
-
-
-def test_normal_l40s_pr_ssim_install_uses_docker_prebuilt_without_build(monkeypatch, tmp_path, capsys) -> None:
-    prebuilt_root = tmp_path / "prebuilt"
-    _write_prebuilt_artifact(prebuilt_root, "9.0a", "hopper-key")
-    l40s_wheel = _write_prebuilt_artifact(prebuilt_root, "8.9", "l40s-key")
-    installed = []
-    monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata("l40s-key"))
-    monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append)
-    monkeypatch.setattr(
-        kernel_build_cache,
-        "_build_and_install_local",
-        lambda repo_root, metadata: pytest.fail("L40S prebuilt hit unexpectedly called build.sh"),
-    )
-
-    kernel_build_cache.install_cached_or_build(
-        tmp_path,
-        cache_root=None,
-        prebuilt_info_path=prebuilt_root,
-    )
-
-    output = capsys.readouterr().out
-    assert installed == [l40s_wheel]
-    assert "Docker-prebuilt cache hit for key l40s-key" in output
-    assert "build.sh" not in output
-
-
-def test_prebuilt_requires_validated_artifact(monkeypatch, tmp_path) -> None:
-    wheel = _write_wheel(tmp_path / WHEEL_NAME)
-    prebuilt_info = tmp_path / "prebuilt.json"
-    payload = {
-        **_metadata(),
-        "artifact": kernel_build_cache._wheel_artifact(wheel),
-        "wheel_path": str(wheel),
-    }
-    prebuilt_info.write_text(json.dumps(payload), encoding="utf-8")
-    installed = []
-    monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append)
-
-    assert kernel_build_cache._try_install_prebuilt(_metadata(), prebuilt_info)
-    assert installed == [wheel]
-
-    installed.clear()
-    wheel.write_bytes(b"corrupt")
-    assert not kernel_build_cache._try_install_prebuilt(_metadata(), prebuilt_info)
-    assert installed == []
-
-
-def test_producer_repairs_invalid_entry_and_stores_valid_wheel(monkeypatch, tmp_path) -> None:
-    cache_root = tmp_path / "cache"
-    cache_root.mkdir()
-    invalid_entry = cache_root / "cache-key"
-    invalid_entry.mkdir()
-    (invalid_entry / kernel_build_cache.METADATA_FILE).write_text("not json", encoding="utf-8")
-    built_dir = tmp_path / "built"
-    built_dir.mkdir()
-    built_wheel = _write_wheel(built_dir / WHEEL_NAME)
-    installed = []
-    monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata())
-    monkeypatch.setattr(kernel_build_cache, "_build_wheel", lambda repo_root, metadata: built_wheel)
-    monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append)
-
-    kernel_build_cache.produce_cached_wheel(tmp_path, cache_root)
-
-    cached_wheel, reason = kernel_build_cache._cache_entry_wheel(cache_root / "cache-key", "cache-key")
-    assert installed == [built_wheel]
-    assert cached_wheel == cache_root / "cache-key" / WHEEL_NAME
-    assert reason == ""

fastvideo/tests/modal/test_pr_test.py

@@ -265,61 +265,13 @@ def test_run_test_command_composes_valid_post_checkout_shell(
         install_command=install_command,
     )
 
-    assert [event for event, _ in events] == ["checkout", "run", "run"]
-    setup_command = events[1][1][-1]
-    test_command = events[2][1][-1]
-    assert "cd /FastVideo" in setup_command
-    assert "git clone" not in setup_command
-    assert ("kernel_build_cache.py install" in setup_command) is build_kernel
-    assert "pytest fastvideo/tests/api -q" not in setup_command
-    assert "kernel_build_cache.py install" not in test_command
-    assert "pytest fastvideo/tests/api -q" in test_command
+    assert [event for event, _ in events] == ["checkout", "run"]
+    shell_command = events[1][1][-1]
+    assert "cd /FastVideo" in shell_command
+    assert "git clone" not in shell_command
+    assert ("./build.sh" in shell_command) is build_kernel
     if install_command:
-        assert install_command in setup_command
+        assert install_command in shell_command
     else:
-        assert "uv pip install -e" not in setup_command
-    for shell_command in (setup_command, test_command):
-        real_run(["/bin/bash", "-n"], input=shell_command, text=True, check=True)
-
-
-def test_run_test_command_uses_nonshared_kernel_install_before_tests(monkeypatch):
-    module = _load_pr_test_module(monkeypatch)
-    commands = []
-
-    def fake_run(args, **_kwargs):
-        commands.append(args[-1])
-        return types.SimpleNamespace(returncode=0)
-
-    monkeypatch.setenv("BUILDKITE_REPO", "https://example.com/FastVideo.git")
-    monkeypatch.setenv("BUILDKITE_COMMIT", "0123456789abcdef")
-    monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false")
-    monkeypatch.setattr(module, "_checkout_repository",
-                        lambda *_args: None)
-    monkeypatch.setattr(subprocess, "run", fake_run)
-
-    module.run_test_command("pytest fastvideo/tests/api -q", build_kernel=True)
-
-    assert len(commands) == 2
-    setup_command, test_command = commands
-    assert "kernel_build_cache.py install" in setup_command
-    assert "--cache-root" not in setup_command
-    assert "pytest fastvideo/tests/api -q" not in setup_command
-    assert "kernel_build_cache.py install" not in test_command
-    assert "pytest fastvideo/tests/api -q" in test_command
-    assert not hasattr(module, "kernel_cache_vol")
-
-
-def test_run_unit_test_collects_modal_cache_runner_tests(monkeypatch):
-    module = _load_pr_test_module(monkeypatch)
-    commands = []
-    monkeypatch.setattr(module, "run_test", commands.append)
-
-    module.run_unit_test()
-
-    assert len(commands) == 1
-    for test_path in (
-        "./fastvideo/tests/modal/test_kernel_build_cache.py",
-        "./fastvideo/tests/modal/test_pr_test.py",
-        "./fastvideo/tests/modal/test_ssim_test.py",
-    ):
-        assert test_path in commands[0]
+        assert "uv pip install -e" not in shell_command
+    real_run(["/bin/bash", "-n"], input=shell_command, text=True, check=True)

fastvideo/tests/modal/test_ssim_test.py

@@ -99,36 +99,3 @@ def test_spawn_ssim_task_passes_hf_api_key_to_pytest_env(monkeypatch, tmp_path):
     assert env["FASTVIDEO_SSIM_MODEL_ID"] == "model-a"
     assert running_task.process.__class__ is FakeProcess
     running_task.log_handle.close()
-
-
-def test_prepare_workspace_uses_nonshared_kernel_install(monkeypatch):
-    module = _load_ssim_test_module(monkeypatch)
-    captured = {}
-    task = module.SSIMTask(
-        task_id=0,
-        test_file="fastvideo/tests/ssim/test_example.py",
-        required_gpus=1,
-    )
-
-    def fake_run(args, **kwargs):
-        captured["args"] = args
-        captured["kwargs"] = kwargs
-        return types.SimpleNamespace(returncode=0, stdout="", stderr="")
-
-    monkeypatch.setattr(module.subprocess, "run", fake_run)
-    monkeypatch.setattr(module, "_discover_ssim_tasks", lambda *_args, **_kwargs: [task])
-
-    repo_root, tasks = module._prepare_ssim_workspace(
-        git_repo="https://example.com/FastVideo.git",
-        git_commit="0123456789abcdef",
-        pr_number="false",
-        hf_api_key="hf_test_token",
-    )
-
-    setup_command = captured["args"][2]
-    assert "kernel_build_cache.py install" in setup_command
-    assert "--cache-root" not in setup_command
-    assert repo_root == "/FastVideo"
-    assert tasks == [task]
-    assert module.SSIM_COMMON_KWARGS["volumes"] == {"/root/data": module.model_vol}
-    assert not hasattr(module, "kernel_cache_vol")

fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py

@@ -1,523 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Single-sample overfit / small end-to-end run for Kandinsky5 QAD:
-stage-1 Attn-QAT finetune -> stage-2 QAT-aware DMD distillation, on the
-new ``fastvideo/train/`` stack -- validated all the way to the public
-inference artifact this recipe exists to produce.
-
-Unlike ``test_e2e_dmd_t2v_crush_smol.py`` (which drives the legacy
-``fastvideo/training/wan_distillation_pipeline.py`` CLI), this drives the
-new-stack entrypoint (``fastvideo.train.entrypoint.train``) via the same
-YAML configs used for real training
-(``examples/train/configs/fine_tuning/kandinsky5/t2v_480p_qat.yaml`` and
-``examples/train/configs/distribution_matching/kandinsky5/dmd2_t2v_480p_qat.yaml``),
-with a handful of steps and dotted-key overrides pointing at a tiny local
-dataset -- mirroring how a real run would be launched, just shrunk down.
-
-The full validated chain, each arrow a hard assertion:
-
-  synthetic clip -> preprocess -> stage-1 train -> stage-1 DCP
-    -> dcp_to_diffusers --verify (strict reload)
-    -> stage-2 train (student/teacher/critic init_from that export)
-    -> stage-2 student DCP -> dcp_to_diffusers --verify (strict reload)
-    -> VideoGenerator.from_pretrained(export,
-           override_pipeline_cls_name="Kandinsky5DMDPipeline",
-           pipeline_config=Kandinsky5DMDConfig())   # the documented recipe
-    -> deterministic (fixed-seed) 4-step generation
-    -> degeneracy checks + MS-SSIM against the committed reference video.
-
-Every artifact (raw clip, parquet, both stage output dirs, both diffusers
-exports, generated video dir) lives under the single test-owned
-``data/kandinsky5_e2e/`` root, which is deleted up front -- the
-documented user dataset paths (``data/kandinsky5_overfit{,_preprocessed}``)
-are never read, written, or removed. ``training.checkpoint.
-resume_from_checkpoint`` is explicitly disabled for both stages: the
-stage-1 YAML defaults to ``resume_from_checkpoint: latest``, so a rerun
-over a dirty ``data/`` dir could otherwise no-op from an old checkpoint
-and go green on stale artifacts.
-
-Reference bootstrap: the committed reference
-(``reference_video_kandinsky5_dmd_v0.mp4`` next to this file) is produced
-by running this test once on a sanctioned GPU box with
-``KANDINSKY5_E2E_WRITE_REFERENCE=1``, reviewing the written video by eye,
-and committing it. Review bar: expect blurry-but-structured content
-loosely matching the prompt (4-step no-CFG sampling of near-base weights
-cannot look good) -- soft shapes, warm sunflower-ish colors, motion.
-Uniform grey static, solid black, or hard garbage blocks mean a broken
-pipeline: do NOT commit such a reference (probe the export with the plain
-50-step ``Kandinsky5T2VPipeline`` to isolate weights-path vs sampler
-problems). Until a reference exists the test FAILS (it does not skip) --
-an artifact-existence-only pass is exactly the false-green this oracle
-exists to prevent. Run it with::
-
-    pytest fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py -vs
-
-The single training sample is a synthesized noise clip (not a downloaded
-dataset): the exact directory/manifest layout of external raw-video HF
-datasets (e.g. the ``crush-smol`` one Wan's e2e test uses) isn't something
-this change can verify without a runnable environment, so depending on it
-here would risk a silently-wrong assumption. ``preprocess_kandinsky5_overfit.py``
-only needs a ``videos2caption.json`` + ``videos/*.mp4`` pair, both of which
-this test fully controls.
-"""
-
-from __future__ import annotations
-
-import json
-import os
-import shutil
-import subprocess
-import sys
-from pathlib import Path
-
-import cv2
-import numpy as np
-import pytest
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-THIS_FILE = Path(__file__).resolve()
-
-NUM_GPUS = os.environ.get("KANDINSKY5_E2E_NUM_GPUS", "1")
-# Every artifact this test reads or writes lives under this single
-# test-owned root. It must NEVER point at (or contain) the documented
-# default dataset paths -- data/kandinsky5_overfit and
-# data/kandinsky5_overfit_preprocessed -- which a user may have populated
-# with their real videos/captions for the actual recipe:
-# _clean_previous_artifacts() deletes this root wholesale on every run,
-# and the preprocess subprocess is pointed here via the
-# KANDINSKY5_OVERFIT_DATA_DIR / KANDINSKY5_OVERFIT_OUTPUT_DIR env
-# overrides instead of the script's user-facing defaults.
-E2E_ROOT = Path("data") / "kandinsky5_e2e"
-RAW_DATA_DIR = E2E_ROOT / "raw"
-PREPROCESSED_DIR = E2E_ROOT / "preprocessed"
-STAGE1_OUTPUT_DIR = E2E_ROOT / "stage1"
-STAGE2_OUTPUT_DIR = E2E_ROOT / "stage2"
-# Kept out of STAGE{1,2}_OUTPUT_DIR: those directories are scanned with
-# glob("checkpoint-*") below to find the latest raw DCP checkpoint, and a
-# "checkpoint-<N>-diffusers" export dir sitting alongside "checkpoint-<N>"
-# would also match that pattern -- on a rerun that doesn't start from a
-# clean data/ dir, it would sort after "checkpoint-<N>" and get picked as
-# the checkpoint instead, feeding an already-exported diffusers directory
-# back into dcp_to_diffusers as if it were a DCP checkpoint.
-STAGE1_DIFFUSERS_DIR = E2E_ROOT / "stage1_diffusers"
-STAGE2_DIFFUSERS_DIR = E2E_ROOT / "stage2_diffusers"
-GENERATED_VIDEO_DIR = E2E_ROOT / "generated"
-
-STAGE1_CONFIG = (REPO_ROOT / "examples" / "train" / "configs" / "fine_tuning" / "kandinsky5" /
-                 "t2v_480p_qat.yaml")
-STAGE2_CONFIG = (REPO_ROOT / "examples" / "train" / "configs" / "distribution_matching" / "kandinsky5" /
-                 "dmd2_t2v_480p_qat.yaml")
-
-# Deterministic-generation oracle. The reference is generated once (see the
-# module docstring's bootstrap procedure), reviewed by a human, and
-# committed next to this file -- mirroring reference_video_1_sample_v0.mp4
-# in the Wan e2e test.
-REFERENCE_VIDEO = THIS_FILE.parent / "reference_video_kandinsky5_dmd_v0.mp4"
-WRITE_REFERENCE_ENV = "KANDINSKY5_E2E_WRITE_REFERENCE"
-# Deliberately semantic even though the training clip is random noise: the
-# caption's relation to the clip content is irrelevant for plumbing, but
-# the caption doubles as the generation prompt, and a semantic prompt makes
-# the reference video human-reviewable (blurry-but-structured output from
-# 4-step sampling of near-base weights). An earlier "random noise" caption
-# produced a grey-static reference that was visually indistinguishable
-# from a broken pipeline AND decorrelates across runs, making the SSIM
-# oracle both unreviewable and potentially flaky.
-GENERATION_PROMPT = ("A curious raccoon peers through a vibrant field of yellow sunflowers, "
-                     "soft natural light filtering through the petals, mid-shot")
-GENERATION_SEED = 42
-# Media-oracle calibration, measured against the committed reference
-# (2026-07-20, from the generator_update_interval=1 / 2-GPU-sharded
-# bootstrap run; regression-tested in
-# fastvideo/tests/workflow/test_kandinsky5_e2e_media_oracle.py, which
-# recomputes these known-bad clips from whatever reference is currently
-# committed -- so it stays valid across reference swaps even though the
-# numbers below are a point-in-time snapshot):
-#
-#   clip                          spatial-std  temporal-diff  mean MS-SSIM
-#   good reference                  >= 11.42      >= 1.12        1.0000
-#   solid @ reference mean RGB       0.0           0.0           0.9162
-#   frozen (reference frame x121)   12.12          0.0           0.8534
-#
-# Two consequences drive the design below:
-#  - MS-SSIM alone CANNOT separate collapsed clips from good ones with
-#    safe margin (a solid clip scores ~0.92!), so the structural checks
-#    in _assert_video_not_degenerate are the primary defense against
-#    solid/frozen/truncated output; the SSIM floor's job is catching a
-#    *different structured video* (wrong remap / re-noise schedule).
-#  - The floors are set >= 5x below the good reference's observed values
-#    and >= 10x above every known-bad clip's.
-# An independent retrain-from-scratch validation run against this exact
-# reference is the last open verification step (see the module
-# docstring's bootstrap procedure) -- 0.95 is chosen to sit comfortably
-# above the ~0.92 solid-clip ceiling while leaving headroom for ordinary
-# run-to-run drift once that number is in.
-EXPECTED_FRAME_COUNT = 121
-EXPECTED_FRAME_HEIGHT = 512
-EXPECTED_FRAME_WIDTH = 768
-MIN_FRAME_SPATIAL_STD = 2.0
-MIN_FRAME_TEMPORAL_DIFF = 0.2
-MIN_MEAN_MS_SSIM = 0.95
-
-
-def _clean_previous_artifacts() -> None:
-    """Delete the single test-owned root a previous run could have left.
-
-    Outputs matter as much as inputs: stage 1's YAML defaults to
-    ``resume_from_checkpoint: latest``, and the final assertions glob for
-    checkpoints/videos -- stale files from an earlier run could satisfy
-    them without this run doing any work.
-
-    Scoped strictly to ``E2E_ROOT``: an earlier version of this cleanup
-    also removed ``data/kandinsky5_overfit{,_preprocessed}``, the
-    documented default dataset paths of the real recipe -- silently and
-    irreversibly deleting whatever real videos/captions a user had
-    prepared there. Nothing outside the test-owned root may be touched.
-    """
-    if E2E_ROOT.exists():
-        shutil.rmtree(E2E_ROOT)
-
-
-def _synthesize_single_sample() -> None:
-    """Write one tiny synthetic clip + caption manifest in the layout
-    preprocess_kandinsky5_overfit.py expects.
-
-    Matches KANDINSKY5_T2V_LITE_5S's native preset (512x768, 121 frames,
-    24fps -- see preprocess_kandinsky5_overfit.py's own constants).
-    """
-    videos_dir = RAW_DATA_DIR / "videos"
-    videos_dir.mkdir(parents=True, exist_ok=True)
-
-    video_path = videos_dir / "sample_0.mp4"
-    writer = cv2.VideoWriter(
-        str(video_path),
-        cv2.VideoWriter_fourcc(*"mp4v"),
-        24.0,
-        (768, 512),
-    )
-    if not writer.isOpened():
-        raise RuntimeError(
-            f"cv2.VideoWriter failed to open {video_path} (mp4v codec unavailable?) -- "
-            "the fixture clip would be empty and preprocessing would fail on it.")
-    rng = np.random.default_rng(0)
-    for _ in range(121):
-        frame = rng.integers(0, 256, size=(512, 768, 3), dtype=np.uint8)
-        writer.write(frame)
-    writer.release()
-
-    # One physical clip, but max(2, NUM_GPUS) manifest rows all pointing at
-    # it: DP_SP_BatchSampler (fastvideo/dataset/parquet_dataset_map_style.py)
-    # floor-divides the number of batches by the number of data-parallel
-    # groups with drop_last=True, so a dataset with fewer rows than GPUs
-    # yields an EMPTY dataloader on every rank. Duplicate rows keep the
-    # single-sample-overfit semantics (identical latents/caption) while
-    # making KANDINSKY5_E2E_NUM_GPUS=2+ a working escape hatch when one GPU
-    # doesn't have the memory headroom for stage 2 (see _run_stage).
-    num_rows = max(2, int(NUM_GPUS))
-    with open(RAW_DATA_DIR / "videos2caption.json", "w") as f:
-        json.dump([{
-            "path": "sample_0.mp4",
-            "cap": [GENERATION_PROMPT],
-        }] * num_rows, f)
-
-
-def _run_preprocessing() -> None:
-    # Point the preprocessor at the test-owned roots. Without these env
-    # overrides it would read/write its user-facing defaults
-    # (data/kandinsky5_overfit{,_preprocessed}) -- the very directories a
-    # user populates for the real recipe, which this test must never touch.
-    env = dict(os.environ)
-    env["KANDINSKY5_OVERFIT_DATA_DIR"] = str(RAW_DATA_DIR)
-    env["KANDINSKY5_OVERFIT_OUTPUT_DIR"] = str(PREPROCESSED_DIR)
-    cmd = [
-        sys.executable, "-m",
-        "fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit",
-    ]
-    subprocess.run(cmd, cwd=str(REPO_ROOT), env=env, check=True)
-
-
-def _export_dcp_to_diffusers(checkpoint_dir: Path, output_dir: Path) -> None:
-    """Convert a stage's DCP checkpoint to a diffusers-style directory.
-
-    ``models.*.init_from`` / ``VideoGenerator.from_pretrained`` (like real
-    training YAML configs, see ``dmd2_t2v_480p_qat.yaml``) expect a
-    diffusers model directory (``model_index.json`` + component
-    subfolders), not the raw ``checkpoint-N`` DCP layout
-    ``CheckpointManager`` writes (``dcp/`` + metadata/RNG state only).
-    ``--verify`` strictly reloads the exported transformer immediately so a
-    key-mapping bug fails here, at the export boundary, rather than deep
-    inside the next launch that loads the directory.
-    """
-    cmd = [
-        sys.executable, "-m",
-        "fastvideo.train.entrypoint.dcp_to_diffusers",
-        "--checkpoint", str(checkpoint_dir),
-        "--output-dir", str(output_dir),
-        "--role", "student",
-        "--overwrite",
-        "--verify",
-    ]
-    subprocess.run(cmd, cwd=str(REPO_ROOT), check=True)
-
-
-def _latest_checkpoint(output_dir: Path) -> Path:
-    checkpoints = sorted(output_dir.glob("checkpoint-*"))
-    assert checkpoints, f"no checkpoint produced under {output_dir}"
-    return checkpoints[-1]
-
-
-def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_overrides: list[str],
-              env_overrides: dict[str, str]) -> None:
-    cmd = [
-        sys.executable, "-m", "torch.distributed.run",
-        "--nnodes", "1",
-        "--nproc_per_node", NUM_GPUS,
-        "-m", "fastvideo.train.entrypoint.train",
-        "--config", str(config),
-        "--training.data.data_path", str(PREPROCESSED_DIR),
-        "--training.distributed.num_gpus", NUM_GPUS,
-        "--training.distributed.hsdp_shard_dim", NUM_GPUS,
-        "--training.loop.max_train_steps", str(max_train_steps),
-        "--training.checkpoint.output_dir", str(output_dir),
-        "--training.checkpoint.training_state_checkpointing_steps", str(max_train_steps),
-        # Explicitly disable resume for BOTH stages, even though
-        # _clean_previous_artifacts() already removed the output dirs: the
-        # stage-1 YAML defaults to resume_from_checkpoint: latest, and this
-        # test's correctness must not silently depend on that default (or a
-        # partially-failed cleanup) ever changing.
-        "--training.checkpoint.resume_from_checkpoint", "none",
-        "--callbacks.validation.every_steps", str(max_train_steps),
-        "--callbacks.validation.dataset_file", str(PREPROCESSED_DIR / "validation_prompts.json"),
-        *extra_overrides,
-    ]
-    env = dict(os.environ)
-    # Stage 2 holds three 2B-param models plus TWO full Adam states (the
-    # student's only exists because generator_update_interval is overridden
-    # to 1 above -- a recorded run peaked at ~76 GiB and died allocating
-    # the last MiBs on an 80 GiB device with ~1.9 GiB lost to
-    # fragmentation). Expandable segments reclaims that headroom; if a
-    # single device still can't fit, shard with KANDINSKY5_E2E_NUM_GPUS=2+
-    # (the synthesized manifest guarantees >= one row per rank).
-    env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-    env.update(env_overrides)
-    subprocess.run(cmd, cwd=str(REPO_ROOT), env=env, check=True)
-
-
-def _generate_from_export(export_dir: Path, output_dir: Path) -> Path:
-    """Deterministically generate one video from an exported student via
-    the documented DMD inference path, in a fresh subprocess.
-
-    A subprocess (re-running this file with ``--generate``, see
-    ``__main__``) keeps VideoGenerator's distributed/CUDA init isolated
-    from the two torchrun stages that already ran from this test process.
-    FASTVIDEO_ATTENTION_BACKEND is removed from the child env: generation
-    must exercise the default dense-attention inference path the README
-    documents for non-sm_120 GPUs, not inherit a QAT training backend.
-    """
-    env = dict(os.environ)
-    env.pop("FASTVIDEO_ATTENTION_BACKEND", None)
-    cmd = [
-        sys.executable,
-        str(THIS_FILE), "--generate",
-        str(export_dir),
-        str(output_dir),
-    ]
-    subprocess.run(cmd, cwd=str(REPO_ROOT), env=env, check=True)
-
-    videos = sorted(Path(output_dir).glob("**/*.mp4"))
-    assert len(videos) == 1, (
-        f"expected exactly one generated video under {output_dir}, found {videos}")
-    return videos[0]
-
-
-def _generate_main(export_dir: str, output_dir: str) -> None:
-    """Subprocess body for _generate_from_export.
-
-    This is the exact inference recipe the QAD README documents for a
-    stage-2 export (minus the optional NVFP4/FP8 weight quantization,
-    which needs flashinfer + sm_120 hardware): the export's
-    model_index.json still names the base T2V pipeline, so
-    override_pipeline_cls_name + Kandinsky5DMDConfig are required to get
-    the 4-step re-noise sampler -- see
-    fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py.
-    """
-    from fastvideo import VideoGenerator
-    from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig
-
-    generator = VideoGenerator.from_pretrained(
-        export_dir,
-        num_gpus=1,
-        override_pipeline_cls_name="Kandinsky5DMDPipeline",
-        pipeline_config=Kandinsky5DMDConfig(),
-        use_fsdp_inference=False,
-    )
-    generator.generate_video(
-        GENERATION_PROMPT,
-        output_path=output_dir,
-        save_video=True,
-        seed=GENERATION_SEED,
-        height=512,
-        width=768,
-        num_frames=121,
-    )
-
-
-def _decode_video(video_path: Path) -> np.ndarray:
-    cap = cv2.VideoCapture(str(video_path))
-    frames = []
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frames.append(frame)
-    cap.release()
-    assert frames, f"video {video_path} decoded to zero frames"
-    return np.stack(frames)
-
-
-def _assert_video_not_degenerate(video_path: Path) -> None:
-    """Structural referee that needs no reference: rejects the collapsed
-    outputs (solid, frozen, truncated, unreadable) a NaN run or broken
-    export/decode path produces.
-
-    This is the primary defense against collapsed clips -- the MS-SSIM
-    floor cannot be (see the calibration table above: a solid clip at the
-    reference's mean color scores 0.9189 against it). Every check here is
-    designed against a measured false-green:
-    - exact frame count/geometry: compute_video_ssim_torchvision
-      truncates both clips to the shorter length, so a 1-frame video
-      could otherwise sail through the comparison;
-    - per-frame *luma* spatial std: the previous global RGB std counted
-      differences between channel means as "variance", scoring a solid
-      [106, 98, 70] clip at 15.4;
-    - consecutive-frame temporal diff: a frozen (single repeated frame)
-      clip has real spatial content and passes the spatial check.
-    """
-    frames = _decode_video(video_path).astype(np.float64)
-    expected_shape = (EXPECTED_FRAME_COUNT, EXPECTED_FRAME_HEIGHT, EXPECTED_FRAME_WIDTH, 3)
-    assert frames.shape == expected_shape, (
-        f"generated video {video_path} has shape {frames.shape}, expected {expected_shape} -- "
-        "wrong frame count or geometry (and the SSIM helper would silently truncate to the "
-        "shorter clip)")
-
-    luma = frames.mean(axis=-1)
-    spatial_std = luma.reshape(len(luma), -1).std(axis=1)
-    assert float(spatial_std.min()) >= MIN_FRAME_SPATIAL_STD, (
-        f"generated video {video_path} has a frame with luma spatial std "
-        f"{spatial_std.min():.3f} < {MIN_FRAME_SPATIAL_STD} (reference min: ~11.4) -- "
-        "solid/near-solid output from a NaN or collapsed run")
-
-    temporal_diff = np.abs(np.diff(luma, axis=0)).mean(axis=(1, 2))
-    assert float(temporal_diff.min()) >= MIN_FRAME_TEMPORAL_DIFF, (
-        f"generated video {video_path} has consecutive frames with mean |luma diff| "
-        f"{temporal_diff.min():.4f} < {MIN_FRAME_TEMPORAL_DIFF} (reference min: ~1.1) -- "
-        "frozen output")
-
-
-def _assert_matches_reference(video_path: Path) -> None:
-    if not REFERENCE_VIDEO.exists():
-        if os.environ.get(WRITE_REFERENCE_ENV) == "1":
-            shutil.copy2(video_path, REFERENCE_VIDEO)
-            print(f"\nWrote new reference video to {REFERENCE_VIDEO} -- "
-                  "review it by eye and commit it. This bootstrap run "
-                  "compares the video against itself (SSIM=1) below.")
-        else:
-            pytest.fail(
-                f"reference video missing at {REFERENCE_VIDEO}. This test "
-                "must not pass on artifact existence alone -- run once on a "
-                f"sanctioned GPU box with {WRITE_REFERENCE_ENV}=1, review "
-                "the written video, and commit it (see module docstring).")
-
-    from fastvideo.tests.utils import compute_video_ssim_torchvision
-
-    mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
-        str(REFERENCE_VIDEO),
-        str(video_path),
-        use_ms_ssim=True,
-    )
-    print("\n===== MS-SSIM vs committed Kandinsky5 DMD reference =====")
-    print(f"Mean MS-SSIM: {mean_ssim:.4f}")
-    print(f"Min MS-SSIM: {min_ssim:.4f}")
-    print(f"Max MS-SSIM: {max_ssim:.4f}")
-    assert mean_ssim >= MIN_MEAN_MS_SSIM, (
-        f"mean MS-SSIM {mean_ssim:.4f} below {MIN_MEAN_MS_SSIM} -- the exported/reloaded "
-        "stage-2 student no longer generates what the reviewed reference recorded "
-        "(wrong weight remap, re-noise schedule, or attention path?)")
-
-
-@pytest.mark.nightly
-def test_e2e_kandinsky5_dmd_overfit_single_sample():
-    if not STAGE1_CONFIG.exists() or not STAGE2_CONFIG.exists():
-        pytest.skip("Kandinsky5 QAT configs not found -- see examples/train/configs/{fine_tuning,distribution_matching}/kandinsky5/")
-
-    os.environ.setdefault("WANDB_MODE", "offline")
-
-    _clean_previous_artifacts()
-    _synthesize_single_sample()
-    _run_preprocessing()
-
-    _run_stage(
-        STAGE1_CONFIG,
-        STAGE1_OUTPUT_DIR,
-        max_train_steps=3,
-        extra_overrides=[
-            # The recipe's full 5e-5 is calibrated for a real dataset; 3
-            # steps of it toward this test's random-noise clip measurably
-            # degrades the model -- empirically enough to collapse the
-            # fragile 4-step no-CFG DMD sampling below into grey static
-            # (50-step T2V sampling of the same export still recovered).
-            # A tiny LR keeps the full optimizer/backward path exercised
-            # while the final generation stays structured: reviewable by
-            # eye at bootstrap, and stable across runs for the SSIM oracle
-            # (noise output decorrelates under any training
-            # nondeterminism; structured output does not).
-            "--training.optimizer.learning_rate", "1e-6",
-        ],
-        env_overrides={"FASTVIDEO_ATTENTION_BACKEND": "ATTN_QAT_TRAIN"},
-    )
-    stage1_ckpt = _latest_checkpoint(STAGE1_OUTPUT_DIR)
-
-    stage1_diffusers_dir = STAGE1_DIFFUSERS_DIR / stage1_ckpt.name
-    _export_dcp_to_diffusers(stage1_ckpt, stage1_diffusers_dir)
-    assert (stage1_diffusers_dir / "model_index.json").exists(), (
-        f"dcp_to_diffusers export did not produce a diffusers model dir at {stage1_diffusers_dir}")
-
-    _run_stage(
-        STAGE2_CONFIG,
-        STAGE2_OUTPUT_DIR,
-        max_train_steps=3,
-        extra_overrides=[
-            "--models.student.init_from", str(stage1_diffusers_dir),
-            "--models.teacher.init_from", str(stage1_diffusers_dir),
-            "--models.critic.init_from", str(stage1_diffusers_dir),
-            # The recipe's generator_update_interval of 5 would mean ZERO
-            # student updates in a 3-step run (the trainer iterates 1..3
-            # and DMD2Method updates the generator only when
-            # iteration % interval == 0) -- the exported "stage-2 student"
-            # would just be the stage-1 weights, and this test could not
-            # catch a broken student backward/optimizer path. Update the
-            # generator every step instead.
-            "--method.generator_update_interval", "1",
-        ],
-        env_overrides={"FASTVIDEO_ATTENTION_BACKEND": "ATTN_QAT_TRAIN"},
-    )
-    assert any(STAGE2_OUTPUT_DIR.glob("*.mp4")), (
-        f"no stage-2 validation video produced under {STAGE2_OUTPUT_DIR}")
-
-    # The actual deliverable of this recipe is the exported stage-2
-    # student: export it, strict-reload it (--verify), then instantiate it
-    # through the documented Kandinsky5DMDPipeline override and generate.
-    stage2_ckpt = _latest_checkpoint(STAGE2_OUTPUT_DIR)
-    stage2_diffusers_dir = STAGE2_DIFFUSERS_DIR / stage2_ckpt.name
-    _export_dcp_to_diffusers(stage2_ckpt, stage2_diffusers_dir)
-    assert (stage2_diffusers_dir / "model_index.json").exists(), (
-        f"dcp_to_diffusers export did not produce a diffusers model dir at {stage2_diffusers_dir}")
-
-    generated_video = _generate_from_export(stage2_diffusers_dir, GENERATED_VIDEO_DIR)
-    _assert_video_not_degenerate(generated_video)
-    _assert_matches_reference(generated_video)
-
-
-if __name__ == "__main__":
-    if len(sys.argv) >= 2 and sys.argv[1] == "--generate":
-        _generate_main(sys.argv[2], sys.argv[3])
-    else:
-        test_e2e_kandinsky5_dmd_overfit_single_sample()

fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py

@@ -1,18 +1,17 @@
 """E2E overfit tests for LTX-2 fine-tuning on the modular trainer (fastvideo/train).
 
-Parametrized over dense LTX-2.0 and LTX-2.3 distilled checkpoints plus
-an LTX-2.0 NVFP4 linear-QAT case. Each case downloads the single-sample
-cats dataset, preprocesses it with the matching LTX-2 VAE + Gemma text
-encoder into the t2v parquet format, trains LTX2Model with FineTuneMethod
-for 300 steps on 4 GPUs, and checks that the final validation video
-reproduces the training clip (MS-SSIM against the preprocessed ground-truth
-clip) better than the step-0 baseline.
+Parametrized over LTX-2.0 and LTX-2.3 distilled checkpoints. Each case
+downloads the single-sample cats dataset, preprocesses it with the
+matching LTX-2 VAE + Gemma text encoder into the t2v parquet format,
+trains LTX2Model with FineTuneMethod for 300 steps on 4 GPUs, and
+checks that the final validation video reproduces the training clip
+(MS-SSIM against the preprocessed ground-truth clip) better than the
+step-0 baseline.
 
 GPU assumptions: 4 x large-memory GPUs (developed on 4x GB200 192GB;
 the 18.9B-param DiT trains FSDP-sharded with fp32 master weights).
 Requires the FastVideo LTX-2 distilled checkpoints (set HF_HOME to a
 cache that contains them, or allow ~60GB of downloads per version).
-The NVFP4-QAT case additionally requires FlashInfer FP4 support and SM100+.
 
 Guarded by FASTVIDEO_NIGHTLY=1 so the default test suite stays fast.
 """
@@ -26,7 +25,6 @@ import sys
 from pathlib import Path
 
 import pytest
-import torch
 
 from huggingface_hub import snapshot_download
 
@@ -54,16 +52,6 @@ _CASES = {
         "prep_dir": Path(DATA_DIR) / "ltx2_3_overfit_preprocessed",
         "out_dir": Path(DATA_DIR) / "outputs_ltx2_3_overfit",
     },
-    # NVFP4 quantization-aware training: same LTX-2.0 data, attention/FFN
-    # linears forward through real FP4 GEMMs with an STE backward.
-    # Requires flashinfer FP4 kernels (sm_100-class GPU).
-    "ltx2_nvfp4_qat": {
-        "model": "FastVideo/LTX2-Distilled-Diffusers",
-        "config": "examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml",
-        "prep_dir": Path(DATA_DIR) / "ltx2_overfit_preprocessed",
-        "out_dir": Path(DATA_DIR) / "outputs_ltx2_nvfp4_qat_overfit",
-        "requires_nvfp4": True,
-    },
 }
 
 
@@ -105,26 +93,10 @@ def run_training(case: dict):
         "--callbacks.validation.dataset_file",
         str(case["prep_dir"] / "validation_prompts.json"),
     ]
-    if (case.get("requires_nvfp4")
-            and torch.cuda.get_device_capability(0) != (12, 0)):
-        cmd.extend(["--callbacks.validation.attn_qat_infer", "false"])
     print(f"Running training: {cmd}")
     subprocess.run(cmd, check=True, env=env)
 
 
-def test_run_training_disables_sm120_only_validation_on_gb200(monkeypatch):
-    commands = []
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (10, 0))
-    monkeypatch.setattr(subprocess, "run",
-                        lambda cmd, **_: commands.append(cmd))
-
-    run_training(_CASES["ltx2_nvfp4_qat"])
-
-    assert commands[0][-2:] == [
-        "--callbacks.validation.attn_qat_infer", "false"
-    ]
-
-
 def _validation_videos_by_step(out_dir: Path) -> dict[int, str]:
     pattern = os.path.join(str(out_dir), "validation_step_*_video_0.mp4")
     videos: dict[int, str] = {}
@@ -142,18 +114,6 @@ def _validation_videos_by_step(out_dir: Path) -> dict[int, str]:
 @pytest.mark.parametrize("case_id", _CASES.keys())
 def test_e2e_ltx2_overfit_new_stack(case_id: str):
     case = _CASES[case_id]
-    if case.get("requires_nvfp4"):
-        if (not torch.cuda.is_available() or
-                torch.cuda.get_device_capability(0) < (10, 0)):
-            pytest.skip("NVFP4 QAT requires an SM100+ GPU")
-        try:
-            from flashinfer import (  # noqa: F401
-                SfLayout,
-                mm_fp4,
-                nvfp4_quantize,
-            )
-        except ImportError:
-            pytest.skip("NVFP4 QAT requires flashinfer with FP4 kernels")
     download_data()
     run_preprocessing(case)
     run_training(case)

fastvideo/tests/ops/quantization/test_nvfp4_ltx2_wiring.py

@@ -18,10 +18,6 @@ from fastvideo.layers.quantization.nvfp4_config import (
     NVFP4Config,
     NVFP4QuantizeMethod,
 )
-from fastvideo.layers.quantization.nvfp4_qat_train_config import (
-    NVFP4QATTrainConfig,
-    NVFP4QATTrainQuantizeMethod,
-)
 from fastvideo.models.dits.ltx2 import (
     BasicAVTransformerBlock,
     FeedForward,
@@ -187,46 +183,3 @@ def test_basic_av_block_propagates_quant_config_to_all_children() -> None:
     for linear in unquantized_projections:
         assert isinstance(linear, ReplicatedLinear)
         assert isinstance(linear.quant_method, UnquantizedLinearMethod)
-
-
-def test_qat_train_uses_the_ltx2_deployment_target_profile() -> None:
-    kwargs = {
-        "idx": 3,
-        "video": TransformerConfig(dim=64, heads=2, d_head=32, context_dim=128),
-        "audio": TransformerConfig(dim=64, heads=2, d_head=32, context_dim=128),
-        "rope_type": LTXRopeType.INTERLEAVED,
-        "norm_eps": 1e-6,
-        "use_distributed_attention": False,
-        "prefix": "ltx2",
-    }
-    deployed = BasicAVTransformerBlock(
-        **kwargs,
-        quant_config=NVFP4Config(),
-    )
-    qat = BasicAVTransformerBlock(
-        **kwargs,
-        quant_config=NVFP4QATTrainConfig(),
-    )
-
-    deployed_names = {
-        name for name, module in deployed.named_modules()
-        if isinstance(getattr(module, "quant_method", None),
-                      NVFP4QuantizeMethod)
-    }
-    qat_names = {
-        name for name, module in qat.named_modules()
-        if isinstance(getattr(module, "quant_method", None),
-                      NVFP4QATTrainQuantizeMethod)
-    }
-    assert qat_names == deployed_names
-    assert len(qat_names) == 12
-
-
-def test_qat_train_preserves_wan_style_targeting() -> None:
-    linear = ReplicatedLinear(
-        16,
-        16,
-        quant_config=NVFP4QATTrainConfig(),
-        prefix="blocks.0.attn1.to_k",
-    )
-    assert isinstance(linear.quant_method, NVFP4QATTrainQuantizeMethod)

fastvideo/tests/ops/quantization/test_nvfp4_purge.py

@@ -1,127 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""CPU construction-level tests for the NVFP4 original-weight purge.
-
-``convert_model_to_nvfp4`` keeps the packed FP4 buffers and, by default
-(auto), purges the bf16 ``layer.weight`` of always-FP4 layers while
-retaining refine-only layers (which still take the stage-profile dense
-path). ``retain_original_weights`` on ``NVFP4Config`` overrides in both
-directions. flashinfer and the FP4 GEMM are monkeypatched, so this runs
-on any host.
-"""
-from __future__ import annotations
-
-import logging
-import types
-
-import pytest
-import torch
-import torch.nn as nn
-
-import fastvideo.layers.quantization.nvfp4_config as nv
-
-_ALWAYS_FP4_PREFIX = "ltx2.blocks.0.attn1.to_q"
-_REFINE_ONLY_PREFIX = "ltx2.blocks.0.audio_to_video_attn.to_q"
-
-
-def _method(prefix: str, retain: bool | None) -> nv.NVFP4QuantizeMethod:
-    # NVFP4QuantizeMethod.__init__ allocates x_global_sf on cuda; build the
-    # object without it so the purge tests run on CPU-only hosts.
-    m = object.__new__(nv.NVFP4QuantizeMethod)
-    m.weight_fp4 = None
-    m.weight_scale = None
-    m.x_global_sf = torch.tensor(1.0, dtype=torch.float32)
-    m.layer_prefix = prefix
-    m._is_refine_only_layer = nv._is_ltx2_refine_only_prefix(prefix)
-    m._retain_original_weights = retain
-    return m
-
-
-class _FakeLinear(nn.Module):
-
-    def __init__(self, prefix: str, retain: bool | None, out_dim: int = 8, in_dim: int = 16) -> None:
-        super().__init__()
-        self.weight = nn.Parameter(torch.randn(out_dim, in_dim, dtype=torch.bfloat16), requires_grad=False)
-        self.quant_method = _method(prefix, retain)
-
-
-def _fake_quantize(weight, global_sf, sfLayout=None, do_shuffle=False):
-    out_dim, in_dim = weight.shape[0], weight.shape[-1]
-    return (torch.zeros(out_dim, in_dim // 2, dtype=torch.int8), torch.zeros(out_dim, in_dim // 16,
-                                                                             dtype=torch.uint8))
-
-
-@pytest.fixture(autouse=True)
-def _patch_flashinfer(monkeypatch):
-    fake_sf_layout = types.SimpleNamespace(layout_128x4=None)
-    monkeypatch.setattr(nv, "_require_flashinfer", lambda: (fake_sf_layout, None, None))
-    monkeypatch.setattr(nv, "_nvfp4_quantize", _fake_quantize)
-
-
-def _model(retain: bool | None) -> nn.Module:
-    root = nn.Module()
-    root.always_fp4 = _FakeLinear(_ALWAYS_FP4_PREFIX, retain)
-    root.refine_only = _FakeLinear(_REFINE_ONLY_PREFIX, retain)
-    return root
-
-
-def test_prefix_classification_sanity() -> None:
-    assert not nv._is_ltx2_refine_only_prefix(_ALWAYS_FP4_PREFIX)
-    assert nv._is_ltx2_refine_only_prefix(_REFINE_ONLY_PREFIX)
-
-
-def test_auto_purges_always_fp4_and_retains_refine_only(caplog) -> None:
-    model = _model(retain=None)
-    with caplog.at_level(logging.INFO):
-        nv.convert_model_to_nvfp4(model)
-    assert model.always_fp4.weight is None
-    assert model.refine_only.weight is not None
-    assert model.always_fp4._nvfp4_weight is not None
-    receipt = [r.message for r in caplog.records if "NVFP4 weight purge receipt" in r.message]
-    assert receipt and "purged 1" in receipt[0] and "retained 1" in receipt[0]
-
-
-def test_retain_true_keeps_everything() -> None:
-    model = _model(retain=True)
-    nv.convert_model_to_nvfp4(model)
-    assert model.always_fp4.weight is not None
-    assert model.refine_only.weight is not None
-
-
-def test_retain_false_still_retains_dense_capable_layers() -> None:
-    """Refine-only layers run dense under the base stage profile by
-    deployment contract (single-stage deploys included), so no flag value
-    may purge them."""
-    model = _model(retain=False)
-    nv.convert_model_to_nvfp4(model)
-    assert model.always_fp4.weight is None
-    assert model.refine_only.weight is not None
-
-
-def test_apply_out_dim_survives_purge(monkeypatch) -> None:
-    model = _model(retain=False)
-    nv.convert_model_to_nvfp4(model)
-    layer = model.always_fp4
-    captured = {}
-
-    def fake_mm_fp4(x_fp4, w_t, x_scale, w_scale_t, alpha, out_dtype, out, backend):
-        captured["out_dim"] = w_t.shape[-1] if w_t.dim() == 2 else None
-        return torch.zeros(x_fp4.shape[0], w_t.shape[-1], dtype=torch.bfloat16)
-
-    monkeypatch.setattr(nv, "_mm_fp4", fake_mm_fp4)
-    monkeypatch.setattr(nv, "_get_ltx2_fp4_stage_profile", lambda default="refine": "refine")
-    layer._weight_global_sf = torch.tensor(1.0, dtype=torch.bfloat16)
-    x = torch.randn(2, 3, 16, dtype=torch.bfloat16)
-    out = layer.quant_method.apply(layer, x)
-    assert out.shape == (2, 3, 8)
-
-
-def test_dense_path_after_purge_raises_with_flag_named(monkeypatch) -> None:
-    """Defensive guard: convert never purges dense-capable layers anymore,
-    but a hand-purged module hitting the dense path must fail loudly."""
-    model = _model(retain=False)
-    nv.convert_model_to_nvfp4(model)
-    model.refine_only.register_parameter("weight", None)  # simulate misuse
-    monkeypatch.setattr(nv, "_get_ltx2_fp4_stage_profile", lambda default="refine": "base")
-    x = torch.randn(2, 3, 16, dtype=torch.bfloat16)
-    with pytest.raises(RuntimeError, match="retain_original_weights"):
-        model.refine_only.quant_method.apply(model.refine_only, x)

fastvideo/tests/ops/quantization/test_nvfp4_qat_receipt.py

@@ -1,70 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""CPU-only test for the NVFP4 QAT attachment receipt.
-
-``_maybe_quantize_model`` is the single post-load walk over the model,
-so it is where the definitive "how many linears carry the QAT-train
-method" count exists. This test builds a tiny module tree — one linear
-whose prefix matches the QAT target set, one that doesn't — and asserts
-the receipt reports the counts derived from that walk.
-"""
-from __future__ import annotations
-
-import logging
-from collections.abc import Iterator
-from contextlib import contextmanager
-
-import torch.nn as nn
-
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.layers.quantization.nvfp4_qat_train_config import (
-    NVFP4QATTrainConfig, )
-from fastvideo.models.loader.fsdp_load import _maybe_quantize_model
-
-
-# ``fastvideo.logger.init_logger`` sets ``propagate=False`` on its
-# loggers, so the standard ``caplog`` fixture cannot observe them.
-# Attach a temporary handler directly to the target logger instead
-# (same pattern as fastvideo/tests/train/callbacks/test_callback.py).
-@contextmanager
-def _capture_logger(name: str, level: int = logging.INFO) -> Iterator[list[logging.LogRecord]]:
-    logger = logging.getLogger(name)
-    records: list[logging.LogRecord] = []
-
-    class _Handler(logging.Handler):
-
-        def emit(self, record: logging.LogRecord) -> None:
-            records.append(record)
-
-    handler = _Handler(level=level)
-    prev_level = logger.level
-    logger.addHandler(handler)
-    logger.setLevel(level)
-    try:
-        yield records
-    finally:
-        logger.removeHandler(handler)
-        logger.setLevel(prev_level)
-
-
-def _model(quant_config: NVFP4QATTrainConfig | None) -> nn.Module:
-    model = nn.Module()
-    # Prefix in DEFAULT_FP4_LAYERS -> gets the QAT-train method attached.
-    model.matched = ReplicatedLinear(16, 16, quant_config=quant_config, prefix="blocks.0.attn1.to_k")
-    # Prefix outside the target set -> UnquantizedLinearMethod fallback.
-    model.unmatched = ReplicatedLinear(16, 16, quant_config=quant_config, prefix="blocks.0.proj_mlp")
-    return model
-
-
-def test_qat_train_receipt_reports_attached_and_skipped_counts() -> None:
-    model = _model(NVFP4QATTrainConfig())
-    with _capture_logger("fastvideo.models.loader.fsdp_load") as records:
-        _maybe_quantize_model(model)
-    messages = [r.getMessage() for r in records]
-    assert "NVFP4 QAT: attached 1 linears (1 skipped by prefix filter)" in messages
-
-
-def test_no_receipt_without_qat_train_layers() -> None:
-    model = _model(quant_config=None)
-    with _capture_logger("fastvideo.models.loader.fsdp_load") as records:
-        _maybe_quantize_model(model)
-    assert not any("NVFP4 QAT" in r.getMessage() for r in records)

fastvideo/tests/stages/test_kandinsky5_attention_backend_guard.py

@@ -1,83 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regression tests for ``Kandinsky5DenoisingStage._assert_local_attention_backend_engaged``.
-
-The stage guard must be strict for ``ATTN_QAT_TRAIN`` only:
-``fastvideo.platforms.cuda`` raises ``ImportError`` at backend selection if
-the training kernel isn't built (training must never silently
-de-quantize), so an exact match is guaranteed whenever module construction
-succeeded -- a mismatch there really does mean a silent SDPA fallback.
-
-``ATTN_QAT_INFER`` is different: ``fastvideo.platforms.cuda`` deliberately
-resolves it to FlashAttention when the sm_120 kernel isn't built, and the
-QAD recipe README (examples/train/configs/fine_tuning/kandinsky5/README.md)
-documents that fallback as the supported path on non-sm_120 GPUs ("the
-attention falls back to a supported dense backend while the FP4 linear
-layers still run"). Treating it as exact-match-only would let module
-construction succeed with FlashAttention and then abort denoising on every
-machine relying on the documented fallback.
-
-Pure logic tests on stand-in modules -- no GPU, no optional kernels, no
-model load, no distributed init needed.
-"""
-from __future__ import annotations
-
-import pytest
-import torch
-
-from fastvideo.attention import LocalAttention
-from fastvideo.pipelines.stages.kandinsky5 import Kandinsky5DenoisingStage
-from fastvideo.platforms import AttentionBackendEnum
-
-
-class _ResolvedLocalAttention(LocalAttention):
-    """``LocalAttention`` stand-in with a pre-resolved backend.
-
-    Bypasses ``LocalAttention.__init__`` (which runs real backend selection
-    -- unavailable without the optional kernels these tests are about)
-    while still passing the stage guard's
-    ``isinstance(module, LocalAttention)`` check.
-    """
-
-    def __init__(self, backend: AttentionBackendEnum) -> None:
-        torch.nn.Module.__init__(self)
-        self.backend = backend
-
-
-def _make_stage(resolved_backend: AttentionBackendEnum) -> Kandinsky5DenoisingStage:
-    """Bypass __init__: only set the field the guard reads."""
-    stage = Kandinsky5DenoisingStage.__new__(Kandinsky5DenoisingStage)
-    transformer = torch.nn.Module()
-    transformer.attn = _ResolvedLocalAttention(resolved_backend)
-    stage.transformer = transformer
-    return stage
-
-
-def test_attn_qat_infer_flash_fallback_is_allowed(monkeypatch):
-    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER")
-    stage = _make_stage(AttentionBackendEnum.FLASH_ATTN)
-
-    # Must not raise: FlashAttention is the documented ATTN_QAT_INFER
-    # fallback on GPUs without the sm_120 kernel.
-    stage._assert_local_attention_backend_engaged()
-
-
-def test_attn_qat_infer_exact_match_is_allowed(monkeypatch):
-    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER")
-    stage = _make_stage(AttentionBackendEnum.ATTN_QAT_INFER)
-
-    stage._assert_local_attention_backend_engaged()
-
-
-def test_attn_qat_train_mismatch_raises(monkeypatch):
-    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN")
-    stage = _make_stage(AttentionBackendEnum.TORCH_SDPA)
-
-    with pytest.raises(AssertionError, match="ATTN_QAT_TRAIN"):
-        stage._assert_local_attention_backend_engaged()
-
-
-def test_attn_qat_train_exact_match_is_allowed(monkeypatch):
-    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN")
-    stage = _make_stage(AttentionBackendEnum.ATTN_QAT_TRAIN)
-
-    stage._assert_local_attention_backend_engaged()

fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py

@@ -1,207 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Stage-boundary regression: ``Kandinsky5DmdDenoisingStage.forward`` must
-invoke the selected attention backend on every denoising step.
-
-``Kandinsky5Attention.forward`` catches the ``AssertionError``
-``LocalAttention`` raises when no pipeline forward context is set and
-silently falls back to plain ``F.scaled_dot_product_attention`` -- so if
-the ``set_forward_context`` wrapper inside the stage's denoising loop were
-ever removed, generation would still "work" while skipping whichever
-kernel ``FASTVIDEO_ATTENTION_BACKEND`` selected. The GPU test
-``fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py``
-installs ``set_forward_context`` itself around a raw transformer call, so
-it cannot catch that removal; this test drives the *production stage* and
-spies on the resolved backend impl (``LocalAttention.attn_impl.forward``,
-which is only reachable once ``get_forward_context()`` succeeds inside
-``LocalAttention.forward``).
-
-CPU-only: a tiny randomly-initialized Kandinsky5 transformer
-(~127K params) driven through the real stage code. TORCH_SDPA is pinned
-via env var so backend resolution is identical on CPU-only and CUDA
-machines. Output *numerics* are deliberately not asserted -- with random
-weights the 4-step predict-x0/re-noise feedback loop amplifies
-magnitudes without bound; the contract under test is backend routing,
-not sample quality (the nightly e2e owns that).
-"""
-from __future__ import annotations
-
-import types
-
-import pytest
-import torch
-
-from fastvideo.attention import LocalAttention
-from fastvideo.attention.selector import _cached_get_attn_backend
-from fastvideo.configs.models.dits.kandinsky5 import (
-    Kandinsky5ArchConfig,
-    Kandinsky5VideoConfig,
-)
-from fastvideo.forward_context import set_forward_context
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler, )
-
-DMD_STEPS = [1000, 750, 500, 250]
-TEXT_SEQ_LEN = 6
-
-
-def _tiny_arch() -> Kandinsky5ArchConfig:
-    # model_dim must be divisible by head_dim = sum(axes_dims) = 32.
-    return Kandinsky5ArchConfig(
-        in_visual_dim=4,
-        in_text_dim=32,
-        in_text_dim2=16,
-        time_dim=32,
-        out_visual_dim=4,
-        patch_size=(1, 2, 2),
-        model_dim=64,
-        ff_dim=128,
-        num_text_blocks=1,
-        num_visual_blocks=1,
-        axes_dims=(8, 12, 12),
-        visual_cond=False,
-        attention_type="regular",
-    )
-
-
-def _build_stage_and_spy(monkeypatch):
-    """Tiny transformer + real DMD stage on CPU, with every resolved
-    backend impl wrapped in a call counter."""
-    # Pin the backend so resolution is identical on CPU-only and CUDA
-    # machines, and clear the process-wide selector cache keyed on
-    # (head_size, dtype, supported_backends) -- NOT on the env var (same
-    # defensive pattern as test_kandinsky5_qat_attention_engages.py).
-    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-    _cached_get_attn_backend.cache_clear()
-
-    from fastvideo.models.dits.kandinsky5 import Kandinsky5Transformer3DModel
-    import fastvideo.pipelines.stages.kandinsky5 as k5_stage_mod
-
-    torch.manual_seed(0)
-    arch = _tiny_arch()
-    transformer = Kandinsky5Transformer3DModel(Kandinsky5VideoConfig(arch_config=arch), hf_config={})
-    transformer.eval()
-
-    # get_local_torch_device() never returns "cpu" (cuda -> mps fallback);
-    # pin the stage to CPU so this test runs identically everywhere.
-    monkeypatch.setattr(k5_stage_mod, "get_local_torch_device", lambda: torch.device("cpu"))
-
-    stage = k5_stage_mod.Kandinsky5DmdDenoisingStage(transformer, FlowMatchEulerDiscreteScheduler(shift=5.0))
-
-    backend_calls: list[int] = []
-    for module in transformer.modules():
-        if isinstance(module, LocalAttention):
-
-            def _spy(*args, __orig=module.attn_impl.forward, **kwargs):
-                backend_calls.append(1)
-                return __orig(*args, **kwargs)
-
-            module.attn_impl.forward = _spy
-
-    return stage, transformer, arch, backend_calls
-
-
-def _make_batch() -> types.SimpleNamespace:
-    # 5 frames / 64x64 with the real VAE ratios (4 temporal, 8 spatial)
-    # -> latents [1, 2, 8, 8, 4], patchified to a 2x4x4 grid.
-    return types.SimpleNamespace(
-        latents=torch.randn(1, 2, 8, 8, 4),
-        prompt_embeds=[torch.randn(1, TEXT_SEQ_LEN, 32), torch.randn(1, 16)],
-        prompt_attention_mask=[torch.ones(1, TEXT_SEQ_LEN, dtype=torch.long)],
-        height=64,
-        width=64,
-        num_frames=5,
-        image_latent=None,
-        generator=None,
-    )
-
-
-def _make_fastvideo_args(arch: Kandinsky5ArchConfig) -> types.SimpleNamespace:
-    return types.SimpleNamespace(
-        model_loaded={"transformer": True},
-        disable_autocast=True,
-        pipeline_config=types.SimpleNamespace(
-            dit_precision="fp32",
-            dmd_denoising_steps=list(DMD_STEPS),
-            dit_config=types.SimpleNamespace(arch_config=arch),
-            vae_config=types.SimpleNamespace(arch_config=types.SimpleNamespace(
-                temporal_compression_ratio=4,
-                spatial_compression_ratio=8,
-            )),
-        ),
-    )
-
-
-def test_dmd_stage_invokes_selected_backend_every_step(monkeypatch):
-    stage, _, arch, backend_calls = _build_stage_and_spy(monkeypatch)
-    try:
-        batch = _make_batch()
-
-        result = stage.forward(batch, _make_fastvideo_args(arch))
-
-        # 3 LocalAttention modules (text self-attn, visual self-attn, visual
-        # cross-attn) x one forward per DMD step. Any silent
-        # missing-forward-context fallback to F.scaled_dot_product_attention
-        # bypasses attn_impl.forward entirely and shows up here as a shortfall.
-        expected = 3 * len(DMD_STEPS)
-        assert len(backend_calls) == expected, (
-            f"selected attention backend invoked {len(backend_calls)} times, expected {expected} -- "
-            "the stage's set_forward_context wrapper is no longer covering the transformer call, "
-            "so Kandinsky5Attention silently fell back to raw SDPA")
-        assert result.latents.shape == (1, 2, 8, 8, 4)
-    finally:
-        _cached_get_attn_backend.cache_clear()
-
-
-def test_raw_transformer_without_context_bypasses_backend(monkeypatch):
-    """Documents the trap the stage wrapper exists to prevent: without a
-    forward context the transformer still returns output, but the selected
-    backend is never invoked."""
-    stage, transformer, _, backend_calls = _build_stage_and_spy(monkeypatch)
-    del stage  # only the spied transformer is needed here
-    try:
-        with torch.no_grad():
-            out = transformer(
-                hidden_states=torch.randn(1, 2, 8, 8, 4),
-                encoder_hidden_states=torch.randn(1, TEXT_SEQ_LEN, 32),
-                pooled_projections=torch.randn(1, 16),
-                timestep=torch.tensor([500.0]),
-                visual_rope_pos=[torch.arange(2), torch.arange(4), torch.arange(4)],
-                text_rope_pos=torch.arange(TEXT_SEQ_LEN),
-                scale_factor=(1.0, 2.0, 2.0),
-                sparse_params=None,
-                return_dict=True,
-            ).sample
-
-        assert out.shape == (1, 2, 8, 8, 4)
-        assert len(backend_calls) == 0, (
-            "expected the missing-forward-context fallback to bypass the backend impl; "
-            "if this now fails, Kandinsky5Attention's fallback behavior changed and the "
-            "stage guard docs/tests should be revisited")
-    finally:
-        _cached_get_attn_backend.cache_clear()
-
-
-def test_dmd_stage_with_context_and_raw_without_share_one_spy(monkeypatch):
-    """Same spy, both paths, one process: proves the counter difference
-    between the two tests above is the forward-context wrapper itself,
-    not some environmental difference."""
-    stage, transformer, arch, backend_calls = _build_stage_and_spy(monkeypatch)
-    try:
-        with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None):
-            transformer(
-                hidden_states=torch.randn(1, 2, 8, 8, 4),
-                encoder_hidden_states=torch.randn(1, TEXT_SEQ_LEN, 32),
-                pooled_projections=torch.randn(1, 16),
-                timestep=torch.tensor([500.0]),
-                visual_rope_pos=[torch.arange(2), torch.arange(4), torch.arange(4)],
-                text_rope_pos=torch.arange(TEXT_SEQ_LEN),
-                scale_factor=(1.0, 2.0, 2.0),
-                sparse_params=None,
-                return_dict=True,
-            )
-        assert len(backend_calls) == 3, "context-wrapped raw call should hit all 3 attention layers"
-
-        stage.forward(_make_batch(), _make_fastvideo_args(arch))
-        assert len(backend_calls) == 3 + 3 * len(DMD_STEPS)
-    finally:
-        _cached_get_attn_backend.cache_clear()

fastvideo/tests/stages/test_kandinsky5_target_dtype.py

@@ -1,139 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regression tests for ``Kandinsky5DenoisingStage._resolve_target_dtype``.
-
-Plain (non-FSDP) transformer: the stage must honor
-``pipeline_config.dit_precision`` exactly -- ``TransformerLoader.load()``
-asserts every parameter matches it for a non-FSDP load, so an explicit
-fp32 pipeline must not be silently cast to bf16 (an earlier
-parameter-scanning version of this helper did exactly that: all-fp32
-params fell through to a bf16 "safe default").
-
-FSDP2-wrapped transformer: the stage must read the active
-``MixedPrecisionPolicy`` recorded by ``set_mixed_precision_policy`` (which
-``maybe_load_fsdp_model`` always calls before sharding -- hardcoded to
-``param_dtype=torch.bfloat16``, ``cast_forward_inputs=False`` today), NOT
-the parameter storage dtypes. FSDP2 computes in the policy's
-``param_dtype`` regardless of the dtype parameters are stored in, so an
-fp16-loaded FSDP model still runs its forward in bf16; a parameter scan
-resolved fp16 there and, with autocast keyed off the wrong dtype and
-input casting disabled, handed bf16-computing modules fp16 inputs.
-
-These are pure logic tests -- the FSDP "wrap" only swaps ``__class__`` the
-same way ``fully_shard`` does (so the ``isinstance(..., FSDPModule)``
-check fires) without any process group, GPU, or model load.
-"""
-from __future__ import annotations
-
-import types
-
-import pytest
-import torch
-
-from fastvideo import utils as fastvideo_utils
-from fastvideo.pipelines.stages.kandinsky5 import Kandinsky5DenoisingStage
-from fastvideo.utils import set_mixed_precision_policy
-
-
-def _make_stage(transformer: torch.nn.Module) -> Kandinsky5DenoisingStage:
-    """Bypass __init__: only set the field _resolve_target_dtype reads."""
-    stage = Kandinsky5DenoisingStage.__new__(Kandinsky5DenoisingStage)
-    stage.transformer = transformer
-    return stage
-
-
-def _fastvideo_args(dit_precision: str) -> types.SimpleNamespace:
-    return types.SimpleNamespace(pipeline_config=types.SimpleNamespace(dit_precision=dit_precision))
-
-
-def _fsdp_wrap(module: torch.nn.Module) -> torch.nn.Module:
-    """Make ``isinstance(module, FSDPModule)`` true the way ``fully_shard``
-    does -- by swapping ``__class__`` to a dynamic ``(FSDPModule, orig)``
-    subclass -- without any actual sharding/distributed setup."""
-    from torch.distributed.fsdp import FSDPModule
-
-    orig_cls = module.__class__
-    module.__class__ = type(f"FSDP{orig_cls.__name__}", (FSDPModule, orig_cls), {})
-    return module
-
-
-@pytest.fixture()
-def mixed_precision_state_reset():
-    """Snapshot/restore the process's thread-local mixed-precision state so
-    these tests neither depend on nor leak state across the pytest run."""
-    state_holder = fastvideo_utils._mixed_precision_state
-    had_state = hasattr(state_holder, "state")
-    prev_state = getattr(state_holder, "state", None)
-    if had_state:
-        del state_holder.state
-    yield
-    if had_state:
-        state_holder.state = prev_state
-    elif hasattr(state_holder, "state"):
-        del state_holder.state
-
-
-def test_resolve_target_dtype_honors_explicit_fp32_for_plain_transformer():
-    transformer = torch.nn.Linear(4, 4).to(torch.float32)
-    stage = _make_stage(transformer)
-
-    resolved = stage._resolve_target_dtype(_fastvideo_args("fp32"))
-
-    assert resolved == torch.float32, (
-        "an explicit fp32 pipeline_config must not be silently cast to bf16 for a "
-        "plain (non-FSDP) transformer")
-
-
-def test_resolve_target_dtype_honors_explicit_fp16_for_plain_transformer():
-    transformer = torch.nn.Linear(4, 4).to(torch.float16)
-    stage = _make_stage(transformer)
-
-    resolved = stage._resolve_target_dtype(_fastvideo_args("fp16"))
-
-    assert resolved == torch.float16
-
-
-def test_resolve_target_dtype_honors_explicit_bf16_for_plain_transformer():
-    transformer = torch.nn.Linear(4, 4).to(torch.bfloat16)
-    stage = _make_stage(transformer)
-
-    resolved = stage._resolve_target_dtype(_fastvideo_args("bf16"))
-
-    assert resolved == torch.bfloat16
-
-
-def test_resolve_target_dtype_fsdp_reads_policy_not_parameter_storage(mixed_precision_state_reset):
-    """An fp16-*stored* FSDP model still computes in the policy's bf16 --
-    the exact mismatch a parameter scan gets wrong."""
-    set_mixed_precision_policy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
-    transformer = _fsdp_wrap(torch.nn.Linear(4, 4).to(torch.float16))
-    stage = _make_stage(transformer)
-
-    resolved = stage._resolve_target_dtype(_fastvideo_args("fp16"))
-
-    assert resolved == torch.bfloat16, (
-        "FSDP compute dtype comes from the MixedPrecisionPolicy param_dtype, not from "
-        "the dtype the parameters happen to be stored in")
-
-
-def test_resolve_target_dtype_fsdp_follows_a_non_default_policy(mixed_precision_state_reset):
-    """Proves the policy is actually read (not a hardcoded bf16)."""
-    set_mixed_precision_policy(param_dtype=torch.float16, reduce_dtype=torch.float32)
-    transformer = _fsdp_wrap(torch.nn.Linear(4, 4).to(torch.float32))
-    stage = _make_stage(transformer)
-
-    resolved = stage._resolve_target_dtype(_fastvideo_args("fp32"))
-
-    assert resolved == torch.float16
-
-
-def test_resolve_target_dtype_fsdp_defaults_to_bf16_without_policy_state(mixed_precision_state_reset):
-    """maybe_load_fsdp_model always records the policy before sharding, but
-    if the thread-local state is somehow unset the stage must still match
-    the loader's hardcoded bf16 param_dtype rather than crash or resolve
-    the storage dtype."""
-    transformer = _fsdp_wrap(torch.nn.Linear(4, 4).to(torch.float16))
-    stage = _make_stage(transformer)
-
-    resolved = stage._resolve_target_dtype(_fastvideo_args("fp16"))
-
-    assert resolved == torch.bfloat16

fastvideo/tests/train/callbacks/test_validation.py

@@ -82,7 +82,6 @@ class TestConstructor:
         assert cb.overlay_actions is False
         assert cb.offload_training_state is False
         assert cb.unload_pipeline_after_validation is False
-        assert cb.attn_qat_infer is False
         # Lazy fields not yet populated.
         assert cb._pipeline is None
         assert cb._sampling_param is None
@@ -103,7 +102,6 @@ class TestConstructor:
             overlay_actions=1,  # type: ignore[arg-type]
             offload_training_state="1",  # type: ignore[arg-type]
             unload_pipeline_after_validation="false",  # type: ignore[arg-type]
-            attn_qat_infer="true",  # type: ignore[arg-type]
         )
         assert cb.every_steps == 50
         assert cb.sampling_steps == [20, 40]
@@ -113,7 +111,6 @@ class TestConstructor:
         assert cb.overlay_actions is True
         assert cb.offload_training_state is True
         assert cb.unload_pipeline_after_validation is False
-        assert cb.attn_qat_infer is True
 
     def test_pipeline_kwargs_collected(self) -> None:
         cb = ValidationCallback(
@@ -222,57 +219,6 @@ class TestOnValidationBegin:
         assert cb.run_calls == [0]
 
 
-class TestAttnQatInferValidation:
-
-    def test_context_swaps_and_restores_qat_attention(
-        self,
-        monkeypatch: pytest.MonkeyPatch,
-    ) -> None:
-        from fastvideo.attention.backends import attn_qat_infer
-        from fastvideo.attention.backends.attn_qat_infer import (
-            AttnQatInferImpl, )
-        from fastvideo.attention.backends.attn_qat_train import (
-            AttnQatTrainImpl, )
-        from fastvideo.platforms import AttentionBackendEnum
-
-        monkeypatch.setattr(attn_qat_infer, "is_attn_qat_infer_available", lambda: True)
-        impl = AttnQatTrainImpl(
-            num_heads=1,
-            head_size=128,
-            causal=False,
-            softmax_scale=128**-0.5,
-        )
-
-        holder = torch.nn.Module()
-        holder.attn_impl = impl
-        holder.num_heads = 1
-        holder.head_size = 128
-        holder.num_kv_heads = 1
-        holder.softmax_scale = 128**-0.5
-        holder.backend = AttentionBackendEnum.ATTN_QAT_TRAIN
-        transformer = torch.nn.Module()
-        transformer.add_module("attention", holder)
-        cb = _make_callback()
-        cb.attn_qat_infer = True
-
-        with pytest.raises(RuntimeError, match="stop"):
-            with cb._attn_qat_infer_context(transformer):
-                assert isinstance(holder.attn_impl, AttnQatInferImpl)
-                assert holder.backend is AttentionBackendEnum.ATTN_QAT_INFER
-                raise RuntimeError("stop")
-
-        assert holder.attn_impl is impl
-        assert holder.backend is AttentionBackendEnum.ATTN_QAT_TRAIN
-
-    def test_context_requires_qat_training_attention(self) -> None:
-        cb = _make_callback()
-        cb.attn_qat_infer = True
-
-        with pytest.raises(RuntimeError, match="no ATTN_QAT_TRAIN layers"):
-            with cb._attn_qat_infer_context(torch.nn.Linear(1, 1)):
-                pass
-
-
 # ---------------------------------------------------------------------------
 # C. _find_ema_callback
 # ---------------------------------------------------------------------------

fastvideo/tests/train/fixtures/kandinsky5_t2v_min.yaml

@@ -1,17 +0,0 @@
-# Minimum config to instantiate Kandinsky5Model for the loading + forward
-# smoke test. Real Kandinsky-5.0 Lite checkpoint loaded via
-# load_module_from_path.
-
-models:
-  student:
-    _target_: fastvideo.train.models.kandinsky5.Kandinsky5Model
-    init_from: kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers
-    trainable: false
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  dit_precision: bf16
-
-pipeline: {}

fastvideo/tests/train/fixtures/ltx2_t2v_qat_finetune_min.yaml

@@ -1,51 +0,0 @@
-# Minimum config to run a single NVFP4-QAT training step of
-# FineTuneMethod on LTX2Model for the per-method smoke test. Identical
-# to ltx2_t2v_finetune_min.yaml except pipeline.dit_config.quant_config
-# selects the LTX-2 deployment-matched NVFP4 QAT-train profile (STE
-# through real FP4 GEMMs). Requires flashinfer FP4 kernels (sm_100-class GPU); the test
-# additionally skips below 60GB GPU memory.
-
-models:
-  student:
-    _target_: fastvideo.train.models.ltx2.LTX2Model
-    init_from: FastVideo/LTX2-Distilled-Diffusers
-    trainable: true
-    attention_backend: ATTN_QAT_TRAIN
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
-
-training:
-  dit_precision: bf16
-
-  distributed:
-    num_gpus: 1
-    sp_size: 1
-    tp_size: 1
-
-  data:
-    seed: 42
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    num_latent_t: 3
-    num_height: 256
-    num_width: 256
-    num_frames: 17
-
-  optimizer:
-    learning_rate: 1.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 1
-    gradient_accumulation_steps: 1
-
-  model:
-    precondition_outputs: false
-
-pipeline:
-  dit_config:
-    quant_config: nvfp4_qat_train

fastvideo/tests/train/methods/grad_norm_refs.json

@@ -9,9 +9,6 @@
   "test_ltx2_finetune": {
     "GB200": 0.3127
   },
-  "test_ltx2_nvfp4_qat_finetune": {
-    "GB200": 0.2421
-  },
   "test_matrixgame2_finetune": {
     "H200": 4.8247
   },

fastvideo/tests/train/methods/test_ltx2_finetune.py

@@ -2,12 +2,12 @@
 """Per-method GPU smoke test: ``LTX2Model`` + ``FineTuneMethod``.
 
 Mirrors ``test_wan_finetune.py`` for the LTX-2 plugin, parametrized
-over dense LTX-2.0/LTX-2.3 and LTX-2.0 NVFP4-QAT checkpoints.
-LTX2-specific differences: the synthetic ``raw_batch`` carries a
-single post-connector Gemma embedding (3840-d for 2.0, 4096-d for 2.3
-which has no in-DiT caption projection) and 128-channel VAE latents;
-the DiT is 18.9B params, so the test skips on GPUs with less than 60GB
-memory (e.g. the L40S CI runner).
+over LTX-2.0 and LTX-2.3 checkpoints. LTX2-specific differences: the
+synthetic ``raw_batch`` carries a single post-connector Gemma
+embedding (3840-d for 2.0, 4096-d for 2.3 which has no in-DiT caption
+projection) and 128-channel VAE latents; the DiT is 18.9B params, so
+the test skips on GPUs with less than 60GB memory (e.g. the L40S CI
+runner).
 """
 
 from __future__ import annotations
@@ -34,17 +34,10 @@ from .grad_norm_regression import (
 
 _FIXTURE_DIR = Path(__file__).resolve().parent.parent / "fixtures"
 
-# id -> (fixture, post-connector text embedding width, grad-norm ref name,
-#        expected NVFP4-QAT attachment count or None for dense training)
+# id -> (fixture, post-connector text embedding width, grad-norm ref name)
 _CASES = {
-    "ltx2": ("ltx2_t2v_finetune_min.yaml", 3840, "test_ltx2_finetune", None),
-    "ltx2_3": ("ltx2_3_t2v_finetune_min.yaml", 4096, "test_ltx2_3_finetune", None),
-    "ltx2_nvfp4_qat": (
-        "ltx2_t2v_qat_finetune_min.yaml",
-        3840,
-        "test_ltx2_nvfp4_qat_finetune",
-        576,
-    ),
+    "ltx2": ("ltx2_t2v_finetune_min.yaml", 3840, "test_ltx2_finetune"),
+    "ltx2_3": ("ltx2_3_t2v_finetune_min.yaml", 4096, "test_ltx2_3_finetune"),
 }
 
 _LTX2_TEXT_LEN = 1024
@@ -59,18 +52,6 @@ def _gpu_too_small() -> bool:
     return total < _MIN_GPU_MEMORY_GB * 1024**3
 
 
-def _flashinfer_fp4_available() -> bool:
-    try:
-        from flashinfer import (  # noqa: F401
-            SfLayout,
-            mm_fp4,
-            nvfp4_quantize,
-        )
-    except ImportError:
-        return False
-    return True
-
-
 def _build_synthetic_batch(
     device: torch.device,
     dtype: torch.dtype,
@@ -104,21 +85,8 @@ def test_ltx2_finetune_single_train_step(
         pytest.skip(f"requires a CUDA GPU with >= {_MIN_GPU_MEMORY_GB}GB "
                     "memory (LTX-2 DiT is 18.9B params)")
 
-    fixture_name, text_dim, ref_name, expected_qat_linears = _CASES[case]
-    if expected_qat_linears is not None:
-        if torch.cuda.get_device_capability(0) < (10, 0):
-            pytest.skip("NVFP4 QAT requires an SM100+ GPU")
-        if not _flashinfer_fp4_available():
-            pytest.skip("requires flashinfer with FP4 kernels")
-
+    fixture_name, text_dim, ref_name = _CASES[case]
     cfg = load_run_config(str(_FIXTURE_DIR / fixture_name))
-    if expected_qat_linears is not None:
-        from fastvideo.layers.quantization.nvfp4_qat_train_config import (
-            NVFP4QATTrainConfig, )
-        assert isinstance(
-            cfg.training.pipeline_config.dit_config.quant_config,
-            NVFP4QATTrainConfig,
-        )
 
     device = torch.device("cuda:0")
     dtype = torch.bfloat16
@@ -137,30 +105,9 @@ def test_ltx2_finetune_single_train_step(
         init_from=cfg.models["student"]["init_from"],
         training_config=cfg.training,
         trainable=True,
-        attention_backend=cfg.models["student"].get("attention_backend"),
     )
     model.transformer = model.transformer.to(device=device, dtype=dtype)
 
-    if expected_qat_linears is not None:
-        from fastvideo.layers.quantization.nvfp4_qat_train_config import (
-            NVFP4QATTrainQuantizeMethod, )
-        quantized = sum(
-            isinstance(getattr(module, "quant_method", None),
-                       NVFP4QATTrainQuantizeMethod)
-            for module in model.transformer.modules())
-        assert quantized == expected_qat_linears, (
-            f"expected {expected_qat_linears} deployment-matched NVFP4-QAT "
-            f"linears on the LTX-2 DiT, found {quantized}")
-
-        from fastvideo.attention.backends.attn_qat_train import (
-            AttnQatTrainImpl, )
-        qat_attention = sum(
-            isinstance(getattr(module, "attn_impl", None), AttnQatTrainImpl)
-            for module in model.transformer.modules())
-        assert qat_attention == 96, (
-            "expected ATTN_QAT_TRAIN on attn1 and attn2 in all 48 LTX-2 "
-            f"blocks, found {qat_attention} implementations")
-
     method = FineTuneMethod(
         cfg=cfg,
         role_models={"student": model},

fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py

@@ -1,162 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Forward+backward smoke test confirming ATTN_QAT_TRAIN actually engages
-for Kandinsky5's dense/local attention path, rather than silently falling
-back to plain SDPA.
-
-Two silent-fallback traps this guards against:
-  - ``Kandinsky5Attention.forward`` catches ``AssertionError`` from
-    ``LocalAttention`` when no forward context is set and falls back to
-    ``F.scaled_dot_product_attention`` (a real fallback for standalone
-    parity tests, not relevant here since we always set a forward context,
-    but worth asserting past explicitly).
-  - ``_cached_get_attn_backend`` is process-wide ``@cache``d on
-    ``(head_size, dtype, supported_attention_backends)`` -- NOT on the env
-    var -- so a stale backend selection from an earlier test in the same
-    process can silently linger. Cleared here before selecting, matching
-    the defensive pattern in
-    ``fastvideo/models/loader/component_loader.py``.
-
-This test itself must not become that stale-selection source for whichever
-model test runs next in the same pytest process: ``FASTVIDEO_ATTENTION_BACKEND``
-is set via ``monkeypatch`` (auto-restored at teardown, including on
-failure) and the cache is cleared again in a ``finally`` block, since
-``monkeypatch`` only restores the env var, not the separate process-wide
-cache keyed on it.
-
-Unlike ``fastvideo.platforms.cuda``'s other backends, ATTN_QAT_TRAIN has no
-silent-fallback path at all if the kernel isn't built -- backend selection
-raises ``ImportError`` loudly. This test skips (does not fail) when the
-kernel isn't available, since it's an optional build artifact from
-fastvideo-kernel/.
-
-Runs the *full* transformer forward (mirroring test_load_kandinsky5.py)
-rather than invoking an attention submodule directly. Two earlier versions
-of this test tried isolating a single Kandinsky5Attention module -- first a
-freshly-random-initialized one (the fake-quant kernel produced NaN/Inf on
-untrained weights), then a real submodule pulled out of the loaded model
-(FSDP2's fully_shard() registers its unshard/reshard hooks on the top-level
-transformer's __call__, not on submodules individually, so invoking a
-submodule directly left its Linear layers' parameters in sharded DTensor
-form -- incompatible with a plain-Tensor input regardless of the
-`trainable` flag). Going through the properly-hooked top-level forward call
-sidesteps both problems.
-"""
-
-from __future__ import annotations
-
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29519")
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.attention import LocalAttention
-from fastvideo.attention.selector import _cached_get_attn_backend
-from fastvideo.forward_context import set_forward_context
-from fastvideo.platforms import AttentionBackendEnum
-
-_FIXTURE = str(
-    Path(__file__).resolve().parent.parent / "fixtures" /
-    "kandinsky5_t2v_min.yaml")
-
-
-@pytest.mark.usefixtures("distributed_setup")
-def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch):
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-
-    from fastvideo.attention.backends.attn_qat_train import (
-        is_attn_qat_train_available, )
-    if not is_attn_qat_train_available():
-        pytest.skip("fastvideo_kernel ATTN_QAT_TRAIN kernel is not built")
-
-    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN")
-    _cached_get_attn_backend.cache_clear()
-    try:
-        from fastvideo.train.models.kandinsky5 import Kandinsky5Model
-        from fastvideo.train.utils.config import load_run_config
-
-        cfg = load_run_config(_FIXTURE)
-        model = Kandinsky5Model(
-            init_from=cfg.models["student"]["init_from"],
-            training_config=cfg.training,
-            trainable=True,
-        )
-
-        device = torch.device("cuda:0")
-        dtype = torch.bfloat16
-        transformer = model.transformer.to(device=device, dtype=dtype)
-
-        attn = transformer.visual_transformer_blocks[0].self_attention
-        assert isinstance(attn.local_attention, LocalAttention)
-        assert attn.local_attention.backend == AttentionBackendEnum.ATTN_QAT_TRAIN, (
-            f"expected ATTN_QAT_TRAIN, got {attn.local_attention.backend} -- "
-            "backend selection silently fell back")
-
-        arch = transformer.config.arch_config
-        patch_size = arch.patch_size
-        in_visual_dim = arch.in_visual_dim
-        in_text_dim = arch.in_text_dim
-        in_text_dim2 = arch.in_text_dim2
-
-        grid_t, grid_h, grid_w = 2, 4, 4
-        latent_t = grid_t * patch_size[0]
-        latent_h = grid_h * patch_size[1]
-        latent_w = grid_w * patch_size[2]
-
-        latents = torch.randn(
-            1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype,
-            requires_grad=True)
-        if bool(getattr(transformer, "visual_cond", False)):
-            # See Kandinsky5Model._build_distill_input_kwargs /
-            # Kandinsky5LatentPreparationStage: visual_cond=True checkpoints
-            # always expect [real | zero_cond | zero_mask] concatenated on the
-            # channel dim.
-            cond = torch.zeros_like(latents)
-            mask = torch.zeros(*latents.shape[:-1], 1, device=device, dtype=dtype)
-            hidden_states = torch.cat([latents, cond, mask], dim=-1)
-        else:
-            hidden_states = latents
-        encoder_hidden_states = torch.randn(1, 8, in_text_dim, device=device, dtype=dtype)
-        pooled_projections = torch.randn(1, in_text_dim2, device=device, dtype=dtype)
-        timestep = torch.tensor([500], device=device, dtype=dtype)
-        visual_rope_pos = [
-            torch.arange(grid_t, device=device),
-            torch.arange(grid_h, device=device),
-            torch.arange(grid_w, device=device),
-        ]
-        text_rope_pos = torch.arange(encoder_hidden_states.shape[1], device=device)
-
-        with set_forward_context(current_timestep=0, attn_metadata=None):
-            out = transformer(
-                hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                pooled_projections=pooled_projections,
-                timestep=timestep,
-                visual_rope_pos=visual_rope_pos,
-                text_rope_pos=text_rope_pos,
-                scale_factor=(1.0, 2.0, 2.0),
-                sparse_params=None,
-                return_dict=True,
-            ).sample
-
-        assert torch.isfinite(out).all().item(), "output contains NaN/Inf"
-
-        out.sum().backward()
-        assert latents.grad is not None
-        assert torch.isfinite(latents.grad).all().item(), "input grad contains NaN/Inf"
-        assert attn.to_query.weight.grad is not None
-        assert torch.isfinite(attn.to_query.weight.grad.to_local()
-                              if hasattr(attn.to_query.weight.grad, "to_local") else
-                              attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf"
-    finally:
-        # monkeypatch restores FASTVIDEO_ATTENTION_BACKEND itself, but the
-        # selector cache is a separate process-wide functools.cache keyed on
-        # (head_size, dtype, supported_backends) -- not on the env var -- so
-        # it must be cleared again here or a later model test in this same
-        # pytest process could reuse this test's ATTN_QAT_TRAIN selection.
-        _cached_get_attn_backend.cache_clear()

fastvideo/tests/train/models/test_load_kandinsky5.py

@@ -1,112 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""GPU loading + forward smoke test for ``Kandinsky5Model``.
-
-Loads the real Kandinsky-5.0 Lite checkpoint via ``Kandinsky5Model.__init__``
-and runs one transformer forward pass on synthetic inputs. Kandinsky5's
-transformer takes a structurally different forward signature than Wan/Hunyuan
--- dual text conditioning (``encoder_hidden_states`` + ``pooled_projections``),
-explicit RoPE position tensors, a ``scale_factor``, and channel-last
-``[B, T, H, W, C]`` hidden_states with a dict return (``.sample``) -- this
-test mirrors the kwargs in ``Kandinsky5Model._build_distill_input_kwargs``
-and ``Kandinsky5DenoisingStage``.
-"""
-
-from __future__ import annotations
-
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29518")
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.forward_context import set_forward_context
-from fastvideo.train.models.kandinsky5 import Kandinsky5Model
-from fastvideo.train.utils.config import load_run_config
-
-_FIXTURE = str(
-    Path(__file__).resolve().parent.parent / "fixtures" /
-    "kandinsky5_t2v_min.yaml")
-
-
-@pytest.mark.usefixtures("distributed_setup")
-def test_kandinsky5_model_loads_and_forwards():
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-
-    cfg = load_run_config(_FIXTURE)
-    model = Kandinsky5Model(
-        init_from=cfg.models["student"]["init_from"],
-        training_config=cfg.training,
-        trainable=False,
-    )
-
-    transformer = model.transformer
-    assert isinstance(transformer, torch.nn.Module)
-    assert sum(p.numel() for p in transformer.parameters()) > 0
-
-    device = torch.device("cuda:0")
-    dtype = torch.bfloat16
-    transformer = transformer.to(device=device, dtype=dtype).eval()
-
-    arch = transformer.config.arch_config
-    patch_size = arch.patch_size
-    in_visual_dim = arch.in_visual_dim
-    out_visual_dim = arch.out_visual_dim
-    in_text_dim = arch.in_text_dim
-    in_text_dim2 = arch.in_text_dim2
-
-    # Small patch grid (T, H, W) so this fits next to the real checkpoint on
-    # a single GPU. scale_factor matches the 480p band Kandinsky5Model
-    # asserts on -- see fastvideo/train/models/kandinsky5/kandinsky5.py.
-    grid_t, grid_h, grid_w = 2, 4, 4
-    latent_t = grid_t * patch_size[0]
-    latent_h = grid_h * patch_size[1]
-    latent_w = grid_w * patch_size[2]
-
-    latents = torch.randn(
-        1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype)
-    if bool(getattr(transformer, "visual_cond", False)):
-        # The shipped checkpoint has visual_cond=True (unified T2V/I2V):
-        # Kandinsky5VisualEmbeddings expects [real | zero_cond | zero_mask]
-        # concatenated on the channel dim, even for pure T2V. Mirrors
-        # Kandinsky5Model._build_distill_input_kwargs / Kandinsky5LatentPreparationStage.
-        cond = torch.zeros_like(latents)
-        mask = torch.zeros(*latents.shape[:-1], 1, device=device, dtype=dtype)
-        hidden_states = torch.cat([latents, cond, mask], dim=-1)
-    else:
-        hidden_states = latents
-    encoder_hidden_states = torch.randn(1, 8, in_text_dim, device=device, dtype=dtype)
-    pooled_projections = torch.randn(1, in_text_dim2, device=device, dtype=dtype)
-    timestep = torch.tensor([500], device=device, dtype=dtype)
-    visual_rope_pos = [
-        torch.arange(grid_t, device=device),
-        torch.arange(grid_h, device=device),
-        torch.arange(grid_w, device=device),
-    ]
-    text_rope_pos = torch.arange(encoder_hidden_states.shape[1], device=device)
-
-    with torch.no_grad(), set_forward_context(
-            current_timestep=0,
-            attn_metadata=None,
-    ):
-        out = transformer(
-            hidden_states=hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            pooled_projections=pooled_projections,
-            timestep=timestep,
-            visual_rope_pos=visual_rope_pos,
-            text_rope_pos=text_rope_pos,
-            scale_factor=(1.0, 2.0, 2.0),
-            sparse_params=None,
-            return_dict=True,
-        ).sample
-
-    expected_shape = (1, latent_t, latent_h, latent_w, out_visual_dim)
-    assert tuple(out.shape) == expected_shape, (
-        f"output shape {tuple(out.shape)} != expected {expected_shape}")
-    assert torch.isfinite(out).all().item(), "output contains NaN/Inf"

fastvideo/tests/train/models/test_load_ltx2.py

@@ -34,7 +34,6 @@ from fastvideo.forward_context import (
     set_forward_context,
 )
 from fastvideo.pipelines import ForwardBatch
-from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.train.models.ltx2 import LTX2Model
 from fastvideo.train.utils.config import load_run_config
 
@@ -110,31 +109,6 @@ def test_ltx2_rejects_audio_training_before_loading(
         )
 
 
-def test_ltx2_forwards_role_attention_backend_to_loader(
-        monkeypatch: pytest.MonkeyPatch) -> None:
-    cfg = load_run_config(str(_FIXTURE_DIR / _CASES["ltx2"][0]))
-    transformer = _TinyLTX2Transformer(is_ltx2_3=False)
-    captured: list[AttentionBackendEnum | None] = []
-
-    def fake_load_module_from_path(**kwargs):
-        captured.append(kwargs.get("attention_backend"))
-        return transformer
-
-    monkeypatch.setattr(
-        "fastvideo.train.models.ltx2.ltx2.load_module_from_path",
-        fake_load_module_from_path,
-    )
-    model = LTX2Model(
-        init_from=cfg.models["student"]["init_from"],
-        training_config=cfg.training,
-        trainable=False,
-        attention_backend="TORCH_SDPA",
-    )
-
-    assert model.attention_backend is AttentionBackendEnum.TORCH_SDPA
-    assert captured == [AttentionBackendEnum.TORCH_SDPA]
-
-
 @pytest.mark.parametrize("case", _CASES.keys())
 def test_ltx2_wrapper_contract_runs_on_cpu(
         case: str, monkeypatch: pytest.MonkeyPatch) -> None:

fastvideo/tests/train/utils/test_config.py

@@ -245,46 +245,6 @@ def test_betas_parses_list_and_string_forms(
     assert cfg.training.optimizer.betas == expected
 
 
-def test_pipeline_quant_config_resolves_for_load_and_export(tmp_path: Path) -> None:
-    from fastvideo.layers.quantization.nvfp4_qat_train_config import (
-        NVFP4QATTrainConfig, )
-    from fastvideo.train.entrypoint.dcp_to_diffusers import (
-        _run_config_from_raw, )
-
-    data = _minimal_yaml()
-    data["models"]["student"]["init_from"] = (
-        "FastVideo/LTX2-Distilled-Diffusers")
-    data["pipeline"] = {
-        "dit_config": {
-            "quant_config": "nvfp4_qat_train"
-        }
-    }
-
-    cfg = load_run_config(_write_yaml(tmp_path, data))
-    assert isinstance(cfg.training.pipeline_config.dit_config.quant_config,
-                      NVFP4QATTrainConfig)
-
-    export_cfg = _run_config_from_raw(cfg.raw)
-    assert isinstance(
-        export_cfg.training.pipeline_config.dit_config.quant_config,
-        NVFP4QATTrainConfig,
-    )
-
-
-def test_pipeline_quant_config_rejects_unknown_name(tmp_path: Path) -> None:
-    data = _minimal_yaml()
-    data["models"]["student"]["init_from"] = (
-        "FastVideo/LTX2-Distilled-Diffusers")
-    data["pipeline"] = {
-        "dit_config": {
-            "quant_config": "not_a_quantization_method"
-        }
-    }
-
-    with pytest.raises(ValueError, match="Invalid quantization method"):
-        load_run_config(_write_yaml(tmp_path, data))
-
-
 def test_data_path_mapping_parses_repeat_counts(tmp_path: Path) -> None:
     # Config loading should preserve structured multi-dataset paths so the
     # dataset layer can interpret repeat counts later.

fastvideo/tests/train/utils/test_finetune_config_construction.py

@@ -1,44 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""CPU-only construction check for every shipped fine-tuning recipe.
-
-Parametrizes over ``examples/train/configs/fine_tuning/**/*.yaml`` and
-asserts each loads through :func:`load_run_config` with a resolved
-pipeline config and an importable-shaped student target — the
-construction proof (parse-green is not construct-green) and a
-regression net for every family's recipes.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-from fastvideo.train.utils.config import load_run_config
-
-_REPO_ROOT = Path(__file__).resolve().parents[4]
-_RECIPE_DIR = _REPO_ROOT / "examples" / "train" / "configs" / "fine_tuning"
-
-_RECIPES = sorted(_RECIPE_DIR.glob("**/*.yaml"))
-
-
-def _recipe_id(path: Path) -> str:
-    return str(path.relative_to(_RECIPE_DIR))
-
-
-@pytest.mark.parametrize("recipe", _RECIPES, ids=_recipe_id)
-def test_fine_tuning_recipe_constructs(recipe: Path) -> None:
-    cfg = load_run_config(str(recipe))
-
-    assert cfg.training.pipeline_config is not None, (
-        f"{_recipe_id(recipe)} did not resolve a pipeline config; "
-        "check its `pipeline:` key and init_from registry match")
-
-    target = cfg.models["student"]["_target_"]
-    assert isinstance(target, str) and target and "." in target, (
-        f"{_recipe_id(recipe)} student _target_ is not importable-shaped: "
-        f"{target!r}")
-
-
-def test_recipe_glob_found_recipes() -> None:
-    # Guard against the glob silently matching nothing after a move.
-    assert len(_RECIPES) >= 20

fastvideo/tests/workflow/test_kandinsky5_e2e_media_oracle.py

@@ -1,106 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Known-bad calibration tests for the Kandinsky5 e2e media oracle.
-
-A review pass demonstrated a concrete false-green: a 121-frame solid clip
-at the committed reference's mean RGB passed the old degeneracy check
-(global RGB std counts channel-mean differences as "variance": 15.4) AND
-the old MS-SSIM floor (scoring a mean of ~0.92 against the low-contrast
-reference, with the SSIM helper additionally truncating both clips to the
-shorter frame count). These tests turn that exploit -- and its frozen and
-truncated siblings -- into permanent regressions against the *currently
-committed* reference video, so the oracle's thresholds can never silently
-drift below a known-bad clip again: if a future reference changes the
-similarity landscape, the floor assertions here fail and force
-recalibration.
-
-CPU-only; skips when the committed reference is absent (pre-bootstrap
-checkouts).
-"""
-from __future__ import annotations
-
-import cv2
-import numpy as np
-import pytest
-
-from fastvideo.tests.nightly.test_e2e_kandinsky5_dmd_t2v_overfit import (
-    EXPECTED_FRAME_COUNT,
-    MIN_MEAN_MS_SSIM,
-    REFERENCE_VIDEO,
-    _assert_video_not_degenerate,
-    _decode_video,
-)
-
-pytestmark = pytest.mark.skipif(
-    not REFERENCE_VIDEO.exists(),
-    reason="committed Kandinsky5 e2e reference video not present",
-)
-
-
-def _write_clip(path, frames) -> None:
-    height, width = frames[0].shape[:2]
-    writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 24.0, (width, height))
-    assert writer.isOpened(), "mp4v VideoWriter unavailable in this OpenCV build"
-    for frame in frames:
-        writer.write(frame)
-    writer.release()
-
-
-def _mean_ms_ssim_vs_reference(path) -> float:
-    pytest.importorskip("pytorch_msssim")
-    pytest.importorskip("av")
-    from fastvideo.tests.utils import compute_video_ssim_torchvision
-
-    mean_ssim, _, _ = compute_video_ssim_torchvision(str(REFERENCE_VIDEO), str(path), use_ms_ssim=True)
-    return float(mean_ssim)
-
-
-def test_committed_reference_passes_the_oracle():
-    """Self-consistency: thresholds must never be tightened past what the
-    reviewed reference itself exhibits."""
-    _assert_video_not_degenerate(REFERENCE_VIDEO)
-
-
-def test_solid_clip_at_reference_mean_color_is_rejected(tmp_path):
-    """The review's exact exploit: solid clip at the reference's own mean
-    RGB. Global-std passed it at 15.4; per-frame luma spatial std is 0."""
-    reference = _decode_video(REFERENCE_VIDEO)
-    mean_color = reference.astype(np.float64).mean(axis=(0, 1, 2)).round().astype(np.uint8)
-    solid_frame = np.full(reference.shape[1:], mean_color, dtype=np.uint8)
-    solid_path = tmp_path / "solid.mp4"
-    _write_clip(solid_path, [solid_frame] * EXPECTED_FRAME_COUNT)
-
-    with pytest.raises(AssertionError, match="solid"):
-        _assert_video_not_degenerate(solid_path)
-
-    # The similarity floor must ALSO sit above this clip's score (measured
-    # ~0.9189 against the 2026-07-18 reference): the structural check is
-    # the primary defense, but the floor must not regress into known-bad
-    # territory either.
-    assert _mean_ms_ssim_vs_reference(solid_path) < MIN_MEAN_MS_SSIM
-
-
-def test_frozen_clip_is_rejected(tmp_path):
-    """A single real reference frame repeated 121x has genuine spatial
-    content (luma std ~12) -- only the temporal check catches it."""
-    reference = _decode_video(REFERENCE_VIDEO)
-    frozen_path = tmp_path / "frozen.mp4"
-    _write_clip(frozen_path, [reference[0]] * EXPECTED_FRAME_COUNT)
-
-    with pytest.raises(AssertionError, match="frozen"):
-        _assert_video_not_degenerate(frozen_path)
-
-    # Measured ~0.8597 against the 2026-07-18 reference.
-    assert _mean_ms_ssim_vs_reference(frozen_path) < MIN_MEAN_MS_SSIM
-
-
-def test_truncated_clip_is_rejected(tmp_path):
-    """compute_video_ssim_torchvision truncates both clips to the shorter
-    frame count, so a partial video could score arbitrarily well -- the
-    oracle must reject on frame count before similarity is even
-    consulted."""
-    reference = _decode_video(REFERENCE_VIDEO)
-    truncated_path = tmp_path / "truncated.mp4"
-    _write_clip(truncated_path, list(reference[:EXPECTED_FRAME_COUNT // 2]))
-
-    with pytest.raises(AssertionError, match="frame count or geometry"):
-        _assert_video_not_degenerate(truncated_path)

fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py

@@ -1,45 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regression test: preprocess_kandinsky5_overfit.py must not silently
-truncate a string caption to its first character.
-
-The documented ``videos2caption.json`` schema
-(``docs/training/data_preprocess.md``) stores ``cap`` as a plain string
-(e.g. ``"Ocean waves at sunset..."``), while some producers (this repo's
-own e2e test fixtures, mirroring ``preprocess_hunyuan_overfit.py``) store a
-non-empty list of caption variants instead. Indexing unconditionally with
-``item["cap"][0]`` silently takes the first *character* of a string caption
-("O") instead of erroring -- ``get_caption`` accepts and validates both
-forms once at ingestion.
-
-Pure logic test -- no GPU, no model load.
-"""
-from __future__ import annotations
-
-import pytest
-
-from fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit import get_caption
-
-
-def test_get_caption_accepts_plain_string():
-    assert get_caption({"path": "a.mp4", "cap": "Ocean waves at sunset..."}) == "Ocean waves at sunset..."
-
-
-def test_get_caption_accepts_nonempty_list():
-    assert get_caption({"path": "a.mp4", "cap": ["a synthetic test clip"]}) == "a synthetic test clip"
-
-
-def test_get_caption_uses_first_list_element():
-    assert get_caption({"path": "a.mp4", "cap": ["first", "second"]}) == "first"
-
-
-@pytest.mark.parametrize("bad_cap", [
-    "",
-    [],
-    [""],
-    None,
-    123,
-    {"nested": "dict"},
-])
-def test_get_caption_rejects_invalid_values(bad_cap):
-    with pytest.raises(ValueError):
-        get_caption({"path": "a.mp4", "cap": bad_cap})

fastvideo/tests/workflow/test_kandinsky5_overfit_load_video.py

@@ -1,49 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regression test: preprocess_kandinsky5_overfit.load_video must fail with
-a path-specific error on an unreadable input, not an opaque IndexError.
-
-``cv2.VideoCapture`` doesn't raise on a missing or corrupt file -- it just
-decodes zero frames, and the repeat-last-frame fill then crashed on
-``frames[-1]`` with an ``IndexError`` that named neither the file nor the
-cause.
-
-CPU-only -- decodes a tiny real clip written with cv2, no GPU or model
-load.
-"""
-from __future__ import annotations
-
-import cv2
-import numpy as np
-import pytest
-
-from fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit import load_video
-
-
-def test_load_video_missing_file_raises_path_specific_error(tmp_path):
-    missing = tmp_path / "does_not_exist.mp4"
-
-    with pytest.raises(ValueError, match="does_not_exist.mp4"):
-        load_video(str(missing), num_frames=5, height=64, width=64)
-
-
-def test_load_video_corrupt_file_raises_path_specific_error(tmp_path):
-    corrupt = tmp_path / "corrupt.mp4"
-    corrupt.write_bytes(b"this is not an mp4 container")
-
-    with pytest.raises(ValueError, match="corrupt.mp4"):
-        load_video(str(corrupt), num_frames=5, height=64, width=64)
-
-
-def test_load_video_short_clip_still_fills_by_repeating_last_frame(tmp_path):
-    """The zero-frame guard must not break the documented short-clip fill."""
-    clip = tmp_path / "short.mp4"
-    writer = cv2.VideoWriter(str(clip), cv2.VideoWriter_fourcc(*"mp4v"), 24.0, (64, 64))
-    assert writer.isOpened(), "mp4v VideoWriter unavailable in this OpenCV build"
-    rng = np.random.default_rng(0)
-    for _ in range(3):
-        writer.write(rng.integers(0, 256, size=(64, 64, 3), dtype=np.uint8))
-    writer.release()
-
-    video = load_video(str(clip), num_frames=8, height=64, width=64)
-
-    assert video.shape == (1, 3, 8, 64, 64)

fastvideo/train/callbacks/validation.py

@@ -133,7 +133,6 @@ class ValidationCallback(Callback):
         keyboard_value_scale: float = 1.0,
         offload_training_state: bool = False,
         unload_pipeline_after_validation: bool = False,
-        attn_qat_infer: bool = False,
         **pipeline_kwargs: Any,
     ) -> None:
         self.pipeline_target = str(pipeline_target)
@@ -151,7 +150,6 @@ class ValidationCallback(Callback):
         self.metrics_config = self._parse_metrics_config(metrics_config)
         self.offload_training_state = self._coerce_bool(offload_training_state)
         self.unload_pipeline_after_validation = self._coerce_bool(unload_pipeline_after_validation)
-        self.attn_qat_infer = self._coerce_bool(attn_qat_infer)
         self.pipeline_kwargs = dict(pipeline_kwargs)
 
         # Set after on_train_start.
@@ -275,7 +273,7 @@ class ValidationCallback(Callback):
                 # EMA weights during validation.
                 ema_cb = self._find_ema_callback()
                 ctx = ema_cb.ema_context(transformer) if ema_cb is not None else contextlib.nullcontext(transformer)
-                with ctx as t, self._attn_qat_infer_context(t):
+                with ctx as t:
                     self._run_validation_inner(
                         method,
                         step,
@@ -285,48 +283,6 @@ class ValidationCallback(Callback):
             if self.unload_pipeline_after_validation:
                 self._clear_pipeline_cache()
 
-    @contextlib.contextmanager
-    def _attn_qat_infer_context(self, transformer: torch.nn.Module):
-        if not self.attn_qat_infer:
-            yield
-            return
-
-        from fastvideo.attention.backends.attn_qat_infer import (AttnQatInferImpl, is_attn_qat_infer_available)
-        from fastvideo.attention.backends.attn_qat_train import (
-            AttnQatTrainImpl, )
-        from fastvideo.platforms import AttentionBackendEnum
-
-        layers = [
-            module for module in transformer.modules()
-            if isinstance(getattr(module, "attn_impl", None), AttnQatTrainImpl)
-        ]
-        if not layers:
-            raise RuntimeError("attn_qat_infer validation requested, but the transformer has no ATTN_QAT_TRAIN layers")
-        if not is_attn_qat_infer_available():
-            from fastvideo.attention.backends.attn_qat_infer import (
-                attn_qat_infer_receipt, )
-            raise RuntimeError("attn_qat_infer validation requested but no ATTN_QAT_INFER kernel serves "
-                               f"this device ({attn_qat_infer_receipt()}). Set "
-                               "callbacks.validation.attn_qat_infer=false to validate with ATTN_QAT_TRAIN.")
-
-        previous = [(layer, layer.attn_impl, layer.backend) for layer in layers]
-        try:
-            for layer, impl, _ in previous:
-                layer.attn_impl = AttnQatInferImpl(
-                    num_heads=layer.num_heads,
-                    head_size=layer.head_size,
-                    num_kv_heads=layer.num_kv_heads,
-                    causal=impl.causal,
-                    softmax_scale=layer.softmax_scale,
-                )
-                layer.backend = AttentionBackendEnum.ATTN_QAT_INFER
-            logger.info("Enabled ATTN_QAT_INFER for %d validation attention layers.", len(layers))
-            yield
-        finally:
-            for layer, impl, backend in previous:
-                layer.attn_impl = impl
-                layer.backend = backend
-
     @contextlib.contextmanager
     def _validation_memory_context(
         self,
@@ -1220,15 +1176,6 @@ class ValidationCallback(Callback):
             VSA_sparsity=tc.vsa_sparsity,
             timesteps=sampling_timesteps_tensor,
         )
-        # shallow_asdict(sampling_param) copies list-typed fields by
-        # reference. sampling_param is cached and reused across every
-        # validation sample/step, so without this reset every ForwardBatch
-        # would share (and keep appending to) the same
-        # prompt_attention_mask/negative_attention_mask list forever --
-        # index [0] would then hold the *first-ever* validation sample's
-        # mask instead of the current one, mismatching prompt_embeds.
-        batch.prompt_attention_mask = []
-        batch.negative_attention_mask = []
         batch._inference_args = inference_args  # type: ignore[attr-defined]
 
         # Conditionally set I2V fields.

fastvideo/train/entrypoint/dcp_to_diffusers.py

@@ -20,10 +20,6 @@ The checkpoint must contain ``metadata.json`` (written by
 ``CheckpointManager``).  If the checkpoint predates metadata
 support, pass ``--config`` explicitly to provide the training
 YAML.
-
-Pass ``--verify`` to strictly reload the exported transformer immediately
-after writing it, so a key-mapping bug fails here instead of deep inside a
-later training/inference launch that loads the exported directory.
 """
 
 from __future__ import annotations
@@ -105,15 +101,10 @@ def _save_role_pretrained(
                                       "Pass --overwrite to replace it.")
 
         def _copy_or_link(src: str, dest: str) -> None:
-            # Resolve symlinks ourselves: os.link's follow_symlinks=True
-            # default isn't honored on all filesystems (e.g. some
-            # network/overlay mounts), which can silently hard-link to
-            # the symlink itself instead of its target.
-            real_src = os.path.realpath(src)
             try:
-                os.link(real_src, dest)
+                os.link(src, dest)
             except OSError:
-                shutil.copy2(real_src, dest)
+                shutil.copy2(src, dest)
 
         logger.info(
             "Creating pretrained export dir at %s "
@@ -124,7 +115,7 @@ def _save_role_pretrained(
         shutil.copytree(
             local_base,
             dst,
-            symlinks=False,
+            symlinks=True,
             copy_function=_copy_or_link,
         )
 
@@ -204,27 +195,6 @@ def _save_role_pretrained(
     return str(dst)
 
 
-def _strict_reload_verify(*, output_dir: str, training_config: Any) -> None:
-    """Reload the just-exported transformer from disk and fail loudly on
-    any key mismatch.
-
-    ``TransformerLoader.load()`` already loads strictly (``strict=True``)
-    for every non-Cosmos25 model, so this doesn't add leniency -- it moves
-    the failure to right after export, with a clear "the export is broken"
-    error, instead of surfacing deep inside a later training/inference
-    launch that happens to load this directory.
-    """
-    from fastvideo.train.utils.moduleloader import load_module_from_path
-
-    logger.info("Verifying export: strictly reloading transformer from %s", output_dir)
-    load_module_from_path(
-        model_path=output_dir,
-        module_type="transformer",
-        training_config=training_config,
-    )
-    logger.info("Strict reload verification passed.")
-
-
 def convert(
     *,
     checkpoint_dir: str,
@@ -232,7 +202,6 @@ def convert(
     config_path: str | None = None,
     role: str = "student",
     overwrite: bool = False,
-    verify: bool = False,
 ) -> str:
     """Load a DCP checkpoint and export as a diffusers model.
 
@@ -324,10 +293,6 @@ def convert(
         model=model,
     )
     logger.info("Export complete: %s", result)
-
-    if verify:
-        _strict_reload_verify(output_dir=result, training_config=tc)
-
     return result
 
 
@@ -436,13 +401,6 @@ def main() -> None:
         action="store_true",
         help="Overwrite output-dir if it exists.",
     )
-    parser.add_argument(
-        "--verify",
-        action="store_true",
-        help=("After exporting, strictly reload the transformer from "
-              "the exported directory to catch key-mapping bugs "
-              "immediately."),
-    )
     args = parser.parse_args(sys.argv[1:])
 
     convert(
@@ -451,7 +409,6 @@ def main() -> None:
         config_path=args.config,
         role=args.role,
         overwrite=args.overwrite,
-        verify=args.verify,
     )
 
 

fastvideo/train/models/kandinsky5/__init__.py

@@ -1,5 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Kandinsky5 model plugin package."""
-
-from fastvideo.train.models.kandinsky5.kandinsky5 import (
-    Kandinsky5Model as Kandinsky5Model, )

fastvideo/train/models/kandinsky5/kandinsky5.py

@@ -1,773 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Kandinsky5 model plugin (per-role instance).
-
-Written directly against ``ModelBase`` (not a ``WanModel`` subclass) because
-Kandinsky5 differs structurally in ways that don't compose cleanly with
-Wan's implementation:
-
-  - Dual text encoders: Qwen/Reason1 sequence embeddings *and* a CLIP pooled
-    projection, vs. Wan's single encoder. The pair is packed into the
-    existing single ``text_embedding`` parquet field by zero-padding the
-    CLIP pooled vector into a row and prepending it to the Qwen sequence
-    (the same trick already used by ``HunyuanModel``/
-    ``preprocess_hunyuan_overfit.py`` for LLaMA+CLIP).
-  - The transformer forward signature needs RoPE position tensors
-    (``visual_rope_pos``, ``text_rope_pos``) and a ``scale_factor`` that Wan
-    never has.
-  - Kandinsky5's native hidden_states layout is channel-last
-    ``[B, T, H, W, C]``; the common ModelBase convention used at the
-    predict_noise/decode_latents boundary is ``[B, T, C, H, W]``.
-  - VAE denorm uses only ``scaling_factor`` (Kandinsky5 shares Hunyuan's
-    VAE, which has no ``latents_mean``/``latents_std``/
-    ``handles_latent_denorm``).
-  - flow_shift default is 5.0 (vs. Wan's 3.0).
-
-Scope: T2V at 480p only, dense/local attention only (NABLA sparse attention
-is never engaged -- ``sparse_params`` is always ``None``).
-"""
-
-from __future__ import annotations
-
-import copy
-import os
-from typing import Any, Literal, TYPE_CHECKING
-
-import torch
-
-from fastvideo.api.sampling_param import SamplingParam
-from fastvideo.distributed import (
-    get_sp_group,
-    get_world_group,
-)
-from fastvideo.forward_context import set_forward_context
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler, )
-from fastvideo.pipelines import TrainingBatch
-from fastvideo.training.activation_checkpoint import (
-    apply_activation_checkpointing, )
-from fastvideo.training.training_utils import (
-    compute_density_for_timestep_sampling,
-    get_sigmas,
-    normalize_dit_input,
-    shift_timestep,
-)
-
-from fastvideo.train.models.base import ModelBase
-from fastvideo.train.utils.module_state import (
-    apply_trainable, )
-from fastvideo.train.utils.moduleloader import (
-    load_module_from_path,
-    make_inference_args,
-)
-
-if TYPE_CHECKING:
-    from fastvideo.train.utils.training_config import (
-        TrainingConfig, )
-    from fastvideo.train.utils.lora import LoraConfig
-
-# 480p is the only supported resolution band for this recipe. Matches
-# Kandinsky5DenoisingStage._scale_factor's low-resolution branch exactly;
-# outside this band Kandinsky5 uses a different visual RoPE scale_factor
-# ((1.0, 3.16, 3.16)) that this wrapper does not implement.
-_SCALE_FACTOR_480P = (1.0, 2.0, 2.0)
-_MIN_480P_SIDE = 480
-_MAX_480P_SIDE = 854
-
-
-class Kandinsky5Model(ModelBase):
-    """Kandinsky5 per-role model: owns transformer + noise_scheduler."""
-
-    _transformer_cls_name: str = "Kandinsky5Transformer3DModel"
-
-    def __init__(
-        self,
-        *,
-        init_from: str,
-        training_config: TrainingConfig,
-        trainable: bool = True,
-        disable_custom_init_weights: bool = False,
-        flow_shift: float = 5.0,
-        enable_gradient_checkpointing_type: str
-        | None = None,
-        transformer_override_safetensor: str
-        | None = None,
-        lora: LoraConfig | dict[str, Any] | None = None,
-    ) -> None:
-        super().__init__(
-            trainable=trainable,
-            lora=lora,
-        )
-        self._init_from = str(init_from)
-
-        self.transformer = self._load_transformer(
-            init_from=self._init_from,
-            trainable=self._trainable,
-            disable_custom_init_weights=(disable_custom_init_weights),
-            enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
-            training_config=training_config,
-            transformer_override_safetensor=(transformer_override_safetensor),
-        )
-
-        self.noise_scheduler = (FlowMatchEulerDiscreteScheduler(shift=float(flow_shift)))
-
-        # Filled by init_preprocessors (student only).
-        self.vae: Any = None
-        self.training_config: TrainingConfig = training_config
-        self.dataloader: Any = None
-        self.validator: Any = None
-        self.start_step: int = 0
-
-        self.world_group: Any = None
-        self.sp_group: Any = None
-
-        # Qwen sequence embeds + mask, and CLIP pooled projection.
-        self.negative_prompt_embeds: (torch.Tensor | None) = None
-        self.negative_prompt_attention_mask: (torch.Tensor | None) = None
-        self.negative_pooled_embeds: (torch.Tensor | None) = None
-        self._requires_negative_conditioning = True
-
-        # Timestep mechanics.
-        self.timestep_shift: float = float(flow_shift)
-        self.num_train_timestep: int = int(self.noise_scheduler.num_train_timesteps)
-        self.min_timestep: int = 0
-        self.max_timestep: int = self.num_train_timestep
-
-    def _load_transformer(
-        self,
-        *,
-        init_from: str,
-        trainable: bool,
-        disable_custom_init_weights: bool,
-        enable_gradient_checkpointing_type: str | None,
-        training_config: TrainingConfig,
-        transformer_override_safetensor: str | None = None,
-    ) -> torch.nn.Module:
-        transformer = load_module_from_path(
-            model_path=init_from,
-            module_type="transformer",
-            training_config=training_config,
-            disable_custom_init_weights=(disable_custom_init_weights),
-            override_transformer_cls_name=(self._transformer_cls_name),
-            transformer_override_safetensor=(transformer_override_safetensor),
-        )
-        ckpt_type = (enable_gradient_checkpointing_type or getattr(
-            getattr(training_config, "model", None),
-            "enable_gradient_checkpointing_type",
-            None,
-        ))
-        if trainable and ckpt_type:
-            transformer = apply_activation_checkpointing(
-                transformer,
-                checkpointing_type=ckpt_type,
-            )
-        if self._enable_lora_if_configured(transformer):
-            return transformer
-        transformer = apply_trainable(transformer, trainable=trainable)
-        return transformer
-
-    # ------------------------------------------------------------------
-    # Lifecycle
-    # ------------------------------------------------------------------
-
-    def init_preprocessors(self, training_config: TrainingConfig) -> None:
-        self.vae = load_module_from_path(
-            model_path=str(training_config.model_path),
-            module_type="vae",
-            training_config=training_config,
-        )
-
-        self.world_group = get_world_group()
-        self.sp_group = get_sp_group()
-
-        self._init_timestep_mechanics()
-
-        from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
-        from fastvideo.train.utils.dataloader import (
-            build_parquet_t2v_train_dataloader, )
-
-        preprocessed_data_type = str(getattr(
-            training_config.data,
-            "preprocessed_data_type",
-            "t2v",
-        )).strip().lower()
-        if preprocessed_data_type != "t2v":
-            raise ValueError("Unsupported Kandinsky5 preprocessed_data_type: "
-                             f"{preprocessed_data_type!r}")
-
-        # Qwen's usable (post-template-trim) embedding length, +1 for the
-        # prepended CLIP pooled row (see module docstring / prepare_batch).
-        qwen_text_len = int(training_config.pipeline_config.text_encoder_configs[  # type: ignore[union-attr]
-            0].arch_config.text_len)
-        self.dataloader = build_parquet_t2v_train_dataloader(
-            training_config.data,
-            text_len=qwen_text_len + 1,
-            parquet_schema=pyarrow_schema_t2v,
-        )
-        self.start_step = 0
-
-    @property
-    def num_train_timesteps(self) -> int:
-        return int(self.num_train_timestep)
-
-    def set_requires_negative_conditioning(self, requires: bool) -> None:
-        self._requires_negative_conditioning = bool(requires)
-
-    def shift_and_clamp_timestep(self, timestep: torch.Tensor) -> torch.Tensor:
-        timestep = shift_timestep(
-            timestep,
-            self.timestep_shift,
-            self.num_train_timestep,
-        )
-        return timestep.clamp(self.min_timestep, self.max_timestep)
-
-    def on_train_start(self) -> None:
-        if self._requires_negative_conditioning:
-            self.ensure_negative_conditioning()
-
-    @torch.no_grad()
-    def decode_latents(
-        self,
-        latents_b_t_c_h_w: torch.Tensor,
-    ) -> torch.Tensor:
-        if self.vae is None:
-            raise RuntimeError("Kandinsky5 VAE is not initialized")
-        latents = latents_b_t_c_h_w.permute(0, 2, 1, 3, 4).float()
-        # Kandinsky5 shares Hunyuan's VAE: only a scalar scaling_factor, no
-        # latents_mean/latents_std, no handles_latent_denorm. Inverse of the
-        # encode-side `normalize_dit_input("hunyuan", ...)` multiply.
-        denorm = latents / self.vae.scaling_factor
-        media = self.vae.to(latents.device).decode(denorm)
-        return (media / 2 + 0.5).clamp(0, 1)
-
-    # ------------------------------------------------------------------
-    # Runtime primitives
-    # ------------------------------------------------------------------
-
-    def prepare_batch(
-        self,
-        raw_batch: dict[str, Any],
-        *,
-        generator: torch.Generator,
-        latents_source: Literal["data", "zeros"] = "data",
-    ) -> TrainingBatch:
-        if self._requires_negative_conditioning:
-            self.ensure_negative_conditioning()
-        assert self.training_config is not None
-        tc = self.training_config
-
-        dtype = self._get_training_dtype()
-        device = self.device
-
-        training_batch = TrainingBatch()
-
-        # Unpack the CLIP-pooled-row-prepended text_embedding field: row 0
-        # is the zero-padded CLIP pooled vector, rows 1: are the Qwen
-        # sequence embeddings. See module docstring.
-        packed_embeds = raw_batch["text_embedding"]
-        packed_mask = raw_batch["text_attention_mask"]
-        pooled_dim = int(tc.pipeline_config.dit_config.arch_config.in_text_dim2  # type: ignore[union-attr]
-                         )
-        pooled_projections = packed_embeds[:, 0, :pooled_dim]
-        qwen_embeds = packed_embeds[:, 1:, :]
-        qwen_mask = packed_mask[:, 1:]
-        infos = raw_batch.get("info_list")
-
-        if latents_source == "zeros":
-            batch_size = packed_embeds.shape[0]
-            vae_config = (
-                tc.pipeline_config.vae_config.arch_config  # type: ignore[union-attr]
-            )
-            # Read off the loaded transformer module, not
-            # tc.pipeline_config.dit_config.arch_config: the latter carries
-            # YAML/dataclass defaults and isn't synced with the actual
-            # checkpoint's transformer/config.json (e.g. the shipped
-            # Kandinsky5-Lite checkpoint has in_visual_dim=16, vs. the
-            # dataclass default of 4).
-            num_channels = int(self.transformer.in_visual_dim)
-            spatial_compression_ratio = (vae_config.spatial_compression_ratio)
-            latent_height = (tc.data.num_height // spatial_compression_ratio)
-            latent_width = (tc.data.num_width // spatial_compression_ratio)
-            latents = torch.zeros(
-                batch_size,
-                num_channels,
-                tc.data.num_latent_t,
-                latent_height,
-                latent_width,
-                device=device,
-                dtype=dtype,
-            )
-        elif latents_source == "data":
-            if "vae_latent" not in raw_batch:
-                raise ValueError("vae_latent not found in batch "
-                                 "and latents_source='data'")
-            latents = raw_batch["vae_latent"]
-            latents = latents[:, :, :tc.data.num_latent_t]
-            latents = latents.to(device, dtype=dtype)
-        else:
-            raise ValueError(f"Unknown latents_source: "
-                             f"{latents_source!r}")
-
-        training_batch.latents = latents
-        training_batch.encoder_hidden_states = qwen_embeds.to(device, dtype=dtype)
-        training_batch.encoder_attention_mask = qwen_mask.to(device, dtype=dtype)
-        training_batch.infos = infos
-        pooled_projections = pooled_projections.to(device, dtype=dtype)
-
-        training_batch.latents = normalize_dit_input("hunyuan", training_batch.latents, self.vae)
-        training_batch = self._prepare_dit_inputs(training_batch, generator, pooled_projections)
-        training_batch = self._build_attention_metadata(training_batch)
-
-        training_batch.attn_metadata_vsa = copy.copy(training_batch.attn_metadata)
-        if training_batch.attn_metadata is not None:
-            training_batch.attn_metadata.VSA_sparsity = 0.0  # type: ignore[attr-defined]
-
-        return training_batch
-
-    def add_noise(
-        self,
-        clean_latents: torch.Tensor,
-        noise: torch.Tensor,
-        timestep: torch.Tensor,
-    ) -> torch.Tensor:
-        b, t = clean_latents.shape[:2]
-        noisy = self.noise_scheduler.add_noise(
-            clean_latents.flatten(0, 1),
-            noise.flatten(0, 1),
-            timestep,
-        ).unflatten(0, (b, t))
-        return noisy
-
-    def predict_noise(
-        self,
-        noisy_latents: torch.Tensor,
-        timestep: torch.Tensor,
-        batch: TrainingBatch,
-        *,
-        conditional: bool,
-        cfg_uncond: dict[str, Any] | None = None,
-        attn_kind: Literal["dense", "vsa"] = "dense",
-    ) -> torch.Tensor:
-        device_type = self.device.type
-        dtype = self._get_training_dtype()
-        if conditional:
-            text_dict = batch.conditional_dict
-            if text_dict is None:
-                raise RuntimeError("Missing conditional_dict in "
-                                   "TrainingBatch")
-        else:
-            text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond)
-
-        if attn_kind == "dense":
-            attn_metadata = batch.attn_metadata
-        elif attn_kind == "vsa":
-            attn_metadata = batch.attn_metadata_vsa
-        else:
-            raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
-
-        if noisy_latents.is_floating_point():
-            noisy_latents = noisy_latents.to(dtype=dtype)
-
-        with torch.autocast(device_type, dtype=dtype), set_forward_context(
-                current_timestep=batch.timesteps,
-                attn_metadata=attn_metadata,
-        ):
-            input_kwargs = (self._build_distill_input_kwargs(noisy_latents, timestep, text_dict))
-            transformer = self._get_transformer(timestep)
-            out = transformer(**input_kwargs)
-            sample = out.sample if hasattr(out, "sample") else out
-            # Kandinsky5's native hidden_states layout is channel-last
-            # [B, T, H, W, C]; convert to the common ModelBase convention
-            # [B, T, C, H, W].
-            pred_noise = sample.permute(0, 1, 4, 2, 3)
-        return pred_noise
-
-    def backward(
-        self,
-        loss: torch.Tensor,
-        ctx: Any,
-        *,
-        grad_accum_rounds: int,
-    ) -> None:
-        timesteps, attn_metadata = ctx
-        with set_forward_context(
-                current_timestep=timesteps,
-                attn_metadata=attn_metadata,
-        ):
-            (loss / max(1, int(grad_accum_rounds))).backward()
-
-    # ------------------------------------------------------------------
-    # Internal helpers
-    # ------------------------------------------------------------------
-
-    def _get_training_dtype(self) -> torch.dtype:
-        return torch.bfloat16
-
-    def _init_timestep_mechanics(self) -> None:
-        assert self.training_config is not None
-        tc = self.training_config
-        self.timestep_shift = float(tc.pipeline_config.flow_shift  # type: ignore[union-attr]
-                                    )
-        self.num_train_timestep = int(self.noise_scheduler.num_train_timesteps)
-        self.min_timestep = 0
-        self.max_timestep = self.num_train_timestep
-
-    def ensure_negative_conditioning(self) -> None:
-        """Encode the negative prompt with dual text encoders (Qwen + CLIP).
-
-        Every rank encodes independently (same rationale as Hunyuan's
-        override: avoids NCCL deadlocks from asymmetric rank-0-only
-        encoding). Cannot reuse ``encode_negative_prompt`` unmodified for
-        the Qwen encoder: ``kandinsky5_qwen_postprocess_text`` requires the
-        attention mask as a second positional argument and returns a
-        ``(embeds, trimmed_mask)`` tuple, unlike the single-arg
-        postprocess-func contract that helper assumes.
-        """
-        if self.negative_prompt_embeds is not None:
-            return
-
-        assert self.training_config is not None
-        tc = self.training_config
-        device = self.device
-        dtype = self._get_training_dtype()
-
-        from transformers import AutoTokenizer
-
-        from fastvideo.configs.pipelines.base import preprocess_text
-        from fastvideo.configs.pipelines.kandinsky5 import (
-            kandinsky5_clip_postprocess_text,
-            kandinsky5_qwen_postprocess_text,
-            kandinsky5_qwen_preprocess_text,
-        )
-        from fastvideo.models.loader.component_loader import TextEncoderLoader
-        from fastvideo.utils import maybe_download_model
-
-        pipeline_config = tc.pipeline_config
-        assert pipeline_config is not None
-        model_path = maybe_download_model(tc.model_path)
-        inference_args = make_inference_args(tc, model_path=model_path)
-        inference_args.text_encoder_cpu_offload = False
-
-        sampling_param = SamplingParam.from_pretrained(tc.model_path)
-        negative_prompt = sampling_param.negative_prompt
-
-        qwen_cfg = pipeline_config.text_encoder_configs[0]
-        clip_cfg = pipeline_config.text_encoder_configs[1]
-        loader = TextEncoderLoader()
-
-        # --- Qwen / Reason1 ---
-        qwen_enc = loader.load(
-            os.path.join(model_path, "text_encoder"),
-            inference_args,
-        ).to(device).eval()
-        qwen_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
-        qwen_tok_kwargs = dict(qwen_cfg.tokenizer_kwargs)
-        qwen_text = kandinsky5_qwen_preprocess_text(negative_prompt)
-
-        with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None):
-            qwen_inputs = qwen_tok(qwen_text, **qwen_tok_kwargs).to(device)
-            qwen_out = qwen_enc(
-                input_ids=qwen_inputs.input_ids,
-                attention_mask=qwen_inputs.attention_mask,
-                output_hidden_states=True,
-            )
-            qwen_embeds, qwen_mask = kandinsky5_qwen_postprocess_text(qwen_out, qwen_inputs.attention_mask)
-
-        del qwen_enc, qwen_tok
-
-        # --- CLIP ---
-        clip_enc = loader.load(
-            os.path.join(model_path, "text_encoder_2"),
-            inference_args,
-        ).to(device).eval()
-        clip_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer_2"))
-        clip_tok_kwargs = dict(clip_cfg.tokenizer_kwargs)
-        clip_text = preprocess_text(negative_prompt)
-
-        with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None):
-            clip_inputs = clip_tok(clip_text, **clip_tok_kwargs).to(device)
-            clip_out = clip_enc(
-                input_ids=clip_inputs.input_ids,
-                attention_mask=clip_inputs.attention_mask,
-            )
-            clip_pooled = kandinsky5_clip_postprocess_text(clip_out)
-
-        del clip_enc, clip_tok
-
-        self.negative_prompt_embeds = qwen_embeds.to(device=device, dtype=dtype)
-        self.negative_prompt_attention_mask = qwen_mask.to(device=device, dtype=dtype)
-        self.negative_pooled_embeds = clip_pooled.to(device=device, dtype=dtype)
-
-    def _sample_timesteps(
-        self,
-        batch_size: int,
-        device: torch.device,
-        generator: torch.Generator,
-    ) -> torch.Tensor:
-        assert self.training_config is not None
-        tc = self.training_config
-
-        u = compute_density_for_timestep_sampling(
-            weighting_scheme=tc.model.weighting_scheme,
-            batch_size=batch_size,
-            generator=generator,
-            device=device,
-            logit_mean=tc.model.logit_mean,
-            logit_std=tc.model.logit_std,
-            mode_scale=tc.model.mode_scale,
-        )
-        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
-        return self.noise_scheduler.timesteps[indices.cpu()].to(device=device)
-
-    def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch:
-        # Scope is 480p T2V with dense/local attention only -- VSA/VMOBA
-        # sparse attention is never engaged for Kandinsky5 in this recipe.
-        training_batch.attn_metadata = None
-        return training_batch
-
-    def _prepare_dit_inputs(
-        self,
-        training_batch: TrainingBatch,
-        generator: torch.Generator,
-        pooled_projections: torch.Tensor,
-    ) -> TrainingBatch:
-        assert self.training_config is not None
-        tc = self.training_config
-        latents = training_batch.latents
-        assert isinstance(latents, torch.Tensor)
-        batch_size = latents.shape[0]
-        device = latents.device
-
-        num_height = int(tc.data.num_height)
-        num_width = int(tc.data.num_width)
-        if not (_MIN_480P_SIDE <= num_height <= _MAX_480P_SIDE and _MIN_480P_SIDE <= num_width <= _MAX_480P_SIDE):
-            raise ValueError("Kandinsky5Model only supports 480p training "
-                             f"(height/width in [{_MIN_480P_SIDE}, {_MAX_480P_SIDE}]); "
-                             f"got num_height={num_height}, num_width={num_width}. "
-                             "A larger resolution needs a different visual RoPE "
-                             "scale_factor, which this wrapper does not implement.")
-
-        noise = torch.randn(
-            latents.shape,
-            generator=generator,
-            device=device,
-            dtype=latents.dtype,
-        )
-        timesteps = self._sample_timesteps(batch_size, device, generator)
-        if int(tc.distributed.sp_size or 1) > 1:
-            self.sp_group.broadcast(timesteps, src=0)
-
-        sigmas = get_sigmas(
-            self.noise_scheduler,
-            device,
-            timesteps,
-            n_dim=latents.ndim,
-            dtype=latents.dtype,
-        )
-        noisy_model_input = ((1.0 - sigmas) * latents + sigmas * noise)
-
-        training_batch.noisy_model_input = noisy_model_input
-        training_batch.timesteps = timesteps
-        training_batch.sigmas = sigmas
-        training_batch.noise = noise
-        training_batch.raw_latent_shape = latents.shape
-
-        # Trim the Qwen sequence to the longest real (non-pad) sample in
-        # this batch, mirroring Kandinsky5DenoisingStage._text_rope_pos.
-        # Kandinsky5's cross-attention has no attention-mask plumbing at
-        # all -- it always attends to the full encoder_hidden_states
-        # sequence -- so any padding left in the tensor beyond this point
-        # would silently leak into cross-attention. Batched *inference*
-        # already relies on the same property (tokenizer padding="True"
-        # pads only to the batch's longest real sample).
-        qwen_embeds = training_batch.encoder_hidden_states
-        qwen_mask = training_batch.encoder_attention_mask
-        assert qwen_embeds is not None and qwen_mask is not None
-        # bf16 cannot exactly accumulate sums beyond ~256, so count valid
-        # tokens in float32 to avoid an off-by-a-few error on long prompts.
-        valid_len = max(1, int(qwen_mask.float().sum(dim=1).max().item()))
-        qwen_embeds = qwen_embeds[:, :valid_len]
-        qwen_mask = qwen_mask[:, :valid_len]
-        text_rope_pos = torch.arange(valid_len, device=device)
-
-        patch_size = (
-            tc.pipeline_config.dit_config.arch_config.patch_size  # type: ignore[union-attr]
-        )
-        spatial_compression_ratio = (
-            tc.pipeline_config.vae_config.arch_config.spatial_compression_ratio  # type: ignore[union-attr]
-        )
-        latent_h = num_height // spatial_compression_ratio
-        latent_w = num_width // spatial_compression_ratio
-        visual_rope_pos = [
-            torch.arange(int(tc.data.num_latent_t), device=device),
-            torch.arange(latent_h // patch_size[1], device=device),
-            torch.arange(latent_w // patch_size[2], device=device),
-        ]
-
-        training_batch.conditional_dict = {
-            "encoder_hidden_states": qwen_embeds,
-            "encoder_attention_mask": qwen_mask,
-            "pooled_projections": pooled_projections,
-            "text_rope_pos": text_rope_pos,
-            "visual_rope_pos": visual_rope_pos,
-            "scale_factor": _SCALE_FACTOR_480P,
-            "sparse_params": None,
-        }
-
-        if (self.negative_prompt_embeds is not None and self.negative_prompt_attention_mask is not None
-                and self.negative_pooled_embeds is not None):
-            neg_embeds = self.negative_prompt_embeds
-            neg_mask = self.negative_prompt_attention_mask
-            neg_pooled = self.negative_pooled_embeds
-            if neg_embeds.shape[0] == 1 and batch_size > 1:
-                neg_embeds = neg_embeds.expand(batch_size, *neg_embeds.shape[1:]).contiguous()
-            if neg_mask.shape[0] == 1 and batch_size > 1:
-                neg_mask = neg_mask.expand(batch_size, *neg_mask.shape[1:]).contiguous()
-            if neg_pooled.shape[0] == 1 and batch_size > 1:
-                neg_pooled = neg_pooled.expand(batch_size, *neg_pooled.shape[1:]).contiguous()
-            neg_text_rope_pos = torch.arange(neg_embeds.shape[1], device=device)
-            training_batch.unconditional_dict = {
-                "encoder_hidden_states": neg_embeds,
-                "encoder_attention_mask": neg_mask,
-                "pooled_projections": neg_pooled,
-                "text_rope_pos": neg_text_rope_pos,
-                "visual_rope_pos": visual_rope_pos,
-                "scale_factor": _SCALE_FACTOR_480P,
-                "sparse_params": None,
-            }
-
-        training_batch.latents = (training_batch.latents.permute(0, 2, 1, 3, 4))
-        return training_batch
-
-    def _build_distill_input_kwargs(
-        self,
-        noise_input: torch.Tensor,
-        timestep: torch.Tensor,
-        text_dict: dict[str, torch.Tensor] | None,
-    ) -> dict[str, Any]:
-        if text_dict is None:
-            raise ValueError("text_dict cannot be None for "
-                             "Kandinsky5 distillation")
-        # noise_input is [B, T, C, H, W] (common ModelBase convention);
-        # Kandinsky5's transformer expects channel-last [B, T, H, W, C].
-        hidden_states = noise_input.permute(0, 1, 3, 4, 2)
-        if getattr(self.transformer, "visual_cond", False):
-            # The shipped Kandinsky5-Lite checkpoint has visual_cond=True
-            # (unified T2V/I2V conditioning): Kandinsky5VisualEmbeddings
-            # always expects [real_latent | zero_cond | zero_mask]
-            # concatenated on the channel dim (2*in_visual_dim+1 total),
-            # even for pure T2V with no actual image conditioning. Mirrors
-            # Kandinsky5LatentPreparationStage's inference-time padding
-            # (fastvideo/pipelines/stages/kandinsky5.py).
-            cond = torch.zeros_like(hidden_states)
-            mask = torch.zeros(*hidden_states.shape[:-1], 1, device=hidden_states.device, dtype=hidden_states.dtype)
-            hidden_states = torch.cat([hidden_states, cond, mask], dim=-1)
-        return {
-            "hidden_states": hidden_states,
-            "encoder_hidden_states": text_dict["encoder_hidden_states"],
-            "pooled_projections": text_dict["pooled_projections"],
-            "timestep": timestep,
-            "visual_rope_pos": text_dict["visual_rope_pos"],
-            "text_rope_pos": text_dict["text_rope_pos"],
-            "scale_factor": text_dict["scale_factor"],
-            "sparse_params": text_dict["sparse_params"],
-            "return_dict": True,
-        }
-
-    def _get_transformer(self, timestep: torch.Tensor) -> torch.nn.Module:
-        return self.transformer
-
-    def _get_uncond_text_dict(
-        self,
-        batch: TrainingBatch,
-        *,
-        cfg_uncond: dict[str, Any] | None,
-    ) -> dict[str, torch.Tensor]:
-        if cfg_uncond is None:
-            text_dict = getattr(batch, "unconditional_dict", None)
-            if text_dict is None:
-                raise RuntimeError("Missing unconditional_dict; "
-                                   "ensure_negative_conditioning() "
-                                   "may have failed")
-            return text_dict
-
-        on_missing_raw = cfg_uncond.get("on_missing", "error")
-        if not isinstance(on_missing_raw, str):
-            raise ValueError("method_config.cfg_uncond.on_missing "
-                             "must be a string, got "
-                             f"{type(on_missing_raw).__name__}")
-        on_missing = on_missing_raw.strip().lower()
-        if on_missing not in {"error", "ignore"}:
-            raise ValueError("method_config.cfg_uncond.on_missing "
-                             "must be one of {error, ignore}, got "
-                             f"{on_missing_raw!r}")
-
-        for channel, policy_raw in cfg_uncond.items():
-            if channel in {"on_missing", "text"}:
-                continue
-            if policy_raw is None:
-                continue
-            if not isinstance(policy_raw, str):
-                raise ValueError("method_config.cfg_uncond values "
-                                 "must be strings, got "
-                                 f"{channel}="
-                                 f"{type(policy_raw).__name__}")
-            policy = policy_raw.strip().lower()
-            if policy == "keep":
-                continue
-            if on_missing == "ignore":
-                continue
-            raise ValueError("Kandinsky5Model does not support "
-                             "cfg_uncond channel "
-                             f"{channel!r} (policy={policy!r}). "
-                             "Set cfg_uncond.on_missing=ignore or "
-                             "remove the channel.")
-
-        text_policy_raw = cfg_uncond.get("text", None)
-        if text_policy_raw is None:
-            text_policy = "negative_prompt"
-        elif not isinstance(text_policy_raw, str):
-            raise ValueError("method_config.cfg_uncond.text must be "
-                             "a string, got "
-                             f"{type(text_policy_raw).__name__}")
-        else:
-            text_policy = (text_policy_raw.strip().lower())
-
-        if text_policy in {"negative_prompt"}:
-            text_dict = getattr(batch, "unconditional_dict", None)
-            if text_dict is None:
-                raise RuntimeError("Missing unconditional_dict; "
-                                   "ensure_negative_conditioning() "
-                                   "may have failed")
-            return text_dict
-        if text_policy == "keep":
-            if batch.conditional_dict is None:
-                raise RuntimeError("Missing conditional_dict in "
-                                   "TrainingBatch")
-            return batch.conditional_dict
-        if text_policy == "zero":
-            if batch.conditional_dict is None:
-                raise RuntimeError("Missing conditional_dict in "
-                                   "TrainingBatch")
-            cond = batch.conditional_dict
-            enc = cond["encoder_hidden_states"]
-            mask = cond["encoder_attention_mask"]
-            pooled = cond["pooled_projections"]
-            if not torch.is_tensor(enc) or not torch.is_tensor(mask) or not torch.is_tensor(pooled):
-                raise TypeError("conditional_dict must contain "
-                                "tensor text inputs")
-            return {
-                "encoder_hidden_states": torch.zeros_like(enc),
-                "encoder_attention_mask": torch.zeros_like(mask),
-                "pooled_projections": torch.zeros_like(pooled),
-                "text_rope_pos": cond["text_rope_pos"],
-                "visual_rope_pos": cond["visual_rope_pos"],
-                "scale_factor": cond["scale_factor"],
-                "sparse_params": cond["sparse_params"],
-            }
-        if text_policy == "drop":
-            raise ValueError("cfg_uncond.text=drop is not supported "
-                             "for Kandinsky5. Use "
-                             "{negative_prompt, keep, zero}.")
-        raise ValueError("cfg_uncond.text must be one of "
-                         "{negative_prompt, keep, zero, drop}, got "
-                         f"{text_policy_raw!r}")

fastvideo/train/models/ltx2/ltx2.py

@@ -38,7 +38,6 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.ltx2 import VideoLatentShape
 from fastvideo.pipelines import ForwardBatch, TrainingBatch
-from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.training.activation_checkpoint import (
     apply_activation_checkpointing, )
 
@@ -95,7 +94,6 @@ class LTX2Model(WanModel):
         transformer_override_safetensor: str
         | None = None,
         lora: LoraConfig | dict[str, Any] | None = None,
-        attention_backend: AttentionBackendEnum | str | None = None,
         train_audio: bool = False,
         timestep_uniform_prob: float = 0.1,
     ) -> None:
@@ -128,7 +126,6 @@ class LTX2Model(WanModel):
             enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
             transformer_override_safetensor=(transformer_override_safetensor),
             lora=lora,
-            attention_backend=attention_backend,
         )
 
         if trainable:
@@ -152,7 +149,6 @@ class LTX2Model(WanModel):
         enable_gradient_checkpointing_type: str | None,
         training_config: TrainingConfig,
         transformer_override_safetensor: str | None = None,
-        attention_backend: AttentionBackendEnum | str | None = None,
     ) -> torch.nn.Module:
         transformer = load_module_from_path(
             model_path=init_from,
@@ -161,7 +157,6 @@ class LTX2Model(WanModel):
             disable_custom_init_weights=(disable_custom_init_weights),
             override_transformer_cls_name=(self._transformer_cls_name),
             transformer_override_safetensor=(transformer_override_safetensor),
-            attention_backend=attention_backend,
         )
         ckpt_type = (enable_gradient_checkpointing_type or getattr(
             getattr(training_config, "model", None),

fastvideo/train/utils/config.py

@@ -282,28 +282,9 @@ def _parse_pipeline_config(
 
     pipeline_config = PipelineConfig.from_kwargs(kwargs)
     _apply_training_dit_arch_overrides(pipeline_config, dit_arch_overrides)
-    _resolve_dit_quant_config(pipeline_config)
     return pipeline_config
 
 
-def _resolve_dit_quant_config(pipeline_config: Any) -> None:
-    """Resolve a string ``pipeline.dit_config.quant_config`` (e.g.
-    ``nvfp4_qat_train``) into its registered QuantizationConfig instance.
-
-    Model construction calls ``quant_config.get_quant_method()`` inside
-    ``LinearBase.__init__``, so a bare YAML string would crash there.
-    This must live in ``_parse_pipeline_config`` (not ``load_run_config``)
-    because ``dcp_to_diffusers`` rebuilds models from a checkpoint's raw
-    config by calling ``_parse_pipeline_config`` directly.
-    """
-    dit_config = getattr(pipeline_config, "dit_config", None)
-    quant = getattr(dit_config, "quant_config", None)
-    if isinstance(quant, str):
-        from fastvideo.layers.quantization import (
-            get_quantization_config, )
-        dit_config.quant_config = get_quantization_config(quant)()
-
-
 def _split_training_dit_arch_overrides(pipeline_raw: Any) -> tuple[Any, dict[str, Any]]:
     """Remove train-only DiT arch overrides before generic config parsing."""
     if not isinstance(pipeline_raw, dict) or not isinstance(pipeline_raw.get("dit_config"), dict):

fastvideo/training/activation_checkpoint.py

@@ -12,8 +12,6 @@ TRANSFORMER_BLOCK_NAMES = [
     "temporal_transformer_blocks",
     "transformer_double_blocks",
     "transformer_single_blocks",
-    "text_transformer_blocks",
-    "visual_transformer_blocks",
 ]
 
 

pyproject.toml

@@ -35,10 +35,7 @@ dependencies = [
     "torch==2.12.0",
     "torchvision",
     "torchaudio",
-    # CUDA-only: no wheels for non-Linux platforms, and the sdist build shells
-    # out to nvcc unconditionally. Gate on sys_platform so installs on macOS
-    # (MPS) and other non-Linux platforms don't try to build it from source.
-    "flashinfer-python; sys_platform == 'linux'",
+    "flashinfer-python",
 
     # Acceleration & Optimization
     "accelerate==1.0.1",
@@ -76,12 +73,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    # CUDA-only: PyPI only publishes manylinux wheels, and the sdist build
-    # requires nvcc via CMake. Gate on sys_platform so installs on macOS
-    # (MPS) and other non-Linux platforms don't try to build it from source;
-    # every fastvideo_kernel import elsewhere is already guarded to degrade
-    # gracefully when the package is absent.
-    "fastvideo-kernel==0.3.2; sys_platform == 'linux'",
+    "fastvideo-kernel==0.3.2",
     "wheel",
 
     # Training Dependencies
@@ -122,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 a cutlass-4.5-compatible revision. torch.compile support
+# FA4 cute, pinned to the CuTe DSL 4.6.0.dev0-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" }
 

tests/local_tests/test_nvfp4_fa4.py

@@ -149,57 +149,3 @@ class TestNVFP4FA4:
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
-    reason="Requires Blackwell GPU (sm100a or sm103a)"
-)
-class TestAttnQatInferFA4Route:
-    """ATTN_QAT_INFER resolves to the FP4 FA4 kernel on sm_100/sm_103.
-
-    Hardware-validated on both GB200 (sm_100a) and GB300 (sm_103a): route
-    resolution, SDPA parity, and cross-attention lengths pass on real
-    silicon for each capability in the FA4-FP4 set.
-    """
-
-    def test_resolves_to_fa4(self):
-        from fastvideo.attention.backends.attn_qat_infer import (
-            _resolved_kernel, attn_qat_infer_receipt, is_attn_qat_infer_available)
-        assert is_attn_qat_infer_available()
-        assert _resolved_kernel() == "fa4_fp4"
-        receipt = attn_qat_infer_receipt()
-        assert "qk_mode=nvfp4(per-16-e4m3-sf)" in receipt
-        assert "pv_mode=bf16" in receipt
-
-    def test_forward_parity_vs_sdpa(self):
-        from fastvideo.attention.backends.attn_qat_infer import AttnQatInferImpl
-        torch.manual_seed(0)
-        b, s, h, d = 1, 4096, 12, 128
-        q = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16)
-        k = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16)
-        v = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16)
-
-        impl = AttnQatInferImpl(num_heads=h, head_size=d, causal=False, softmax_scale=d**-0.5)
-        out = impl.forward(q, k, v, attn_metadata=None)
-
-        ref = torch.nn.functional.scaled_dot_product_attention(
-            q.transpose(1, 2).float(), k.transpose(1, 2).float(),
-            v.transpose(1, 2).float()).transpose(1, 2)
-
-        cos = torch.nn.functional.cosine_similarity(out.float().flatten(), ref.flatten(), dim=0).item()
-        # Same bound the sm_120 CUTLASS kernel test uses; FA4 NVFP4 QK
-        # measures cos_sim ~0.99 vs BF16 (kernel repo README precision table).
-        assert cos >= 0.97, f"cos_sim={cos:.4f} < 0.97"
-
-    def test_cross_attention_lengths(self):
-        from fastvideo.attention.backends.attn_qat_infer import AttnQatInferImpl
-        torch.manual_seed(0)
-        b, h, d = 1, 12, 128
-        q = torch.randn(b, 384, h, d, device="cuda", dtype=torch.bfloat16)
-        k = torch.randn(b, 512, h, d, device="cuda", dtype=torch.bfloat16)
-        v = torch.randn(b, 512, h, d, device="cuda", dtype=torch.bfloat16)
-        impl = AttnQatInferImpl(num_heads=h, head_size=d, causal=False, softmax_scale=d**-0.5)
-        out = impl.forward(q, k, v, attn_metadata=None)
-        assert out.shape == (b, 384, h, d)