[misc]: re-land doc fixes orphaned by the #1655 merge point

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

.buildkite/pipeline.yml

@@ -1,6 +1,8 @@
 env:
   IMAGE_VERSION: "py3.12-latest"
   BUILDKITE_CLEAN_CHECKOUT: true
+  # Buildkite only launches Modal; remote jobs initialize their own submodules.
+  BUILDKITE_GIT_SUBMODULES: false
 
 notify:
   - github_commit_status:
@@ -402,7 +404,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile"
               config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                command: "timeout 25m .buildkite/scripts/pr_test.sh"
                 label: ":test_tube: Training Tests"
                 env:
                   - TEST_TYPE=training
@@ -413,7 +415,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile"
               config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                command: "timeout 25m .buildkite/scripts/pr_test.sh"
                 label: ":test_tube: Distillation DMD Tests"
                 env:
                   - TEST_TYPE=distillation_dmd

.github/workflows/ci-precommit.yml

@@ -27,14 +27,25 @@ jobs:
         ref: ${{ inputs.ref || '' }}
     # For PR events, lint the PR head — but keep the hook definitions from
     # the base branch so an untrusted PR cannot alter what gets executed.
-    - name: Save trusted hook config
+    # The gate scripts are saved too: the self-test step below executes them,
+    # so it must run the base-branch copies, not the PR head's.
+    - name: Save trusted hook config and gate scripts
       if: github.event_name == 'pull_request_target'
-      run: cp .pre-commit-config.yaml "$RUNNER_TEMP/trusted-pre-commit-config.yaml"
-    - uses: actions/checkout@v4
+      run: |
+        cp .pre-commit-config.yaml "$RUNNER_TEMP/trusted-pre-commit-config.yaml"
+        cp -a .github/scripts "$RUNNER_TEMP/trusted-scripts"
+        echo "GATE_SCRIPTS_DIR=$RUNNER_TEMP/trusted-scripts" >> "$GITHUB_ENV"
+    # allow-unsafe-pr-checkout acknowledges checkout's pull_request_target
+    # guard: the head is data for the trusted hooks to lint; nothing from it
+    # is executed (config and gate scripts are pinned to the base branch
+    # above) and credentials are not persisted. SHA-pinned to v4.4.0 because
+    # actionlint's action schema does not know the new input yet.
+    - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
       if: github.event_name == 'pull_request_target'
       with:
         ref: ${{ github.event.pull_request.head.sha }}
         persist-credentials: false
+        allow-unsafe-pr-checkout: true
     - name: Restore trusted hook config
       if: github.event_name == 'pull_request_target'
       run: cp "$RUNNER_TEMP/trusted-pre-commit-config.yaml" .pre-commit-config.yaml
@@ -48,5 +59,7 @@ jobs:
       with:
         extra_args: --all-files --hook-stage manual
     # After pre-commit so a self-test failure cannot mask lint failures.
+    # GATE_SCRIPTS_DIR points at the base-branch copy on fork PRs (set above);
+    # push / workflow_call runs use the checked-out tree directly.
     - name: Full-suite gate self-test
-      run: bash .github/scripts/test_gate_full_suite.sh
+      run: bash "${GATE_SCRIPTS_DIR:-.github/scripts}/test_gate_full_suite.sh"

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

@@ -13,16 +13,23 @@ on:
         required: false
         default: false
         type: boolean
-  # 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.
+  # 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.
   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:
@@ -50,7 +57,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 a push that changed docker/Dockerfile (inputs are null on push). The
+    # on an in-scope main push (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' }}

.github/workflows/infra-docs.yml

@@ -11,9 +11,7 @@ on:
       - 'requirements-mkdocs.txt'
       - 'scripts/check_docs_links.py'
       - '.github/workflows/infra-docs.yml'
-  # Run the trusted base-branch workflow so fork PRs can be skipped without
-  # waiting for maintainer approval.
-  pull_request_target:
+  pull_request:
     branches: [ main ]
     paths:
       - 'docs/**'
@@ -26,19 +24,21 @@ on:
 
 permissions:
   contents: read
+  pages: write
+  id-token: write
+
+concurrency:
+  group: "pages"
+  cancel-in-progress: false
 
 jobs:
   build:
-    # MkDocs executes repository code; only trusted same-repository PRs run it.
-    if: github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
         uses: actions/checkout@v4
         with:
-          ref: ${{ github.event.pull_request.head.sha || github.sha }}
           fetch-depth: 0
-          persist-credentials: false
 
       - name: Setup Python
         uses: actions/setup-python@v5
@@ -52,7 +52,6 @@ jobs:
         run: uv pip install --system -r requirements-mkdocs.txt
 
       - name: Setup Pages
-        if: github.event_name == 'push'
         uses: actions/configure-pages@v4
 
       - name: Build documentation
@@ -62,22 +61,17 @@ jobs:
         run: python scripts/check_docs_links.py
 
       - name: Upload artifact
-        if: github.event_name == 'push'
         uses: actions/upload-pages-artifact@v3
         with:
           path: ./site
 
   deploy:
-    permissions:
-      pages: write
-      id-token: write
-    concurrency: pages
     environment:
       name: github-pages
       url: ${{ steps.deployment.outputs.page_url }}
     runs-on: ubuntu-latest
     needs: build
-    if: github.event_name == 'push'
+    if: github.ref == 'refs/heads/main'
     steps:
       - name: Deploy to GitHub Pages
         id: deployment

.gitignore

@@ -6,6 +6,7 @@ results/
 wandb/
 *.ipynb
 *.jpg
+!examples/dataset/lingbotworld2/image.jpg
 *.safetensors
 *.mp4
 *.png
@@ -34,6 +35,7 @@ env
 *.log
 weights/
 logs/
+/Z-Image/
 official_weights/
 converted_weights/
 
@@ -73,6 +75,7 @@ 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

README.md

@@ -9,7 +9,7 @@
 **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.**
 
 ## NEWS
-- `2026/06/23`: Release FastWan-QAD: 5s of Video generated in 1.8s E2E. [FastWan-QAD models](https://huggingface.co/FastVideo/FastWan-QAD-FP8-1.3B), check out the [Blog](https://haoailab.com/blogs/fastwan-qad/).
+- `2026/06/23`: Release FastWan-QAD: 5s of Video generated in 1.8s E2E. See the [FastWan-QAD models](https://huggingface.co/FastVideo/FastWan-QAD-FP8-1.3B), [Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/), and [blog](https://haoailab.com/blogs/fastwan-qad/).
 - `2026/03/17`: Release demo: Into the Dreamverse: Vibe Directing in FastVideo, check out the [Blog](https://haoailab.com/blogs/dreamverse/).
 - `2026/03/13`: Release demo: Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU, check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/).
 - `2025/11/19`: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py).

apps/dreamverse/scripts/launch/launch_backend_dreamverse.sh

@@ -39,6 +39,17 @@ export FASTVIDEO_GENERATION_SEGMENT_CAP="${FASTVIDEO_GENERATION_SEGMENT_CAP:-6}"
 export FASTVIDEO_PROMPT_AUTO_SLEEP_MS="${FASTVIDEO_PROMPT_AUTO_SLEEP_MS:-120}"
 export FASTVIDEO_PROMPT_AUTO_TIMEOUT_MS="${FASTVIDEO_PROMPT_AUTO_TIMEOUT_MS:-1800}"
 
+if [[ "${ENABLE_TORCH_COMPILE}" == "1" ]]; then
+  # Persist Inductor, AOTAutograd, and Triton artifacts across launches.
+  export DREAMVERSE_TORCH_COMPILE_CACHE_ROOT="${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT:-${HOME}/.cache/dreamverse/torch_compile}"
+  export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}/inductor}"
+  export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}/triton}"
+  export TORCHINDUCTOR_FX_GRAPH_CACHE="${TORCHINDUCTOR_FX_GRAPH_CACHE:-1}"
+  export TORCHINDUCTOR_AUTOGRAD_CACHE="${TORCHINDUCTOR_AUTOGRAD_CACHE:-1}"
+  mkdir -p "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}"
+  echo "[launch-demo] torch.compile cache: ${DREAMVERSE_TORCH_COMPILE_CACHE_ROOT}"
+fi
+
 cd "${DREAMVERSE_ROOT}"
 
 if ! command -v dreamverse-server >/dev/null 2>&1; then

docker/Dockerfile

@@ -186,15 +186,40 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 
 COPY . .
 
-# 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.
+# 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
 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_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
-    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} ./build.sh
+    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
 
 # Install FastVideo itself (editable) now that the source is present, and set up
 # shell configuration. Dependencies and the local kernel are already installed,

docs/attention/index.md

@@ -8,13 +8,15 @@ FastVideo provides highly optimized custom attention kernels to accelerate video
 * **[Sliding Tile Attention (STA)](sta/index.md)**: STA kernel support is kept in
   `fastvideo-kernel`; full FastVideo STA pipeline workflow is archived in
   `sta_do_not_delete`.
+* **[Attn-QAT Training](../training/attn_qat.md)**: Runtime-JIT Triton forward
+  and backward kernels for role-local quantization-aware training.
 * **Backend development guide**: See the developer guide at
   [Attention Backend Development](../contributing/attention_backend.md).
 
 ## General Build Instructions
 
 These instructions apply to building the `fastvideo-kernel` package from
-source, which includes both STA and VSA kernels.
+source, which includes STA, VSA, and Attn-QAT kernels.
 
 ### Prerequisites
 

docs/contributing/ci_architecture.md

@@ -255,10 +255,9 @@ If you add a new CI test category:
 
 ### Documentation
 
-`.github/workflows/infra-docs.yml` builds documentation for same-repository PRs
-that touch `docs/**`, `mkdocs.yml`, `requirements-mkdocs.txt`, or the workflow
-itself. Fork PRs skip this executable build instead of waiting for maintainer
-approval. On pushes to `main`, it also deploys the built site to GitHub Pages.
+`.github/workflows/infra-docs.yml` builds documentation for PRs that touch
+`docs/**`, `mkdocs.yml`, `requirements-mkdocs.txt`, or the workflow itself. On
+pushes to `main`, it also deploys the built site to GitHub Pages.
 
 The docs job:
 
@@ -270,16 +269,26 @@ The docs job:
 ### Docker Images
 
 `.github/workflows/infra-build-image.yml` supports manual `workflow_dispatch`
-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`.
+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.
 
 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/design/inference_schema_parity_inventory.yaml

@@ -333,15 +333,25 @@ surfaces:
       use_distill:
         sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
       scheduler_arch:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       text_encoder_archs:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       tokenizer_archs:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       transformer_arch:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       vae_arch:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       expand_timesteps:
         sources:
           - fastvideo.configs.pipelines.wan.FastWan2_2_TI2V_5B_Config
@@ -421,6 +431,8 @@ surfaces:
       frame_receptive_field: "MagiHuman internal data-proxy receptive-field setting."
       image_conditioning: "MagiHuman preset variant marker for reference-image conditioning."
       ref_audio_offset: "MagiHuman internal data-proxy audio alignment offset."
+      scheduler_sigma_min: "Z-Image scheduler parity invariant; not part of the public typed inference API."
+      scheduler_use_reference_discrete_timesteps: "Z-Image scheduler parity invariant; not part of the public typed inference API."
       sr_local_attn_layers: "MagiHuman SR internal sparse-attention layer selection."
       text_offset: "MagiHuman internal data-proxy text alignment offset."
       vae_stride: "MagiHuman internal VAE/data-proxy stride setting."
@@ -438,6 +450,7 @@ surfaces:
       grid_sizes: request.inputs.grid_sizes
       pose: request.inputs.pose
       c2ws_plucker_emb: request.inputs.c2ws_plucker_emb
+      action_path: request.inputs.action_path
       refine_from: request.inputs.refine_from
       stage1_video: request.inputs.stage1_video
       prompt: request.prompt
@@ -447,6 +460,7 @@ surfaces:
       output_video_name: request.output.output_video_name
       num_videos_per_prompt: request.sampling.num_videos_per_prompt
       seed: request.sampling.seed
+      max_sequence_length: request.sampling.max_sequence_length
       num_frames: request.sampling.num_frames
       height: request.sampling.height
       width: request.sampling.width
@@ -456,7 +470,10 @@ surfaces:
       num_inference_steps: request.sampling.num_inference_steps
       num_inference_steps_sr: request.sampling.num_inference_steps_sr
       guidance_scale: request.sampling.guidance_scale
+      batch_cfg: request.sampling.batch_cfg
       guidance_scale_2: request.sampling.guidance_scale_2
+      cfg_normalization: request.sampling.cfg_normalization
+      cfg_truncation: request.sampling.cfg_truncation
       guidance_rescale: request.sampling.guidance_rescale
       use_embedded_guidance: request.sampling.use_embedded_guidance
       true_cfg_scale: request.sampling.true_cfg_scale
@@ -508,7 +525,6 @@ surfaces:
     internal_only:
       data_type: "Derived from the request shape and not a public input."
       latents: "Pre-generated diffusion latents supplied by parity/debug harnesses; not a public input."
-      max_sequence_length: "Model-specific text-encoder sequence cap; not part of the public typed inference API."
 
   sampling_param_extensions: {}
 

docs/distillation/dmd.md

@@ -2,6 +2,12 @@
 
 We introduce a new finetuning strategy - **Sparse-distill**, which jointly integrates **[DMD](https://arxiv.org/abs/2405.14867)** and **[VSA](https://arxiv.org/abs/2505.13389)** in a single training process. This approach combines the benefits of both distillation to shorten diffusion steps and sparse attention to reduce attention computation, enabling much faster video generation.
 
+!!! tip "Attn-QAT DMD2 workflow"
+    The modular trainer also provides a Wan2.1 MixKit recipe that first
+    fine-tunes with fake-quantized attention, then distills the student to
+    timesteps `[1000, 757, 522]` while teacher and critic remain on Flash
+    Attention. See [Attn-QAT Training](../training/attn_qat.md).
+
 ## 📊 Model Overview
 
 We provide two distilled models:

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-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ) 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-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ) 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-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ)
+our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ)
 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-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ).
 
 ## Step 1: Pipeline Modules
 

docs/inference/inference_quick_start.md

@@ -134,5 +134,4 @@ 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 Discord](https://discord.gg/JA7cksDz86).
-- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ).

docs/inference/optimizations.md

@@ -103,7 +103,7 @@ 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 or 3.11
+- **Python**: 3.10 – 3.12
 
 #### Installation
 
@@ -112,13 +112,36 @@ 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+ssh://git@github.com/hao-ai-lab/flash-attention-fp4.git@fix/cutlass-dsl-4.5#subdirectory=flash_attn/cute"
+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
 ```
 
 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.
+
 #### Usage
 
 Enable FP4 attention via the `--nvfp4_fa4` flag:

docs/inference/support_matrix.md

@@ -13,6 +13,92 @@ 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
@@ -25,6 +111,9 @@ 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 {
@@ -56,7 +145,7 @@ pipeline initialization and sampling.
 | Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sliding Tile Attn (Legacy Branch) | Sage Attn | VSA | BSA |
 |------------|---------------------|-------------|----------|-------------------|-----------|-----|-----|
 | FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
-| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
+| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
 | DreamX-World 5B Cam | `FastVideo/DreamX-World-5B-Cam-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | DreamX-World 5B AR | `FastVideo/DreamX-World-5B-Diffusers` | 704px1280p | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
@@ -65,14 +154,14 @@ pipeline initialization and sampling.
 | 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 | See note** | 480P<br>720P | ❌ | ❌ | ❌ | ⭕ | ✅ |
+| LongCat T2V 13.6B | `FastVideo/LongCat-Video-T2V-Diffusers` | 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 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
@@ -98,6 +187,21 @@ 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/attn_qat.md

@@ -0,0 +1,147 @@
+# Attn-QAT Training
+
+Attn-QAT simulates low-bit attention during training while keeping the rest of
+the training method unchanged. In the modular `fastvideo/train` framework it is
+a per-role model option, not a separate training method: supervised fine-tuning
+and DMD2 still own their losses and optimizer cadence.
+
+This guide covers the QAD Wan2.1-T2V-1.3B MixKit workflow:
+
+1. run a 4,000-step supervised Attn-QAT fine-tune;
+2. export the stage-1 DCP checkpoint to Diffusers format; and
+3. distill the student to three denoising steps with DMD2.
+
+The ready-to-run configs and wrappers are in
+`examples/train/scenario/qad_wan2_1_mixkit/`.
+
+## Role-local attention backends
+
+A DMD2 run owns three independent model roles. Configure the attention backend
+on each role so fake quantization is applied only to the student:
+
+```yaml
+models:
+  student:
+    attention_backend: ATTN_QAT_TRAIN
+  teacher:
+    attention_backend: FLASH_ATTN
+  critic:
+    attention_backend: FLASH_ATTN
+```
+
+The override is active only while that role's transformer is constructed, then
+the previous process-wide backend is restored. This lets student, teacher, and
+critic use different implementations in one process. Invalid role-level names
+fail during configuration instead of silently selecting another backend.
+
+See [Training Infrastructure](train_infra.md) for the complete model-role
+configuration reference.
+
+## Prerequisites
+
+- Install FastVideo and make the `fastvideo-kernel` Python package importable.
+  `ATTN_QAT_TRAIN` intentionally fails instead of falling back to dense
+  attention when its kernel cannot be loaded.
+- Prepare the precomputed MixKit VAE latents and text embeddings.
+- Run the commands below from the repository root. The supplied recipe expects
+  four GPUs by default; set `NUM_GPUS` to override it.
+
+Download the published preprocessed dataset:
+
+```bash
+bash examples/training/finetune/wan_t2v_1.3B/mixkit/download_mixkit_data.sh
+```
+
+## Stage 1: supervised Attn-QAT fine-tuning
+
+The stage-1 config uses `ATTN_QAT_TRAIN` on the student, sequence parallelism
+across four GPUs, FP32 master weights, and 4,000 optimizer steps:
+
+```bash
+NUM_GPUS=4 \
+  bash examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh
+```
+
+Pass a dataset directory as the first positional argument when it differs from
+the default:
+
+```bash
+NUM_GPUS=4 \
+  bash examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh \
+  /path/to/combined_parquet_dataset
+```
+
+The wrapper calls `examples/train/run.sh`; the YAML file remains the source of
+truth for optimizer, validation, checkpointing, and distributed settings.
+
+## Export the stage-1 checkpoint
+
+Modular training checkpoints use Distributed Checkpoint (DCP) format. Export
+the student before using it to initialize stage 2:
+
+```bash
+bash examples/train/scenario/qad_wan2_1_mixkit/export_stage1.sh \
+  checkpoints/wan_t2v_qat_finetune/checkpoint-4000 \
+  checkpoints/wan_t2v_qat_finetune/diffusers
+```
+
+Both arguments are optional; the command above shows their defaults.
+
+## Stage 2: three-step DMD2 distillation
+
+Stage 2 loads the exported student weights, keeps Attn-QAT on the student, and
+uses Flash Attention for the teacher and critic:
+
+```bash
+NUM_GPUS=4 \
+  bash examples/train/scenario/qad_wan2_1_mixkit/run_stage2.sh \
+  data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset \
+  checkpoints/wan_t2v_qat_finetune/diffusers/transformer/model.safetensors
+```
+
+The migrated recipe preserves these behaviors:
+
+| Behavior | Modular configuration |
+|---|---|
+| Student fake-quantized attention | `models.student.attention_backend: ATTN_QAT_TRAIN` |
+| Teacher and critic full-precision attention | Role-local `FLASH_ATTN` |
+| Generator update every five critic steps | `method.generator_update_interval: 5` |
+| Three-step rollout | `method.dmd_denoising_steps: [1000, 757, 522]` |
+| Score timestep range | `method.min_timestep_ratio: 0.02`, `max_timestep_ratio: 0.98` |
+| Legacy guidance `cond + 2(cond - uncond)` | Standard CFG scale `3.0` |
+| Stage handoff | DCP checkpoint to Diffusers export to student override weights |
+
+The timestep ratios apply to randomly sampled teacher and critic score
+timesteps; `dmd_denoising_steps` separately controls the student rollout. See
+[DMD Distillation](../distillation/dmd.md) for general DMD concepts.
+
+## Architecture-specific Triton routing
+
+The training kernel is runtime-JIT-compiled Triton code and selects its route on
+every call. It supports different query and key/value sequence lengths for
+cross-attention; key and value must have the same sequence length.
+
+| Hardware/configuration | Route |
+|---|---|
+| SM100, validated non-causal BF16 QAT configuration with head dimension 128 | Large-tile forward and split 64x64 backward; optimized backward requires a 16-aligned KV length |
+| SM120, including RTX 5090 | Previous forward tiling with joined quantized/STE P@V operations and a shallower backward pipeline for long sequences |
+| Unsupported configurations | Previous Triton implementation |
+
+Warp specialization is disabled automatically on SM100 and SM120 because the
+Triton 3.7 NVWS compiler pass aborts for this kernel on Blackwell. No user
+setting is required.
+
+The available tuning and comparison controls are:
+
+| Environment variable | Default | Effect |
+|---|---|---|
+| `FASTVIDEO_ATTN_QAT_FWD_MODE` | `fast` | Selects `fast`, `balanced`, or `reference` forward tiling on the SM100 optimized route |
+| `FASTVIDEO_ATTN_QAT_FWD_EXACT_M` | `0` | Set to `1` to recompute reference-order softmax statistics and keep `dV` bitwise-compatible on the SM100 optimized route |
+| `FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED` | `1` | Set to `0` to force the previous SM100 forward and backward for comparison |
+| `FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV` | `1` | Set to `0` to compare SM120 against the split P@V path |
+
+The first invocation JIT-compiles the selected configuration; later calls reuse
+the Triton cache. To measure the production shape, run
+`python benchmarks/benchmark_attn_qat_train.py` from `fastvideo-kernel/`.
+
+For import and backend-selection failures, see [Debugging](../utilities/debugging.md).

docs/training/finetune.md

@@ -62,6 +62,18 @@ bash examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
 - Gradient checkpointing: `--enable_gradient_checkpointing_type "full"`
 - Memory scaling: Increase `--sp_size` or reduce `--num_latent_t` to fit in memory
 
+## Attention Quantization-Aware Training
+
+Attn-QAT fine-tunes a model while simulating low-bit attention in the forward
+and backward passes. The modular trainer can select the backend per model role,
+so a later DMD2 stage can keep fake quantization on the student while the
+teacher and critic use Flash Attention.
+
+The ready-to-run Wan2.1 MixKit workflow includes supervised fine-tuning,
+checkpoint export, and three-step DMD2 distillation:
+
+**→ [Follow the Attn-QAT training guide](attn_qat.md)**
+
 ## LoRA Finetuning
 
 LoRA (Low-Rank Adaptation) trains lightweight adapters while keeping the base model frozen. This significantly reduces memory usage and training time.
@@ -166,6 +178,7 @@ Ready-to-run training scripts are available for multiple models:
 | Wan2.1 I2V 14B | I2V | `examples/training/finetune/wan_i2v_14B_480p/crush_smol/` |
 | Wan2.1-Fun 1.3B InP | I2V | `examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/` |
 | Wan2.1 VSA | T2V/I2V | `examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/` |
+| Wan2.1 T2V 1.3B Attn-QAT | QAT SFT + DMD2 | `examples/train/scenario/qad_wan2_1_mixkit/` |
 
 Each example includes:
 

docs/training/overview.md

@@ -43,6 +43,9 @@ Ready-to-run examples with preprocessing scripts, training launchers, and valida
 
 **→ [Browse all training examples](examples/examples_training_index.md)**
 
+For the complete two-stage Wan2.1 MixKit quantization-aware workflow, see
+**[Attn-QAT Training](attn_qat.md)**.
+
 Each example includes:
 
 - `download_dataset.sh` — download sample data
@@ -59,9 +62,11 @@ FastVideo supports several training approaches:
 | **Full finetune** | Adapt entire model to a new domain or style |
 | **LoRA finetune** | Lightweight adaptation with frozen base weights |
 | **VSA finetune** | Finetune with Variable Sparse Attention for efficiency |
+| **Attn-QAT** | Train with fake-quantized attention, optionally followed by DMD2 distillation |
 
 ## Next Steps
 
 1. **Get started**: Pick an example from the [training examples index](examples/examples_training_index.md)
 2. **Prepare data**: Follow [data preprocessing](data_preprocess.md) for your own dataset
-3. **Run inference**: After training, see [inference examples](../inference/examples/examples_inference_index.md)
+3. **Train with quantized attention**: Follow the [Attn-QAT two-stage recipe](attn_qat.md)
+4. **Run inference**: After training, see [inference examples](../inference/examples/examples_inference_index.md)

docs/training/train_infra.md

@@ -81,6 +81,7 @@ Common model parameters:
 | `disable_custom_init_weights` | `false` | Skip custom weight initialization (use for teacher/critic) |
 | `flow_shift` | `3.0` | Timestep shifting factor |
 | `enable_gradient_checkpointing_type` | `null` | Gradient checkpointing (`"full"` or `null`) |
+| `attention_backend` | `null` | Optional role-local backend for Wan models (for example `ATTN_QAT_TRAIN`); overrides the process default only while this role's transformer is built |
 
 Which roles are needed depends on the training method:
 
@@ -211,6 +212,29 @@ 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
@@ -298,6 +322,8 @@ method:
 | `dmd_denoising_steps` | *(required)* | Timestep schedule for student rollout |
 | `generator_update_interval` | `1` | Update student every N critic steps |
 | `real_score_guidance_scale` | `1.0` | CFG scale for teacher predictions |
+| `min_timestep_ratio` | `0.0` | Lower bound for randomly sampled teacher/critic score timesteps |
+| `max_timestep_ratio` | `1.0` | Upper bound for randomly sampled teacher/critic score timesteps |
 | `fake_score_learning_rate` | *(required)* | Critic optimizer learning rate |
 | `fake_score_betas` | *(required)* | Critic optimizer Adam betas |
 | `fake_score_lr_scheduler` | *(required)* | Critic LR scheduler type |

docs/utilities/debugging.md

@@ -77,8 +77,13 @@ 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`
-- `ATTN_QAT_TRAIN`: `fastvideo-kernel` install exposing `fastvideo_kernel`
+  `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_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
+  [Attn-QAT Training](../training/attn_qat.md) for architecture controls.
 
 As a fallback, use:
 

examples/dataset/lingbotworld2/README.md

@@ -0,0 +1,17 @@
+# LingBot World 2 Example Dataset
+
+These files were copied unchanged from the LingBot World 2 repository for the
+FastVideo causal-fast inference example.
+
+- Repository: `https://github.com/Robbyant/lingbot-world-v2.git`
+- Source commit: `94f43115de8d4a4f9f282126528c300a0b232c5f`
+- Source directory: `examples/03`
+
+## Files
+
+- `image.jpg`: source image for image-to-video generation. SHA-256:
+  `6ee3dacfef32cfef504dd698adb8a660cf15f686535c52fed4903fef27c0edd0`
+- `poses.npy`: camera-to-world trajectory matrices. SHA-256:
+  `bd0a23a696e184b0b43e7767eb432bfe644690560fe327fa96961affc941c404`
+- `intrinsics.npy`: camera intrinsic parameters. SHA-256:
+  `821fca6cf957ae8fbb1181307f02479efb1705e04c9e05734cd02fb43462e082`

examples/inference/basic/basic_lingbot_video.py

@@ -0,0 +1,52 @@
+"""Generate a five-second Dense LingBot-Video clip with the official defaults."""
+
+import argparse
+from pathlib import Path
+
+from fastvideo import VideoGenerator
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse the converted checkpoint and output paths for the sample."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--model-path",
+        type=Path,
+        required=True,
+        help="Path to a converted Dense LingBot-Video checkpoint.",
+    )
+    parser.add_argument(
+        "--output-path",
+        type=Path,
+        default=Path("outputs/lingbot-video/dense-t2v"),
+        help="Directory for the generated video.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    """Load the converted Dense checkpoint and generate the default T2V sample."""
+    args = parse_args()
+    generator = VideoGenerator.from_pretrained(
+        str(args.model_path),
+        num_gpus=1,
+        use_fsdp_inference=False,
+        text_encoder_cpu_offload=True,
+        vae_cpu_offload=False,
+        pin_cpu_memory=True,
+    )
+    try:
+        generator.generate({
+            "prompt": "A red fox runs through fresh snow at sunrise.",
+            "output": {
+                "output_path": str(args.output_path),
+                "save_video": True,
+                "return_frames": False,
+            },
+        })
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_lingbotworld2_causal_fast.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run LingBot World 2 14B causal-fast I2V generation with FastVideo."""
+
+import os
+from pathlib import Path
+
+from fastvideo import VideoGenerator
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DATASET_DIR = REPO_ROOT / "examples" / "dataset" / "lingbotworld2"
+OUTPUT_PATH = REPO_ROOT / "outputs" / "lingbotworld2_causal_fast.mp4"
+
+
+def main() -> None:
+    """Load the native FastVideo LingBot World 2 causal-fast pipeline and generate one video."""
+    generator = VideoGenerator.from_pretrained(
+        os.environ["LINGBOTWORLD2_MODEL_PATH"],
+        num_gpus=8,
+        sp_size=8,
+        hsdp_shard_dim=8,
+        use_fsdp_inference=True,
+        dit_layerwise_offload=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        pin_cpu_memory=True,
+        override_pipeline_cls_name="LingBotWorld2CausalFastPipeline",
+    )
+
+    try:
+        generator.generate_video(
+            "A serene lakeside scene with a lone tree standing in calm water, surrounded by distant snow-capped mountains under a bright blue sky with drifting white clouds; gentle ripples reflect the tree and sky, creating a tranquil, meditative atmosphere.",
+            image_path=str(DATASET_DIR / "image.jpg"),
+            action_path=str(DATASET_DIR),
+            output_path=str(OUTPUT_PATH),
+            save_video=True,
+            height=480,
+            width=832,
+            num_frames=65,
+            num_inference_steps=4,
+            guidance_scale=1.0,
+            negative_prompt="",
+            fps=16,
+            seed=42,
+        )
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_ltx2_distilled_fast_profile.py

@@ -191,6 +191,22 @@ 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/basic/basic_zimage.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run Z-Image-Turbo text-to-image generation through FastVideo.
+
+User story:
+    "I want the official Z-Image-Turbo defaults and a deterministic PNG from
+    a local or Hugging Face checkpoint."
+"""
+
+import argparse
+from pathlib import Path
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+DEFAULT_PROMPT = (
+    "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. "
+    "Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. "
+    "Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, "
+    "silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."
+)
+DEFAULT_REVISION = "f332072aa78be7aecdf3ee76d5c247082da564a6"
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Z-Image-Turbo text-to-image generation.")
+    parser.add_argument("--model-path", default="Tongyi-MAI/Z-Image-Turbo")
+    parser.add_argument("--revision", default=DEFAULT_REVISION)
+    parser.add_argument("--output", default="outputs/zimage/zimage_turbo.png")
+    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
+    parser.add_argument("--negative-prompt", default="")
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=8)
+    parser.add_argument("--guidance-scale", type=float, default=0.0)
+    parser.add_argument("--max-sequence-length", type=int, default=512)
+    parser.add_argument("--cfg-normalization", action=argparse.BooleanOptionalAction, default=False)
+    parser.add_argument("--cfg-truncation", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=42)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+
+    generator = VideoGenerator.from_config(
+        GeneratorConfig(
+            model_path=args.model_path,
+            revision=args.revision,
+            engine=EngineConfig(
+                num_gpus=1,
+                parallelism=ParallelismConfig(tp_size=1, sp_size=1),
+                use_fsdp_inference=False,
+            ),
+            # The model registry selects the native zimage_turbo preset.
+            pipeline=PipelineSelection(workload_type="t2i"),
+        ))
+    try:
+        generator.generate(
+            GenerationRequest(
+                prompt=args.prompt,
+                negative_prompt=args.negative_prompt,
+                sampling=SamplingConfig(
+                    height=args.height,
+                    width=args.width,
+                    num_frames=1,
+                    fps=1,
+                    num_inference_steps=args.steps,
+                    guidance_scale=args.guidance_scale,
+                    max_sequence_length=args.max_sequence_length,
+                    cfg_normalization=args.cfg_normalization,
+                    cfg_truncation=args.cfg_truncation,
+                    seed=args.seed,
+                ),
+                output=OutputConfig(
+                    output_path=str(output),
+                    save_video=True,
+                    return_frames=False,
+                ),
+            ))
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

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

@@ -6,10 +6,6 @@ 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 |
@@ -38,27 +34,21 @@ 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.
 
-## Blocking prereqs (why this draft PR cannot be merged yet)
+## Upstream surfaces the demo relies on
 
-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:
+Everything the demo needs ships in the upstream `fastvideo` package:
 
-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.
+- **`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).
 
 ## 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.fp4_config import FP4Config
+from fastvideo.layers.quantization.nvfp4_config import NVFP4Config
 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 = FP4Config()
+        pipeline_config.dit_config.quant_config = NVFP4Config()
         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.configs.sample.base import SamplingParam
+from fastvideo 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.configs.sample.base import SamplingParam
+from fastvideo import SamplingParam
 from fastvideo.entrypoints.video_generator import VideoGenerator
 
 from .config import (

examples/inference/ltx2_3/optimized_nvfp4_t2v.py

@@ -0,0 +1,176 @@
+# 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/README.md

@@ -91,3 +91,7 @@ examples/train/
 ```
 
 See `configs/README.md` and `scenario/README.md` for details.
+
+The featured QAD Wan2.1 MixKit scenario runs Attn-QAT supervised fine-tuning,
+checkpoint export, and three-step DMD2. See the
+[Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/).

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

@@ -0,0 +1,70 @@
+#!/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

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

@@ -25,6 +25,7 @@ models:
     disable_custom_init_weights: false                # default: false
     flow_shift: 3.0                                   # default: 3.0
     enable_gradient_checkpointing_type: null           # default: null (falls back to training.model)
+    attention_backend: null                            # default: null (global/default); role-local when set
 
   teacher:
     _target_: fastvideo.train.models.wan.WanModel
@@ -52,6 +53,8 @@ method:
   rollout_mode: simulate              # required: "simulate" or "data_latent"
   generator_update_interval: 5        # default: 1
   dmd_denoising_steps: [1000, 750, 500, 250]  # SDE timestep schedule
+  min_timestep_ratio: 0.0            # score-model timestep lower bound
+  max_timestep_ratio: 1.0            # score-model timestep upper bound
 
   # Critic optimizer (all required — no fallback)
   fake_score_learning_rate: 8.0e-6
@@ -192,6 +195,10 @@ 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

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

@@ -0,0 +1,33 @@
+#!/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

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

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

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

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

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

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

@@ -0,0 +1,116 @@
+# 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_3_t2v.yaml

@@ -0,0 +1,95 @@
+# LTX-2.3 T2V overfitting test config.
+#
+# Overfits on a single short video (480x832, 81 frames @ 24fps) to
+# verify the LTX-2 training plugin works end-to-end on an LTX-2.3
+# checkpoint (gated attention, cross-attention AdaLN, 4096-d
+# post-connector text embeddings, no in-DiT caption projection).
+# Uses the distilled checkpoint (validation is 8 sampling steps).
+#
+# Preprocess data first (writes data/ltx2_3_overfit_preprocessed):
+#   CUDA_VISIBLE_DEVICES=0 \
+#   LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers \
+#   LTX2_OVERFIT_OUTPUT_DIR=data/ltx2_3_overfit_preprocessed \
+#   python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
+#
+# Run:
+#   NUM_GPUS=4 bash examples/train/run.sh examples/train/configs/overfit_ltx2_3_t2v.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    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:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/ltx2_3_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_3_overfit
+    # 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_3_overfit
+
+  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_3_overfit_preprocessed/validation_prompts.json
+    every_steps: 50
+    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.yaml

@@ -0,0 +1,90 @@
+# LTX-2 T2V overfitting test config.
+#
+# Overfits on a single short video (480x832, 81 frames @ 24fps) to
+# verify the LTX-2 training plugin works end-to-end. Uses the
+# distilled checkpoint (validation is 8 sampling steps, single pass).
+#
+# Preprocess data first (writes data/ltx2_overfit_preprocessed):
+#   CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
+#
+# Run:
+#   NUM_GPUS=4 bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX2-Distilled-Diffusers
+    trainable: true
+    enable_gradient_checkpointing_type: full
+
+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
+    # 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
+
+  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
+
+# 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

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

examples/train/scenario/README.md

@@ -5,9 +5,12 @@ all configs, scripts, and data needed to run a complete workflow.
 
 ```
 scenario/
-└── ode_init_self_forcing_wan_causal/   # KD → export → Self-Forcing
+├── ode_init_self_forcing_wan_causal/   # KD → export → Self-Forcing
+└── qad_wan2_1_mixkit/                  # Attn-QAT SFT → export → DMD2
 ```
 
-See the `usage.md` inside each scenario for step-by-step instructions.
+See the `usage.md` inside each scenario for step-by-step instructions. The QAD
+workflow is also documented in the website-visible
+[Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/).
 
 For single-step configs, see `examples/train/configs/`.

examples/train/scenario/qad_wan2_1_mixkit/export_stage1.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+# Export a modular-trainer DCP checkpoint for stage-2 initialization.
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
+cd "${REPO_ROOT}"
+
+CHECKPOINT_DIR=${1:-checkpoints/wan_t2v_qat_finetune/checkpoint-4000}
+OUTPUT_DIR=${2:-checkpoints/wan_t2v_qat_finetune/diffusers}
+
+python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+    --role student \
+    --checkpoint "${CHECKPOINT_DIR}" \
+    --output-dir "${OUTPUT_DIR}"

examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+# Run the modular Attn-QAT finetune recipe.
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
+cd "${REPO_ROOT}"
+
+DATA_DIR=${1:-data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset}
+NUM_GPUS=${NUM_GPUS:-4}
+export NUM_GPUS
+export FASTVIDEO_ATTN_QAT_FWD_EXACT_M=${FASTVIDEO_ATTN_QAT_FWD_EXACT_M:-0}
+
+bash "${REPO_ROOT}/examples/train/run.sh" \
+    "${SCRIPT_DIR}/stage1_attn_qat_finetune.yaml" \
+    --training.data.data_path "${DATA_DIR}" \
+    --training.distributed.num_gpus "${NUM_GPUS}" \
+    --training.distributed.sp_size "${NUM_GPUS}" \
+    --training.distributed.hsdp_replicate_dim 1 \
+    --training.distributed.hsdp_shard_dim "${NUM_GPUS}"

examples/train/scenario/qad_wan2_1_mixkit/run_stage2.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Run modular DMD2 with Attn-QAT on the student only.
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
+cd "${REPO_ROOT}"
+
+DATA_DIR=${1:-data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset}
+INIT_WEIGHTS=${2:-checkpoints/wan_t2v_qat_finetune/diffusers/transformer/model.safetensors}
+NUM_GPUS=${NUM_GPUS:-4}
+export NUM_GPUS
+
+if [[ ! -f "${INIT_WEIGHTS}" ]]; then
+    echo "Missing exported stage-1 weights: ${INIT_WEIGHTS}" >&2
+    echo "Run export_stage1.sh before stage 2." >&2
+    exit 1
+fi
+
+bash "${REPO_ROOT}/examples/train/run.sh" \
+    "${SCRIPT_DIR}/stage2_attn_qat_dmd.yaml" \
+    --models.student.transformer_override_safetensor "${INIT_WEIGHTS}" \
+    --training.data.data_path "${DATA_DIR}" \
+    --training.distributed.num_gpus "${NUM_GPUS}" \
+    --training.distributed.sp_size 1 \
+    --training.distributed.hsdp_replicate_dim "${NUM_GPUS}" \
+    --training.distributed.hsdp_shard_dim 1

examples/train/scenario/qad_wan2_1_mixkit/stage1_attn_qat_finetune.yaml

@@ -0,0 +1,72 @@
+# QAD stage 1: Attn-QAT finetune of Wan2.1-T2V-1.3B on MixKit.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    # Role-local: does not change the backend used by other models loaded in
+    # this process.
+    attention_backend: ATTN_QAT_TRAIN
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset
+    dataloader_num_workers: 1
+    train_batch_size: 1
+    training_cfg_rate: 0.1
+    seed: 1000
+    num_latent_t: 20
+    num_height: 480
+    num_width: 832
+    num_frames: 77
+
+  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: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: checkpoints/wan_t2v_qat_finetune
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 0
+
+  tracker:
+    project_name: wan_t2v_qat_finetune
+    run_name: wan_t2v_qat_finetune
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+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.wan.wan_pipeline.WanPipeline
+    dataset_file: examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json
+    every_steps: 50
+    sampling_steps: [50]
+    guidance_scale: 5.0
+
+pipeline:
+  flow_shift: 1

examples/train/scenario/qad_wan2_1_mixkit/stage2_attn_qat_dmd.yaml

@@ -0,0 +1,101 @@
+# QAD stage 2: distill the Attn-QAT student to three sampling steps.
+#
+# Export the stage-1 DCP checkpoint first, then point
+# models.student.transformer_override_safetensor at the exported weight file.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    transformer_override_safetensor: checkpoints/wan_t2v_qat_finetune/diffusers/transformer/model.safetensors
+    trainable: true
+    attention_backend: ATTN_QAT_TRAIN
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: data_latent
+  generator_update_interval: 5
+  dmd_denoising_steps: [1000, 757, 522]
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+  # The modular DMD method uses standard CFG:
+  # uncond + scale * (cond - uncond). This is equivalent to the legacy
+  # recipe's cond + 2.0 * (cond - uncond).
+  real_score_guidance_scale: 3.0
+
+  # The legacy recipe inherited these values from its global optimizer.
+  fake_score_learning_rate: 2.0e-6
+  fake_score_betas: [0.9, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 4
+    hsdp_shard_dim: 1
+
+  data:
+    data_path: data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 480
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: checkpoints/wan_t2v_distill_dmd_qat
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: wan_t2v_distill_dmd_qat
+    run_name: wan_t2v_distill_dmd_qat
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+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.wan.wan_dmd_pipeline.WanDMDPipeline
+    dataset_file: examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json
+    every_steps: 200
+    sampling_steps: [3]
+    sampling_timesteps: [1000, 757, 522]
+    guidance_scale: 6.0
+
+pipeline:
+  flow_shift: 3

examples/train/scenario/qad_wan2_1_mixkit/usage.md

@@ -0,0 +1,26 @@
+# QAD Wan2.1 MixKit Attn-QAT
+
+This scenario runs entirely on the modular `fastvideo/train` stack. The
+student's attention backend is configured per role, so DMD2 can keep the
+teacher and critic on Flash Attention while the student uses the fake-quantized
+Attn-QAT kernel.
+
+From the repository root:
+
+```bash
+# 1. Download the preprocessed MixKit data.
+bash examples/training/finetune/wan_t2v_1.3B/mixkit/download_mixkit_data.sh
+
+# 2. Stage 1: Attn-QAT supervised finetune.
+NUM_GPUS=4 bash examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh
+
+# 3. Export the stage-1 DCP checkpoint to a Diffusers weight file.
+bash examples/train/scenario/qad_wan2_1_mixkit/export_stage1.sh
+
+# 4. Stage 2: three-step Attn-QAT DMD2 distillation.
+NUM_GPUS=4 bash examples/train/scenario/qad_wan2_1_mixkit/run_stage2.sh
+```
+
+The two YAML configs are also directly runnable through `examples/train/run.sh`.
+The wrapper scripts only provide dataset/checkpoint paths and derive distributed
+dimensions from `NUM_GPUS`.

examples/training/finetune/wan_t2v_1.3B/mixkit/README.md

@@ -52,7 +52,7 @@ for the full parameter reference.
 ## Train (QAT finetune)
 
 With the data in place, run the quantization-aware finetune. The 4-bit attention
-path is **config-driven** — selected purely by an env var, no monkey-patching:
+path is **config-driven** and selected by an environment variable:
 
 ```bash
 bash examples/training/finetune/wan_t2v_1.3B/mixkit/finetune_qat.sh
@@ -60,10 +60,19 @@ bash examples/training/finetune/wan_t2v_1.3B/mixkit/finetune_qat.sh
 NUM_GPUS=4 bash .../mixkit/finetune_qat.sh data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset/
 ```
 
-`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` routes attention through the
-fake-quantized Triton kernel (straight-through estimator), so the DiT learns to
-absorb FP4 attention error. This kernel is Triton, so it runs on both `sm_100`
-(B200/GB200) and `sm_120` (RTX 5090).
+`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` keeps the fake-quantized Triton
+forward and backward (straight-through estimator). Both kernels ship in
+`fastvideo-kernel`: SM100 automatically uses the optimized Triton path for the
+production non-causal, head-dimension-128 configuration, while SM120 GPUs such
+as RTX 5090 join the quantized and STE P@V operations and use a shallower
+backward pipeline for long sequences. Set
+`FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV=0` to compare against the split P@V path.
+The script defaults
+`FASTVIDEO_ATTN_QAT_FWD_EXACT_M=0` for throughput; set it to `1` to reproduce
+the previous forward softmax statistic and bitwise-compatible `dV`.
+
+For the website-visible modular SFT-to-DMD2 workflow, see the
+[Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/).
 
 ## Train stage 2 (QAT DMD distillation to 3 steps)
 
@@ -71,8 +80,8 @@ Distill the QAT-finetuned generator down to **3 sampling steps**. Only the
 generator is quantized (Attn-QAT); the teacher (`real_score`) and critic
 (`fake_score`) stay full precision. This is enforced in the loader
 (`component_loader.py`, via the `_loading_teacher_critic_model` flag), so the
-same global `ATTN_QAT_TRAIN` env reaches **only** the generator — no per-model
-flags or monkey-patching.
+same global `ATTN_QAT_TRAIN` env reaches **only** the generator, with no
+per-model flags.
 
 ```bash
 # generator init = the stage-1 finetune checkpoint

examples/training/finetune/wan_t2v_1.3B/mixkit/finetune_qat.sh

@@ -2,20 +2,23 @@
 # QAD recipe — quantization-aware finetune of Wan2.1-T2V-1.3B with fake-quant
 # (Attn-QAT) attention.
 #
-# The 4-bit attention path is selected purely by env var (config-driven, no
-# monkey-patching): FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN routes attention
-# through the fake-quantized Triton kernel (straight-through estimator), so the
-# DiT learns to absorb FP4 attention error instead of fighting it.
+# The 4-bit attention path is selected by env var:
+# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN keeps the fake-quantized Triton
+# forward and backward, so the DiT learns to absorb FP4 attention error instead
+# of fighting it. SM100 selects the optimized kernels; RTX 5090 keeps the
+# previous Triton implementation.
 #
 # Data: run download_mixkit_data.sh first (preprocessed Parquet).
 #
 # Verified end-to-end on Blackwell (GB200/sm_100): the ATTN_QAT_TRAIN backend is
 # selected (not a fallback), forward+backward run, loss/grad are healthy, and
-# validation generates videos. The kernel is Triton so it runs on sm_100 and
-# sm_120 alike (the FP4 inference kernel, by contrast, is sm_120-only).
+# validation generates videos.
 set -euo pipefail
 
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../../.." && pwd)"
+export PYTHONPATH="${REPO_ROOT}/fastvideo-kernel/python${PYTHONPATH:+:${PYTHONPATH}}"
 export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN   # <-- enables Attn-QAT training
+export FASTVIDEO_ATTN_QAT_FWD_EXACT_M=${FASTVIDEO_ATTN_QAT_FWD_EXACT_M:-0}
 export WANDB_MODE=${WANDB_MODE:-online}
 export TOKENIZERS_PARALLELISM=false
 
@@ -23,6 +26,8 @@ MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=${1:-"data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset/"}
 VALIDATION_FILE="$(dirname "$0")/../crush_smol/validation.json"
 NUM_GPUS=${NUM_GPUS:-4}
+MAX_TRAIN_STEPS=${MAX_TRAIN_STEPS:-4000}
+VALIDATION_SAMPLING_STEPS=${VALIDATION_SAMPLING_STEPS:-50}
 
 torchrun --nnodes 1 --nproc_per_node "${NUM_GPUS}" \
     fastvideo/training/wan_training_pipeline.py \
@@ -30,15 +35,15 @@ torchrun --nnodes 1 --nproc_per_node "${NUM_GPUS}" \
     --hsdp_replicate_dim 1 --hsdp_shard_dim "${NUM_GPUS}" \
     --model_path "${MODEL_PATH}" --pretrained_model_name_or_path "${MODEL_PATH}" \
     --data_path "${DATA_DIR}" --dataloader_num_workers 1 \
-    --max_train_steps 2000 --train_batch_size 1 --train_sp_batch_size 1 \
+    --max_train_steps "${MAX_TRAIN_STEPS}" --train_batch_size 1 --train_sp_batch_size 1 \
     --gradient_accumulation_steps 1 \
     --num_latent_t 20 --num_height 480 --num_width 832 --num_frames 77 \
     --enable_gradient_checkpointing_type full \
     --log_validation --validation_dataset_file "${VALIDATION_FILE}" \
-    --validation_steps 200 --validation_sampling_steps 50 --validation_guidance_scale 3.0 \
-    --learning_rate 5e-5 --mixed_precision bf16 --weight_decay 1e-4 --max_grad_norm 1.0 \
+    --validation_steps 50 --validation_sampling_steps "${VALIDATION_SAMPLING_STEPS}" --validation_guidance_scale 5.0 \
+    --learning_rate 1e-6 --mixed_precision bf16 --weight_decay 0.01 --max_grad_norm 1.0 \
     --weight_only_checkpointing_steps 500 --training_state_checkpointing_steps 500 \
     --tracker_project_name wan_t2v_qat_finetune --output_dir checkpoints/wan_t2v_qat_finetune \
     --inference_mode False --training_cfg_rate 0.1 --not_apply_cfg_solver \
-    --dit_precision fp32 --num_euler_timesteps 50 --ema_start_step 0 \
+    --dit_precision fp32 --num_euler_timesteps 50 --ema_start_step 0 --flow_shift 1 \
     --multi_phased_distill_schedule "4000-1"

fastvideo-kernel/README.md

@@ -108,6 +108,33 @@ out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)
 
 ## Benchmark
 
+### Attn-QAT training
+
+The default shape matches one sequence-parallel rank of the 4-GPU
+Wan2.1-T2V-1.3B MixKit recipe (`B=1, H=3, L=31200, D=128`):
+
+```bash
+cd fastvideo-kernel
+python benchmarks/benchmark_attn_qat_train.py
+```
+
+The benchmark reports both conventional attention FLOPs and the extra matrix
+multiplications executed by the QAT straight-through path. Override
+`--peak-tflops` when running on a GPU other than RTX 5090.
+
+The QAT kernel is entirely Triton and routes by architecture at runtime. SM100
+uses a large-tile forward and split 64x64 backward for the production
+non-causal, head-dimension-128 configuration with a 16-aligned KV length. SM120
+(including RTX 5090) keeps the previous tiling but joins the quantized and STE
+P@V operations and uses a shallower backward pipeline for long sequences. Set
+`FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV=0` to compare against the split P@V path.
+Unsupported configurations retain the previous implementation. Set
+`FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED=0` to benchmark that previous path on SM100. Forward tuning is available through
+`FASTVIDEO_ATTN_QAT_FWD_MODE=fast|balanced|reference`; exact reference-order
+softmax statistics are controlled by `FASTVIDEO_ATTN_QAT_FWD_EXACT_M` and are
+disabled by default for maximum throughput. Set it to `1` for reference-order
+statistics and bitwise-compatible `dV`.
+
 ### VSA (block-sparse) TFLOPs
 
 After building/installing `fastvideo-kernel`, run:

fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""Benchmark the Attn-QAT training kernel on a single GPU.
+
+Defaults model one rank of the 4-GPU Wan2.1-T2V-1.3B MixKit recipe:
+``B=1, H=12/4, L=20*30*52, D=128``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import statistics
+import time
+from collections.abc import Callable
+
+import torch
+
+from fastvideo_kernel.triton_kernels.attn_qat_train import attention
+
+
+RTX_5090_DENSE_BF16_TFLOPS = 209.5
+
+
+def _qat_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    consumer_blackwell = torch.cuda.get_device_capability()[0] == 12
+    return attention(
+        q,
+        k,
+        v,
+        False,
+        q.shape[-1]**-0.5,
+        True,  # use_qat_qkv_backward
+        False,  # smooth_k
+        not consumer_blackwell,  # warp_specialize
+        True,  # IS_QAT
+        False,  # two_level_quant_P
+        True,  # fake_quant_P
+        True,  # use_high_prec_o
+        False,  # smooth_q
+        False,  # use_global_sf_P
+        False,  # use_global_sf_QKV
+    )
+
+
+def _measure_ms(fn: Callable[[], object], warmup: int, repeat: int) -> tuple[float, float, float]:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+
+    samples = []
+    for _ in range(repeat):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end))
+    return statistics.median(samples), min(samples), max(samples)
+
+
+def _format_result(
+    label: str,
+    timing_ms: tuple[float, float, float],
+    algorithmic_flops: int,
+    executed_matmul_flops: int,
+    peak_tflops: float,
+) -> str:
+    median_ms, min_ms, max_ms = timing_ms
+    algorithmic_tflops = algorithmic_flops / (median_ms * 1e9)
+    executed_tflops = executed_matmul_flops / (median_ms * 1e9)
+    return (
+        f"{label}: {median_ms:.3f} ms (min={min_ms:.3f}, max={max_ms:.3f}), "
+        f"algorithmic={algorithmic_tflops:.2f} TFLOPS/{100 * algorithmic_tflops / peak_tflops:.2f}% MFU, "
+        f"executed_matmul={executed_tflops:.2f} TFLOPS/{100 * executed_tflops / peak_tflops:.2f}% MFU"
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--heads", type=int, default=3, help="Heads per SP rank; Wan 1.3B has 12 total.")
+    parser.add_argument("--query-length", type=int, default=31_200)
+    parser.add_argument("--kv-length", type=int, help="Defaults to --query-length.")
+    parser.add_argument("--head-dim", type=int, default=128)
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--repeat", type=int, default=10)
+    parser.add_argument(
+        "--peak-tflops",
+        type=float,
+        default=RTX_5090_DENSE_BF16_TFLOPS,
+        help="Dense BF16 Tensor TFLOPS with FP32 accumulation; default is RTX 5090 boost-clock peak.",
+    )
+    args = parser.parse_args()
+    kv_length = args.kv_length or args.query_length
+
+    torch.manual_seed(0)
+    q_shape = (args.batch_size, args.heads, args.query_length, args.head_dim)
+    kv_shape = (args.batch_size, args.heads, kv_length, args.head_dim)
+    q = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    k = torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    v = torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    grad_out = torch.randn_like(q)
+
+    compile_start = time.perf_counter()
+    output = _qat_attention(q, k, v)
+    torch.cuda.synchronize()
+    compile_seconds = time.perf_counter() - compile_start
+
+    forward_ms = _measure_ms(lambda: _qat_attention(q, k, v), args.warmup, args.repeat)
+    backward_ms = _measure_ms(
+        lambda: torch.autograd.grad(output, (q, k, v), grad_out, retain_graph=True),
+        args.warmup,
+        args.repeat,
+    )
+
+    base_flops = args.batch_size * args.heads * args.query_length * kv_length * args.head_dim
+    # Conventional attention FLOPs are 4*base forward and 10*base backward.
+    # QAT additionally computes the STE high-precision P@V path in forward and
+    # the quantized-P dV path in backward, for 6*base and 14*base matmul FLOPs.
+    print(f"device: {torch.cuda.get_device_name()}")
+    print(f"q: {q_shape}; k/v: {kv_shape}; compile+first-forward: {compile_seconds:.3f} s")
+    print(_format_result("forward", forward_ms, 4 * base_flops, 6 * base_flops, args.peak_tflops))
+    print(_format_result("backward", backward_ms, 10 * base_flops, 14 * base_flops, args.peak_tflops))
+
+
+if __name__ == "__main__":
+    main()

fastvideo-kernel/build.sh

@@ -4,6 +4,7 @@ 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)
@@ -67,12 +68,25 @@ fi
 # Install build dependencies
 uv pip install scikit-build-core cmake ninja
 
-RELEASE=0
 GPU_BACKEND=CUDA
-for arg in "$@"; do
-    case "$arg" in
+WHEEL_DIR=""
+while [[ $# -gt 0 ]]; do
+    case "$1" 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
@@ -165,4 +179,15 @@ echo "CMAKE_ARGS: ${CMAKE_ARGS:-<unset>}"
 echo "GPU_BACKEND: ${GPU_BACKEND:-<unset>}"
 # Build and install
 # Use -v for verbose output
-uv pip install . -v --no-build-isolation
+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

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py

@@ -20,14 +20,68 @@ def supports_host_descriptor():
     return is_cuda() and torch.cuda.get_device_capability()[0] >= 9
 
 
+def is_sm100(device=None):
+    return is_cuda() and torch.cuda.get_device_capability(device) == (10, 0)
+
+
 def is_blackwell():
     return is_cuda() and torch.cuda.get_device_capability()[0] == 10
 
 
+def is_consumer_blackwell():
+    return is_cuda() and torch.cuda.get_device_capability()[0] == 12
+
+
 def is_hopper():
     return is_cuda() and torch.cuda.get_device_capability()[0] == 9
 
 
+def _sm100_optimization_enabled():
+    return os.environ.get("FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED", "1") != "0"
+
+
+def _sm100_exact_m_enabled():
+    return os.environ.get("FASTVIDEO_ATTN_QAT_FWD_EXACT_M", "0") != "0"
+
+
+def _consumer_blackwell_join_qat_pv_enabled():
+    """Return whether SM120 uses the joined quantized/STE P@V path."""
+    return os.environ.get("FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV", "1") != "0"
+
+
+def _use_sm100_optimized_qat(
+    device,
+    head_dim: int,
+    causal: bool,
+    is_qat: bool,
+    fake_quant_p: bool,
+    two_level_quant_p: bool,
+    use_global_sf_p: bool,
+) -> bool:
+    """Return whether this call matches the validated SM100 fast path."""
+    return (
+        _sm100_optimization_enabled()
+        and is_sm100(device)
+        and head_dim == 128
+        and not causal
+        and is_qat
+        and fake_quant_p
+        and not two_level_quant_p
+        and not use_global_sf_p
+    )
+
+
+def _select_sm100_forward_config(n_ctx_q: int, n_ctx_kv: int, mode: str):
+    n_ctx = max(n_ctx_q, n_ctx_kv)
+    if mode == "reference":
+        return 32, 32, 4, 4 if n_ctx >= 16_384 else 5
+    if n_ctx <= 2_048:
+        return 32, 32, 4, 5
+    if mode == "balanced":
+        return 64, 32, 4, 4
+    return 128, 128, 8, 3
+
+
 @triton.jit
 def _mul_alpha(acc, alpha, BM: tl.constexpr, BN: tl.constexpr):
     acc0, acc1 = acc.reshape([BM, 2, BN // 2]).permute(0, 2, 1).split()
@@ -47,7 +101,8 @@ def _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid,
                     IS_QAT: tl.constexpr,
                     fake_quant_P: tl.constexpr = True,
                     two_level_quant_P: tl.constexpr = False,
-                    use_global_sf_P: tl.constexpr = True):
+                    use_global_sf_P: tl.constexpr = True,
+                    JOIN_QAT_PV: tl.constexpr = False):
     # range of values handled by this stage (kv blocks)
     if STAGE == 1:
         lo, hi = 0, start_m * BLOCK_M
@@ -113,10 +168,19 @@ def _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid,
             v = desc_v.load([offsetv_y, 0])
             v = tl.where(kv_valid[:, None], v, 0.0)
         p = p.to(dtype)
-        # note that this non transposed v for FP8 is only supported on Blackwell
-        acc = tl.dot(p, v.to(dtype), acc)
-        if IS_QAT:
-            high_prec_acc = tl.dot(high_prec_p, v, high_prec_acc)
+        # Keep the quantized and STE paths in one tensor-core operation. They
+        # share V, so joining along M avoids issuing two small, independent
+        # dot operations for every KV tile.
+        if IS_QAT and JOIN_QAT_PV:
+            joined_p = tl.join(p, high_prec_p).permute(2, 0, 1).reshape([2 * BLOCK_M, BLOCK_N])
+            joined_acc = tl.join(acc, high_prec_acc).permute(2, 0, 1).reshape([2 * BLOCK_M, HEAD_DIM])
+            joined_acc = tl.dot(joined_p, v.to(dtype), joined_acc)
+            acc, high_prec_acc = joined_acc.reshape([2, BLOCK_M, HEAD_DIM]).permute(1, 2, 0).split()
+        else:
+            # note that this non transposed v for FP8 is only supported on Blackwell
+            acc = tl.dot(p, v.to(dtype), acc)
+            if IS_QAT:
+                high_prec_acc = tl.dot(high_prec_p, v, high_prec_acc)
         # update m_i and l_i
         # place this at the end of the loop to reduce register pressure
         l_i = l_i * alpha + l_ij
@@ -204,6 +268,7 @@ def _attn_fwd(sm_scale, M,
               fake_quant_P: tl.constexpr = True,
               two_level_quant_P: tl.constexpr = False,
               use_global_sf_P: tl.constexpr = True,
+              JOIN_QAT_PV: tl.constexpr = False,
               ):
     dtype = tl.float8e5 if FP8_OUTPUT else tl.bfloat16
     tl.static_assert(BLOCK_N <= HEAD_DIM)
@@ -268,7 +333,7 @@ def _attn_fwd(sm_scale, M,
             offset_y_kv, dtype, start_m, qk_scale,
             BLOCK_M, HEAD_DIM, BLOCK_N,
             4 - STAGE, offs_m, offs_n, N_CTX_KV,
-            warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P
+            warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P, JOIN_QAT_PV
         )
     # stage 2: on-band
     if STAGE & 2:
@@ -278,7 +343,7 @@ def _attn_fwd(sm_scale, M,
             offset_y_kv, dtype, start_m, qk_scale,
             BLOCK_M, HEAD_DIM, BLOCK_N,
             2, offs_m, offs_n, N_CTX_KV,
-            warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P
+            warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P, JOIN_QAT_PV
         )
     # epilogue
     m_i += tl.math.log2(l_i)
@@ -292,6 +357,52 @@ def _attn_fwd(sm_scale, M,
         desc_high_prec_o.store([off_hz, start_m * BLOCK_M, 0], high_prec_acc[None, :, :])
 
 
+@triton.jit
+def _attn_fwd_exact_m(
+    desc_q,
+    desc_k,
+    M,
+    sm_scale,
+    N_CTX_Q,
+    N_CTX_KV,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Reproduce the legacy 32x32 forward softmax statistic exactly."""
+    start_m = tl.program_id(0) * BLOCK_M
+    off_hz = tl.program_id(1)
+    offs_m = start_m + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    q_valid = offs_m < N_CTX_Q
+
+    q_base = off_hz * N_CTX_Q
+    kv_base = off_hz * N_CTX_KV
+    q = desc_q.load([q_base + start_m, 0])
+    q = tl.where(q_valid[:, None], q, 0.0)
+
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.full([BLOCK_M], 1.0, tl.float32)
+    qk_scale = sm_scale * 1.44269504
+
+    for start_n in tl.range(0, N_CTX_KV, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        kv_valid = start_n + offs_n < N_CTX_KV
+        k = desc_k.load([kv_base + start_n, 0])
+        k = tl.where(kv_valid[:, None], k, 0.0)
+        qk = tl.dot(q, tl.trans(k))
+        qk = tl.where(kv_valid[None, :], qk, -1.0e6)
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1) * qk_scale)
+        p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
+        l_ij = tl.sum(p.to(tl.bfloat16), axis=1)
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        m_i = m_ij
+
+    m_i += tl.math.log2(l_i)
+    tl.store(M + off_hz * N_CTX_Q + offs_m, m_i, mask=q_valid)
+
+
 @triton.jit
 def _attn_bwd_preprocess(O, DO,
                          Delta,
@@ -835,10 +946,33 @@ class _attention(torch.autograd.Function):
         assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
         assert HEAD_DIM_K in {16, 32, 64, 128, 256}
 
+        # Triton 3.7's NVWS pass aborts for this kernel on Blackwell. Keep the
+        # architecture guard next to the kernel so direct callers and the
+        # FastVideo backend follow the same supported path on sm_100/sm_120.
+        consumer_blackwell = is_consumer_blackwell()
+        blackwell = is_blackwell() or consumer_blackwell
+        warp_specialize = warp_specialize and not blackwell
+
         # Support different sequence lengths for q and k/v (needed for cross attention)
         N_CTX_Q = q.shape[2]  # Query sequence length
         N_CTX_KV = k.shape[2]  # Key/Value sequence length (may differ from query)
         assert k.shape[2] == v.shape[2], "k and v must have the same sequence length"
+        sm100_optimized = (
+            q.dtype == torch.bfloat16
+            and k.dtype == q.dtype
+            and v.dtype == q.dtype
+            and k.device == q.device
+            and v.device == q.device
+            and _use_sm100_optimized_qat(
+                q.device,
+                HEAD_DIM_K,
+                causal,
+                IS_QAT,
+                fake_quant_P,
+                two_level_quant_P,
+                use_global_sf_P,
+            )
+        )
 
         # smoothing k from SageAttn
         ctx.k_mean = None
@@ -924,7 +1058,20 @@ class _attention(torch.autograd.Function):
             else:
                 extra_kern_args["maxnreg"] = 80
 
-        BLOCK_M, BLOCK_N = 32, 32
+        qkv_block_m, qkv_block_n = 32, 32
+        fwd_block_m, fwd_block_n = 32, 32
+        fwd_num_warps, fwd_num_stages = 4, 2
+        fwd_mode = "legacy"
+        if sm100_optimized:
+            fwd_mode = os.environ.get("FASTVIDEO_ATTN_QAT_FWD_MODE", "fast").lower()
+            if fwd_mode not in {"fast", "balanced", "reference"}:
+                raise ValueError(
+                    f"FASTVIDEO_ATTN_QAT_FWD_MODE={fwd_mode!r} "
+                    "(want fast|balanced|reference)"
+                )
+            fwd_block_m, fwd_block_n, fwd_num_warps, fwd_num_stages = _select_sm100_forward_config(
+                N_CTX_Q, N_CTX_KV, fwd_mode
+            )
         if IS_QAT:
             fake_q = torch.empty_like(q)
             fake_k = torch.empty_like(k)
@@ -942,8 +1089,8 @@ class _attention(torch.autograd.Function):
                 desc_v = fake_v
 
             H = q.shape[1]
-            grid_1 = (triton.cdiv(q.shape[2], BLOCK_M), q.shape[0] * q.shape[1], 1)
-            grid_2 = (triton.cdiv(k.shape[2], BLOCK_N), q.shape[0] * q.shape[1], 1)
+            grid_1 = (triton.cdiv(q.shape[2], qkv_block_m), q.shape[0] * q.shape[1], 1)
+            grid_2 = (triton.cdiv(k.shape[2], qkv_block_n), q.shape[0] * q.shape[1], 1)
 
             fake_quantize_q[grid_1](
                 q, fake_q,
@@ -952,7 +1099,7 @@ class _attention(torch.autograd.Function):
                 fake_q.stride(0), fake_q.stride(1),
                 fake_q.stride(2), fake_q.stride(3),
                 H, N_CTX_Q,
-                BLOCK_M=BLOCK_M, HEAD_DIM=HEAD_DIM_K,
+                BLOCK_M=qkv_block_m, HEAD_DIM=HEAD_DIM_K,
                 use_global_sf=use_global_sf_QKV,
             )
             fake_quantize_kv[grid_2](
@@ -962,14 +1109,14 @@ class _attention(torch.autograd.Function):
                 fake_k.stride(0), fake_k.stride(1),
                 fake_k.stride(2), fake_k.stride(3),
                 H, N_CTX_KV,
-                BLOCK_N=BLOCK_N, HEAD_DIM=HEAD_DIM_K, 
+                BLOCK_N=qkv_block_n, HEAD_DIM=HEAD_DIM_K,
                 use_global_sf=use_global_sf_QKV,
             )
 
         # Apply pre-hook to set block shapes on tensor descriptors
         _host_descriptor_pre_hook({
-            "BLOCK_M": BLOCK_M,
-            "BLOCK_N": BLOCK_N,
+            "BLOCK_M": fwd_block_m,
+            "BLOCK_N": fwd_block_n,
             "HEAD_DIM": HEAD_DIM_K,
             "desc_q": desc_q,
             "desc_k": desc_k,
@@ -986,7 +1133,7 @@ class _attention(torch.autograd.Function):
             N_CTX_Q=N_CTX_Q,
             N_CTX_KV=N_CTX_KV,
             HEAD_DIM=HEAD_DIM_K,
-            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
+            BLOCK_M=fwd_block_m, BLOCK_N=fwd_block_n,
             FP8_OUTPUT=q.dtype == torch.float8_e5m2,
             STAGE=stage,
             warp_specialize=warp_specialize,
@@ -995,10 +1142,40 @@ class _attention(torch.autograd.Function):
             fake_quant_P=fake_quant_P,
             two_level_quant_P=two_level_quant_P,
             use_global_sf_P=use_global_sf_P,
-            num_warps=4,
-            num_stages=2,
+            JOIN_QAT_PV=(consumer_blackwell and _consumer_blackwell_join_qat_pv_enabled()),
+            num_warps=fwd_num_warps,
+            num_stages=fwd_num_stages,
             **extra_kern_args
         )
+
+        exact_m = _sm100_exact_m_enabled()
+        if (
+            sm100_optimized
+            and fwd_mode != "reference"
+            and exact_m
+            and (fwd_block_m != 32 or fwd_block_n != 32)
+        ):
+            # The large forward tile changes a legal reduction order. Restore
+            # the legacy statistic so dV remains bitwise-compatible while the
+            # two output paths retain the faster large-tile PV computation.
+            assert isinstance(desc_q, TensorDescriptor)
+            assert isinstance(desc_k, TensorDescriptor)
+            desc_q.block_shape = [32, HEAD_DIM_K]
+            desc_k.block_shape = [32, HEAD_DIM_K]
+            stats_grid = (triton.cdiv(N_CTX_Q, 32), q.shape[0] * q.shape[1])
+            _attn_fwd_exact_m[stats_grid](
+                desc_q,
+                desc_k,
+                M,
+                sm_scale,
+                N_CTX_Q,
+                N_CTX_KV,
+                HEAD_DIM=HEAD_DIM_K,
+                BLOCK_M=32,
+                BLOCK_N=32,
+                num_warps=8,
+                num_stages=4,
+            )
         o_for_bwd = high_prec_o if IS_QAT and use_high_prec_o else o
 
         if IS_QAT:
@@ -1018,6 +1195,7 @@ class _attention(torch.autograd.Function):
         ctx.smooth_q = smooth_q
         ctx.use_global_sf_P = use_global_sf_P
         ctx.warp_specialize = warp_specialize
+        ctx.sm100_optimized = sm100_optimized
         return o
 
     @staticmethod
@@ -1032,7 +1210,10 @@ class _attention(torch.autograd.Function):
         N_CTX_KV = k.shape[2]
         assert k.shape[2] == v.shape[2], "k and v must have the same sequence length"
         PRE_BLOCK = 128
-        NUM_STAGES = 3
+        # Long video sequences are occupancy-bound on consumer Blackwell: a
+        # third software-pipeline stage consumes shared memory without hiding
+        # additional latency. Shorter sequences retain the deeper pipeline.
+        NUM_STAGES = 2 if is_consumer_blackwell() and max(N_CTX_Q, N_CTX_KV) >= 8192 else 3
         NUM_WARPS = 4
         BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 32, 32, 32
         if not ctx.use_qat_qkv_backward:
@@ -1057,7 +1238,85 @@ class _attention(torch.autograd.Function):
             # _, q_m = triton_group_mean(q)
             q_m = q_m.repeat_interleave(q.shape[2] // q_m.shape[2], dim=2)  # B,H,L,D
 
-        if N_CTX_Q == N_CTX_KV:
+        sm100_optimized_backward = (
+            getattr(ctx, "sm100_optimized", False)
+            and ctx.use_qat_qkv_backward
+            and not ctx.smooth_k
+            and not ctx.smooth_q
+            and N_CTX_KV % 16 == 0
+        )
+        if sm100_optimized_backward:
+            # Keeping dQ and dK/dV in separate programs allows 64x64 tiles
+            # without carrying all three fp32 accumulators at once. On SM100
+            # this is substantially faster than the legacy 32x32 combined
+            # self-attention program with the same math and BF16 parity bounds.
+            block_m, block_n = 64, 64
+            grid_dq = ((N_CTX_Q + block_m - 1) // block_m, 1, BATCH * N_HEAD)
+            _attn_bwd_dq_cross[grid_dq](
+                q,
+                arg_k,
+                v,
+                ctx.sm_scale,
+                do,
+                dq,
+                M,
+                delta,
+                q.stride(0),
+                k.stride(0),
+                q.stride(1),
+                k.stride(1),
+                q.stride(2),
+                k.stride(2),
+                q.stride(3),
+                k.stride(3),
+                N_HEAD,
+                N_CTX_Q,
+                N_CTX_KV,
+                ctx.k_mean,
+                BLOCK_M2=block_m,
+                BLOCK_N2=block_n,
+                HEAD_DIM=ctx.HEAD_DIM,
+                SMOOTH_K=False,
+                warp_specialize=False,
+                num_warps=8,
+                num_stages=2,
+            )
+            grid_dkdv = ((N_CTX_KV + block_n - 1) // block_n, 1, BATCH * N_HEAD)
+            _attn_bwd_dkdv_cross[grid_dkdv](
+                q,
+                arg_k,
+                v,
+                ctx.sm_scale,
+                do,
+                dk,
+                dv,
+                M,
+                delta,
+                q_m,
+                q.stride(0),
+                k.stride(0),
+                q.stride(1),
+                k.stride(1),
+                q.stride(2),
+                k.stride(2),
+                q.stride(3),
+                k.stride(3),
+                N_HEAD,
+                N_CTX_Q,
+                N_CTX_KV,
+                BLOCK_M1=block_m,
+                BLOCK_N1=block_n,
+                HEAD_DIM=ctx.HEAD_DIM,
+                IS_QAT=True,
+                two_level_quant_P=False,
+                fake_quant_P=True,
+                SMOOTH_Q=False,
+                use_global_sf_P=False,
+                warp_specialize=False,
+                num_warps=8,
+                num_stages=3,
+            )
+        elif N_CTX_Q == N_CTX_KV:
             # Use existing kernel for self-attention (same sequence lengths)
             grid = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD)
             _attn_bwd[grid](
@@ -1074,10 +1333,10 @@ class _attention(torch.autograd.Function):
                 IS_QAT=ctx.IS_QAT,
                 SMOOTH_K=ctx.smooth_k,
                 two_level_quant_P=ctx.two_level_quant_P,
-            fake_quant_P=ctx.fake_quant_P,
-            SMOOTH_Q=ctx.smooth_q,
-            use_global_sf_P=ctx.use_global_sf_P,
-            warp_specialize=ctx.warp_specialize,
+                fake_quant_P=ctx.fake_quant_P,
+                SMOOTH_Q=ctx.smooth_q,
+                use_global_sf_P=ctx.use_global_sf_P,
+                warp_specialize=ctx.warp_specialize,
                 num_warps=NUM_WARPS,
                 num_stages=NUM_STAGES
             )

fastvideo-kernel/tests/test_attn_qat_train_routing.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+
+import pytest
+import torch
+
+from fastvideo_kernel.triton_kernels import attn_qat_train as kernel
+
+
+def _production_route_kwargs():
+    return {
+        "device": torch.device("cuda"),
+        "head_dim": 128,
+        "causal": False,
+        "is_qat": True,
+        "fake_quant_p": True,
+        "two_level_quant_p": False,
+        "use_global_sf_p": False,
+    }
+
+
+def test_sm100_production_configuration_uses_optimized_route(monkeypatch):
+    monkeypatch.setattr(kernel, "is_sm100", lambda device=None: True)
+    monkeypatch.delenv("FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED", raising=False)
+
+    assert kernel._use_sm100_optimized_qat(**_production_route_kwargs())
+
+
+@pytest.mark.parametrize(
+    ("override", "value"),
+    [
+        ("head_dim", 64),
+        ("causal", True),
+        ("is_qat", False),
+        ("fake_quant_p", False),
+        ("two_level_quant_p", True),
+        ("use_global_sf_p", True),
+    ],
+)
+def test_unsupported_configuration_keeps_legacy_route(monkeypatch, override, value):
+    monkeypatch.setattr(kernel, "is_sm100", lambda device=None: True)
+    kwargs = _production_route_kwargs()
+    kwargs[override] = value
+
+    assert not kernel._use_sm100_optimized_qat(**kwargs)
+
+
+def test_non_sm100_and_debug_switch_keep_legacy_route(monkeypatch):
+    kwargs = _production_route_kwargs()
+    monkeypatch.setattr(kernel, "is_sm100", lambda device=None: False)
+    assert not kernel._use_sm100_optimized_qat(**kwargs)
+
+    monkeypatch.setattr(kernel, "is_sm100", lambda device=None: True)
+    monkeypatch.setenv("FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED", "0")
+    assert not kernel._use_sm100_optimized_qat(**kwargs)
+
+
+def test_exact_m_is_opt_in(monkeypatch):
+    monkeypatch.delenv("FASTVIDEO_ATTN_QAT_FWD_EXACT_M", raising=False)
+    assert not kernel._sm100_exact_m_enabled()
+
+    monkeypatch.setenv("FASTVIDEO_ATTN_QAT_FWD_EXACT_M", "1")
+    assert kernel._sm100_exact_m_enabled()
+
+
+def test_sm120_joined_pv_is_enabled_by_default_and_can_be_disabled(monkeypatch):
+    monkeypatch.delenv("FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV", raising=False)
+    assert kernel._consumer_blackwell_join_qat_pv_enabled()
+
+    monkeypatch.setenv("FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV", "0")
+    assert not kernel._consumer_blackwell_join_qat_pv_enabled()
+
+
+@pytest.mark.parametrize(
+    ("n_ctx", "mode", "expected"),
+    [
+        (2_048, "fast", (32, 32, 4, 5)),
+        (4_096, "fast", (128, 128, 8, 3)),
+        (4_096, "balanced", (64, 32, 4, 4)),
+        (4_096, "reference", (32, 32, 4, 5)),
+        (31_200, "reference", (32, 32, 4, 4)),
+    ],
+)
+def test_sm100_forward_config_selection(n_ctx, mode, expected):
+    assert kernel._select_sm100_forward_config(n_ctx, n_ctx, mode) == expected
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0),
+    reason="SM100 parity test",
+)
+@pytest.mark.parametrize(("q_length", "kv_length"), [(2_112, 2_112), (2_112, 2_080)])
+def test_sm100_optimized_forward_backward_matches_legacy(monkeypatch, q_length, kv_length):
+    torch.manual_seed(7)
+    q_shape = (1, 1, q_length, 128)
+    kv_shape = (1, 1, kv_length, 128)
+    inputs = [
+        torch.randn(q_shape, device="cuda", dtype=torch.bfloat16),
+        torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16),
+        torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16),
+    ]
+    grad_out = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16)
+    flags = (
+        True,  # use_qat_qkv_backward
+        False,  # smooth_k
+        True,  # warp_specialize (disabled internally on Blackwell)
+        True,  # IS_QAT
+        False,  # two_level_quant_P
+        True,  # fake_quant_P
+        True,  # use_high_prec_o
+        False,  # smooth_q
+        False,  # use_global_sf_P
+        False,  # use_global_sf_QKV
+    )
+
+    def run(optimized: bool):
+        monkeypatch.setenv("FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED", "1" if optimized else "0")
+        monkeypatch.setenv("FASTVIDEO_ATTN_QAT_FWD_MODE", "fast")
+        monkeypatch.setenv("FASTVIDEO_ATTN_QAT_FWD_EXACT_M", "1")
+        q, k, v = [tensor.clone().requires_grad_(True) for tensor in inputs]
+        output = kernel.attention(
+            q,
+            k,
+            v,
+            False,
+            1.0 / math.sqrt(q_shape[-1]),
+            *flags,
+        )
+        output.backward(grad_out)
+        return output.detach(), q.grad, k.grad, v.grad
+
+    legacy = run(False)
+    optimized = run(True)
+
+    assert (optimized[0].float() - legacy[0].float()).abs().max().item() <= 1e-2
+    assert (optimized[1].float() - legacy[1].float()).abs().max().item() <= 4e-3
+    assert (optimized[2].float() - legacy[2].float()).abs().max().item() <= 4e-3
+    assert torch.equal(optimized[3], legacy[3])
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12,
+    reason="SM120 parity test",
+)
+@pytest.mark.parametrize(("q_length", "kv_length"), [(2_112, 2_112), (2_112, 2_080)])
+def test_sm120_joined_pv_forward_backward_matches_split_path(monkeypatch, q_length, kv_length):
+    torch.manual_seed(11)
+    q_shape = (1, 1, q_length, 128)
+    kv_shape = (1, 1, kv_length, 128)
+    inputs = [
+        torch.randn(q_shape, device="cuda", dtype=torch.bfloat16),
+        torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16),
+        torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16),
+    ]
+    grad_out = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16)
+    flags = (
+        True,  # use_qat_qkv_backward
+        False,  # smooth_k
+        True,  # warp_specialize (disabled internally on Blackwell)
+        True,  # IS_QAT
+        False,  # two_level_quant_P
+        True,  # fake_quant_P
+        True,  # use_high_prec_o
+        False,  # smooth_q
+        False,  # use_global_sf_P
+        False,  # use_global_sf_QKV
+    )
+
+    def run(joined_pv: bool):
+        monkeypatch.setenv("FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV", "1" if joined_pv else "0")
+        q, k, v = [tensor.clone().requires_grad_(True) for tensor in inputs]
+        output = kernel.attention(
+            q,
+            k,
+            v,
+            False,
+            1.0 / math.sqrt(q_shape[-1]),
+            *flags,
+        )
+        output.backward(grad_out)
+        return output.detach(), q.grad, k.grad, v.grad
+
+    split = run(False)
+    joined = run(True)
+
+    assert torch.equal(joined[0], split[0])
+    assert torch.equal(joined[1], split[1])
+    assert torch.equal(joined[2], split[2])
+    assert torch.equal(joined[3], split[3])

fastvideo/api/compat.py

@@ -304,6 +304,8 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An
         for key in _LTX2_REFINE_FLAT_KEYS:
             if key in refine:
                 kwargs[f"ltx2_refine_{key}"] = refine[key]
+        if "enabled" in refine:
+            kwargs["refine_enabled"] = refine["enabled"]
     kwargs.update(preset_overrides)
     kwargs.update(deepcopy(normalized.pipeline.experimental))
     return FastVideoArgs.from_kwargs(**kwargs)

fastvideo/api/sampling_param.py

@@ -51,8 +51,9 @@ class SamplingParam:
     gt_latents: Any | None = None  # Ground truth latents [B, 16, T, H, W]
     conditioning_mask: Any | None = None  # Mask [B, 1, T, H, W]
 
-    # Camera control inputs (LingBotWorld)
+    # Camera control inputs (LingBotWorld and LingBotWorld2)
     c2ws_plucker_emb: Any | None = None  # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
+    action_path: str | None = None  # Directory containing poses.npy and intrinsics.npy
 
     # Refine inputs (LongCat 480p->720p upscaling)
     # Path-based refine (load stage1 video from disk, e.g. MP4)
@@ -89,7 +90,13 @@ class SamplingParam:
     num_inference_steps: int = 50
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
+    batch_cfg: bool = False
     guidance_scale_2: float | None = None
+    # Z-Image CFG controls. ``cfg_normalization=True`` caps the guided
+    # prediction norm at the positive-prediction norm; ``cfg_truncation``
+    # disables CFG above the normalized-noise threshold.
+    cfg_normalization: bool = False
+    cfg_truncation: float | None = 1.0
     # Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG.
     use_embedded_guidance: bool = False
     # Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding).
@@ -323,6 +330,24 @@ class SamplingParam:
             default=SamplingParam.guidance_scale,
             help="Classifier-free guidance scale",
         )
+        parser.add_argument(
+            "--cfg-normalization",
+            action=StoreBoolean,
+            default=SamplingParam.cfg_normalization,
+            help="Cap Z-Image CFG prediction norm to the positive-prediction norm",
+        )
+        parser.add_argument(
+            "--cfg-truncation",
+            type=float,
+            default=SamplingParam.cfg_truncation,
+            help="Disable Z-Image CFG above this normalized-noise threshold",
+        )
+        parser.add_argument(
+            "--batch-cfg",
+            action=StoreBoolean,
+            default=SamplingParam.batch_cfg,
+            help="Evaluate conditional and unconditional CFG branches in one batch",
+        )
         parser.add_argument(
             "--guidance-rescale",
             type=float,

fastvideo/api/schema.py

@@ -130,6 +130,7 @@ class InputConfig:
     keyboard_cond: Any | None = None
     grid_sizes: Any | None = None
     c2ws_plucker_emb: Any | None = None
+    action_path: str | None = None
     refine_from: str | None = None
     stage1_video: Any | None = None
 
@@ -138,6 +139,7 @@ class InputConfig:
 class SamplingConfig:
     num_videos_per_prompt: int = 1
     seed: int = 1024
+    max_sequence_length: int | None = None
     num_frames: int = 125
     height: int = 720
     width: int = 1280
@@ -147,7 +149,10 @@ class SamplingConfig:
     num_inference_steps: int = 50
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
+    batch_cfg: bool = False
     guidance_scale_2: float | None = None
+    cfg_normalization: bool = False
+    cfg_truncation: float | None = 1.0
     guidance_rescale: float = 0.0
     true_cfg_scale: float | None = None
     use_embedded_guidance: bool | None = None

fastvideo/attention/backends/attn_qat_infer.py

@@ -50,8 +50,143 @@ 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:
-    return _get_attn_qat_infer() is not None
+    """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
 
 
 class AttnQatInferBackend(AttentionBackend):
@@ -97,6 +232,12 @@ 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,
@@ -105,10 +246,18 @@ 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()
-        if attn_qat_infer is None:
-            raise ImportError("attn_qat_infer is not available. Please ensure the "
-                              "attn_qat_infer kernel package is installed.")
+        assert attn_qat_infer is not None  # kernel == "cutlass_sm12x" implies the import succeeded
 
         query = query.transpose(1, 2).contiguous()
         key = key.transpose(1, 2).contiguous()
@@ -123,3 +272,39 @@ 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/attn_qat_train.py

@@ -18,14 +18,14 @@ from fastvideo.logger import init_logger
 logger = init_logger(__name__)
 
 _project_root = Path(__file__).resolve().parent.parent.parent.parent
-_kernel_root = _project_root / "fastvideo-kernel"
-_kernel_python_root = _kernel_root / "python"
+_kernel_python_root = _project_root / "fastvideo-kernel" / "python"
 _attn_qat_train_attention: Callable[..., torch.Tensor] | None = None
 _attn_qat_train_import_attempted = False
+_attn_qat_train_import_error: ImportError | None = None
 
 
 def _ensure_kernel_paths() -> None:
-    for path in (_project_root, _kernel_root, _kernel_python_root):
+    for path in (_project_root, _kernel_python_root):
         path_str = str(path)
         if path_str not in sys.path:
             sys.path.insert(0, path_str)
@@ -34,6 +34,7 @@ def _ensure_kernel_paths() -> None:
 def _get_attn_qat_train_attention() -> Callable[..., torch.Tensor] | None:
     global _attn_qat_train_attention
     global _attn_qat_train_import_attempted
+    global _attn_qat_train_import_error
 
     if _attn_qat_train_import_attempted:
         return _attn_qat_train_attention
@@ -42,8 +43,11 @@ def _get_attn_qat_train_attention() -> Callable[..., torch.Tensor] | None:
     _ensure_kernel_paths()
 
     try:
-        _attn_qat_train_attention = importlib.import_module("fastvideo_kernel.triton_kernels.attn_qat_train").attention
-    except ImportError:
+        triton_qat = importlib.import_module("fastvideo_kernel.triton_kernels.attn_qat_train")
+        _attn_qat_train_attention = triton_qat.attention
+        logger.info("ATTN_QAT_TRAIN loaded FastVideo's architecture-optimized Triton kernel")
+    except ImportError as exc:
+        _attn_qat_train_import_error = exc
         _attn_qat_train_attention = None
 
     return _attn_qat_train_attention
@@ -60,8 +64,9 @@ def attn_qat_train(q_BLHD: torch.Tensor,
                    sm_scale: float | None = None) -> torch.Tensor:
     attention = _get_attn_qat_train_attention()
     if attention is None:
-        raise ImportError("fastvideo_kernel.triton_kernels.attn_qat_train is not available. "
-                          "Please ensure the FastVideo kernel package is installed.")
+        detail = f" Original import error: {_attn_qat_train_import_error}" if _attn_qat_train_import_error else ""
+        raise ImportError("ATTN_QAT_TRAIN requires FastVideo's fastvideo-kernel package. Install it or make "
+                          f"fastvideo-kernel/python importable.{detail}")
 
     q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
     k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -69,7 +74,11 @@ def attn_qat_train(q_BLHD: torch.Tensor,
 
     use_qat_qkv_backward = True
     smooth_k = False
-    warp_specialize = True
+    # Triton 3.7's NVWS pass aborts while compiling this kernel on Blackwell.
+    # The kernel has a supported non-warp-specialized path, so use it on both
+    # datacenter (sm_100) and consumer (sm_120) Blackwell GPUs.
+    capability_major = torch.cuda.get_device_capability()[0]
+    warp_specialize = capability_major not in (10, 12)
     is_qat = True
     two_level_quant_p_sage3 = False
     fake_quant_p_bwd = True
@@ -106,7 +115,7 @@ class AttnQatTrainBackend(AttentionBackend):
 
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
-        return [64, 96, 128, 160, 192, 224, 256]
+        return [128]
 
     @staticmethod
     def get_name() -> str:

fastvideo/attention/backends/flash_attn.py

@@ -32,8 +32,30 @@ except ImportError:
     flash_attn_fp4_func = None
     _FA4_FP4_AVAILABLE = False
 
+_FA4_QUANT_OPS: tuple | None = None
 
-def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]:
+
+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]:
     """Quantize a (batch, seqlen, nheads, headdim) BF16 tensor to FP4.
 
     Returns:
@@ -42,7 +64,7 @@ def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor, ) -> tuple[torch.Tensor, to
             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
     """
-    from flashinfer.quantization import nvfp4_quantize, SfLayout
+    nvfp4_quantize, SfLayout = _resolve_fa4_quant_ops()
 
     batch, seqlen, nheads, headdim = tensor_4d.shape
     sf_vec_size = 16
@@ -82,6 +104,43 @@ def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor, ) -> tuple[torch.Tensor, to
     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
 
@@ -180,8 +239,11 @@ class FlashAttentionImpl(AttentionImpl):
         # SP). Cast through bf16 and restore, matching TORCH_SDPA's tolerance.
         orig_dtype = query.dtype
         if orig_dtype not in (torch.float16, torch.bfloat16):
-            logger.warning_once(f"FLASH_ATTN received {orig_dtype} inputs; casting to "
-                                f"bfloat16 for the kernel and restoring on output.")
+            # Keep logging (and its Python lru_cache) out of compiled traces
+            # so that the AOTAutograd cache can serialize.
+            if not torch.compiler.is_compiling():
+                logger.warning_once(f"FLASH_ATTN received {orig_dtype} inputs; casting to "
+                                    f"bfloat16 for the kernel and restoring on output.")
             query = query.to(torch.bfloat16)
             key = key.to(torch.bfloat16)
             value = value.to(torch.bfloat16)

fastvideo/attention/selector.py

@@ -32,6 +32,27 @@ def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None:
           None
 
 
+def coerce_attn_backend(attn_backend: AttentionBackendEnum | str | None, ) -> AttentionBackendEnum | None:
+    """Normalize an explicit backend selection.
+
+    Environment-variable parsing remains permissive via
+    :func:`backend_name_to_enum`, but typed/config-driven call sites should
+    fail fast on typos instead of silently falling back to another backend.
+    """
+    if attn_backend is None or isinstance(attn_backend, AttentionBackendEnum):
+        return attn_backend
+    if not isinstance(attn_backend, str) or not attn_backend.strip():
+        raise ValueError("attention backend must be a non-empty string, "
+                         f"an AttentionBackendEnum, or None; got {attn_backend!r}")
+
+    backend_name = attn_backend.strip().upper()
+    backend = backend_name_to_enum(backend_name)
+    if backend is None:
+        raise ValueError(f"Unknown attention backend {attn_backend!r}. "
+                         f"Expected one of {sorted(AttentionBackendEnum.__members__)}")
+    return backend
+
+
 def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
     '''
     Get the backend override specified by the FastVideo attention
@@ -69,6 +90,11 @@ def global_force_attn_backend(attn_backend: AttentionBackendEnum | None) -> None
     '''
     global forced_attn_backend
     forced_attn_backend = attn_backend
+    # Backend selection is cached by tensor shape/dtype, while the global
+    # override is intentionally not part of that cache key. Invalidate cached
+    # resolutions whenever the override changes so independently constructed
+    # role models can bind different attention implementations.
+    _cached_get_attn_backend.cache_clear()
 
 
 def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
@@ -124,14 +150,15 @@ 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 automatic "
-            "selection.",
+            "layer; supported backends are %s. Falling back to %s.",
             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 = None
+        selected_backend = fallback_backend
     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/attention/utils/flash_attn_cute.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import functools
 from collections.abc import Callable
 
 import torch
@@ -43,19 +42,26 @@ except ImportError:
     _flash_attn_2_func = None
     _flash_attn_2_varlen_func = None
 
+# Dynamo unwraps functools caches, so keep this cache inside the opaque helper.
+_SM90_OR_NEWER_BY_DEVICE: dict[int, bool] = {}
+
 
 def _check_dropout(dropout_p: float) -> None:
     if dropout_p != 0.0:
         raise NotImplementedError(f"flash_attn.cute does not support dropout (got dropout_p={dropout_p})")
 
 
-@functools.cache
-def _sm90_or_newer() -> bool:
-    return current_platform.has_device_capability(90)
+@torch.compiler.assume_constant_result
+def _sm90_or_newer(device_id: int) -> bool:
+    if device_id not in _SM90_OR_NEWER_BY_DEVICE:
+        _SM90_OR_NEWER_BY_DEVICE[device_id] = (current_platform.has_device_capability(90, device_id))
+    return _SM90_OR_NEWER_BY_DEVICE[device_id]
 
 
 def _use_fa2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool:
-    if _sm90_or_newer():
+    device_id = q.device.index
+    assert device_id is not None
+    if _sm90_or_newer(device_id):
         return False
     # Pre-sm90 FA4 cute limitations, both served by FA2 (deterministic
     # capability gate, not a runtime fallback):

fastvideo/configs/models/dits/__init__.py

@@ -12,12 +12,16 @@ from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
+from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
+from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig
+from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
     "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
     "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
-    "StableAudioConfig", "GlmImageDiTConfig"
+    "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig",
+    "ZImageDiTConfig"
 ]

fastvideo/configs/models/dits/kandinsky5.py

@@ -13,12 +13,15 @@ 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, plus
-    # the dense backends every DiT supports.
+    # 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/).
     _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/dits/lingbot_video.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Architecture configuration for LingBot-Video Dense and MoE DiTs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def _is_lingbot_video_block(name: str, module: object) -> bool:
+    """Select top-level transformer blocks for FSDP and compilation."""
+    del module
+    parts = name.split(".")
+    return len(parts) == 2 and parts[0] == "blocks" and parts[1].isdigit()
+
+
+@dataclass
+class LingBotVideoArchConfig(DiTArchConfig):
+    """One-to-one representation of the released transformer config JSON."""
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_lingbot_video_block])
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
+        AttentionBackendEnum.TORCH_SDPA,
+        AttentionBackendEnum.FLASH_ATTN,
+    )
+    param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"})
+
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    in_channels: int = 16
+    out_channels: int = 16
+    hidden_size: int = 2048
+    num_attention_heads: int = 16
+    depth: int = 24
+    intermediate_size: int = 6144
+    text_dim: int = 2560
+    freq_dim: int = 256
+    norm_eps: float = 1e-6
+    rope_theta: float = 256.0
+    axes_dims: tuple[int, int, int] = (32, 48, 48)
+    axes_lens: tuple[int, int, int] = (8192, 1024, 1024)
+    qkv_bias: bool = False
+    out_bias: bool = True
+    patch_embed_bias: bool = True
+    timestep_mlp_bias: bool = True
+    num_experts: int = 0
+    num_experts_per_tok: int = 8
+    moe_intermediate_size: int = 512
+    decoder_sparse_step: int = 1
+    mlp_only_layers: tuple[int, ...] = ()
+    n_shared_experts: int | None = None
+    score_func: str = "sigmoid"
+    norm_topk_prob: bool = True
+    n_group: int | None = None
+    topk_group: int | None = None
+    routed_scaling_factor: float = 1.0
+
+    def __post_init__(self) -> None:
+        """Populate FastVideo loader fields from the released architecture."""
+        super().__post_init__()
+        self.num_channels_latents = self.in_channels
+        self.attention_head_dim = self.hidden_size // self.num_attention_heads
+
+
+@dataclass
+class LingBotVideoConfig(DiTConfig):
+    """FastVideo component configuration for LingBot-Video transformers."""
+
+    arch_config: DiTArchConfig = field(default_factory=LingBotVideoArchConfig)

fastvideo/configs/models/dits/lingbotworld2.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class LingBotWorld2CausalFastArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+    param_names_mapping: dict = field(default_factory=dict)
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    model_type: str = "i2v"
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len: int = 512
+    in_dim: int = 36
+    dim: int = 5120
+    ffn_dim: int = 13824
+    freq_dim: int = 256
+    text_dim: int = 4096
+    out_dim: int = 16
+    num_heads: int = 40
+    num_layers: int = 40
+    qk_norm: bool = True
+    cross_attn_norm: bool = True
+    eps: float = 1e-6
+
+    local_attn_size: int = 18
+    sink_size: int = 6
+    chunk_size: int = 4
+    sample_shift: float = 10.0
+    num_train_timesteps: int = 1000
+    timesteps_index: tuple[int, int, int, int] = (0, 250, 500, 750)
+    max_area: int = 480 * 832
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.dim
+        self.num_attention_heads = self.num_heads
+        self.attention_head_dim = self.dim // self.num_heads
+        self.in_channels = self.in_dim
+        self.out_channels = self.out_dim
+        self.num_channels_latents = self.out_dim
+
+
+@dataclass
+class LingBotWorld2CausalFastVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=LingBotWorld2CausalFastArchConfig)
+
+    prefix: str = "Wan"

fastvideo/configs/models/dits/zimage.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def is_zimage_block(name: str, module) -> bool:
+    parts = name.split(".")
+    return len(parts) >= 2 and parts[-2] in {"noise_refiner", "context_refiner", "layers"} and parts[-1].isdigit()
+
+
+@dataclass
+class ZImageDiTArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_block])
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.TORCH_SDPA, )
+
+    all_patch_size: tuple[int, ...] = (2, )
+    all_f_patch_size: tuple[int, ...] = (1, )
+    in_channels: int = 16
+    dim: int = 3840
+    n_layers: int = 30
+    n_refiner_layers: int = 2
+    n_heads: int = 30
+    n_kv_heads: int = 30
+    norm_eps: float = 1e-5
+    qk_norm: bool = True
+    cap_feat_dim: int = 2560
+    rope_theta: float = 256.0
+    t_scale: float = 1000.0
+    axes_dims: tuple[int, ...] = (32, 48, 48)
+    axes_lens: tuple[int, ...] = (1536, 512, 512)
+
+    adaln_embed_dim: int = 256
+    frequency_embedding_size: int = 256
+    timestep_mid_size: int = 1024
+    max_period: int = 10000
+    seq_multi_of: int = 32
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if len(self.all_patch_size) != len(self.all_f_patch_size):
+            raise ValueError("all_patch_size and all_f_patch_size must have equal length")
+        if self.dim % self.n_heads:
+            raise ValueError("dim must be divisible by n_heads")
+        if self.dim // self.n_heads != sum(self.axes_dims):
+            raise ValueError("attention head dimension must equal sum(axes_dims)")
+        if len(self.axes_dims) != len(self.axes_lens) or any(dim % 2 for dim in self.axes_dims):
+            raise ValueError("RoPE axes require matching lengths and even dimensions")
+
+        self.hidden_size = self.dim
+        self.num_attention_heads = self.n_heads
+        self.num_channels_latents = self.in_channels
+        self.out_channels = self.in_channels
+
+
+@dataclass
+class ZImageDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=ZImageDiTArchConfig)
+    prefix: str = "ZImage"

fastvideo/configs/models/encoders/__init__.py

@@ -2,6 +2,7 @@ from fastvideo.configs.models.encoders.base import (BaseEncoderOutput, EncoderCo
                                                     TextEncoderConfig)
 from fastvideo.configs.models.encoders.clip import (CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig)
 from fastvideo.configs.models.encoders.llama import LlamaConfig
+from fastvideo.configs.models.encoders.lingbotworld2_t5 import LingBotWorld2UMT5ArchConfig, LingBotWorld2UMT5Config
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
 from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
 from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig
@@ -9,6 +10,7 @@ from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1
 from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
 from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
 from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
+from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig
 from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig,
                                                                         StableAudioConditionerConfig)
 from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig
@@ -17,5 +19,6 @@ __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
     "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
     "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
-    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig"
+    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig",
+    "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config", "LingBotVideoQwen3VLTextConfig"
 ]

fastvideo/configs/models/encoders/base.py

@@ -83,6 +83,7 @@ class TextEncoderConfig(EncoderConfig):
     arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
     is_chat_model: bool = False
     treat_empty_as_dot: bool = False
+    chat_template_enable_thinking: bool = field(default=False, kw_only=True)
 
 
 @dataclass

fastvideo/configs/models/encoders/lingbot_video.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3-VL text-only encoder configuration used by LingBot-Video."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.base import TextEncoderArchConfig
+from fastvideo.configs.models.encoders.qwen3 import Qwen3TextArchConfig, Qwen3TextConfig
+
+
+@dataclass
+class LingBotVideoQwen3VLTextArchConfig(Qwen3TextArchConfig):
+    """Exact Qwen3-VL language-model architecture released with LingBot-Video."""
+
+    architectures: list[str] = field(default_factory=lambda: ["LingBotVideoQwen3VLTextModel"])
+    vocab_size: int = 151936
+    hidden_size: int = 2560
+    intermediate_size: int = 9728
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    max_position_embeddings: int = 262144
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 5000000.0
+    rope_scaling: dict | None = None
+    mrope_interleaved: bool = True
+    mrope_section: tuple[int, int, int] = (24, 20, 20)
+    bos_token_id: int = 151643
+    eos_token_id: int = 151645
+    pad_token_id: int = 151643
+    text_len: int = 37698
+    output_hidden_states: bool = True
+    require_processor: bool = True
+
+    def __post_init__(self) -> None:
+        """Match the official processor call used by LingBotVideoPipeline."""
+        self.tokenizer_kwargs = {
+            "truncation": True,
+            "max_length": self.text_len,
+            "padding": "longest",
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class LingBotVideoQwen3VLTextConfig(Qwen3TextConfig):
+    """FastVideo loader config for the LingBot-Video text-only Qwen3-VL path."""
+
+    arch_config: TextEncoderArchConfig = field(default_factory=LingBotVideoQwen3VLTextArchConfig)
+    prefix: str = "language_model"
+    is_chat_model: bool = False

fastvideo/configs/models/encoders/lingbotworld2_t5.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+
+@dataclass
+class LingBotWorld2UMT5ArchConfig(TextEncoderArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["LingBotWorld2T5EncoderModel"])
+    vocab_size: int = 256384
+    dim: int = 4096
+    dim_attn: int = 4096
+    dim_ffn: int = 10240
+    num_heads: int = 64
+    num_layers: int = 24
+    num_buckets: int = 32
+    text_len: int = 512
+    hidden_size: int = 4096
+    dropout: float = 0.1
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": self.text_len,
+            "add_special_tokens": True,
+            "return_attention_mask": True,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class LingBotWorld2UMT5Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=LingBotWorld2UMT5ArchConfig)
+
+    prefix: str = "text_encoder"

fastvideo/configs/models/encoders/reason1.py

@@ -7,6 +7,18 @@ 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)."""
@@ -61,6 +73,8 @@ 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,9 @@ 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 Kandinsky5I2VConfig, Kandinsky5T2VConfig
+from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig, Kandinsky5I2VConfig, Kandinsky5T2VConfig
+from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
+from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
@@ -19,5 +21,6 @@ __all__ = [
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
     "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
-    "Kandinsky5I2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig",
+    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/kandinsky5.py

@@ -120,3 +120,21 @@ 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/configs/pipelines/lingbot_video.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dense LingBot-Video T2V pipeline configuration."""
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+PROMPT_CROP_START = 140
+PROMPT_TEMPLATE = ("<|im_start|>system\nGiven a user input that may include a text prompt alone, "
+                   "a text prompt with an image reference, or a text prompt with a video reference "
+                   'or a video reference alone, generate an "Enhanced prompt" that provides detailed '
+                   "visual descriptions suitable for video generation. Evaluate the level of detail "
+                   "in the user's input: if it is simple, enrich it by adding specifics about colors, "
+                   "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
+                   "progression, and spatial relationships to create vivid, concrete, and temporally "
+                   "coherent scenes to create vivid and concrete scenes. Please generate only the "
+                   "enhanced description for the prompt below and avoid including any additional "
+                   "commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n"
+                   "<|im_start|>assistant\n")
+
+
+def preprocess_lingbot_video_prompt(prompt: str) -> str:
+    """Apply the released T2V system/user/assistant prompt template."""
+    return PROMPT_TEMPLATE.format(prompt)
+
+
+def postprocess_lingbot_video_text(
+    outputs: BaseEncoderOutput,
+    attention_mask: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Select the final hidden state, crop the template, and trim batch-one padding."""
+    if outputs.hidden_states is None:
+        raise ValueError("LingBot-Video requires text-encoder hidden states")
+    prompt_embeds = outputs.hidden_states[-1][:, PROMPT_CROP_START:]
+    prompt_mask = attention_mask[:, PROMPT_CROP_START:]
+    if prompt_embeds.shape[0] == 1:
+        true_length = int(prompt_mask[0].sum().item())
+        prompt_embeds = prompt_embeds[:, :true_length]
+        prompt_mask = prompt_mask[:, :true_length]
+    return prompt_embeds, prompt_mask
+
+
+@dataclass
+class LingBotVideoT2VConfig(PipelineConfig):
+    """Released Dense T2V component wiring and numerical precision policy."""
+
+    dit_config: DiTConfig = field(default_factory=LingBotVideoConfig)
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (LingBotVideoQwen3VLTextConfig(), ))
+    preprocess_text_funcs: tuple[Callable, ...] = field(default_factory=lambda: (preprocess_lingbot_video_prompt, ))
+    postprocess_text_funcs: tuple[Callable, ...] = field(default_factory=lambda: (postprocess_lingbot_video_text, ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp32"
+    vae_decode_precision: str | None = "fp32"
+    vae_tiling: bool = False
+    vae_sp: bool = False
+    flow_shift: float | None = 3.0
+    embedded_cfg_scale: float | None = None
+    scheduler_step_in_fp32: bool = True
+
+    def __post_init__(self) -> None:
+        """Load only the VAE decoder for the T2V workload."""
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True

fastvideo/configs/pipelines/lingbotworld2.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+import html
+from dataclasses import dataclass, field
+
+import ftfy
+import torch
+
+from fastvideo.configs.models import DiTConfig
+from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput, LingBotWorld2UMT5Config
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.pipelines.wan import Wan2_2_I2V_A14B_Config
+
+
+def lingbotworld2_whitespace_preprocess(prompt: str) -> str:
+    text = ftfy.fix_text(prompt)
+    text = html.unescape(html.unescape(text))
+    return " ".join(text.strip().split())
+
+
+def lingbotworld2_t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    assert outputs.last_hidden_state is not None
+    return outputs.last_hidden_state
+
+
+@dataclass
+class LingBotWorld2CausalFastI2V480PConfig(Wan2_2_I2V_A14B_Config):
+    dit_config: DiTConfig = field(default_factory=LingBotWorld2CausalFastVideoConfig)
+    vae_config: WanVAEConfig = field(default_factory=WanVAEConfig)
+    text_encoder_configs: tuple = field(default_factory=lambda: (LingBotWorld2UMT5Config(), ))
+    preprocess_text_funcs: tuple = field(default_factory=lambda: (lingbotworld2_whitespace_preprocess, ))
+    postprocess_text_funcs: tuple = field(default_factory=lambda: (lingbotworld2_t5_postprocess_text, ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+    flow_shift: float | None = 10.0
+    boundary_ratio: float | None = 0.947
+    vae_precision: str = "fp32"
+    vae_decode_precision: str | None = "fp32"
+    vae_tiling: bool = False
+    vae_sp: bool = False
+    dit_precision: str = "bf16"
+    is_causal: bool = False
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True

fastvideo/configs/pipelines/zimage.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import EncoderConfig
+from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
+from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
+
+
+def _zimage_text_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    if outputs.hidden_states is None:
+        raise RuntimeError("Z-Image requires Qwen3 hidden states")
+    return outputs.hidden_states[-2]
+
+
+@dataclass
+class ZImagePipelineConfig(PipelineConfig):
+    """Configuration for the native Z-Image text-to-image pipeline."""
+
+    scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
+    transformer_arch: str = "ZImageTransformer2DModel"
+    vae_arch: str = "AutoencoderKL"
+    text_encoder_archs: tuple[str, ...] = ("Qwen3Model", )
+    tokenizer_archs: tuple[str, ...] = ("Qwen2Tokenizer", )
+
+    dit_config: ZImageDiTConfig = field(default_factory=ZImageDiTConfig)
+    vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Qwen3TextConfig(chat_template_enable_thinking=True), ))
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=lambda: (preprocess_text, ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda: (_zimage_text_postprocess, ))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    embedded_cfg_scale: float = 0.0
+    flow_shift: float | None = 3.0
+    scheduler_step_in_fp32: bool = True
+    scheduler_sigma_min: float = 0.0
+    scheduler_use_reference_discrete_timesteps: bool = True

fastvideo/entrypoints/cli/main.py

@@ -8,6 +8,7 @@ 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]:
@@ -23,7 +24,7 @@ def cmd_init() -> list[CLISubcommand]:
 
 def main() -> None:
     parser = FlexibleArgumentParser(description="FastVideo CLI")
-    parser.add_argument('-v', '--version', action='version', version='0.1.0')
+    parser.add_argument('-v', '--version', action='version', version=__version__)
 
     subparsers = parser.add_subparsers(required=False, dest="subparser")
 

fastvideo/entrypoints/video_generator.py

@@ -112,6 +112,40 @@ def _infer_latent_batch_size(batch: ForwardBatch) -> int:
     return latent_batch_size
 
 
+def _resolve_output_size(
+    samples: torch.Tensor,
+    fallback: tuple[int, int, int],
+    *,
+    pixel_output: bool,
+) -> tuple[int, int, int]:
+    """Report the final decoded video's `(height, width, frames)`.
+
+    Refiner stages can produce a different resolution from the base request, so
+    pixel outputs use the final `[batch, channels, frames, height, width]` tensor.
+    Latent and audio outputs keep the requested fallback because their tensor
+    dimensions do not describe decoded pixels.
+    """
+    if pixel_output and samples.ndim == 5:
+        return (int(samples.shape[-2]), int(samples.shape[-1]), int(samples.shape[-3]))
+    return fallback
+
+
+def _validate_request_stage_overrides(model_path: str, request: GenerationRequest) -> None:
+    """Validate typed stage overrides against the model's registered preset."""
+    if not request.stage_overrides:
+        return
+    from fastvideo.api.presets import validate_preset_selection
+    from fastvideo.registry import get_preset_selection
+    preset_name, model_family = get_preset_selection(model_path)
+    if preset_name is None or model_family is None:
+        raise ValueError(f"Model {model_path!r} has no preset for stage override validation")
+    validate_preset_selection(
+        preset_name,
+        model_family,
+        stage_overrides=request.stage_overrides,
+    )
+
+
 class VideoGenerator:
     """
     A unified class for generating videos using diffusion models.
@@ -443,6 +477,7 @@ class VideoGenerator:
         self,
         request: GenerationRequest,
     ) -> GenerationResult | list[GenerationResult]:
+        _validate_request_stage_overrides(self.fastvideo_args.model_path, request)
         if isinstance(request.prompt, list):
             if request.inputs.prompt_path is not None:
                 raise ValueError("request.prompt list cannot be combined with request.inputs.prompt_path")
@@ -808,6 +843,15 @@ class VideoGenerator:
         #   2. Audio-only workload — `samples` is a 1×3×1×8×8 placeholder
         #      no caller will use; skip the grid loop and save a `.wav`.
         #   3. Pixel video / image — the historical happy path.
+        # `GenerationResult.size` describes the produced media, not only the
+        # base-stage request. Refiner pipelines can change the final pixel
+        # dimensions, so derive this result metadata from the decoded output.
+        output_size = _resolve_output_size(
+            samples,
+            (target_height, target_width, batch.num_frames),
+            pixel_output=not is_latent_output and not audio_only,
+        )
+
         postprocess_start = time.perf_counter()
         frames: list[np.ndarray] | None
         if is_latent_output or audio_only:
@@ -916,7 +960,7 @@ class VideoGenerator:
             "audio": output_batch.extra.get("audio"),
             "audio_sample_rate": output_batch.extra.get("audio_sample_rate"),
             "ltx2_audio_latents": output_batch.extra.get("ltx2_audio_latents"),
-            "size": (target_height, target_width, batch.num_frames),
+            "size": output_size,
             "generation_time": gen_time,
             "e2e_latency": e2e_time,
             "logging_info": logging_info,

fastvideo/layers/quantization/fp8_config.py

@@ -26,6 +26,9 @@ 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",
@@ -33,6 +36,11 @@ _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,6 +58,29 @@ _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)
@@ -286,6 +309,11 @@ 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):
@@ -326,7 +354,11 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
         | None = None,
     ) -> torch.Tensor:
         SfLayout, _, _ = _require_flashinfer()
-        out_dim = layer.weight.shape[0]
+        # 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]
         original_shape = x.shape
 
         # Stage-aware profile: keep refine-only FP4 layers in dense mode
@@ -334,8 +366,13 @@ 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":
-            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)))
+            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)))
             return out.view(*original_shape[:-1], out_dim)
         if pre_quantized is not None:
             x_fp4, x_scale, x_global_sf = pre_quantized
@@ -393,11 +430,19 @@ class NVFP4Config(QuantizationConfig):
     instead of hardcoding it here.
     """
 
-    def __init__(self, layer_profile: str = "refine"):
+    def __init__(self, layer_profile: str = "refine", retain_original_weights: bool | None = None):
         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"
@@ -415,32 +460,20 @@ class NVFP4Config(QuantizationConfig):
 
     @classmethod
     def from_config(cls, config: dict[str, Any]) -> NVFP4Config:
-        return cls(layer_profile=config.get("layer_profile", "refine"))
+        return cls(
+            layer_profile=config.get("layer_profile", "refine"),
+            retain_original_weights=config.get("retain_original_weights"),
+        )
 
     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.
-        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)
+        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
         return None
 
 
@@ -448,6 +481,9 @@ 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):
@@ -480,9 +516,40 @@ 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,7 +46,11 @@ 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").
+# 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.
 DEFAULT_FP4_LAYERS = (
     "ffn.fc_in",
     "ffn.fc_out",
@@ -54,6 +58,11 @@ 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,10 +9,11 @@ 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, 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.
+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.
 """
 import logging
 from typing import Any
@@ -21,6 +22,8 @@ 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__)
@@ -72,7 +75,10 @@ class NVFP4QATTrainConfig(QuantizationConfig):
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
         from fastvideo.layers.linear import LinearBase
-        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):
+        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:
             return NVFP4QATTrainQuantizeMethod()
         return None

fastvideo/models/dits/kandinsky5.py

@@ -15,6 +15,7 @@ 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
@@ -285,6 +286,7 @@ 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
@@ -293,20 +295,24 @@ 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,
@@ -413,9 +419,11 @@ class Kandinsky5Attention(nn.Module):
 
 class Kandinsky5FeedForward(nn.Module):
 
-    def __init__(self, dim: int, ff_dim: int):
+    def __init__(self, dim: int, ff_dim: int, prefix: str = "",
+                 quant_config: QuantizationConfig | None = None):
         super().__init__()
-        self.mlp = MLP(dim, ff_dim, bias=False, act_type="gelu")
+        self.mlp = MLP(dim, ff_dim, bias=False, act_type="gelu",
+                      quant_config=quant_config, prefix=f"{prefix}.mlp")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.mlp(x)
@@ -467,7 +475,8 @@ class Kandinsky5TransformerEncoderBlock(nn.Module):
                  head_dim: int,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
-                 prefix: str = ""):
+                 prefix: str = "",
+                 quant_config: QuantizationConfig | None = None):
         super().__init__()
         self.text_modulation = Kandinsky5Modulation(time_dim, model_dim, 6)
 
@@ -482,7 +491,8 @@ class Kandinsky5TransformerEncoderBlock(nn.Module):
             model_dim,
             head_dim,
             supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.self_attention")
+            prefix=f"{prefix}.self_attention",
+            quant_config=quant_config)
 
         self.feed_forward_norm = LayerNormScaleShift(
             model_dim,
@@ -491,7 +501,9 @@ class Kandinsky5TransformerEncoderBlock(nn.Module):
             elementwise_affine=False,
             dtype=torch.float32,
             compute_dtype=torch.float32)
-        self.feed_forward = Kandinsky5FeedForward(model_dim, ff_dim)
+        self.feed_forward = Kandinsky5FeedForward(
+            model_dim, ff_dim, prefix=f"{prefix}.feed_forward",
+            quant_config=quant_config)
 
     def forward(self, x: torch.Tensor, time_embed: torch.Tensor,
                 rope: torch.Tensor) -> torch.Tensor:
@@ -523,7 +535,8 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
                  prefix: str = "",
-                 use_nabla: bool = False):
+                 use_nabla: bool = False,
+                 quant_config: QuantizationConfig | None = None):
         super().__init__()
         self.visual_modulation = Kandinsky5Modulation(time_dim, model_dim, 9)
 
@@ -539,7 +552,8 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
             head_dim,
             supported_attention_backends=supported_attention_backends,
             prefix=f"{prefix}.self_attention",
-            use_nabla=use_nabla)
+            use_nabla=use_nabla,
+            quant_config=quant_config)
 
         self.cross_attention_norm = LayerNormScaleShift(
             model_dim,
@@ -552,7 +566,8 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
             model_dim,
             head_dim,
             supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.cross_attention")
+            prefix=f"{prefix}.cross_attention",
+            quant_config=quant_config)
 
         self.feed_forward_norm = LayerNormScaleShift(
             model_dim,
@@ -561,7 +576,9 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
             elementwise_affine=False,
             dtype=torch.float32,
             compute_dtype=torch.float32)
-        self.feed_forward = Kandinsky5FeedForward(model_dim, ff_dim)
+        self.feed_forward = Kandinsky5FeedForward(
+            model_dim, ff_dim, prefix=f"{prefix}.feed_forward",
+            quant_config=quant_config)
 
     def forward(self, visual_embed: torch.Tensor, text_embed: torch.Tensor,
                 time_embed: torch.Tensor, rope: torch.Tensor,
@@ -636,6 +653,8 @@ 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
@@ -664,7 +683,8 @@ class Kandinsky5Transformer3DModel(BaseDiT):
                                               arch.ff_dim,
                                               head_dim,
                                               self._supported_attention_backends,
-                                              prefix=f"{config.prefix}.text_transformer_blocks.{i}")
+                                              prefix=f"{config.prefix}.text_transformer_blocks.{i}",
+                                              quant_config=quant_config)
             for i in range(arch.num_text_blocks)
         ])
         self.visual_transformer_blocks = nn.ModuleList([
@@ -673,7 +693,8 @@ class Kandinsky5Transformer3DModel(BaseDiT):
                                               head_dim,
                                               self._supported_attention_backends,
                                               prefix=f"{config.prefix}.visual_transformer_blocks.{i}",
-                                              use_nabla=arch.attention_type == "nabla")
+                                              use_nabla=arch.attention_type == "nabla",
+                                              quant_config=quant_config)
             for i in range(arch.num_visual_blocks)
         ])
 

fastvideo/models/dits/lingbot_video.py

@@ -0,0 +1,811 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-native LingBot-Video Dense and MoE diffusion transformers."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig, _is_lingbot_video_block
+from fastvideo.distributed.communication_op import (
+    sequence_model_parallel_all_gather_with_unpad,
+    sequence_model_parallel_all_to_all_4D,
+    sequence_model_parallel_shard,
+)
+from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.visual_embedding import Timesteps
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.platforms import AttentionBackendEnum
+
+
+@dataclass
+class LingBotVideoTransformerOutput:
+    """Output container matching the released transformer contract."""
+
+    sample: torch.Tensor
+
+
+_FP32_MODULE_NAMES = (
+    "time_embedder",
+    "time_modulation",
+    "scale_shift_table",
+    "norm",
+    "norm1",
+    "norm2",
+    "norm_q",
+    "norm_k",
+    "norm_post_attn",
+    "norm_post_ffn",
+    "norm_out",
+    "norm_out_modulation",
+    "router",
+)
+
+
+def _keep_in_fp32(name: str) -> bool:
+    """Return whether a released checkpoint module keeps fp32 parameters."""
+    return any(module_name in name.split(".") for module_name in _FP32_MODULE_NAMES)
+
+
+def _sequence_parallel_world_size() -> int:
+    """Use standalone single-rank behavior before distributed initialization."""
+    return get_sp_world_size() if model_parallel_is_initialized() else 1
+
+
+class LingBotVideoLinear(ReplicatedLinear):
+    """Replicated FastVideo linear with the tensor-only official call contract."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Return the projected tensor while preserving the normal weight surface."""
+        output, _ = super().forward(hidden_states)
+        return output
+
+
+class LingBotVideoRMSNorm(nn.Module):
+    """RMSNorm with fp32 accumulation and input-dtype output."""
+
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Normalize the last dimension using the official accumulation order."""
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)
+
+
+def _apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply complex 3D rotary embeddings to `(B, S, H, D)` tensors."""
+    with torch.amp.autocast("cuda", enabled=False):
+        x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        output = torch.view_as_real(x_complex * freqs_cis.unsqueeze(2)).flatten(3)
+        return output.type_as(x)
+
+
+class LingBotVideoRotaryEmbedding(nn.Module):
+    """Complex64 rotary table indexed by temporal and spatial positions."""
+
+    def __init__(self, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...], theta: float) -> None:
+        super().__init__()
+        self.axes_dims = tuple(axes_dims)
+        self.axes_lens = list(axes_lens)
+        self.theta = theta
+        self.freqs_cis: list[torch.Tensor] | None = None
+
+    @staticmethod
+    def _precompute(dims: tuple[int, ...], lengths: tuple[int, ...], theta: float) -> list[torch.Tensor]:
+        """Build the per-axis complex frequency tables on CPU."""
+        tables: list[torch.Tensor] = []
+        for dim, length in zip(dims, lengths, strict=True):
+            frequencies = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float64, device="cpu") / dim))
+            positions = torch.arange(length, device=frequencies.device, dtype=torch.float64)
+            phases = torch.outer(positions, frequencies).float()
+            tables.append(torch.polar(torch.ones_like(phases), phases).to(torch.complex64))
+        return tables
+
+    def forward(
+        self,
+        position_ids: torch.Tensor,
+        maxima: tuple[int, ...] | None = None,
+    ) -> torch.Tensor:
+        """Gather and concatenate rotary frequencies for `(S, 3)` positions."""
+        device = position_ids.device
+        if maxima is None:
+            maxima = tuple(int(value) for value in position_ids.max(dim=0).values.tolist())
+        rebuild = self.freqs_cis is None or any(maximum >= length
+                                                for maximum, length in zip(maxima, self.axes_lens, strict=True))
+        if rebuild:
+            for index, maximum in enumerate(maxima):
+                if maximum >= self.axes_lens[index]:
+                    self.axes_lens[index] = int(maximum * 1.5) + 1
+            self.freqs_cis = self._precompute(self.axes_dims, tuple(self.axes_lens), self.theta)
+            self.freqs_cis = [table.to(device) for table in self.freqs_cis]
+        elif self.freqs_cis[0].device != device:
+            self.freqs_cis = [table.to(device) for table in self.freqs_cis]
+        return torch.cat(
+            [self.freqs_cis[index][position_ids[:, index]] for index in range(len(self.axes_dims))],
+            dim=-1,
+        )
+
+
+def _make_joint_position_ids(
+    text_len: int,
+    grid_t: int,
+    grid_h: int,
+    grid_w: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Create official `[video; text]` 3D positions for one sample."""
+    temporal = torch.arange(grid_t, device=device, dtype=torch.int32) + text_len + 1
+    height = torch.arange(grid_h, device=device, dtype=torch.int32)
+    width = torch.arange(grid_w, device=device, dtype=torch.int32)
+    video_positions = torch.stack(torch.meshgrid(temporal, height, width, indexing="ij"), dim=-1).flatten(0, 2)
+    text_temporal = torch.arange(text_len, device=device, dtype=torch.int32) + 1
+    text_positions = torch.stack(
+        [text_temporal, torch.zeros_like(text_temporal),
+         torch.zeros_like(text_temporal)], dim=-1)
+    return torch.cat([video_positions, text_positions], dim=0)
+
+
+class LingBotVideoTextEmbedder(nn.Module):
+    """Project Qwen3-VL hidden states into the DiT hidden dimension."""
+
+    def __init__(self, text_dim: int, hidden_size: int) -> None:
+        super().__init__()
+        self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6)
+        self.linear_1 = LingBotVideoLinear(text_dim, hidden_size, bias=True)
+        self.linear_2 = LingBotVideoLinear(hidden_size, hidden_size, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Apply RMSNorm followed by the released two-layer SiLU projection."""
+        hidden_states = self.norm(hidden_states)
+        return self.linear_2(F.silu(self.linear_1(hidden_states)))
+
+
+class LingBotVideoAttention(nn.Module):
+    """Joint video-text attention shared by the Dense and MoE variants."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        norm_eps: float,
+        qkv_bias: bool,
+        out_bias: bool,
+    ) -> None:
+        """Create released QKV projections, per-head norms, and output projection."""
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.to_q = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias)
+        self.to_k = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias)
+        self.to_v = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias)
+        self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps)
+        self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps)
+        self.to_out = LingBotVideoLinear(hidden_size, hidden_size, bias=out_bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rotary_emb: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        original_seq_len: int | None = None,
+    ) -> torch.Tensor:
+        """Project QKV, apply rotary embeddings, and run non-causal SDPA."""
+        batch, sequence, _ = hidden_states.shape
+        query = self.to_q(hidden_states).unflatten(2, (self.num_heads, self.head_dim))
+        key = self.to_k(hidden_states).unflatten(2, (self.num_heads, self.head_dim))
+        value = self.to_v(hidden_states).unflatten(2, (self.num_heads, self.head_dim))
+        query = _apply_rotary_emb(self.norm_q(query), rotary_emb)
+        key = _apply_rotary_emb(self.norm_k(key), rotary_emb)
+
+        if original_seq_len is not None:
+            # Attention needs the full unpadded joint sequence while projections
+            # and residual blocks stay sharded over tokens.
+            qkv = torch.cat([query, key, value], dim=0)
+            qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)
+            padded_seq_len = qkv.shape[1]
+            query, key, value = qkv[:, :original_seq_len].chunk(3, dim=0)
+        output = F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+        ).transpose(1, 2)
+        if original_seq_len is not None:
+            output = F.pad(output, (0, 0, 0, 0, 0, padded_seq_len - original_seq_len))
+            output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)
+        return self.to_out(output.reshape(batch, sequence, -1).type_as(hidden_states))
+
+
+class LingBotVideoMLP(nn.Module):
+    """Dense SwiGLU feed-forward network."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int) -> None:
+        super().__init__()
+        self.gate_proj = LingBotVideoLinear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = LingBotVideoLinear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = LingBotVideoLinear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Apply the released SiLU-gated MLP ordering."""
+        return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+
+
+class LingBotVideoRouter(nn.Module):
+    """Released token-choice router with bias-only expert selection correction."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_experts: int,
+        top_k: int,
+        score_func: str,
+        norm_topk_prob: bool,
+        n_group: int | None,
+        topk_group: int | None,
+        route_scale: float,
+    ) -> None:
+        """Create fp32-routed expert scores with the released persistent bias."""
+        super().__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.score_func = score_func
+        self.norm_topk_prob = norm_topk_prob
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.route_scale = route_scale
+        self.weight = nn.Parameter(torch.empty(num_experts, hidden_size))
+        self.register_buffer("e_score_correction_bias", torch.zeros(num_experts), persistent=True)
+
+    def _group_limited_topk(self, scores_for_choice: torch.Tensor) -> torch.Tensor:
+        """Restrict token choices to groups with the two strongest expert scores."""
+        sequence_length = scores_for_choice.shape[0]
+        experts_per_group = self.num_experts // self.n_group
+        grouped = scores_for_choice.view(sequence_length, self.n_group, experts_per_group)
+        group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1)
+        group_indices = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_indices, 1)
+        score_mask = (group_mask.unsqueeze(-1).expand(sequence_length, self.n_group,
+                                                      experts_per_group).reshape(sequence_length, -1))
+        masked_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
+        return torch.topk(masked_scores, k=self.top_k, dim=-1, sorted=False)[1]
+
+    def forward(self,
+                tokens: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Score in fp32, select with correction bias, and weight without it."""
+        with torch.amp.autocast(tokens.device.type, enabled=False):
+            logits = F.linear(tokens.float(), self.weight.float())
+        scores = F.softmax(logits, dim=-1) if self.score_func == "softmax" else logits.sigmoid()
+        scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)
+        if self.n_group is not None and self.n_group > 1:
+            top_indices = self._group_limited_topk(scores_for_choice)
+        else:
+            top_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
+        top_scores = scores.gather(1, top_indices)
+        if self.top_k > 1 and self.norm_topk_prob:
+            top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20)
+        top_scores = top_scores * self.route_scale
+        return top_indices, top_scores.to(tokens.dtype), logits, scores, scores_for_choice
+
+
+class LingBotVideoGroupedExperts(nn.Module):
+    """Released grouped-expert parameter layout: w1/w3 `[E,I,H]`, w2 `[E,H,I]`."""
+
+    def __init__(self, num_experts: int, hidden_size: int, intermediate_size: int) -> None:
+        super().__init__()
+        self.num_experts = num_experts
+        self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size))
+        self.w2 = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size))
+        self.w3 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size))
+
+
+def _round_up_to_multiple(value: int, multiple: int) -> int:
+    """Round an integer up to the next multiple."""
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+class LingBotVideoSparseMoeBlock(nn.Module):
+    """Token-choice sparse feed-forward block matching the released state surface."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        num_experts: int,
+        top_k: int,
+        moe_intermediate_size: int,
+        score_func: str,
+        norm_topk_prob: bool,
+        n_group: int | None,
+        topk_group: int | None,
+        routed_scaling_factor: float,
+        n_shared_experts: int | None,
+    ) -> None:
+        """Create routed and optional shared experts with released parameter names."""
+        super().__init__()
+        del intermediate_size
+        self.hidden_size = hidden_size
+        self.num_experts = num_experts
+        self.router = LingBotVideoRouter(
+            hidden_size,
+            num_experts,
+            top_k,
+            score_func,
+            norm_topk_prob,
+            n_group,
+            topk_group,
+            routed_scaling_factor,
+        )
+        self.experts = LingBotVideoGroupedExperts(num_experts, hidden_size, moe_intermediate_size)
+        self.shared_experts = None
+        if n_shared_experts is not None and n_shared_experts > 0:
+            self.shared_experts = LingBotVideoMLP(hidden_size, moe_intermediate_size * n_shared_experts)
+
+    @staticmethod
+    def _reorder_tokens(
+        tokens: torch.Tensor,
+        top_scores: torch.Tensor,
+        top_indices: torch.Tensor,
+        num_experts: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+        """Pack active token choices into stable expert-major order."""
+        num_tokens = tokens.shape[0]
+        top_k = top_indices.shape[1]
+        flat_scores = top_scores.reshape(-1)
+        flat_indices = top_indices.reshape(-1)
+        active_positions = torch.where(flat_scores != 0)[0]
+        active_experts = flat_indices[active_positions]
+        counts = torch.zeros(num_experts, device=tokens.device, dtype=torch.int64)
+        counts.scatter_add_(0, active_experts, torch.ones_like(active_experts, dtype=torch.int64))
+        sort_order = torch.argsort(active_experts, stable=True)
+        sorted_positions = active_positions[sort_order]
+        sorted_scores = flat_scores[sorted_positions]
+        original_token_indices = sorted_positions // top_k
+        permuted_tokens = tokens[original_token_indices]
+        return permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k
+
+    @staticmethod
+    def _pad_grouped_tokens(tokens: torch.Tensor,
+                            counts: torch.Tensor,
+                            align: int = 8) -> tuple[torch.Size, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Align each expert segment for `torch._grouped_mm` and retain unpad indices."""
+        num_tokens = tokens.shape[0]
+        num_experts = int(counts.shape[0])
+        max_length = _round_up_to_multiple(num_tokens + num_experts * align, align)
+        # Build the small padding index on CPU after one synchronization instead
+        # of reading three CUDA scalars for every expert.
+        counts_cpu = counts.to(device="cpu", dtype=torch.int64)
+        total_per_expert = torch.clamp_min(counts_cpu, align)
+        aligned_counts_cpu = ((total_per_expert + align - 1) // align * align).to(torch.int32)
+        write_offsets = torch.cumsum(aligned_counts_cpu, dim=0) - aligned_counts_cpu
+        start_indices = torch.cumsum(counts_cpu, dim=0) - counts_cpu
+        permuted_indices_cpu = torch.full((max_length, ), num_tokens, dtype=torch.int64, device="cpu")
+        for expert_index in range(num_experts):
+            length = int(counts_cpu[expert_index])
+            if length == 0:
+                continue
+            write_start = int(write_offsets[expert_index])
+            start = int(start_indices[expert_index])
+            permuted_indices_cpu[write_start:write_start + length] = torch.arange(start,
+                                                                                  start + length,
+                                                                                  device="cpu",
+                                                                                  dtype=torch.int64)
+        permuted_indices = permuted_indices_cpu.to(tokens.device)
+        aligned_counts = aligned_counts_cpu.to(tokens.device)
+        tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1], ))))
+        input_shape = tokens_with_pad.shape
+        return input_shape, tokens_with_pad[permuted_indices], permuted_indices, aligned_counts
+
+    @staticmethod
+    def _unpad_grouped_tokens(output: torch.Tensor, input_shape: torch.Size,
+                              permuted_indices: torch.Tensor) -> torch.Tensor:
+        """Undo per-expert alignment while dropping the shared padding row."""
+        unpermuted = output.new_empty(input_shape)
+        unpermuted[permuted_indices, :] = output
+        return unpermuted[:-1]
+
+    def _run_grouped_experts(self, tokens: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
+        """Use the released bf16 grouped matmuls on CUDA and an eager CPU fallback."""
+        if tokens.device.type == "cpu" or not hasattr(torch, "_grouped_mm"):
+            return self._run_experts_for_loop(tokens, counts)
+        input_shape, padded_tokens, permuted_indices, aligned_counts = self._pad_grouped_tokens(tokens, counts)
+        offsets = torch.cumsum(aligned_counts, dim=0, dtype=torch.int32)
+        hidden = F.silu(
+            torch._grouped_mm(
+                padded_tokens.bfloat16(),
+                self.experts.w1.bfloat16().transpose(-2, -1),
+                offs=offsets,
+            ))
+        hidden = hidden * torch._grouped_mm(
+            padded_tokens.bfloat16(),
+            self.experts.w3.bfloat16().transpose(-2, -1),
+            offs=offsets,
+        )
+        output = torch._grouped_mm(
+            hidden,
+            self.experts.w2.bfloat16().transpose(-2, -1),
+            offs=offsets,
+        ).type_as(padded_tokens)
+        return self._unpad_grouped_tokens(output, input_shape, permuted_indices)
+
+    def _run_experts_for_loop(self, tokens: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
+        """Evaluate contiguous expert segments eagerly for CPU correctness tests."""
+        splits = torch.split(tokens, counts.tolist(), dim=0)
+        outputs: list[torch.Tensor] = []
+        for expert_index, expert_tokens in enumerate(splits):
+            if expert_tokens.numel() == 0:
+                continue
+            hidden = F.silu(expert_tokens @ self.experts.w1[expert_index].transpose(-2, -1))
+            hidden = hidden * (expert_tokens @ self.experts.w3[expert_index].transpose(-2, -1))
+            outputs.append(hidden @ self.experts.w2[expert_index].transpose(-2, -1))
+        if not outputs:
+            return tokens.new_zeros(tokens.shape)
+        return torch.cat(outputs, dim=0)
+
+    @staticmethod
+    def _restore_tokens(
+        expert_output: torch.Tensor,
+        sorted_positions: torch.Tensor,
+        sorted_scores: torch.Tensor,
+        num_tokens: int,
+        top_k: int,
+    ) -> torch.Tensor:
+        """Restore token order and combine expert outputs with fp32 weighted sums."""
+        hidden_size = expert_output.shape[-1]
+        unsorted = torch.zeros(
+            (num_tokens * top_k, hidden_size),
+            dtype=expert_output.dtype,
+            device=expert_output.device,
+        )
+        unsorted[sorted_positions] = expert_output
+        unsorted = unsorted.reshape(num_tokens, top_k, hidden_size)
+        scores_unsorted = torch.zeros(
+            num_tokens * top_k,
+            dtype=sorted_scores.dtype,
+            device=sorted_scores.device,
+        )
+        scores_unsorted[sorted_positions] = sorted_scores
+        scores_unsorted = scores_unsorted.reshape(num_tokens, top_k, 1)
+        return (unsorted.float() * scores_unsorted).sum(dim=1).to(expert_output.dtype)
+
+    def _run_selected_experts(
+        self,
+        tokens: torch.Tensor,
+        top_scores: torch.Tensor,
+        top_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        """Dispatch routed choices, execute experts, and restore token-major order."""
+        permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k = self._reorder_tokens(
+            tokens, top_scores, top_indices, self.router.num_experts)
+        expert_output = self._run_grouped_experts(permuted_tokens, counts)
+        return self._restore_tokens(expert_output, sorted_positions, sorted_scores, num_tokens, top_k)
+
+    def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor:
+        """Route token choices, zero padded routes, and add optional shared experts."""
+        batch = hidden_states.shape[0]
+        tokens = hidden_states.view(-1, self.hidden_size)
+        top_indices, top_scores, logits, scores, scores_for_choice = self.router(tokens)
+        del logits, scores, scores_for_choice
+        if padding_mask is not None:
+            mask = padding_mask.unsqueeze(-1).to(top_scores.dtype)
+            top_scores = top_scores * mask
+            top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9)
+            top_scores = top_scores * self.router.route_scale
+        output = self._run_selected_experts(tokens, top_scores, top_indices)
+        output = output.view(batch, -1, self.hidden_size)
+        if self.shared_experts is not None:
+            output = output + self.shared_experts(hidden_states)
+        return output
+
+
+class LingBotVideoBlock(nn.Module):
+    """One Dense or sparse LingBot-Video transformer block."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        intermediate_size: int,
+        norm_eps: float,
+        qkv_bias: bool,
+        out_bias: bool,
+        num_experts: int,
+        num_experts_per_tok: int,
+        moe_intermediate_size: int,
+        decoder_sparse_step: int,
+        mlp_only_layers: tuple[int, ...] | list[int],
+        n_shared_experts: int | None,
+        score_func: str,
+        norm_topk_prob: bool,
+        n_group: int | None,
+        topk_group: int | None,
+        routed_scaling_factor: float,
+        layer_idx: int,
+    ) -> None:
+        """Select the released Dense or sparse feed-forward structure for one layer."""
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * hidden_size))
+        self.norm1 = LingBotVideoRMSNorm(hidden_size, norm_eps)
+        self.attn = LingBotVideoAttention(hidden_size, num_attention_heads, norm_eps, qkv_bias, out_bias)
+        self.norm_post_attn = LingBotVideoRMSNorm(hidden_size, norm_eps)
+        self.norm2 = LingBotVideoRMSNorm(hidden_size, norm_eps)
+        if layer_idx not in mlp_only_layers and (num_experts > 0 and (layer_idx + 1) % decoder_sparse_step == 0):
+            self.ffn = LingBotVideoSparseMoeBlock(
+                hidden_size,
+                intermediate_size,
+                num_experts,
+                num_experts_per_tok,
+                moe_intermediate_size,
+                score_func,
+                norm_topk_prob,
+                n_group,
+                topk_group,
+                routed_scaling_factor,
+                n_shared_experts,
+            )
+        else:
+            self.ffn = LingBotVideoMLP(hidden_size, intermediate_size)
+        self.norm_post_ffn = LingBotVideoRMSNorm(hidden_size, norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb6: torch.Tensor,
+        rotary_emb: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        moe_padding_mask: torch.Tensor | None = None,
+        original_seq_len: int | None = None,
+    ) -> torch.Tensor:
+        """Run attention and configured feed-forward residual branches with fp32 AdaLN."""
+        expected_tokens = hidden_states.shape[0] * hidden_states.shape[1]
+        if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
+            raise ValueError("LingBotVideoBlock expects token-level temb6 with shape "
+                             f"(B*S, 6D); got {tuple(temb6.shape)} for {tuple(hidden_states.shape)}.")
+        modulation = temb6.view(*hidden_states.shape[:2], -1) + self.scale_shift_table.unsqueeze(0)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation.chunk(6, dim=-1)
+        gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+        scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+        bulk_dtype = self.attn.to_q.weight.dtype
+
+        attention_input = (self.norm1(hidden_states) * scale_msa + shift_msa).to(bulk_dtype)
+        attention_output = self.attn(attention_input, rotary_emb, attention_mask, original_seq_len)
+        hidden_states = hidden_states + (gate_msa * self.norm_post_attn(attention_output)).to(hidden_states.dtype)
+        mlp_input = (self.norm2(hidden_states) * scale_mlp + shift_mlp).to(bulk_dtype)
+        if isinstance(self.ffn, LingBotVideoSparseMoeBlock):
+            mlp_output = self.ffn(mlp_input, padding_mask=moe_padding_mask)
+        else:
+            mlp_output = self.ffn(mlp_input)
+        mlp_output = self.norm_post_ffn(mlp_output)
+        return hidden_states + (gate_mlp * mlp_output).to(hidden_states.dtype)
+
+
+class LingBotVideoTimestepEmbedding(nn.Module):
+    """Two-layer timestep embedding with released parameter names."""
+
+    def __init__(self, input_dim: int, hidden_size: int, bias: bool) -> None:
+        super().__init__()
+        self.linear_1 = LingBotVideoLinear(input_dim, hidden_size, bias=bias)
+        self.linear_2 = LingBotVideoLinear(hidden_size, hidden_size, bias=bias)
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        """Apply the official linear-SiLU-linear timestep projection."""
+        return self.linear_2(F.silu(self.linear_1(sample)))
+
+
+class LingBotVideoTransformer3DModel(BaseDiT):
+    """LingBot-Video DiT with a source-compatible Dense or MoE state surface."""
+
+    _fsdp_shard_conditions = [_is_lingbot_video_block]
+    _compile_conditions = [_is_lingbot_video_block]
+    _supported_attention_backends = (
+        AttentionBackendEnum.TORCH_SDPA,
+        AttentionBackendEnum.FLASH_ATTN,
+    )
+    param_names_mapping = {r"^(.*)$": r"\1"}
+    reverse_param_names_mapping = {r"^(.*)$": r"\1"}
+
+    def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -> torch.dtype:
+        """Select the released mixed-precision dtype while loading each parameter."""
+        return torch.float32 if _keep_in_fp32(name) else default_dtype
+
+    def __init__(self, config: LingBotVideoConfig, hf_config: dict[str, Any]) -> None:
+        """Construct the Dense or MoE variant from its released transformer config."""
+        config.update_model_arch(hf_config)
+        super().__init__(config, hf_config)
+        head_dim = config.hidden_size // config.num_attention_heads
+        if head_dim != sum(config.axes_dims):
+            raise ValueError(f"head_dim {head_dim} != sum(axes_dims) {sum(config.axes_dims)}")
+        sp_world_size = _sequence_parallel_world_size()
+        assert config.num_attention_heads % sp_world_size == 0, (
+            f"The number of attention heads ({config.num_attention_heads}) must be divisible by "
+            f"the sequence parallel size ({sp_world_size})")
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.in_channels
+        self.patch_embedder = LingBotVideoLinear(
+            config.in_channels * math.prod(config.patch_size),
+            config.hidden_size,
+            bias=config.patch_embed_bias,
+        )
+        self.time_proj = Timesteps(config.freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embedder = LingBotVideoTimestepEmbedding(config.freq_dim, config.hidden_size,
+                                                           config.timestep_mlp_bias)
+        self.time_modulation = nn.Sequential(nn.SiLU(), LingBotVideoLinear(config.hidden_size, 6 * config.hidden_size))
+        self.text_embedder = LingBotVideoTextEmbedder(config.text_dim, config.hidden_size)
+        self.rope = LingBotVideoRotaryEmbedding(tuple(config.axes_dims), tuple(config.axes_lens), config.rope_theta)
+        self.blocks = nn.ModuleList([
+            LingBotVideoBlock(
+                config.hidden_size,
+                config.num_attention_heads,
+                config.intermediate_size,
+                config.norm_eps,
+                config.qkv_bias,
+                config.out_bias,
+                config.num_experts,
+                config.num_experts_per_tok,
+                config.moe_intermediate_size,
+                config.decoder_sparse_step,
+                config.mlp_only_layers,
+                config.n_shared_experts,
+                config.score_func,
+                config.norm_topk_prob,
+                config.n_group,
+                config.topk_group,
+                config.routed_scaling_factor,
+                layer_index,
+            ) for layer_index in range(config.depth)
+        ])
+        self.norm_out = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=config.norm_eps)
+        self.norm_out_modulation = nn.Sequential(nn.SiLU(),
+                                                 LingBotVideoLinear(config.hidden_size, 2 * config.hidden_size))
+        self.proj_out = LingBotVideoLinear(config.hidden_size, math.prod(config.patch_size) * config.out_channels)
+        self.__post_init__()
+
+    def to(self, *args: Any, **kwargs: Any) -> LingBotVideoTransformer3DModel:
+        """Cast bulk weights while retaining the released fp32-sensitive modules."""
+        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
+        if dtype is None or dtype == torch.float32:
+            return super().to(*args, **kwargs)
+        if not torch.is_floating_point(torch.empty((), dtype=dtype)):
+            return super().to(*args, **kwargs)
+        if device is not None:
+            super().to(device=device, non_blocking=non_blocking)
+        for name, parameter in self.named_parameters():
+            if torch.is_floating_point(parameter):
+                target_dtype = torch.float32 if _keep_in_fp32(name) else dtype
+                parameter.data = parameter.data.to(target_dtype, non_blocking=non_blocking)
+                if parameter.grad is not None:
+                    parameter.grad.data = parameter.grad.data.to(target_dtype, non_blocking=non_blocking)
+        for name, buffer in self.named_buffers():
+            if torch.is_floating_point(buffer):
+                target_dtype = torch.float32 if _keep_in_fp32(name) else dtype
+                buffer.data = buffer.data.to(target_dtype, non_blocking=non_blocking)
+        return self
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        guidance: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        return_dict: bool = True,
+        **kwargs: Any,
+    ) -> LingBotVideoTransformerOutput | tuple[torch.Tensor]:
+        """Denoise video latents with joint video-text attention."""
+        del encoder_hidden_states_image, guidance, kwargs
+        if isinstance(encoder_hidden_states, list):
+            if len(encoder_hidden_states) != 1:
+                raise ValueError("LingBot-Video expects one text-encoder output tensor.")
+            encoder_hidden_states = encoder_hidden_states[0]
+        batch, channels, frames, height, width = hidden_states.shape
+        patch_t, patch_h, patch_w = self.config.patch_size
+        grid_t, grid_h, grid_w = frames // patch_t, height // patch_h, width // patch_w
+        video_tokens = grid_t * grid_h * grid_w
+        text_tokens = encoder_hidden_states.shape[1]
+        device = hidden_states.device
+        if encoder_attention_mask is None:
+            encoder_attention_mask = torch.ones((batch, text_tokens), device=device, dtype=torch.bool)
+        text_lengths = encoder_attention_mask.sum(dim=-1).long()
+
+        patches = hidden_states.reshape(batch, channels, grid_t, patch_t, grid_h, patch_h, grid_w, patch_w)
+        patches = patches.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(batch, video_tokens,
+                                                                  patch_t * patch_h * patch_w * channels)
+        video_hidden = self.patch_embedder(patches)
+        text_hidden = self.text_embedder(encoder_hidden_states)
+        joint = torch.cat([video_hidden, text_hidden], dim=1)
+
+        rotary_parts: list[torch.Tensor] = []
+        for index in range(batch):
+            real_text_length = int(text_lengths[index].item())
+            positions = _make_joint_position_ids(real_text_length, grid_t, grid_h, grid_w, device)
+            maxima = (real_text_length + grid_t, grid_h - 1, grid_w - 1)
+            rotary = self.rope(positions, maxima=maxima)
+            if real_text_length < text_tokens:
+                padding = torch.zeros(
+                    text_tokens - real_text_length,
+                    rotary.shape[-1],
+                    device=device,
+                    dtype=rotary.dtype,
+                )
+                rotary = torch.cat([rotary, padding], dim=0)
+            rotary_parts.append(rotary)
+        rotary_emb = torch.stack(rotary_parts, dim=0)
+
+        attention_mask = None
+        moe_padding_mask = None
+        if bool((text_lengths < text_tokens).any()):
+            key_mask = torch.cat(
+                [
+                    torch.ones(batch, video_tokens, dtype=torch.bool, device=device),
+                    encoder_attention_mask.bool(),
+                ],
+                dim=1,
+            )
+            attention_mask = key_mask[:, None, None, :]
+            moe_padding_mask = key_mask
+        timestep_projection = self.time_proj(timestep.float())
+        timestep_embedding = self.time_embedder(timestep_projection)
+        token_embedding = timestep_embedding.unsqueeze(1).expand(batch, joint.shape[1], -1)
+        original_joint_length: int | None = None
+        if _sequence_parallel_world_size() > 1:
+            # Match the official CP order: project token modulation before
+            # placing every token-aligned tensor on the same padded shard.
+            temb6 = self.time_modulation(token_embedding.reshape(-1,
+                                                                 self.hidden_size)).reshape(batch, joint.shape[1], -1)
+            joint, original_joint_length = sequence_model_parallel_shard(joint, dim=1)
+            rotary_emb, _ = sequence_model_parallel_shard(rotary_emb, dim=1)
+            token_embedding, _ = sequence_model_parallel_shard(token_embedding, dim=1)
+            temb6, _ = sequence_model_parallel_shard(temb6, dim=1)
+            if moe_padding_mask is None:
+                moe_padding_mask = torch.ones(batch, original_joint_length, dtype=torch.bool, device=device)
+            moe_padding_mask, _ = sequence_model_parallel_shard(moe_padding_mask, dim=1)
+            moe_padding_mask = moe_padding_mask.reshape(-1)
+            temb6 = temb6.reshape(-1, 6 * self.hidden_size)
+        else:
+            temb6 = self.time_modulation(token_embedding.reshape(-1, self.hidden_size))
+            if moe_padding_mask is not None:
+                moe_padding_mask = moe_padding_mask.reshape(-1)
+
+        for block in self.blocks:
+            joint = block(
+                joint,
+                temb6,
+                rotary_emb,
+                attention_mask,
+                moe_padding_mask,
+                original_joint_length,
+            )
+
+        final_modulation = self.norm_out_modulation(token_embedding.reshape(-1, self.hidden_size))
+        shift, scale = final_modulation.reshape(*joint.shape[:2], -1).chunk(2, dim=-1)
+        final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
+        projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype))
+        if original_joint_length is not None:
+            projected = sequence_model_parallel_all_gather_with_unpad(projected, original_joint_length, dim=1)
+        projected = projected[:, :video_tokens]
+        output_channels = self.config.out_channels
+        output = projected.reshape(batch, grid_t, grid_h, grid_w, patch_t, patch_h, patch_w, output_channels)
+        output = output.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(batch, output_channels, frames, height, width)
+        if not return_dict:
+            return (output, )
+        return LingBotVideoTransformerOutput(sample=output)
+
+
+EntryClass = LingBotVideoTransformer3DModel

fastvideo/models/dits/lingbotworld2/__init__.py

@@ -0,0 +1,5 @@
+from .causal_fast import LingBotWorld2CausalFastTransformer3DModel
+
+__all__ = ["LingBotWorld2CausalFastTransformer3DModel"]
+
+EntryClass = LingBotWorld2CausalFastTransformer3DModel

fastvideo/models/dits/lingbotworld2/cam_utils.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+import numpy as np
+import os
+import torch
+from scipy.interpolate import interp1d
+from scipy.spatial.transform import Rotation, Slerp
+
+
+# --- Official Code (Leave Unchanged) ---
+ 
+def interpolate_camera_poses(
+    src_indices: np.ndarray, 
+    src_rot_mat: np.ndarray, 
+    src_trans_vec: np.ndarray, 
+    tgt_indices: np.ndarray,
+) -> torch.Tensor:
+    # interpolate translation
+    interp_func_trans = interp1d(
+        src_indices, 
+        src_trans_vec, 
+        axis=0, 
+        kind='linear', 
+        bounds_error=False,
+        fill_value="extrapolate",
+    )
+    interpolated_trans_vec = interp_func_trans(tgt_indices)
+
+    # interpolate rotation
+    src_quat_vec = Rotation.from_matrix(src_rot_mat)
+    # ensure there is no sudden change in qw
+    quats = src_quat_vec.as_quat().copy()  # [N, 4]
+    for i in range(1, len(quats)):
+        if np.dot(quats[i], quats[i-1]) < 0:
+            quats[i] = -quats[i]
+    src_quat_vec = Rotation.from_quat(quats)
+    slerp_func_rot = Slerp(src_indices, src_quat_vec)
+    interpolated_rot_quat = slerp_func_rot(tgt_indices)
+    interpolated_rot_mat = interpolated_rot_quat.as_matrix()
+
+    poses = np.zeros((len(tgt_indices), 4, 4))
+    poses[:, :3, :3] = interpolated_rot_mat
+    poses[:, :3, 3] = interpolated_trans_vec
+    poses[:, 3, 3] = 1.0
+    return torch.from_numpy(poses).float()
+
+
+def SE3_inverse(T: torch.Tensor) -> torch.Tensor:
+    Rot = T[:, :3, :3] # [B,3,3]
+    trans = T[:, :3, 3:] # [B,3,1]
+    R_inv = Rot.transpose(-1, -2)
+    t_inv = -torch.bmm(R_inv, trans)
+    T_inv = torch.eye(4, device=T.device, dtype=T.dtype)[None, :, :].repeat(T.shape[0], 1, 1)
+    T_inv[:, :3, :3] = R_inv
+    T_inv[:, :3, 3:] = t_inv
+    return T_inv
+
+
+def compute_relative_poses(
+    c2ws_mat: torch.Tensor, 
+    framewise: bool = False, 
+    normalize_trans: bool = True, 
+) -> torch.Tensor:
+    ref_w2cs = SE3_inverse(c2ws_mat[0:1])
+    relative_poses = torch.matmul(ref_w2cs, c2ws_mat)
+    # ensure identity matrix for 1st frame
+    relative_poses[0] = torch.eye(4, device=c2ws_mat.device, dtype=c2ws_mat.dtype)
+    if framewise:
+        # compute pose between i and i+1
+        relative_poses_framewise = torch.bmm(SE3_inverse(relative_poses[:-1]), relative_poses[1:])
+        relative_poses[1:] = relative_poses_framewise
+    if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning."
+        translations = relative_poses[:, :3, 3] # [f, 3]
+        max_norm = torch.norm(translations, dim=-1).max()
+        # only normlaize when moving
+        if max_norm > 0:
+            relative_poses[:, :3, 3] = translations / max_norm
+    return relative_poses
+
+
+@torch.no_grad()
+def create_meshgrid(n_frames: int, height: int, width: int, bias: float = 0.5, device='cuda', dtype=torch.float32) -> torch.Tensor:
+    x_range = torch.arange(width, device=device, dtype=dtype)
+    y_range = torch.arange(height, device=device, dtype=dtype)
+    grid_y, grid_x = torch.meshgrid(y_range, x_range, indexing='ij')
+    grid_xy = torch.stack([grid_x, grid_y], dim=-1).view([-1, 2]) + bias # [h*w, 2]
+    grid_xy = grid_xy[None, ...].repeat(n_frames, 1, 1) # [f, h*w, 2]
+    return grid_xy
+
+
+def get_plucker_embeddings(
+    c2ws_mat: torch.Tensor,
+    Ks: torch.Tensor,
+    height: int,
+    width: int,
+):
+    n_frames = c2ws_mat.shape[0]
+    device = c2ws_mat.device
+    dtype = c2ws_mat.dtype
+
+    grid_y, grid_x = torch.meshgrid(
+        torch.arange(height, device=device, dtype=dtype) + 0.5,
+        torch.arange(width, device=device, dtype=dtype) + 0.5,
+        indexing='ij',
+    )
+    x_flat = grid_x.reshape(-1)
+    y_flat = grid_y.reshape(-1)
+
+    fx, fy, cx, cy = Ks[0, 0], Ks[0, 1], Ks[0, 2], Ks[0, 3]
+    dirs = torch.stack([
+        (x_flat - cx) / fx,
+        (y_flat - cy) / fy,
+        torch.ones_like(x_flat),
+    ], dim=-1)
+    dirs = dirs / dirs.norm(dim=-1, keepdim=True)
+
+    rays_d = (c2ws_mat[:, :3, :3] @ dirs.T).transpose(1, 2)
+    rays_o = c2ws_mat[:, :3, 3].unsqueeze(1).expand_as(rays_d)
+    plucker_embeddings = torch.cat([rays_o, rays_d], dim=-1)
+    return plucker_embeddings.view(n_frames, height, width, 6)
+
+
+def get_Ks_transformed(
+    Ks: torch.Tensor,
+    height_org: int,
+    width_org: int,
+    height_resize: int,
+    width_resize: int,
+    height_final: int,
+    width_final: int,
+):
+    fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1]
+
+    scale_x = width_resize / width_org
+    scale_y = height_resize / height_org
+
+    fx_resize = fx * scale_x
+    fy_resize = fy * scale_y
+    cx_resize = cx * scale_x
+    cy_resize = cy * scale_y
+
+    crop_offset_x = (width_resize - width_final) / 2
+    crop_offset_y = (height_resize - height_final) / 2
+
+    cx_final = cx_resize - crop_offset_x
+    cy_final = cy_resize - crop_offset_y
+    
+    Ks_transformed = torch.zeros_like(Ks)
+    Ks_transformed[:, 0:1] = fx_resize
+    Ks_transformed[:, 1:2] = fy_resize
+    Ks_transformed[:, 2:3] = cx_final
+    Ks_transformed[:, 3:4] = cy_final
+
+    return Ks_transformed
+
+
+# --- Custom ---
+
+def prepare_camera_embedding(
+    action_path: str,
+    num_frames: int,
+    height: int,
+    width: int,
+    spatial_scale: int = 8,
+) -> tuple[torch.Tensor, int]:
+    c2ws = np.load(os.path.join(action_path, "poses.npy"))
+    len_c2ws = ((len(c2ws) - 1) // 4) * 4 + 1
+    num_frames = min(num_frames, len_c2ws)
+    c2ws = c2ws[:num_frames]
+
+    Ks = torch.from_numpy(
+        np.load(os.path.join(action_path, "intrinsics.npy"))
+    ).float()
+    Ks = get_Ks_transformed(
+        Ks,
+        height_org=480,
+        width_org=832,
+        height_resize=height,
+        width_resize=width,
+        height_final=height,
+        width_final=width,
+    )
+    Ks = Ks[0]  # use first frame
+
+    len_c2ws = len(c2ws)
+    num_latent_frames = (len_c2ws - 1) // 4 + 1
+    c2ws_infer = interpolate_camera_poses(
+        src_indices=np.linspace(0, len_c2ws - 1, len_c2ws),
+        src_rot_mat=c2ws[:, :3, :3],
+        src_trans_vec=c2ws[:, :3, 3],
+        tgt_indices=np.linspace(0, len_c2ws - 1, num_latent_frames),
+    )
+    c2ws_infer = compute_relative_poses(c2ws_infer, framewise=True)
+    Ks = Ks.repeat(num_latent_frames, 1)
+    plucker = get_plucker_embeddings(c2ws_infer, Ks, height, width)  # [F, H, W, 6]
+
+    # reshpae
+    latent_height = height // spatial_scale
+    latent_width = width // spatial_scale
+    plucker = plucker.view(num_latent_frames, latent_height, spatial_scale, latent_width, spatial_scale, 6)
+    plucker = plucker.permute(0, 1, 3, 5, 2, 4).contiguous()
+    plucker = plucker.view(num_latent_frames, latent_height, latent_width, 6 * spatial_scale * spatial_scale)
+    c2ws_plucker_emb = plucker.permute(3, 0, 1, 2).contiguous().unsqueeze(0)
+
+    return c2ws_plucker_emb, num_frames