Revert redundant fixes

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

.buildkite/pipeline.yml

@@ -217,6 +217,17 @@ steps:
             limit: 2
       agents:
         queue: "default"
+    - label: ":test_tube: Eval Metrics Tests"
+      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "eval"
+      command: "timeout 90m .buildkite/scripts/pr_test.sh"
+      retry:
+        automatic:
+          - exit_status: 128
+            limit: 3
+          - exit_status: -1
+            limit: 2
+      agents:
+        queue: "default"
 
   # ============================================================
   # Fastcheck: Runs on every PR (~10-15 min parallel)
@@ -483,3 +494,15 @@ steps:
                   - TEST_TYPE=train_framework
                 agents:
                   queue: "default"
+            - path:
+                - "fastvideo/eval/**"
+                - "fastvideo/tests/eval/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 90m .buildkite/scripts/pr_test.sh"
+                label: ":test_tube: Eval Metrics Tests"
+                env:
+                  - TEST_TYPE=eval
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -76,7 +76,7 @@ EFFECTIVE_PR=${BUILDKITE_PULL_REQUEST:-false}
 if [ "$EFFECTIVE_PR" = "false" ] && [ -n "${PR_NUMBER:-}" ]; then
     EFFECTIVE_PR=$PR_NUMBER
 fi
-MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} TEST_SCOPE=${TEST_SCOPE:-} IMAGE_VERSION=$IMAGE_VERSION"
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} TEST_SCOPE=${TEST_SCOPE:-} BUILDKITE_BUILD_URL=${BUILDKITE_BUILD_URL:-} BUILDKITE_BUILD_ID=${BUILDKITE_BUILD_ID:-} BUILDKITE_JOB_ID=${BUILDKITE_JOB_ID:-} IMAGE_VERSION=$IMAGE_VERSION"
 
 POST_RUN_HOOK=""
 
@@ -219,6 +219,10 @@ case "$TEST_TYPE" in
         log "Running fastvideo.train framework tests..."
         MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_train_framework_tests"
         ;;
+    "eval")
+        log "Running eval metric tests..."
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_eval_tests"
+        ;;
     "lora_extraction")
         log "Running LoRA extraction tests..."
         MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_lora_extraction_tests"

.github/mergify.yml

@@ -272,7 +272,7 @@ pull_request_rules:
         remove: [needs-rebase]
 
   # ============================================================
-  # Auto-merge and auto-rebase
+  # Auto-merge
   # ============================================================
 
   - name: auto-merge when ready and all checks pass
@@ -290,16 +290,6 @@ pull_request_rules:
       merge:
         method: squash
 
-  - name: auto-update when ready
-    conditions:
-      - label=ready
-      - "#approved-reviews-by>=1"
-      - -conflict
-      - -closed
-      - -draft
-    actions:
-      update: {}
-
   # ============================================================
   # PR title format help
   # ============================================================

.github/workflows/ci-slash-commands.yml

@@ -125,7 +125,7 @@ jobs:
           set -euo pipefail
           TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true)
 
-          VALID="encoder vae transformer kernel unit dreamverse ssim training lora-inference lora-training distillation self-forcing vsa vmoba performance api train-framework full fastcheck pre-commit"
+          VALID="encoder vae transformer kernel unit dreamverse ssim training lora-inference lora-training distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit"
           if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then
             echo "Unknown test: '$TEST_NAME'. Valid: $VALID"
             exit 1
@@ -139,7 +139,7 @@ jobs:
             [distillation]=distillation_dmd [self-forcing]=self_forcing
             [vsa]=training_vsa [vmoba]=inference_vmoba
             [performance]=performance [api]=api_server
-            [train-framework]=train_framework
+            [train-framework]=train_framework [eval]=eval
           )
 
           if [ "$TEST_NAME" = "full" ]; then

.github/workflows/publish-kernel.yml

@@ -53,23 +53,15 @@ jobs:
       fail-fast: false
       matrix:
           os: [ubuntu-22.04]
-          python-version: ['3.10', '3.11', '3.12']
+          python-version: ['3.12']
           torch-cuda:
-            # - torch-version: '2.5.1'
-            #   cuda-version: '12.4.1'
-            #   torch-cuda-short: 'cu124'
-            # - torch-version: '2.6.0'
-            #   cuda-version: '12.6.3'
-            #   torch-cuda-short: 'cu126'
-            # - torch-version: '2.7.1'
-            #   cuda-version: '12.8.0'
-            #   torch-cuda-short: 'cu128'
-            # - torch-version: '2.9.1'
-            #   cuda-version: '12.8.0'
-            #   torch-cuda-short: 'cu128'
-            - torch-version: '2.10.0'
-              cuda-version: '12.8.0'
-              torch-cuda-short: 'cu128'
+            # torch 2.12 dropped cu128; cu130 (CUDA 13) is the default wheel, cu126 covers older drivers.
+            - torch-version: '2.12.0'
+              cuda-version: '12.6.3'
+              torch-cuda-short: 'cu126'
+            - torch-version: '2.12.0'
+              cuda-version: '13.0.0'
+              torch-cuda-short: 'cu130'
 
     steps:
       - name: Free up disk space
@@ -104,7 +96,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }}
-        uses: Jimver/cuda-toolkit@v0.2.21
+        uses: Jimver/cuda-toolkit@v0.2.35
         id: cuda-toolkit
         with:
           cuda: ${{ matrix.torch-cuda.cuda-version }}
@@ -152,9 +144,27 @@ jobs:
 
           cd fastvideo-kernel
           git submodule update --init --recursive  # Ensure ThunderKittens submodule is initialized
-          # Release builds are produced on GPU-less runners, so force-enable TK and target Hopper.
-          export TORCH_CUDA_ARCH_LIST="9.0a"
-          export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+          # Release builds run on GPU-less runners, so force-enable TK and set the arch explicitly.
+          # The CUDA 13 (cu130) leg additionally targets Blackwell sm_120a and builds the
+          # attn_qat_infer FP4 inference kernels (they need CUDA Toolkit 12.8+, so the cu126 leg
+          # stays Hopper-only and ships without them).
+          # Cap Ninja build concurrency to fit the 16 GB standard runner. CUTLASS FP4
+          # and ThunderKittens template TUs each need several GB; Ninja's default
+          # (cores+2 ≈ 6 parallel jobs) exhausts RAM and the host OOM-kills the build
+          # (SIGTERM / exit 143, surfaced as "runner lost communication").
+          if [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then
+            export TORCH_CUDA_ARCH_LIST="9.0a;12.0a"
+            export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+            # A single FP4 TU (attn_qat_infer) can use ~8-12 GB on its own, so serialize.
+            export CMAKE_BUILD_PARALLEL_LEVEL=1
+          else
+            export TORCH_CUDA_ARCH_LIST="9.0a"
+            export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+            # No FP4 here; only the two TK TUs are heavy (single-arch) and they fit
+            # side by side. Uncapped (cores+2=6) has built but occasionally OOMs when
+            # the light TUs pile on; -j4 overlaps them safely (~12-14 GB peak).
+            export CMAKE_BUILD_PARALLEL_LEVEL=4
+          fi
 
           # Build standard wheel (no local version suffix) for PyPI
           python -m build --wheel --outdir dist
@@ -183,8 +193,8 @@ jobs:
           mv fixed_dist/*.whl dist/
 
       - name: Upload wheel artifact
-        # Only upload if it's the "main" CUDA version we want on PyPI
-        # We upload all to artifacts for inspection/GH releases, but give them distinct artifact names
+        # Upload every matrix leg as a distinct artifact (for inspection / GitHub releases).
+        # The publish job below selects which CUDA build is pushed to PyPI.
         uses: actions/upload-artifact@v4
         with:
           name: fastvideo_kernel-py${{ matrix.python-version }}-${{ matrix.torch-cuda.torch-cuda-short }}-torch${{ matrix.torch-cuda.torch-version }}
@@ -204,13 +214,15 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: '3.12'
 
       - name: Download PyPI wheels
+        # Publish only the cu130 (CUDA 13) wheel to PyPI — it includes the Blackwell FP4 kernels.
+        # The cu126 wheel stays available as a build artifact / GitHub-release asset.
         uses: actions/download-artifact@v4
         with:
           path: fastvideo-kernel/dist/
-          pattern: 'fastvideo_kernel-py*'
+          pattern: 'fastvideo_kernel-py*-cu130-*'
           merge-multiple: true
 
       - name: Install uv

.gitignore

@@ -130,3 +130,8 @@ apps/dreamverse/web/.env.production.local
 .sisyphus/
 openspec/
 fastvideo/tests/ssim/reference_videos/**
+
+# Editor logs and local Python version pins (accidentally committed)
+*.nvimlog
+.nvimlog
+.python-version

.nvimlog

@@ -1 +0,0 @@
-WRN 2026-03-26T13:46:33.469 ?.19646    server_start:193: Failed to start server: operation not permitted: /var/folders/z_/h_6myyk14d1b7z87z3vy4mjh0000gn/T/nvim.dsynkd/iSe0el/nvim.19646.0

.python-version

@@ -1 +0,0 @@
-3.12

README.md

@@ -9,8 +9,9 @@
 **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.**
 
 ## NEWS
-- `2026/03/17`: Release Live demo: [Into the Dreamverse: Vibe Directing in FastVideo](https://dreamverse.fastvideo.org/), check out the [Blog](https://haoailab.com/blogs/dreamverse/).
-- `2026/03/13`: Release Live demo: [Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU](https://1080p.fastvideo.org/), check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/).
+- `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/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).
 - `2025/08/04`: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
 

apps/dreamverse/docker/README.md

@@ -32,7 +32,7 @@ checkout with the `dreamverse` extra, installs the FA4
 flash-attention fork, builds native FFmpeg, and installs FlashInfer for NVFP4
 quantization.
 
-FastVideo's pinned `fastvideo-kernel==0.2.6` package is installed by default.
+FastVideo's pinned `fastvideo-kernel==0.3.0` package is installed by default.
 To rebuild `fastvideo-kernel` from this checkout during the image build, set:
 
 ```bash

docs/contributing/ci_architecture.md

@@ -103,6 +103,7 @@ failing test's output.
 | [Performance Tests](performance_benchmarks.md) | `performance` | 30 min |
 | API Server Tests | `api_server` | 30 min |
 | Train Framework Tests | `train_framework` | 30 min |
+| Eval Metrics Tests | `eval` | 90 min |
 
 If a Full Suite test fails, check the Buildkite build log for the failing step's output.
 Fix the regression, push, and comment `/merge` again to re-trigger.
@@ -281,6 +282,7 @@ Triggers a specific Buildkite test or suite on the current PR branch.
 | `/test performance` | Performance benchmarks | `performance` |
 | `/test api` | API server integration tests | `api_server` |
 | `/test train-framework` | `fastvideo.train` GPU model loading + per-method tests | `train_framework` |
+| `/test eval` | Eval metric regression vs upstream reference scores | `eval` |
 | `/test full` | Entire Full Suite | all (with `TEST_SCOPE=full`) |
 | `/test fastcheck` | Entire Fastcheck suite | fastcheck (with `TEST_SCOPE=fastcheck`) |
 | `/test pre-commit` | Pre-commit checks on PR code | — (runs `ci-precommit.yml` via `workflow_call`) |

docs/contributing/performance_benchmarks.md

@@ -21,22 +21,31 @@ It serves three audiences:
 # fastvideo/tests/performance/results/
 pytest fastvideo/tests/performance/ -vs
 
-# Optional: compare against the rolling HF baseline (read-only outside CI).
+# Optional: compare against the rolling HF baseline.
 # PERF_REPORTS_DIR defaults to /root/data/perf_reports for Modal/CI, so
 # override it when running outside the container.
 PERF_REPORTS_DIR=/tmp/fastvideo_perf_reports \
 python fastvideo/tests/performance/compare_baseline.py
 
+# Optional: explicitly upload a passing local/manual run.
+HF_TOKEN=hf_... \
+PERF_RUN_SOURCE=local \
+PERF_UPLOAD_POLICY=pass \
+PERF_REPORTS_DIR=/tmp/fastvideo_perf_reports \
+python fastvideo/tests/performance/compare_baseline.py
+
 # Optional: build the Plotly dashboard locally.
 PERF_REPORTS_DIR=/tmp/fastvideo_perf_reports \
 python fastvideo/tests/performance/dashboard.py
 ```
 
-The pytest run never uploads anything. `compare_baseline.py` only writes to
-the HF dataset when `TEST_SCOPE=full` *and* `BUILDKITE_BRANCH=main`, so local
-runs are always read-only. The report directory default is container-oriented;
-set `PERF_REPORTS_DIR` to a writable local path when generating dashboards or
-when you want local Markdown/normalized-result artifacts from the comparator.
+The pytest run never uploads anything. `compare_baseline.py` uploads only when
+`PERF_UPLOAD_POLICY` is set. Local uploads are explicit opt-in and require HF
+credentials. PR/direct performance runs upload passing records for dashboard
+visibility, while scheduled-main runs upload both pass and fail records. The
+report directory default is container-oriented; set `PERF_REPORTS_DIR` to a
+writable local path when generating dashboards or when you want local
+Markdown/normalized-result artifacts from the comparator.
 `compare_baseline.py` reads every `perf_*.json` currently present in
 `fastvideo/tests/performance/results/`; remove stale result files if you only
 want to compare the latest local run.
@@ -68,7 +77,9 @@ fastvideo/tests/performance/
 
 The HF dataset (`FastVideo/performance-tracking` by default) holds one
 normalized JSON per `(model_id, gpu_type, run)` tuple. The rolling baseline is
-the median of the last 5 successful records for that model+GPU.
+the median of the last 5 successful, baseline-eligible records for that
+model+GPU. PR and local records are visible in the dashboard but are not
+baseline eligible.
 
 ## Planned Coverage
 
@@ -94,11 +105,16 @@ Each benchmark records six metrics:
 
 `test_inference_performance.py` temporarily sets `FASTVIDEO_STAGE_LOGGING=1`
 while it runs so pipeline stage execution times are available in
-`generate_video(...).logging_info`. It maps `TextEncodingStage` to
-`text_encoder_time_s`, `DenoisingStage` and `DmdDenoisingStage` to
-`dit_time_s`, and `DecodingStage` to `vae_decode_time_s`. If a pipeline does
-not report one of those stages, that component metric is stored as `null` and
-is skipped by the static threshold and rolling baseline checks.
+`generate_video(...).logging_info`. Stage logs use pipeline-unique keys such as
+`prompt_encoding_stage` so duplicate stage classes do not collide. For
+`PipelineStage` entries, the extractor maps the `stage_class` field:
+`TextEncodingStage` maps to `text_encoder_time_s`, `DenoisingStage` and
+`DmdDenoisingStage` map to `dit_time_s`, and `DecodingStage` maps to
+`vae_decode_time_s`, with a fallback for older logs that used the class name as
+the stage key. Generator-side timings such as `PostDecodeFrameProcessStage`,
+`VideoSaveStage`, and `AudioMuxStage` are intentionally ignored. If a pipeline
+does not report one of the mapped stages, that component metric is stored as
+`null` and is skipped by the static threshold and rolling baseline checks.
 
 ## The two gates
 
@@ -138,16 +154,16 @@ headroom and almost never need touching.
 
 ### Rolling baseline (per `(model_id, gpu_type)`)
 
-`compare_baseline.py` loads the last 5 successful records for the same
-`(model_id, gpu_type)` from the HF dataset, computes the median for each
-available metric, and fails if the current run regresses by more than
+`compare_baseline.py` loads the last 5 successful, baseline-eligible records
+for the same `(model_id, gpu_type)` from the HF dataset, computes the median
+for each available metric, and fails if the current run regresses by more than
 `PERF_MAX_REGRESSION` (default 5%). For latency, memory, and component times,
 higher values are regressions. For throughput, lower values are regressions.
 
 This is the **drift detector** — it catches sub-threshold regressions that
-slowly add up. It only persists new records when running the full suite on
-`main`. Local and pull-request runs can compare against the HF baseline, but
-they do not update it.
+slowly add up. Only scheduled-main successful records are baseline eligible.
+Local and pull-request runs can upload dashboard-visible records, but they do
+not update future gating baselines.
 
 When the baseline shifts for a legitimate reason (torch upgrade, kernel
 change, etc.) and CI starts failing, use the
@@ -215,6 +231,8 @@ result, used as the rolling-baseline source of truth.
 Older records in the HF dataset may not have component timing fields. The
 comparator ignores missing or `null` metrics when computing a median, and the
 dashboard lists skipped plots for metric series that have no non-null values.
+Records missing both `run_source` and `baseline_eligible` are treated as legacy
+successful main/full-suite uploads and remain eligible for rolling baselines.
 
 ## Environment variable reference
 
@@ -224,8 +242,11 @@ dashboard lists skipped plots for metric series that have no non-null values.
 | `PERFORMANCE_TRACKING_ROOT` | `/tmp/perf-tracking` | `compare_baseline.py`, `dashboard.py` | Local directory the HF dataset is synced to. |
 | `PERF_REPORTS_DIR` | `/root/data/perf_reports` | `compare_baseline.py`, `dashboard.py` | Where the Markdown summary and Plotly HTML get written for Buildkite to pick up. |
 | `HF_REPO_ID` | `FastVideo/performance-tracking` | `hf_store.py` | HF dataset repo holding rolling-baseline records. |
-| `HF_API_KEY` | unset | `hf_store.py` | Required for upload (main-branch full-suite only); reads work without it. |
-| `TEST_SCOPE` | unset | `compare_baseline.py` | Set to `full` together with `BUILDKITE_BRANCH=main` to enable HF persistence. |
+| `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `hf_store.py` | Required for upload or private dataset reads. |
+| `PERF_RUN_SOURCE` | inferred | `compare_baseline.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
+| `PERF_UPLOAD_POLICY` | `never` | `compare_baseline.py` | Upload policy: `never`, `pass`, or `always`. |
+| `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Static-threshold pytest exit code, used so scheduled-main failures can be uploaded with `success=false`. |
+| `TEST_SCOPE` | unset | `compare_baseline.py` | CI context used to infer scheduled-main runs together with `BUILDKITE_BRANCH=main`. |
 | `BUILDKITE_BRANCH`, `BUILDKITE_COMMIT`, `BUILDKITE_PULL_REQUEST` | unset | `compare_baseline.py`, `test_inference_performance.py` | CI metadata stamped into records. |
 | `DASHBOARD_DAYS` | `30` | `dashboard.py` | Lookback window for the Plotly trend pages. |
 | `PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS` | `3600` | `hf_store.py` | Freshness window for reusing an existing HF sync when requested by dashboard consumers. |

docs/inference/optimizations.md

@@ -11,6 +11,9 @@ This page describes the various options for speeding up generation times in Fast
   - [Sliding Tile Attention (Archived)](#sliding-tile-attention-archived)
   - [Sage Attention](#sage-attention)
   - [Sage Attention 3](#sage-attention-3)
+
+- [FP8 Weight Quantization](#fp8-weight-quantization)
+
 - [Adaptive Guidance (CFG gating)](#adaptive-guidance-cfg-gating)
 
 - [torch.compile](#torch-compile)
@@ -218,6 +221,50 @@ These backends are model-specific and require the corresponding kernels and
 dependencies. Use the support matrix and model examples to confirm compatibility
 before enabling them.
 
+## FP8 Weight Quantization
+
+**`transformer_quant="FP8"`**
+
+Quantizes DiT linear layers (attention projections and FFN) to FP8 e4m3.
+
+On GPUs older than sm89, the FP8 matmul falls back to a bf16 dequant path
+automatically.
+
+### Requirements
+
+- **GPU**: sm89+ (H100, L40S, RTX 4090, or newer) for hardware FP8 compute
+- No additional packages required beyond the base FastVideo install
+
+### Usage
+
+```python
+from fastvideo import VideoGenerator
+from fastvideo.layers.quantization import get_quantization_config
+
+gen = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    # Pass an instance — the bare string is not resolved on the from_pretrained path.
+    transformer_quant=get_quantization_config("FP8")(),          # per-tensor (default)
+    # transformer_quant=get_quantization_config("FP8")(granularity="channel"),  # slower, higher accuracy
+)
+gen.generate(request={"prompt": "A raccoon in sunflowers", "output": {"save_video": True}})
+```
+
+Or run the example script:
+
+```bash
+python examples/inference/optimizations/fp8_wan2_1_1_3b.py
+python examples/inference/optimizations/fp8_wan2_1_1_3b.py --granularity channel
+python examples/inference/optimizations/fp8_wan2_1_1_3b.py --bf16  # baseline
+```
+
+### Granularity
+
+| Mode | Weight scales | Activation scales | Speed | Accuracy |
+|------|--------------|-------------------|-------|----------|
+| `tensor` (default) | per-tensor | per-tensor | faster | lower |
+| `channel` | per-output-channel | per-token (rowwise) | slower | higher |
+
 <a id="torch-compile"></a>
 
 ## torch.compile

examples/inference/eval/all_metrics_demo.py

@@ -0,0 +1,159 @@
+"""End-to-end demo: generate one LTX2 video, score it with every
+applicable metric via the top-level Evaluator API.
+
+Two phases:
+
+1. **Generate** one LTX2-Distilled video (single GPU).  Mirrors the
+   pattern in ``examples/inference/basic/basic_ltx2_distilled.py``,
+   forced to single GPU for portability.
+
+2. **Score** the same video duplicated N times into ``gen/`` and ``ref/``
+   directories.  Because gen == ref by construction:
+
+   * Per-sample paired metrics (PSNR / SSIM / LPIPS / gt_optical_flow)
+     return their perfect-match values (PSNR ≈ 100, SSIM ≈ 1, LPIPS ≈ 0).
+   * Set metric ``common.fvd`` returns ≈ 0.
+   * VBench per-sample metrics return the same score on every duplicate
+     (a stable read on the LTX2 sample).
+
+The scoring block is the point of this script: 4 lines to build the
+samples list (``samples_from``) and run the Evaluator
+(``ev.evaluate``).  All input-shape ceremony goes away because each
+metric reads only the keys it needs from the fat sample dict.
+"""
+from __future__ import annotations
+
+import os
+import shutil
+import tempfile
+import time
+from pathlib import Path
+
+import torch
+
+# ---------------------------------------------------------------------
+# Generation — make one LTX2 video to evaluate.
+# ---------------------------------------------------------------------
+
+PROMPT = (
+    "A warm sunny backyard. The camera starts in a tight cinematic close-up "
+    "of a woman and a man in their 30s, facing each other with serious "
+    "expressions. The woman, emotional and dramatic, says softly, \"That's "
+    "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly "
+    "annoyed: \"Stop being so dramatic, Jess.\""
+)
+
+OUTPUT_PATH = "fastvideo/tests/eval/asset/ltx2.mp4"
+N_DUP = 4  # how many times to duplicate the video for the gen/ref corpora
+
+
+def generate_one_ltx2_video() -> str:
+    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
+    from fastvideo import VideoGenerator
+
+    Path(OUTPUT_PATH).parent.mkdir(parents=True, exist_ok=True)
+    # Davids048/LTX2-Base-Diffusers is the audio-capable LTX-2 checkpoint
+    # (the Distilled variant ships without the audio VAE, so its mp4
+    # audio track is silence/noise — unusable for audio.* metrics).
+    generator = VideoGenerator.from_pretrained(
+        "Davids048/LTX2-Base-Diffusers",
+        num_gpus=1,
+    )
+    generator.generate_video(
+        prompt=PROMPT,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+        num_frames=121,  # ~5s @ 24 fps — long enough for audio.desync (Synchformer ≥14 segments)
+        height=480,
+        width=832,
+        fps=24,
+    )
+    generator.shutdown()
+    torch.cuda.empty_cache()
+    return OUTPUT_PATH
+
+
+# ---------------------------------------------------------------------
+# Eval — the point of the script.  4 lines from "two paths" to results.
+# ---------------------------------------------------------------------
+
+def _all_registered_metrics() -> list[str]:
+    """Every metric in the registry, sorted.  Combined with
+    ``skip_missing_deps=True`` this is the "run everything that works in
+    this venv" pattern — missing-dep / setup-failed metrics drop with a
+    warning, runtime-missing inputs (audio, masks) surface as the metric
+    returning a skipped MetricResult."""
+    from fastvideo.eval.registry import list_metrics
+    return list_metrics()
+
+
+def score_all_metrics(video_path: str) -> None:
+    # Duplicate one video into gen/ ref/ — identical content by design.
+    tmp = Path(tempfile.mkdtemp(prefix="all_metrics_demo_"))
+    gen_dir = tmp / "gen"
+    ref_dir = tmp / "ref"
+    gen_dir.mkdir()
+    ref_dir.mkdir()
+    for i in range(N_DUP):
+        shutil.copyfile(video_path, gen_dir / f"clip_{i:03d}.mp4")
+        shutil.copyfile(video_path, ref_dir / f"clip_{i:03d}.mp4")
+
+    # === The whole eval, in 4 lines. ===
+    from fastvideo.eval import create_evaluator, samples_from
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    t_init0 = time.perf_counter()
+    ev = create_evaluator(metrics=_all_registered_metrics(),
+                          device="cuda:0", num_gpus=1, skip_missing_deps=True)
+    t_init1 = time.perf_counter()
+    samples = samples_from(video=gen_dir, reference=ref_dir, text_prompt=PROMPT, fps=24.0,
+                           extract_audio=True)  # auto-extract audio track from videos
+    t_eval0 = time.perf_counter()
+    results = ev.evaluate(samples=samples)
+    torch.cuda.synchronize()
+    t_eval1 = time.perf_counter()
+    # ===================================
+
+    peak_gb = torch.cuda.max_memory_allocated() / 1024**3
+    peak_reserved_gb = torch.cuda.max_memory_reserved() / 1024**3
+    print(f"\n=== Timing & memory ===")
+    print(f"  Evaluator construction + model loads: {t_init1 - t_init0:.1f}s")
+    print(f"  ev.evaluate(samples):                 {t_eval1 - t_eval0:.1f}s")
+    print(f"  total eval block:                      {t_eval1 - t_init0:.1f}s")
+    print(f"  peak GPU mem allocated:                {peak_gb:.2f} GB")
+    print(f"  peak GPU mem reserved (caches):        {peak_reserved_gb:.2f} GB")
+
+    print("\n=== Per-sample scores (sample 0; identical for every duplicate) ===")
+    for name in ev.metric_names:
+        r = results[0].get(name)
+        if r is None:
+            continue  # set metric, will show in corpus section
+        if r.score is None:
+            print(f"  {name}: SKIPPED — {r.details.get('skipped', 'no score')}")
+        else:
+            print(f"  {name}: {r.score:.4f}")
+
+    print("\n=== Corpus scores (set metrics) ===")
+    if not results.corpus:
+        print("  (none)")
+    for name, r in results.corpus.items():
+        if r.score is None:
+            print(f"  {name}: SKIPPED — {r.details.get('skipped', 'no score')}")
+        else:
+            print(f"  {name}: {r.score:.4f}  details={r.details}")
+
+
+def main() -> None:
+    # Reuse an existing generation if one is already on disk so iteration
+    # on the eval block doesn't pay the generation cost every run.
+    if not Path(OUTPUT_PATH).exists():
+        print(f"[gen] No video at {OUTPUT_PATH}; generating with LTX2-Distilled (1 GPU)...")
+        generate_one_ltx2_video()
+    else:
+        print(f"[gen] Reusing existing {OUTPUT_PATH}")
+
+    score_all_metrics(OUTPUT_PATH)
+
+
+if __name__ == "__main__":
+    main()

examples/inference/optimizations/FastWan_QAD_TAEHV.py

@@ -0,0 +1,286 @@
+"""Fast NVFP4 linear inference for Wan2.1-T2V-1.3B with TAEHV decoding.
+
+This is the FP4-linear fast path from ``fp4_linear_wan2_1_1_3b.py`` with the
+heavy Wan VAE swapped out for TAEHV -- a tiny autoencoder that decodes Wan2.1
+latents directly (no denormalization) and is dramatically faster / lighter.
+
+How it works: the generator runs with ``output_type="latent"`` so the pipeline
+returns raw denoised latents instead of pixels (the Wan VAE is offloaded and
+never used). We then decode those latents with TAEHV in this script and save
+the frames ourselves. This mirrors the FastVideo-Quantization
+``quantization_example_taehv.py`` proof-of-concept, but kept clean: TAEHV is a
+pip package (no ``sys.path`` hacks), the latent->uint8 conversion is vectorized,
+and there is no dead profiler / sanitization code.
+
+Requirements:
+    - Blackwell GPU (B200/B300, sm100a/sm103a) for the FP4 linear path
+    - flashinfer (``pip install flashinfer-python``)
+    - TAEHV weights ``taew2_1.pth`` (https://github.com/madebyollin/taehv)
+
+Usage:
+    python fp4_linear_taehv_wan2_1_1_3b.py                    # FP4 + TAEHV + compile
+    python fp4_linear_taehv_wan2_1_1_3b.py --no-taehv         # FP4 + full Wan VAE
+    python fp4_linear_taehv_wan2_1_1_3b.py --no-compile       # eager
+    python fp4_linear_taehv_wan2_1_1_3b.py --baseline         # dense bf16 reference
+    python fp4_linear_taehv_wan2_1_1_3b.py --distilled_model ''  # base Wan2.1 weights
+"""
+
+import argparse
+import contextlib
+import logging
+import os
+import time
+
+import imageio
+import torch
+
+from fastvideo import VideoGenerator
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig
+
+OUTPUT_PATH = "video_samples"
+
+# Distilled, quantization-aware (QAD) transformer for Wan2.1-1.3B (3 steps,
+# guidance 1.0). Loaded on top of the base Wan2.1 pipeline; pass
+# ``--distilled_model ''`` to run the base weights instead.
+DEFAULT_DISTILLED_MODEL = "FastVideo/FastWan-QAD-1.3B"
+DISTILLED_WEIGHTS_FILE = (
+    "generator_inference_transformer/diffusion_pytorch_model.safetensors"
+)
+
+# TAEHV checkpoint for Wan2.1. Clone https://github.com/madebyollin/taehv to get
+# ``taew2_1.pth`` (Wan 2.1 / Wan 2.2-14B / Qwen-Image all use this VAE).
+DEFAULT_TAEHV_CHECKPOINT = "/root/taehv/taew2_1.pth"
+
+PROMPT = (
+    "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+    "wide with interest. The playful yet serene atmosphere is complemented by soft "
+    "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+)
+
+
+class TaehvDecoder:
+    """Thin wrapper around the TAEHV tiny autoencoder for Wan2.1 latents.
+
+    TAEHV consumes the *normalized* latents the diffusion model produces (the
+    same representation FastVideo carries internally), so no denormalization is
+    needed -- unlike the full Wan VAE path.
+    """
+
+    def __init__(self, checkpoint_path: str, device: str = "cuda",
+                 dtype: torch.dtype = torch.float16) -> None:
+        from taehv import TAEHV  # pip-installed; no sys.path manipulation
+        self.device = device
+        self.dtype = dtype
+        print(f"Loading TAEHV from {checkpoint_path} ...")
+        self.model = TAEHV(checkpoint_path=checkpoint_path).to(device, dtype).eval()
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor):
+        """Decode FastVideo latents into uint8 RGB frames.
+
+        Args:
+            latents: ``[B, C, T, H, W]`` (NCTHW) normalized latent tensor.
+
+        Returns:
+            A ``(T, H, W, 3)`` uint8 numpy array ready for ``imageio.mimsave``.
+        """
+        # NCTHW -> NTCHW (TAEHV's expected layout), on the TAEHV device/dtype.
+        latents = latents.permute(0, 2, 1, 3, 4).to(self.device, self.dtype)
+        decoded = self.model.decode_video(
+            latents, parallel=True, show_progress_bar=False)
+        # decoded: [B, T, 3, H, W] in [0, 1]. Take batch 0, vectorize to uint8.
+        frames = (decoded[0].clamp(0, 1) * 255).to(torch.uint8)
+        return frames.permute(0, 2, 3, 1).cpu().numpy()
+
+
+def resolve_distilled_weights(hf_id: str) -> str:
+    """Return a local path to the distilled transformer safetensors."""
+    if os.path.exists(hf_id):
+        return hf_id
+    from huggingface_hub import hf_hub_download
+    return hf_hub_download(repo_id=hf_id, filename=DISTILLED_WEIGHTS_FILE)
+
+
+@contextlib.contextmanager
+def silence_request_log():
+    """Quiet ``VideoGenerator.generate``'s per-request config printout.
+
+    Each ``generate(...)`` call logs a multi-line debug block (height/width/
+    prompt/steps/...) at INFO via ``logger.info`` in
+    ``fastvideo.entrypoints.video_generator``. There is no built-in switch,
+    so this context manager raises that logger's level to WARNING while the
+    warmup calls run, then restores it for the timed run.
+    """
+    vg_logger = logging.getLogger("fastvideo.entrypoints.video_generator")
+    prev_level = vg_logger.level
+    vg_logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        vg_logger.setLevel(prev_level)
+
+
+def resolve_taehv_checkpoint(path: str) -> str:
+    """Validate the TAEHV checkpoint path, with a helpful error if missing."""
+    if os.path.exists(path):
+        return path
+    raise FileNotFoundError(
+        f"TAEHV checkpoint not found at {path!r}. Clone the weights with:\n"
+        "    git clone https://github.com/madebyollin/taehv\n"
+        "and pass --taehv_checkpoint <repo>/taew2_1.pth")
+
+
+def build_generator(args: argparse.Namespace) -> VideoGenerator:
+    model_id = args.model
+
+    # Half precision everywhere; DiT linears are additionally NVFP4-quantized
+    # via dit_config.quant_config below.
+    pipeline_config = PipelineConfig.from_pretrained(model_id)
+    pipeline_config.dit_precision = "bf16"
+    pipeline_config.vae_precision = "bf16"
+    pipeline_config.text_encoder_precisions = ("bf16",)
+
+    if not args.baseline:
+        pipeline_config.dit_config.quant_config = NVFP4QATConfig()
+
+    compile_enabled = not args.no_compile
+
+    extra_kwargs = {}
+    if args.distilled_model:
+        weights_path = resolve_distilled_weights(args.distilled_model)
+        print(f"Using distilled weights: {args.distilled_model} -> {weights_path}")
+        extra_kwargs["init_weights_from_safetensors"] = weights_path
+
+    if args.taehv:
+        # Skip the in-pipeline VAE decode entirely: the pipeline returns raw
+        # latents, the Wan VAE is offloaded to CPU (and not compiled) since we
+        # decode with TAEHV in this script instead.
+        extra_kwargs["output_type"] = "latent"
+
+    generator = VideoGenerator.from_pretrained(
+        model_id,
+        pipeline_config=pipeline_config,
+        num_gpus=args.num_gpus,
+        # Keep everything resident on the GPU -- no offloading, except the
+        # unused Wan VAE when TAEHV handles decoding.
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        vae_cpu_offload=args.taehv,
+        text_encoder_cpu_offload=False,
+        pin_cpu_memory=False,
+        enable_torch_compile=compile_enabled,
+        enable_torch_compile_text_encoder=compile_enabled,
+        enable_torch_compile_vae=compile_enabled and not args.taehv,
+        **extra_kwargs,
+    )
+    return generator
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="FP4 linear Wan2.1-1.3B with TAEHV decoding benchmark")
+    parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+                        help="Model path or HuggingFace ID")
+    parser.add_argument("--baseline", action="store_true",
+                        help="Run dense bf16 instead of FP4 linear")
+    parser.add_argument("--no-compile", action="store_true",
+                        help="Disable torch.compile (eager)")
+    parser.add_argument("--taehv", action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="Decode with TAEHV instead of the full Wan VAE "
+                             "(use --no-taehv for the Wan VAE path)")
+    parser.add_argument("--taehv_checkpoint", default=DEFAULT_TAEHV_CHECKPOINT,
+                        help="Path to the TAEHV taew2_1.pth checkpoint")
+    parser.add_argument("--distilled_model", default=DEFAULT_DISTILLED_MODEL,
+                        help="HuggingFace ID (or local path) of a distilled "
+                             "transformer checkpoint to load on top of --model. "
+                             "Pass '' to use the base --model weights instead.")
+    parser.add_argument("--num_gpus", type=int, default=1)
+    parser.add_argument("--infer_steps", type=int, default=3)
+    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for FP4 inference.")
+
+    cap = torch.cuda.get_device_capability()
+    print(f"GPU: {torch.cuda.get_device_name()} (capability {cap[0]}.{cap[1]})")
+    if not args.baseline and cap[0] < 10:
+        print("Warning: NVFP4 requires Blackwell (capability 10.0+); "
+              "FP4 kernels may be unavailable on this GPU.")
+
+    mode = "bf16" if args.baseline else "fp4_linear"
+    mode += "_taehv" if args.taehv else "_wanvae"
+    if not args.no_compile:
+        mode += "_compile"
+    print(f"Mode: {mode.upper()}")
+
+    # Load TAEHV before the (slow) generator build so a bad checkpoint path
+    # fails fast.
+    taehv = TaehvDecoder(resolve_taehv_checkpoint(args.taehv_checkpoint)) \
+        if args.taehv else None
+
+    generator = build_generator(args)
+
+    os.makedirs(OUTPUT_PATH, exist_ok=True)
+
+    # Warmup: with compile enabled the first call(s) pay the DiT compilation
+    # cost. When using TAEHV we also decode the warmup latents so the timed
+    # decode below is warm -- TAEHV's decoder is all conv/upsample, so the
+    # first call otherwise pays cuDNN algo selection + allocator growth
+    # (~0.2s), which is exactly the cold-start overhead we want to exclude.
+    n_warmup = 2 if not args.no_compile else 1
+    with silence_request_log():
+        for _ in range(n_warmup):
+            warm = generator.generate(request={
+                "prompt": PROMPT,
+                "sampling": {"num_inference_steps": 2, "guidance_scale": args.guidance_scale},
+                "output": {"save_video": False, "return_frames": args.taehv},
+            })
+            if args.taehv:
+                taehv.decode(warm.samples)
+
+    output_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    result = generator.generate(request={
+        "prompt": PROMPT,
+        "sampling": {
+            "num_inference_steps": args.infer_steps,
+            "guidance_scale": args.guidance_scale,
+        },
+        # When using TAEHV we need the latents back and save manually; the Wan
+        # VAE path lets the pipeline decode and save the mp4 itself.
+        "output": {
+            "save_video": not args.taehv,
+            "return_frames": args.taehv,
+            "output_path": output_path,
+        },
+    })
+    torch.cuda.synchronize()
+    denoise_elapsed = result.generation_time
+    
+    if args.taehv:
+        torch.cuda.synchronize()
+        decode_start = time.perf_counter()
+        frames = taehv.decode(result.samples)
+        torch.cuda.synchronize()
+        decode_elapsed = time.perf_counter() - decode_start
+
+        imageio.mimsave(output_path, frames, fps=16, format="mp4")
+        total = denoise_elapsed + decode_elapsed
+        print(f"[{mode.upper()}] denoise {denoise_elapsed:.2f}s + TAEHV decode "
+              f"{decode_elapsed:.2f}s = {total:.2f}s "
+              f"({frames.shape[0]} frames @ {tuple(frames.shape[1:3])})")
+        print(f"Saved video to {output_path}")
+    else:
+        print(f"[{mode.upper()}] {args.infer_steps} steps in {denoise_elapsed:.2f}s "
+              f"({args.infer_steps / denoise_elapsed:.2f} it/s)")
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/optimizations/fp4_qad_wan2_1_1_3b.py

@@ -0,0 +1,150 @@
+"""NVFP4 QAD inference example.
+
+Runs Wan2.1-T2V-1.3B with the FastWan-QAD-1.3B distilled checkpoint and
+NVFP4QATConfig quantization. Uses ATTN_QAT_INFER attention backend.
+
+Requirements:
+    - GPU: Blackwell (B200/B300, sm100a+) for the FP4 linear path
+    - TAEHV (optional): Follow install instructions at https://github.com/madebyollin/taehv
+
+Usage:
+    python fp4_qad_wan2_1_1_3b.py                              # NVFP4 QAD (default)
+    python fp4_qad_wan2_1_1_3b.py --bf16                       # BF16 baseline
+    python fp4_qad_wan2_1_1_3b.py --taehv-checkpoint /path/to/taew2_1.pth
+"""
+
+import argparse
+import os
+import sys
+import time
+
+import torch
+
+OUTPUT_PATH = "video_samples"
+
+
+def load_taehv(checkpoint_path, device="cuda", dtype=torch.float16):
+    repo_dir = os.path.dirname(checkpoint_path)
+    if repo_dir not in sys.path:
+        sys.path.insert(0, repo_dir)
+    from taehv import TAEHV
+    print(f"Loading TAEHV from {checkpoint_path}...")
+    model = TAEHV(checkpoint_path=checkpoint_path).to(device, dtype)
+    print("TAEHV loaded.")
+    return model
+
+
+@torch.no_grad()  # type: ignore[misc]
+def decode_with_taehv(taehv_model, latents):
+    latents = latents.permute(0, 2, 1, 3, 4)
+    latents = latents.to(device=next(taehv_model.parameters()).device,
+                         dtype=next(taehv_model.parameters()).dtype)
+    decoded = taehv_model.decode_video(latents, parallel=False, show_progress_bar=False)
+    frames = []
+    for frame in decoded[0]:
+        frame_np = (frame.clamp(0, 1) * 255).byte().cpu().permute(1, 2, 0).numpy()
+        frames.append(frame_np)
+    return frames
+
+
+def main():
+    parser = argparse.ArgumentParser(description="NVFP4 QAD video generation benchmark")
+    parser.add_argument("--bf16", action="store_true",
+                        help="BF16 baseline (no NVFP4 quantization)")
+    parser.add_argument("--taehv-checkpoint", default=None, metavar="PATH",
+                        help="Path to taew2_1.pth; enables TAEHV tiny autoencoder decoding")
+    parser.add_argument("--model", default="FastVideo/FastWan-QAD-1.3B",
+                        help="Model path or HuggingFace ID")
+    parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile for the DiT")
+    parser.add_argument("--num_gpus", type=int, default=1)
+    parser.add_argument("--infer_steps", type=int, default=3)
+    args = parser.parse_args()
+
+    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER")
+    os.environ["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] = "0"
+    os.environ["FLASHINFER_CUDA_ARCH_LIST"] = "12.0a"
+    # os.environ["FLASHINFER_EXTRA_CFLAGS"] = "-DCCCL_DISABLE_CTK_COMPATIBILITY_CHECK"
+    # os.environ["FLASHINFER_EXTRA_CUDAFLAGS"] = "-DCCCL_DISABLE_CTK_COMPATIBILITY_CHECK"
+    # os.environ["CUDA_HOME"] = "/root/miniconda3/envs/fastvideo/lib/python3.12/site-packages/nvidia/cu13"
+
+    from fastvideo import VideoGenerator
+    from fastvideo.configs.pipelines.base import PipelineConfig
+
+    mode = "bf16" if args.bf16 else "nvfp4_qad"
+    if not args.no_compile:
+        mode += "_compile"
+    use_taehv = args.taehv_checkpoint is not None
+    print(f"Mode: {mode.upper()}" + ("  decoder=TAEHV" if use_taehv else "  decoder=VAE"))
+
+    taehv_model = load_taehv(args.taehv_checkpoint) if use_taehv else None
+
+    pipeline_config = PipelineConfig.from_pretrained(args.model)
+    pipeline_config.text_encoder_precisions = ("bf16",)
+    if not args.bf16:
+        from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig
+        pipeline_config.dit_config.quant_config = NVFP4QATConfig()
+
+    generator = VideoGenerator.from_pretrained(
+        args.model,
+        pipeline_config=pipeline_config,
+        num_gpus=args.num_gpus,
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        vae_cpu_offload=use_taehv,
+        text_encoder_cpu_offload=False,
+        pin_cpu_memory=False,
+        enable_torch_compile=not args.no_compile,
+        enable_torch_compile_text_encoder=not args.no_compile,
+        enable_torch_compile_vae=not args.no_compile and not use_taehv,
+        output_type="latent" if use_taehv else "pil",
+    )
+
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+
+    n_warmup = 2 if not args.no_compile else 0
+    for _ in range(n_warmup):
+        warmup_result = generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 3, "guidance_scale": 1.0},
+                                                    "output": {"save_video": False}})
+        if use_taehv:
+            decode_with_taehv(taehv_model, warmup_result.samples)
+
+    os.makedirs(OUTPUT_PATH, exist_ok=True)
+    video_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
+    if use_taehv:
+        import imageio
+        result = generator.generate(request={
+            "prompt": prompt,
+            "sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
+            "output": {"save_video": False},
+        })
+        denoise_elapsed = result.generation_time
+        torch.cuda.synchronize()
+        t_decode = time.perf_counter()
+        frames = decode_with_taehv(taehv_model, result.samples)
+        torch.cuda.synchronize()
+        decode_elapsed = time.perf_counter() - t_decode
+        total = denoise_elapsed + decode_elapsed
+        imageio.mimsave(video_path, frames, fps=16, format="mp4")
+        print(f"Saved TAEHV-decoded video to: {video_path}")
+        print(f"[{mode.upper()}] {args.infer_steps} steps in {total:.3f}s "
+              f"(denoise {denoise_elapsed:.3f}s + decode {decode_elapsed:.3f}s)")
+    else:
+        result = generator.generate(request={
+            "prompt": prompt,
+            "sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
+            "output": {"save_video": True, "output_path": video_path},
+        })
+        elapsed = result.generation_time
+        print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.3f}s "
+              f"({args.infer_steps / elapsed:.2f} it/s)")
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/optimizations/fp8_wan2_1_1_3b.py

@@ -0,0 +1,140 @@
+"""FP8 weight quantization inference example.
+
+Runs Wan2.1-T2V-1.3B with FP8 e4m3 quantized DiT linear layers (attention
+projections and FFN). Weights are quantized in-place after loading; activations
+are quantized dynamically at runtime. Reduces GPU memory relative to BF16 and
+can improve throughput on sm89+ GPUs.
+
+Requirements:
+    - GPU: sm89+ (H100, L40S, RTX 4090, Ada Lovelace, or newer)
+      Falls back to a bf16 dequant path on older GPUs.
+    - TAEHV (optional): Follow install instructions at https://github.com/madebyollin/taehv
+
+Usage:
+    python fp8_wan2_1_1_3b.py              # FP8 per-tensor (default)
+    python fp8_wan2_1_1_3b.py --bf16       # BF16 baseline
+    python fp8_wan2_1_1_3b.py --granularity channel  # per-channel (higher accuracy but slower)
+    python fp8_wan2_1_1_3b.py --taehv-checkpoint /path/to/taew2_1.pth
+"""
+
+import argparse
+import os
+import sys
+import time
+
+import torch
+
+OUTPUT_PATH = "video_samples"
+
+
+def load_taehv(checkpoint_path, device="cuda", dtype=torch.float16):
+    repo_dir = os.path.dirname(checkpoint_path)
+    if repo_dir not in sys.path:
+        sys.path.insert(0, repo_dir)
+    from taehv import TAEHV
+    print(f"Loading TAEHV from {checkpoint_path}...")
+    model = TAEHV(checkpoint_path=checkpoint_path).to(device, dtype)
+    print("TAEHV loaded.")
+    return model
+
+
+@torch.no_grad()  # type: ignore[misc]
+def decode_with_taehv(taehv_model, latents):
+    latents = latents.permute(0, 2, 1, 3, 4)
+    latents = latents.to(device=next(taehv_model.parameters()).device,
+                         dtype=next(taehv_model.parameters()).dtype)
+    decoded = taehv_model.decode_video(latents, parallel=False, show_progress_bar=False)
+    frames = []
+    for frame in decoded[0]:
+        frame_np = (frame.clamp(0, 1) * 255).byte().cpu().permute(1, 2, 0).numpy()
+        frames.append(frame_np)
+    return frames
+
+
+def main():
+    parser = argparse.ArgumentParser(description="FP8 video generation benchmark")
+    parser.add_argument("--bf16", action="store_true",
+                        help="BF16 baseline (no FP8 quantization)")
+    parser.add_argument("--granularity", choices=["tensor", "channel"], default="tensor",
+                        help="FP8 weight scale granularity: tensor (faster) or channel (more accurate)")
+    parser.add_argument("--taehv-checkpoint", default=None, metavar="PATH",
+                        help="Path to taew2_1.pth; enables TAEHV tiny autoencoder decoding")
+    parser.add_argument("--model", default="FastVideo/FastWan-QAD-FP8-1.3B",
+                        help="Model path or HuggingFace ID")
+    parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile for the DiT")
+    parser.add_argument("--num_gpus", type=int, default=1)
+    parser.add_argument("--infer_steps", type=int, default=3)
+    args = parser.parse_args()
+
+    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
+
+    from fastvideo import VideoGenerator
+    from fastvideo.layers.quantization import get_quantization_config
+
+    mode = "bf16" if args.bf16 else f"fp8_{args.granularity}"
+    if not args.no_compile:
+        mode += "_compile"
+    use_taehv = args.taehv_checkpoint is not None
+    print(f"Mode: {mode.upper()}" + ("  decoder=TAEHV" if use_taehv else "  decoder=VAE"))
+
+    taehv_model = load_taehv(args.taehv_checkpoint) if use_taehv else None
+
+    # transformer_quant needs a QuantizationConfig *instance* — the bare string
+    # is not resolved on the from_pretrained kwarg path.
+    extra = {} if args.bf16 else {
+        "transformer_quant": get_quantization_config("FP8")(granularity=args.granularity)
+    }
+    generator = VideoGenerator.from_pretrained(
+        args.model,
+        num_gpus=args.num_gpus,
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        vae_cpu_offload=use_taehv,
+        text_encoder_cpu_offload=False,
+        pin_cpu_memory=False,
+        enable_torch_compile=not args.no_compile,
+        enable_torch_compile_vae=not args.no_compile and not use_taehv,
+        output_type="latent" if use_taehv else "pil",
+        **extra,
+    )
+
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+
+    n_warmup = 1 if not args.no_compile else 0
+    for _ in range(n_warmup):
+        generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 3, "guidance_scale": 1.0},
+                                    "output": {"save_video": False}})
+
+    os.makedirs(OUTPUT_PATH, exist_ok=True)
+    start = time.time()
+    if use_taehv:
+        result = generator.generate(request={
+            "prompt": prompt,
+            "sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
+            "output": {"save_video": False},
+        })
+        import imageio
+        frames = decode_with_taehv(taehv_model, result.samples)
+        video_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
+        imageio.mimsave(video_path, frames, fps=16, format="mp4")
+        print(f"Saved TAEHV-decoded video to: {video_path}")
+    else:
+        generator.generate(request={
+            "prompt": prompt,
+            "sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
+            "output": {"save_video": True, "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")},
+        })
+    elapsed = time.time() - start
+    print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s "
+          f"({args.infer_steps / elapsed:.2f} it/s)")
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu

@@ -286,24 +286,19 @@
  
    // load input
    const int token_id = token_block_id * BLOCK_SIZE + threadIdx.x / NUM_THREADS_PER_TOKEN;
-   // Permute V rows within each 32-element block so the PV MMA K-indexed
-   // access reads the correct CLayout N-indexed values (Edenzzzz causal fix).
-   const int k_intra = token_id & 31;
-   const int load_token_id = (token_id & ~31)
-       | ((k_intra & 6) << 2) | ((k_intra & 24) >> 2) | (k_intra & 1);
- 
+
    PackedVec in_vec;
-   
+
    #pragma unroll
    for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
      reinterpret_cast<uint32_t&>(in_vec.elts[i]) = 0;
    }
-   
-   if (load_token_id < num_tokens) {
-     in_vec = reinterpret_cast<PackedVec const*>(input + 
+
+   if (token_id < num_tokens) {
+     in_vec = reinterpret_cast<PackedVec const*>(input +
                                            batch_id * stride_bz_input + // batch dim
                                            head_id * stride_h_input +   // head dim
-                                           load_token_id * stride_seq_input + // seq dim (permuted)
+                                           token_id * stride_seq_input + // seq dim
                                            (threadIdx.x % NUM_THREADS_PER_TOKEN) * CVT_FP4_ELTS_PER_THREAD)[0]; // feature dim
    }
  

fastvideo-kernel/benchmarks/bench_fused_compress_topk.py

@@ -0,0 +1,315 @@
+#!/usr/bin/env python3
+"""
+Benchmark fused compress (block mean) and topk mask kernels vs. PyTorch baselines.
+
+Compares:
+  compress:
+    - Old: .view() -> .float() -> .sum(dim=3) -> / vbs -> .to(bf16)
+    - New: fused_block_mean (Triton: bf16 read -> fp32 accumulate -> div -> bf16 write)
+
+  topk:
+    - Old: torch.topk() -> zeros() -> scatter_()
+    - New: fused_topk_mask (Triton: binary-search pivot -> bool mask)
+
+Reports per-kernel latency (ms), speedup, and numerical accuracy (max abs error,
+cosine similarity for compress; mask match rate for topk).
+
+Also benchmarks backward pass of compress (fused Triton bwd kernel vs. PyTorch autograd).
+"""
+
+from __future__ import annotations
+
+import argparse
+import random
+
+import numpy as np
+import torch
+
+try:
+    from triton.testing import do_bench
+except ImportError as e:
+    raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e
+
+from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_block_mean, fused_topk_mask
+
+
+def set_seed(seed: int = 42) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def parse_arguments() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Benchmark fused compress & topk vs. PyTorch baselines")
+    p.add_argument("--batch_size", type=int, default=1)
+    p.add_argument("--num_heads", type=int, default=12)
+    p.add_argument("--head_dim", type=int, default=128, choices=[64, 128])
+    p.add_argument("--seq_lens", type=int, nargs="+", default=[49152],
+                   help="Sequence lengths to benchmark (must be divisible by block_elements)")
+    p.add_argument("--block_elements", type=int, default=64, choices=[64, 256],
+                   help="Tokens per block (64 or 256)")
+    p.add_argument("--topk", type=int, default=None,
+                   help="KV blocks to select per Q block (default: ~10%% of kv_blocks)")
+    p.add_argument("--warmup", type=int, default=10)
+    p.add_argument("--rep", type=int, default=50)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"])
+    return p.parse_args()
+
+
+# ---------------------------------------------------------------------------
+# Old PyTorch baselines (extracted from ops.py before commit d9cd0b9)
+# ---------------------------------------------------------------------------
+
+def pytorch_block_mean(
+    x: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    block_elements: int,
+) -> torch.Tensor:
+    """Old PyTorch compress: .view() -> .float() -> .sum(dim=3) -> / vbs -> .to(dtype)"""
+    B, H, seq_len, D = x.shape
+    num_blocks = seq_len // block_elements
+    x_c = x.view(B, H, num_blocks, block_elements, D)
+    x_c = (x_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(x.dtype)
+    return x_c
+
+
+def pytorch_topk_mask(
+    scores: torch.Tensor,
+    topk: int,
+) -> torch.Tensor:
+    """Old PyTorch topk: torch.topk() -> zeros() -> scatter_()"""
+    topk = min(topk, scores.shape[-1])
+    topk_idx = torch.topk(scores, topk, dim=-1).indices
+    mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
+    return mask
+
+
+# ---------------------------------------------------------------------------
+# Accuracy helpers
+# ---------------------------------------------------------------------------
+
+def accuracy_compress(ref: torch.Tensor, test: torch.Tensor) -> dict:
+    ref_f = ref.float()
+    test_f = test.float()
+    abs_err = (ref_f - test_f).abs()
+    cos_sim = torch.nn.functional.cosine_similarity(ref_f.flatten(), test_f.flatten(), dim=0)
+    return {
+        "max_abs_err": abs_err.max().item(),
+        "mean_abs_err": abs_err.mean().item(),
+        "cosine_sim": cos_sim.item(),
+    }
+
+
+def accuracy_topk(ref_mask: torch.Tensor, test_mask: torch.Tensor) -> dict:
+    match = (ref_mask == test_mask).all(dim=-1).float().mean().item()
+    true_per_row_ref = ref_mask.sum(dim=-1).float()
+    true_per_row_test = test_mask.sum(dim=-1).float()
+    count_match = (true_per_row_ref == true_per_row_test).float().mean().item()
+    return {
+        "row_exact_match": match,
+        "count_match": count_match,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Benchmark runner
+# ---------------------------------------------------------------------------
+
+def bench_compress_fwd(
+    B: int, H: int, seq_len: int, D: int, block_elements: int,
+    dtype: torch.dtype, warmup: int, rep: int,
+) -> None:
+    num_blocks = seq_len // block_elements
+    x = torch.randn(B, H, seq_len, D, dtype=dtype, device="cuda")
+    vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device="cuda")
+    if num_blocks > 4:
+        vbs[1] = block_elements - 2
+        vbs[-2] = block_elements - 5
+
+    # Accuracy
+    ref = pytorch_block_mean(x, vbs, block_elements)
+    fused = fused_block_mean(x, vbs, block_elements)
+    acc = accuracy_compress(ref, fused)
+
+    # Latency
+    old_ms = do_bench(lambda: pytorch_block_mean(x, vbs, block_elements), warmup=warmup, rep=rep)
+    new_ms = do_bench(lambda: fused_block_mean(x, vbs, block_elements), warmup=warmup, rep=rep)
+
+    speedup = old_ms / new_ms if new_ms > 0 else float("inf")
+    print(f"  compress fwd | old: {old_ms:8.3f} ms | new: {new_ms:8.3f} ms | speedup: {speedup:5.2f}x "
+          f"| max_abs_err: {acc['max_abs_err']:.2e} | cos_sim: {acc['cosine_sim']:.8f}")
+
+
+def bench_compress_bwd(
+    B: int, H: int, seq_len: int, D: int, block_elements: int,
+    dtype: torch.dtype, warmup: int, rep: int,
+) -> None:
+    num_blocks = seq_len // block_elements
+    vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device="cuda")
+    if num_blocks > 4:
+        vbs[1] = block_elements - 2
+        vbs[-2] = block_elements - 5
+
+    # --- Accuracy: compare gradients ---
+    x_old = torch.randn(B, H, seq_len, D, dtype=dtype, device="cuda", requires_grad=True)
+    grad_out = torch.randn(B, H, num_blocks, D, dtype=dtype, device="cuda")
+
+    out_old = pytorch_block_mean(x_old, vbs, block_elements)
+    out_old.backward(grad_out)
+    grad_ref = x_old.grad.clone()
+
+    x_new = x_old.detach().clone().requires_grad_(True)
+    out_new = fused_block_mean(x_new, vbs, block_elements)
+    out_new.backward(grad_out)
+    grad_fused = x_new.grad.clone()
+
+    acc = accuracy_compress(grad_ref, grad_fused)
+
+    # --- Latency: isolate backward-only via retain_graph ---
+    x_o = x_old.detach().clone().requires_grad_(True)
+    out_o = pytorch_block_mean(x_o, vbs, block_elements)
+    loss_o = (out_o * grad_out).sum()
+    for _ in range(warmup):
+        torch.autograd.grad(loss_o, x_o, retain_graph=True)
+    old_ms = do_bench(
+        lambda: torch.autograd.grad(loss_o, x_o, retain_graph=True),
+        warmup=0, rep=rep,
+    )
+
+    x_n = x_old.detach().clone().requires_grad_(True)
+    out_n = fused_block_mean(x_n, vbs, block_elements)
+    loss_n = (out_n * grad_out).sum()
+    for _ in range(warmup):
+        torch.autograd.grad(loss_n, x_n, retain_graph=True)
+    new_ms = do_bench(
+        lambda: torch.autograd.grad(loss_n, x_n, retain_graph=True),
+        warmup=0, rep=rep,
+    )
+
+    speedup = old_ms / new_ms if new_ms > 0 else float("inf")
+    print(f"  compress bwd | old: {old_ms:8.3f} ms | new: {new_ms:8.3f} ms | speedup: {speedup:5.2f}x "
+          f"| max_abs_err: {acc['max_abs_err']:.2e} | cos_sim: {acc['cosine_sim']:.8f}")
+
+
+def bench_topk(
+    B: int, H: int, num_blocks: int, topk: int,
+    dtype: torch.dtype, warmup: int, rep: int,
+) -> None:
+    scores = torch.randn(B, H, num_blocks, num_blocks, dtype=dtype, device="cuda")
+
+    # Accuracy
+    ref = pytorch_topk_mask(scores, topk)
+    fused = fused_topk_mask(scores, topk)
+    acc = accuracy_topk(ref, fused)
+
+    # Latency
+    old_ms = do_bench(lambda: pytorch_topk_mask(scores, topk), warmup=warmup, rep=rep)
+    new_ms = do_bench(lambda: fused_topk_mask(scores, topk), warmup=warmup, rep=rep)
+
+    speedup = old_ms / new_ms if new_ms > 0 else float("inf")
+    print(f"  topk      | old: {old_ms:8.3f} ms | new: {new_ms:8.3f} ms | speedup: {speedup:5.2f}x "
+          f"| row_exact_match: {acc['row_exact_match']:.4f} | count_match: {acc['count_match']:.4f}")
+
+
+def bench_topk_neginf(
+    B: int, H: int, num_blocks: int, topk: int,
+    dtype: torch.dtype, inf_ratio: float = 0.3,
+) -> None:
+    """Test topk correctness when a fraction of scores are -inf (masked positions)."""
+    scores = torch.randn(B, H, num_blocks, num_blocks, dtype=dtype, device="cuda")
+    inf_mask = torch.rand(B, H, num_blocks, num_blocks, device="cuda") < inf_ratio
+    scores[inf_mask] = float("-inf")
+
+    ref = pytorch_topk_mask(scores, topk)
+    fused = fused_topk_mask(scores, topk)
+    acc = accuracy_topk(ref, fused)
+
+    status = "PASS" if acc["row_exact_match"] == 1.0 else "FAIL"
+    print(f"  topk -inf | {inf_ratio*100:.0f}% masked | {status} "
+          f"| row_exact_match: {acc['row_exact_match']:.4f} | count_match: {acc['count_match']:.4f}")
+
+
+_MAX_KV_BLOCK_SIZE = 4096
+
+
+def bench_topk_large_kv(
+    B: int, H: int, kv_blocks: int, topk: int,
+    dtype: torch.dtype, warmup: int, rep: int,
+) -> None:
+    """Test topk with large kv_blocks that may exceed Triton block size limit."""
+    import triton
+    kv_block_size = triton.next_power_of_2(kv_blocks)
+    fallback = kv_block_size > _MAX_KV_BLOCK_SIZE
+
+    q_blocks = max(1, kv_blocks // 8)
+    scores = torch.randn(B, H, q_blocks, kv_blocks, dtype=dtype, device="cuda")
+
+    ref = pytorch_topk_mask(scores, topk)
+    fused = fused_topk_mask(scores, topk)
+    acc = accuracy_topk(ref, fused)
+
+    old_ms = do_bench(lambda: pytorch_topk_mask(scores, topk), warmup=warmup, rep=rep)
+    new_ms = do_bench(lambda: fused_topk_mask(scores, topk), warmup=warmup, rep=rep)
+
+    speedup = old_ms / new_ms if new_ms > 0 else float("inf")
+    path = "fallback" if fallback else "triton"
+    print(f"  topk kv={kv_blocks:<5d} | {path:8s} | old: {old_ms:8.3f} ms | new: {new_ms:8.3f} ms "
+          f"| speedup: {speedup:5.2f}x | row_exact_match: {acc['row_exact_match']:.4f}")
+
+
+def main() -> None:
+    args = parse_arguments()
+    set_seed(args.seed)
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+    B, H, D = args.batch_size, args.num_heads, args.head_dim
+    block_elements = args.block_elements
+
+    print("Fused Compress & TopK Benchmark")
+    print(f"device: {torch.cuda.get_device_name(0)}")
+    print(f"batch={B}, heads={H}, head_dim={D}, block_elements={block_elements}, dtype={args.dtype}")
+    print(f"warmup={args.warmup}, rep={args.rep}")
+
+    for seq_len in args.seq_lens:
+        if seq_len % block_elements != 0:
+            print(f"\n[skip] seq_len={seq_len} not divisible by block_elements={block_elements}")
+            continue
+
+        num_blocks = seq_len // block_elements
+        topk = args.topk if args.topk is not None else max(1, num_blocks // 10)
+        topk = min(topk, num_blocks)
+
+        print(f"\n{'=' * 100}")
+        print(f"seq_len={seq_len}, num_blocks={num_blocks}, topk={topk}")
+        print("-" * 100)
+
+        bench_compress_fwd(B, H, seq_len, D, block_elements, dtype, args.warmup, args.rep)
+        bench_compress_bwd(B, H, seq_len, D, block_elements, dtype, args.warmup, args.rep)
+        bench_topk(B, H, num_blocks, topk, dtype, args.warmup, args.rep)
+
+    # --- Edge case: topk with -inf scores ---
+    print(f"\n{'=' * 100}")
+    print("TopK with -inf scores (masked positions)")
+    print("-" * 100)
+    kv = args.seq_lens[0] // block_elements
+    tk = args.topk if args.topk is not None else max(1, kv // 10)
+    tk = min(tk, kv)
+    for inf_ratio in [0.1, 0.3, 0.5, 0.8]:
+        bench_topk_neginf(B, H, kv, tk, dtype, inf_ratio)
+
+    # --- Edge case: topk with large kv_blocks (Triton block size limit) ---
+    print(f"\n{'=' * 100}")
+    print("TopK with large kv_blocks (Triton block size limit test)")
+    print("-" * 100)
+    for kv_blocks in [1024, 2048, 4096, 8192]:
+        tk = max(1, kv_blocks // 10)
+        bench_topk_large_kv(1, 1, kv_blocks, tk, dtype, args.warmup, args.rep)
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark.")
+    main()

fastvideo-kernel/build.sh

@@ -103,6 +103,11 @@ if [ "${GPU_BACKEND}" = "CUDA" ]; then
     if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then
         if [ "${cc_major}" = "9" ] && [ "${cc_minor}" = "0" ]; then
             export TORCH_CUDA_ARCH_LIST="9.0a"
+        elif [ "${cc_major}" = "12" ] && [ "${cc_minor}" = "0" ]; then
+            # Blackwell sm_120 needs the arch-conditional 'a' suffix so CMake's
+            # AUTO gate (matches 12.0a/120a/sm_120a) builds the attn_qat_infer
+            # (modified SageAttention3 FP4) kernels instead of silently skipping.
+            export TORCH_CUDA_ARCH_LIST="12.0a"
         else
             export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}"
         fi

fastvideo-kernel/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "fastvideo-kernel"
-version = "0.2.6"
+version = "0.3.0"
 description = "Unified CUDA kernels for FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -1,11 +1,12 @@
 import math
 import torch
-from .block_sparse_attn import block_sparse_attn, block_sparse_attn_from_indices
+from .block_sparse_attn import block_sparse_attn
 from .block_sparse_attn_256 import (
     block_sparse_attn_256,
     block_sparse_attn_256_bshd,
 )
 from .triton_kernels.st_attn_triton import sliding_tile_attention_triton
+from .triton_kernels.fused_compress_topk import fused_block_mean, fused_topk_mask
 
 # Try to load the C++ extension
 try:
@@ -118,13 +119,10 @@ def video_sparse_attn(
             f"got {q_variable_block_sizes.numel()}"
         )
 
-    # Compression branch (token-level average per block + dense block-level attn).
-    q_c = q.view(batch, heads, q_num_blocks, block_elements, dim)
-    k_c = k.view(batch, heads, kv_num_blocks, block_elements, dim)
-    v_c = v.view(batch, heads, kv_num_blocks, block_elements, dim)
-    q_c = (q_c.float().sum(dim=3) / q_variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype)
-    k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype)
-    v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype)
+    # Compression branch (fused Triton: bf16 read → fp32 accumulate → div → bf16 write)
+    q_c = fused_block_mean(q, q_variable_block_sizes, block_elements)
+    k_c = fused_block_mean(k, variable_block_sizes, block_elements)
+    v_c = fused_block_mean(v, variable_block_sizes, block_elements)
 
     scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim ** 0.5)
     attn = torch.softmax(scores, dim=-1)
@@ -132,25 +130,13 @@ def video_sparse_attn(
     out_c = out_c.view(batch, heads, q_num_blocks, 1, dim)
     out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, q_seq_len, dim)
 
-    # Sparse branch.
-    topk_idx = torch.topk(scores, topk, dim=-1).indices
+    # Sparse branch (fused Triton topk mask)
+    mask = fused_topk_mask(scores, topk)
 
     if block_elements == 256:
-        # CuTe path consumes a bool mask (full/partial split inside the wrapper).
-        mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
         out_s = block_sparse_attn_256(q, k, v, mask, variable_block_sizes)[0]
     else:
-        # Index-native path for 64-block (TK/Triton).
-        q2k_idx = topk_idx.to(torch.int32).contiguous()
-        q2k_num = torch.full(
-            (batch, heads, q_num_blocks),
-            topk,
-            dtype=torch.int32,
-            device=q.device,
-        )
-        out_s = block_sparse_attn_from_indices(
-            q, k, v, q2k_idx, q2k_num, variable_block_sizes
-        )[0]
+        out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]
 
     if compress_attn_weight is not None:
         return out_c * compress_attn_weight + out_s
@@ -236,9 +222,8 @@ def video_sparse_attn_bshd(
     out_c_ch = torch.matmul(attn, v_ch)
     out_c_blk = out_c_ch.permute(0, 2, 1, 3).contiguous()
 
-    # Sparse branch (CuTe BSHD).
-    topk_idx = torch.topk(scores, topk, dim=-1).indices
-    mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
+    # Sparse branch (fused Triton topk mask + CuTe BSHD).
+    mask = fused_topk_mask(scores, topk)
     out_s, _ = block_sparse_attn_256_bshd(q, k, v, mask, variable_block_sizes)
 
     out = out_s

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

@@ -0,0 +1,334 @@
+"""
+Fused Triton kernels for VSA compress (block mean) and topk mask construction.
+
+Replaces the multi-kernel PyTorch pipeline:
+  Original compress: .view() -> .float() -> .sum(dim=3) -> / vbs -> .to(bf16)
+  Original topk:     torch.topk() -> zeros() -> scatter_()
+
+With single-pass fused kernels:
+  fused_block_mean:  read bf16, accumulate fp32, div by vbs, write bf16
+  fused_topk_mask:   read scores, find k-th value, write bool mask
+"""
+
+import logging
+
+import torch
+import triton
+import triton.language as tl
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _fused_block_mean_kernel(
+    X_ptr,
+    Out_ptr,
+    VBS_ptr,
+    stride_x_bh,
+    stride_x_seq,
+    stride_o_bh,
+    stride_o_blk,
+    num_blocks,
+    BLOCK_ELEMENTS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    OUTPUT_DTYPE: tl.constexpr,
+):
+    """Fused block mean: one program computes mean of one block for one (b,h).
+
+    X is viewed as [B*H, num_blocks*BLOCK_ELEMENTS, HEAD_DIM] contiguous.
+    Out is [B*H, num_blocks, HEAD_DIM] contiguous.
+    2D load + parallel tl.sum reduction, accumulates in fp32.
+    """
+    block_idx = tl.program_id(0)
+    bh_idx = tl.program_id(1)
+
+    if block_idx >= num_blocks:
+        return
+
+    vbs = tl.load(VBS_ptr + block_idx).to(tl.float32)
+
+    x_base = X_ptr + bh_idx * stride_x_bh + block_idx * BLOCK_ELEMENTS * stride_x_seq
+
+    row_offsets = tl.arange(0, BLOCK_ELEMENTS)
+    dim_offsets = tl.arange(0, HEAD_DIM)
+    offsets = row_offsets[:, None] * stride_x_seq + dim_offsets[None, :]
+    block_data = tl.load(x_base + offsets).to(tl.float32)
+    acc = tl.sum(block_data, axis=0) / vbs
+
+    out_base = Out_ptr + bh_idx * stride_o_bh + block_idx * stride_o_blk + dim_offsets
+    tl.store(out_base, acc.to(OUTPUT_DTYPE))
+
+
+@triton.jit
+def _fused_block_mean_bwd_kernel(
+    GradOut_ptr,
+    GradX_ptr,
+    VBS_ptr,
+    stride_go_bh,
+    stride_go_blk,
+    stride_gx_bh,
+    stride_gx_seq,
+    num_blocks,
+    BLOCK_ELEMENTS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    OUTPUT_DTYPE: tl.constexpr,
+):
+    """Backward of block mean: broadcast grad_out / vbs to each token in the block.
+
+    Mirrors the forward kernel: one program per (block, bh).
+    GradOut is [B*H, num_blocks, HEAD_DIM].
+    GradX  is [B*H, num_blocks*BLOCK_ELEMENTS, HEAD_DIM].
+    2D store writes all BLOCK_ELEMENTS rows in parallel.
+    """
+    block_idx = tl.program_id(0)
+    bh_idx = tl.program_id(1)
+
+    if block_idx >= num_blocks:
+        return
+
+    vbs = tl.load(VBS_ptr + block_idx).to(tl.float32)
+
+    dim_offsets = tl.arange(0, HEAD_DIM)
+    go_base = GradOut_ptr + bh_idx * stride_go_bh + block_idx * stride_go_blk
+    grad_val = tl.load(go_base + dim_offsets).to(tl.float32) / vbs
+
+    row_offsets = tl.arange(0, BLOCK_ELEMENTS)
+    gx_base = GradX_ptr + bh_idx * stride_gx_bh + block_idx * BLOCK_ELEMENTS * stride_gx_seq
+    offsets = row_offsets[:, None] * stride_gx_seq + dim_offsets[None, :]
+    grad_2d = tl.broadcast_to(grad_val[None, :], [BLOCK_ELEMENTS, HEAD_DIM])
+    tl.store(gx_base + offsets, grad_2d.to(OUTPUT_DTYPE))
+
+
+_TORCH_TO_TRITON_DTYPE = {
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+    torch.float32: tl.float32,
+}
+
+
+def _fused_block_mean_bwd(
+    grad_output: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    block_elements: int,
+) -> torch.Tensor:
+    B, H, num_blocks, D = grad_output.shape
+    seq_len = num_blocks * block_elements
+
+    grad_x = torch.empty(B, H, seq_len, D, dtype=grad_output.dtype, device=grad_output.device)
+
+    go_flat = grad_output.contiguous().view(B * H, num_blocks, D)
+    gx_flat = grad_x.view(B * H, seq_len, D)
+
+    grid = (num_blocks, B * H)
+
+    _fused_block_mean_bwd_kernel[grid](
+        go_flat, gx_flat, variable_block_sizes,
+        go_flat.stride(0), go_flat.stride(1),
+        gx_flat.stride(0), gx_flat.stride(1),
+        num_blocks,
+        BLOCK_ELEMENTS=block_elements,
+        HEAD_DIM=D,
+        OUTPUT_DTYPE=_TORCH_TO_TRITON_DTYPE[grad_output.dtype],
+    )
+
+    return grad_x
+
+
+def _fused_block_mean_fwd(
+    x: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    block_elements: int,
+) -> torch.Tensor:
+    B, H, seq_len, D = x.shape
+    num_blocks = seq_len // block_elements
+    assert seq_len % block_elements == 0
+
+    x = x.contiguous()
+    out = torch.empty(B, H, num_blocks, D, dtype=x.dtype, device=x.device)
+
+    x_flat = x.view(B * H, seq_len, D)
+    out_flat = out.view(B * H, num_blocks, D)
+
+    grid = (num_blocks, B * H)
+
+    _fused_block_mean_kernel[grid](
+        x_flat, out_flat, variable_block_sizes,
+        x_flat.stride(0), x_flat.stride(1),
+        out_flat.stride(0), out_flat.stride(1),
+        num_blocks,
+        BLOCK_ELEMENTS=block_elements,
+        HEAD_DIM=D,
+        OUTPUT_DTYPE=_TORCH_TO_TRITON_DTYPE[x.dtype],
+    )
+
+    return out
+
+
+class _FusedBlockMeanAutograd(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, variable_block_sizes, block_elements):
+        ctx.save_for_backward(variable_block_sizes)
+        ctx.block_elements = block_elements
+        return _fused_block_mean_fwd(x, variable_block_sizes, block_elements)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        variable_block_sizes, = ctx.saved_tensors
+        block_elements = ctx.block_elements
+        return _fused_block_mean_bwd(grad_output, variable_block_sizes, block_elements), None, None
+
+
+def fused_block_mean(
+    x: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    block_elements: int,
+) -> torch.Tensor:
+    """Compute block-wise mean with fp32 accumulation, fused in one kernel.
+
+    Forward: fused Triton kernel (bf16 read → fp32 accumulate → div → bf16 write).
+    Backward: broadcasts grad_output / vbs back to each token position.
+
+    Args:
+        x: [B, H, seq_len, D] in bf16
+        variable_block_sizes: [num_blocks] number of valid tokens per block
+        block_elements: tokens per block (e.g. 64)
+
+    Returns:
+        [B, H, num_blocks, D] in bf16
+    """
+    return _FusedBlockMeanAutograd.apply(x, variable_block_sizes, block_elements)
+
+
+@triton.jit
+def _fused_topk_mask_kernel(
+    Scores_ptr,
+    Mask_ptr,
+    stride_s_bh,
+    stride_s_q,
+    stride_s_kv,
+    stride_m_bh,
+    stride_m_q,
+    stride_m_kv,
+    kv_blocks: tl.constexpr,
+    topk: tl.constexpr,
+    KV_BLOCK_SIZE: tl.constexpr,
+):
+    """Build topk boolean mask via randomized pivot selection (quickselect-style).
+
+    For each (b,h,q_block) row: find the k-th largest score using iterative
+    pivot-based partitioning, then build mask by comparing against threshold.
+
+    Grid: (num_q_blocks, B * H)
+    """
+    q_idx = tl.program_id(0)
+    bh_idx = tl.program_id(1)
+
+    kv_offsets = tl.arange(0, KV_BLOCK_SIZE)
+    score_base = Scores_ptr + bh_idx * stride_s_bh + q_idx * stride_s_q
+    mask_base = Mask_ptr + bh_idx * stride_m_bh + q_idx * stride_m_q
+
+    valid_mask = kv_offsets < kv_blocks
+    scores = tl.load(score_base + kv_offsets * stride_s_kv, mask=valid_mask, other=-float("inf"))
+    scores_f32 = scores.to(tl.float32)
+
+    # Binary search for threshold: find value T such that count(scores > T) <= topk
+    # and count(scores >= T) >= topk
+    # Exclude -inf from lo so the binary search can converge when masked
+    # scores are present (mid = (-inf + hi) * 0.5 = -inf would stall).
+    finite_mask = valid_mask & (scores_f32 > float("-inf"))
+    lo = tl.min(tl.where(finite_mask, scores_f32, float("inf")), axis=0)
+    hi = tl.max(tl.where(valid_mask, scores_f32, float("-inf")), axis=0)
+    # If all valid scores are -inf, lo > hi; threshold stays at -inf and
+    # the tie-breaking logic below selects the first topk positions.
+    lo = tl.minimum(lo, hi)
+
+    # 32 iterations of fp32 bisection give range-relative resolution (hi-lo)/2^32.
+    # For VSA's softmax-input scores (q_c@k_c/sqrt(d), O(1) magnitude, range < ~10),
+    # this resolves to ~2e-9, well below the bf16 ULP at that magnitude (~8e-3),
+    # so the threshold converges exactly to the k-th bf16 score value and the
+    # > / == comparisons below are exact.
+    for _i in range(32):
+        mid = (lo + hi) * 0.5
+        count_ge = tl.sum(((scores_f32 >= mid) & valid_mask).to(tl.int32), axis=0)
+        lo = tl.where(count_ge >= topk, mid, lo)
+        hi = tl.where(count_ge >= topk, hi, mid)
+
+    # lo is our threshold: count(scores >= lo) >= topk
+    threshold = lo
+    above_threshold = scores_f32 > threshold
+    at_threshold = scores_f32 == threshold
+    n_above = tl.sum(above_threshold.to(tl.int32), axis=0)
+    n_needed_at_thresh = topk - n_above
+
+    at_thresh_cumsum = tl.cumsum(at_threshold.to(tl.int32), axis=0)
+    at_thresh_selected = at_threshold & (at_thresh_cumsum <= n_needed_at_thresh)
+
+    final_mask = above_threshold | at_thresh_selected
+
+    tl.store(mask_base + kv_offsets * stride_m_kv, final_mask, mask=valid_mask)
+
+
+# Triton kernel loads the entire kv row into registers via tl.arange(0, KV_BLOCK_SIZE).
+# Each row spawns multiple same-sized register arrays (scores_f32, valid_mask,
+# above_threshold, at_threshold, cumsum, etc.).  GPU SMs have a fixed register file
+# (e.g. 65536 × 32-bit), so the maximum array length per program is bounded.
+# 4096 (2^12) is an empirical power-of-2 cap that avoids register spilling on
+# mainstream GPUs.  Beyond this the kernel either fails to compile or spills to
+# local memory with severe perf regression, so we fall back to torch.topk.
+MAX_KV_BLOCK_SIZE = 4096
+
+
+def _pytorch_topk_mask_fallback(
+    scores: torch.Tensor,
+    topk: int,
+) -> torch.Tensor:
+    topk_idx = torch.topk(scores, topk, dim=-1).indices
+    return torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
+
+
+def fused_topk_mask(
+    scores: torch.Tensor,
+    topk: int,
+) -> torch.Tensor:
+    """Build topk boolean mask from scores using fused Triton kernel.
+
+    Falls back to PyTorch when kv_blocks exceeds the Triton block size
+    limit (MAX_KV_BLOCK_SIZE) to avoid compilation failures.
+
+    Args:
+        scores: [B, H, q_blocks, kv_blocks] block-level attention scores
+        topk: number of top blocks to select per q-block
+
+    Returns:
+        mask: [B, H, q_blocks, kv_blocks] bool tensor with exactly topk True per row
+    """
+    B, H, q_blocks, kv_blocks = scores.shape
+    topk = min(topk, kv_blocks)
+
+    KV_BLOCK_SIZE = triton.next_power_of_2(kv_blocks)
+    if KV_BLOCK_SIZE > MAX_KV_BLOCK_SIZE:
+        logger.debug(
+            "fused_topk_mask: kv_blocks=%d exceeds Triton limit %d, "
+            "falling back to PyTorch topk (slower)",
+            kv_blocks, MAX_KV_BLOCK_SIZE,
+        )
+        return _pytorch_topk_mask_fallback(scores, topk)
+
+    mask = torch.zeros(B, H, q_blocks, kv_blocks, dtype=torch.bool, device=scores.device)
+
+    scores_flat = scores.contiguous().view(B * H, q_blocks, kv_blocks)
+    mask_flat = mask.view(B * H, q_blocks, kv_blocks)
+
+    grid = (q_blocks, B * H)
+
+    _fused_topk_mask_kernel[grid](
+        scores_flat, mask_flat,
+        scores_flat.stride(0), scores_flat.stride(1), scores_flat.stride(2),
+        mask_flat.stride(0), mask_flat.stride(1), mask_flat.stride(2),
+        kv_blocks=kv_blocks,
+        topk=topk,
+        KV_BLOCK_SIZE=KV_BLOCK_SIZE,
+    )
+
+    return mask

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.2.6"
+__version__ = "0.3.0"

fastvideo-kernel/tests/test_fused_compress_topk.py

@@ -0,0 +1,401 @@
+"""Tests for fused_block_mean and fused_topk_mask equivalence against PyTorch reference."""
+
+import torch
+import pytest
+
+
+def _require_cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
+
+# ---------------------------------------------------------------------------
+# fused_topk_mask: tie-boundary correctness
+# ---------------------------------------------------------------------------
+
+
+class TestFusedTopkMaskTies:
+    """Verify that fused_topk_mask selects exactly topk per row, even with ties."""
+
+    def _run_topk_mask(self, scores, topk):
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_topk_mask
+        return fused_topk_mask(scores.cuda(), topk)
+
+    def _ref_topk_mask(self, scores, topk):
+        """PyTorch reference: topk + scatter."""
+        idx = torch.topk(scores, topk, dim=-1).indices
+        mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, idx, True)
+        return mask
+
+    def test_all_equal_scores(self):
+        """All scores identical — must still select exactly topk."""
+        _require_cuda()
+        scores = torch.full((1, 1, 4, 8), 0.5, dtype=torch.bfloat16)
+        for topk in [1, 3, 5, 8]:
+            mask = self._run_topk_mask(scores, topk)
+            row_counts = mask.sum(dim=-1)
+            assert (row_counts == topk).all(), (
+                f"topk={topk}, row_counts={row_counts}"
+            )
+
+    def test_reviewer_example(self):
+        """Exact case from the reviewer: two 0.8875 with topk=1 → must select 1."""
+        _require_cuda()
+        vals = torch.tensor(
+            [0.5222, 0.5222, 0.8875, 0.5222, 0.5222, 0.5222, 0.8875],
+            dtype=torch.bfloat16,
+        )
+        scores = vals.view(1, 1, 1, -1)
+        mask = self._run_topk_mask(scores, topk=1)
+        assert mask.sum().item() == 1, f"Expected 1 selected, got {mask.sum().item()}"
+
+    def test_tie_at_boundary_various_topk(self):
+        """Scores with a large tied cluster at the k-th boundary."""
+        _require_cuda()
+        scores = torch.tensor(
+            [1.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.2],
+            dtype=torch.bfloat16,
+        ).view(1, 1, 1, -1)
+        for topk in [1, 2, 3, 4, 7, 8]:
+            mask = self._run_topk_mask(scores, topk)
+            assert mask.sum().item() == topk, (
+                f"topk={topk}, selected={mask.sum().item()}"
+            )
+
+    def test_batch_heads_ties(self):
+        """Multiple batch/heads with tie-heavy scores."""
+        _require_cuda()
+        B, H, Q, KV = 2, 4, 8, 16
+        scores = torch.zeros(B, H, Q, KV, dtype=torch.bfloat16)
+        scores[:, :, :, :4] = 1.0
+        scores[:, :, :, 4:] = 0.5
+        for topk in [1, 4, 8]:
+            mask = self._run_topk_mask(scores, topk)
+            row_counts = mask.sum(dim=-1)
+            assert (row_counts == topk).all(), (
+                f"topk={topk}, row_counts min={row_counts.min()}, max={row_counts.max()}"
+            )
+
+
+# ---------------------------------------------------------------------------
+# fused_topk_mask: equivalence with PyTorch topk+scatter
+# ---------------------------------------------------------------------------
+
+
+class TestFusedTopkMaskEquivalence:
+    """fused_topk_mask must select the same set as torch.topk (modulo tie order)."""
+
+    def test_random_scores_row_count(self):
+        """Random scores (unlikely ties) — row count must match exactly."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_topk_mask
+
+        B, H, Q, KV = 2, 8, 32, 64
+        topk = 8
+        scores = torch.randn(B, H, Q, KV, dtype=torch.bfloat16, device="cuda")
+        mask = fused_topk_mask(scores, topk)
+        assert (mask.sum(dim=-1) == topk).all()
+
+    def test_random_scores_value_match(self):
+        """With distinct scores, fused and reference must select identical indices."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_topk_mask
+
+        B, H, Q, KV = 1, 2, 16, 32
+        topk = 4
+        scores = torch.arange(KV, dtype=torch.float32).view(1, 1, 1, KV)
+        scores = scores.expand(B, H, Q, KV).contiguous().to(torch.bfloat16).cuda()
+
+        fused_mask = fused_topk_mask(scores, topk)
+
+        ref_idx = torch.topk(scores.cuda(), topk, dim=-1).indices
+        ref_mask = torch.zeros_like(scores, dtype=torch.bool, device="cuda").scatter_(-1, ref_idx, True)
+
+        assert (fused_mask == ref_mask).all()
+
+
+# ---------------------------------------------------------------------------
+# fused_block_mean: equivalence with PyTorch reference
+# ---------------------------------------------------------------------------
+
+
+class TestFusedBlockMeanEquivalence:
+
+    def test_against_pytorch_reference(self):
+        """fused_block_mean must match the original view→float→sum→div→bf16 path."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_block_mean
+
+        B, H, D = 2, 4, 128
+        num_blocks = 16
+        block_elements = 64
+        seq_len = num_blocks * block_elements
+
+        x = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda")
+        vbs = torch.randint(16, 65, (num_blocks,), dtype=torch.int32, device="cuda")
+
+        fused_out = fused_block_mean(x, vbs, block_elements)
+
+        x_blocks = x.view(B, H, num_blocks, block_elements, D)
+        ref_out = (x_blocks.float().sum(dim=3) / vbs.view(1, 1, -1, 1).float()).to(torch.bfloat16)
+
+        assert fused_out.shape == ref_out.shape
+        assert torch.allclose(fused_out, ref_out, atol=1e-2, rtol=1e-2), (
+            f"max diff={( fused_out - ref_out).abs().max().item()}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# fused_block_mean: backward (gradient) parity
+# ---------------------------------------------------------------------------
+
+
+def _ref_block_mean_with_grad(x, vbs, block_elements):
+    """Eager PyTorch block-mean that supports autograd for gradient reference."""
+    B, H, seq_len, D = x.shape
+    num_blocks = seq_len // block_elements
+    x_blocks = x.view(B, H, num_blocks, block_elements, D)
+    return (x_blocks.float().sum(dim=3) / vbs.view(1, 1, -1, 1).float()).to(x.dtype)
+
+
+class TestFusedBlockMeanBackward:
+
+    def test_gradient_parity(self):
+        """Gradient through fused_block_mean must match the eager view→sum→div path."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_block_mean
+
+        B, H, D = 2, 4, 128
+        num_blocks = 16
+        block_elements = 64
+        seq_len = num_blocks * block_elements
+
+        vbs = torch.randint(16, 65, (num_blocks,), dtype=torch.int32, device="cuda")
+        grad_out = torch.randn(B, H, num_blocks, D, dtype=torch.bfloat16, device="cuda")
+
+        x_fused = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda")
+        x_ref = x_fused.clone()
+        x_fused.requires_grad_(True)
+        x_ref.requires_grad_(True)
+
+        out_fused = fused_block_mean(x_fused, vbs, block_elements)
+        out_fused.backward(grad_out)
+
+        out_ref = _ref_block_mean_with_grad(x_ref, vbs, block_elements)
+        out_ref.backward(grad_out)
+
+        assert x_fused.grad is not None, "fused path did not produce gradient"
+        assert x_ref.grad is not None, "ref path did not produce gradient"
+        assert x_fused.grad.shape == x_ref.grad.shape
+
+        max_diff = (x_fused.grad - x_ref.grad).abs().max().item()
+        assert torch.allclose(x_fused.grad, x_ref.grad, atol=1e-2, rtol=1e-2), (
+            f"gradient max diff={max_diff}"
+        )
+
+    def test_gradient_parity_variable_block_sizes(self):
+        """Backward parity with non-uniform variable_block_sizes."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_block_mean
+
+        B, H, D = 1, 8, 64
+        num_blocks = 32
+        block_elements = 64
+        seq_len = num_blocks * block_elements
+
+        vbs = torch.randint(1, 65, (num_blocks,), dtype=torch.int32, device="cuda")
+        grad_out = torch.randn(B, H, num_blocks, D, dtype=torch.bfloat16, device="cuda")
+
+        x_fused = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda")
+        x_ref = x_fused.clone()
+        x_fused.requires_grad_(True)
+        x_ref.requires_grad_(True)
+
+        out_fused = fused_block_mean(x_fused, vbs, block_elements)
+        out_fused.backward(grad_out)
+
+        out_ref = _ref_block_mean_with_grad(x_ref, vbs, block_elements)
+        out_ref.backward(grad_out)
+
+        max_diff = (x_fused.grad - x_ref.grad).abs().max().item()
+        assert torch.allclose(x_fused.grad, x_ref.grad, atol=1e-2, rtol=1e-2), (
+            f"gradient max diff={max_diff}"
+        )
+
+    def test_gradient_through_matmul_chain(self):
+        """End-to-end backward through the compress branch: block_mean → matmul → softmax → matmul."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_block_mean
+
+        B, H, D = 1, 4, 128
+        num_blocks = 16
+        block_elements = 64
+        seq_len = num_blocks * block_elements
+
+        vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device="cuda")
+
+        def _forward_compress(x_q, x_k, x_v, use_fused):
+            mean_fn = fused_block_mean if use_fused else (
+                lambda x, v, be: _ref_block_mean_with_grad(x, v, be)
+            )
+            q_c = mean_fn(x_q, vbs, block_elements)
+            k_c = mean_fn(x_k, vbs, block_elements)
+            v_c = mean_fn(x_v, vbs, block_elements)
+            scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D ** 0.5)
+            attn = torch.softmax(scores, dim=-1)
+            return torch.matmul(attn, v_c)
+
+        base_q = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda")
+        base_k = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda")
+        base_v = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda")
+        grad_out = torch.randn(B, H, num_blocks, D, dtype=torch.bfloat16, device="cuda")
+
+        grads = {}
+        for label, use_fused in [("fused", True), ("ref", False)]:
+            q = base_q.clone().requires_grad_(True)
+            k = base_k.clone().requires_grad_(True)
+            v = base_v.clone().requires_grad_(True)
+            out = _forward_compress(q, k, v, use_fused)
+            out.backward(grad_out)
+            grads[label] = (q.grad, k.grad, v.grad)
+
+        for name, (g_fused, g_ref) in zip(
+            ["dQ", "dK", "dV"],
+            zip(grads["fused"], grads["ref"]),
+        ):
+            max_diff = (g_fused - g_ref).abs().max().item()
+            assert torch.allclose(g_fused, g_ref, atol=1e-2, rtol=1e-2), (
+                f"{name} gradient max diff={max_diff}"
+            )
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: fused video_sparse_attn vs old PyTorch pipeline
+# ---------------------------------------------------------------------------
+
+
+def _old_pytorch_video_sparse_attn(
+    q, k, v, variable_block_sizes, q_variable_block_sizes,
+    topk, block_elements,
+):
+    """Reproduce the pre-optimization PyTorch pipeline from ops.py."""
+    batch, heads, q_seq_len, dim = q.shape
+    kv_seq_len = k.shape[2]
+    q_num_blocks = q_seq_len // block_elements
+    kv_num_blocks = kv_seq_len // block_elements
+
+    q_c = q.view(batch, heads, q_num_blocks, block_elements, dim)
+    k_c = k.view(batch, heads, kv_num_blocks, block_elements, dim)
+    v_c = v.view(batch, heads, kv_num_blocks, block_elements, dim)
+    q_c = (q_c.float().sum(dim=3) / q_variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype)
+    k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype)
+    v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype)
+
+    scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim ** 0.5)
+    attn = torch.softmax(scores, dim=-1)
+    out_c = torch.matmul(attn, v_c)
+    out_c = out_c.view(batch, heads, q_num_blocks, 1, dim)
+    out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, q_seq_len, dim)
+
+    topk_idx = torch.topk(scores, topk, dim=-1).indices
+    mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
+
+    return out_c, scores, mask
+
+
+class TestVideoSparseAttnEquivalence:
+    """End-to-end: fused compress+topk path vs old PyTorch pipeline."""
+
+    def _make_inputs(self, B, H, num_blocks, D, block_elements, device="cuda"):
+        seq_len = num_blocks * block_elements
+        q = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device=device)
+        k = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device=device)
+        v = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device=device)
+        vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device=device)
+        return q, k, v, vbs
+
+    def test_compress_branch_equivalence(self):
+        """Compression branch (block_mean → matmul → softmax → matmul) must match."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import fused_block_mean
+
+        B, H, num_blocks, D = 1, 4, 32, 128
+        block_elements = 64
+        topk = 4
+        q, k, v, vbs = self._make_inputs(B, H, num_blocks, D, block_elements)
+
+        q_c_fused = fused_block_mean(q, vbs, block_elements)
+        k_c_fused = fused_block_mean(k, vbs, block_elements)
+        v_c_fused = fused_block_mean(v, vbs, block_elements)
+        scores_fused = torch.matmul(q_c_fused, k_c_fused.transpose(-2, -1)) / (D ** 0.5)
+        attn_fused = torch.softmax(scores_fused, dim=-1)
+        out_c_fused = torch.matmul(attn_fused, v_c_fused)
+
+        old_out_c, old_scores, _ = _old_pytorch_video_sparse_attn(
+            q, k, v, vbs, vbs, topk, block_elements,
+        )
+
+        assert torch.allclose(scores_fused, old_scores, atol=1e-2, rtol=1e-2), (
+            f"scores max diff={( scores_fused - old_scores).abs().max().item()}"
+        )
+
+    def test_topk_mask_row_count_matches(self):
+        """Fused topk mask must select exactly topk per row, same as torch.topk."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import (
+            fused_block_mean,
+            fused_topk_mask,
+        )
+
+        B, H, num_blocks, D = 1, 8, 48, 128
+        block_elements = 64
+        topk = 6
+        q, k, v, vbs = self._make_inputs(B, H, num_blocks, D, block_elements)
+
+        q_c = fused_block_mean(q, vbs, block_elements)
+        k_c = fused_block_mean(k, vbs, block_elements)
+        scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D ** 0.5)
+
+        fused_mask = fused_topk_mask(scores, topk)
+        row_counts = fused_mask.sum(dim=-1)
+        assert (row_counts == topk).all(), (
+            f"row_counts min={row_counts.min()}, max={row_counts.max()}"
+        )
+
+        _, _, ref_mask = _old_pytorch_video_sparse_attn(
+            q, k, v, vbs, vbs, topk, block_elements,
+        )
+        ref_counts = ref_mask.sum(dim=-1)
+        assert (ref_counts == topk).all()
+
+    def test_e2e_mask_equivalence_realistic_shape(self):
+        """On realistic bf16 q_c@k_c scores, fused and ref mask should agree."""
+        _require_cuda()
+        from fastvideo_kernel.triton_kernels.fused_compress_topk import (
+            fused_block_mean,
+            fused_topk_mask,
+        )
+
+        B, H, num_blocks, D = 1, 16, 32, 128
+        block_elements = 64
+        topk = 4
+        q, k, v, vbs = self._make_inputs(B, H, num_blocks, D, block_elements)
+
+        q_c = fused_block_mean(q, vbs, block_elements)
+        k_c = fused_block_mean(k, vbs, block_elements)
+        scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D ** 0.5)
+
+        fused_mask = fused_topk_mask(scores, topk)
+
+        ref_idx = torch.topk(scores, topk, dim=-1).indices
+        ref_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, ref_idx, True)
+
+        total_rows = B * H * num_blocks
+        matching_rows = (fused_mask == ref_mask).all(dim=-1).sum().item()
+        match_rate = matching_rows / total_rows
+
+        print(f"\n[e2e mask equivalence] {matching_rows}/{total_rows} rows match "
+              f"({match_rate:.2%})")
+
+        assert (fused_mask.sum(dim=-1) == topk).all(), "row count mismatch"
+        assert match_rate > 0.99, f"match rate too low: {match_rate:.2%}"

fastvideo/attention/backends/sage_attn3.py

@@ -55,6 +55,20 @@ class SageAttention3Impl(AttentionImpl):
         self.softmax_scale = softmax_scale
         self.dropout = extra_impl_args.get("dropout_p", 0.0)
 
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """Transpose stacked QKV from [3B, L, H, D] to [3B, H, L, D].
+
+        Single bulk permute+contiguous on the entire stacked tensor rather than
+        three separate transposed views for Q, K, V. The .contiguous() is
+        required: sageattn_blackwell's fake kernel returns empty_like(q), so the
+        op's output strides must match contiguous q under torch.compile.
+        """
+        return qkv.permute(0, 2, 1, 3).contiguous()
+
     def forward(
         self,
         query: torch.Tensor,
@@ -62,9 +76,15 @@ class SageAttention3Impl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
-        query = query.transpose(1, 2)
-        key = key.transpose(1, 2)
-        value = value.transpose(1, 2)
+        """Call sageattn3_blackwell directly. Input is already [B, H, L, D]
+        and contiguous from preprocess_qkv."""
         output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
-        output = output.transpose(1, 2)
         return output
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """Transpose output from [B, H, L, D] back to [B, L, H, D]."""
+        return output.permute(0, 2, 1, 3).contiguous()

fastvideo/attention/layer.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+
 import torch
 import torch.nn as nn
 
@@ -13,6 +15,26 @@ from fastvideo.utils import get_compute_dtype
 from fastvideo.layers.rotary_embedding import _apply_rotary_emb
 
 
+def _attention_compile_disabled() -> bool:
+    """Whether to keep attention ``forward`` out of the torch.compile graph.
+
+    Defaults to ``True`` (the historical behavior: attention runs eager via
+    ``torch.compiler.disable``). Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=0``
+    to let attention be traced/compiled into the surrounding graph.
+    """
+    val = os.environ.get("FASTVIDEO_DISABLE_ATTENTION_COMPILE")
+    if val is None:
+        return True
+    return val.strip().lower() not in ("0", "false", "no", "off", "")
+
+
+def _maybe_compiler_disable(fn):
+    """Apply ``torch.compiler.disable`` unless disabled via env var."""
+    if _attention_compile_disabled():
+        return torch.compiler.disable(fn)
+    return fn
+
+
 class DistributedAttention(nn.Module):
     """Distributed attention layer.
     """
@@ -56,7 +78,7 @@ class DistributedAttention(nn.Module):
         self.backend = backend_name_to_enum(attn_backend.get_name())
         self.dtype = dtype
 
-    @torch.compiler.disable
+    @_maybe_compiler_disable
     def forward(
         self,
         q: torch.Tensor,
@@ -146,7 +168,7 @@ class DistributedAttention_VSA(DistributedAttention):
     """Distributed attention layer with VSA support.
     """
 
-    @torch.compiler.disable
+    @_maybe_compiler_disable
     def forward(
         self,
         q: torch.Tensor,

fastvideo/eval/metrics/audio/desync/metric.py

@@ -77,11 +77,21 @@ def _segment_video(frames: torch.Tensor) -> torch.Tensor:
 
 
 def _load_audio_waveform(audio_path: str, resample_cache: dict[int, Any]) -> torch.Tensor:
-    """Load *audio_path* at 16 kHz mono, native length. Returns ``(N,)``."""
-    import torchaudio
-    waveform, sr = torchaudio.load(audio_path)
-    waveform = waveform.mean(dim=0)
+    """Load *audio_path* at 16 kHz mono, native length. Returns ``(N,)``.
+
+    Decode with librosa (libsndfile) rather than ``torchaudio.load``: recent
+    torchaudio routes ``load`` through torchcodec, whose shared libraries fail
+    to load when the installed build doesn't match the pinned torch. The other
+    audio metrics already decode via librosa, and the asset / PyAV-extracted
+    clips are PCM WAV, so libsndfile reproduces the samples bit-for-bit — the
+    score is unchanged. ``sr=None`` keeps the native rate so the torchaudio
+    resampler below is identical to before.
+    """
+    import librosa
+    waveform_np, sr = librosa.load(audio_path, sr=None, mono=True)
+    waveform = torch.from_numpy(waveform_np)
     if sr != _AUDIO_SR:
+        import torchaudio
         if sr not in resample_cache:
             resample_cache[sr] = torchaudio.transforms.Resample(sr, _AUDIO_SR)
         waveform = resample_cache[sr](waveform)
@@ -134,7 +144,7 @@ class DeSyncMetric(BaseMetric):
     needs_gpu = True
     is_set_metric = False
     backbone = "synchformer"
-    dependencies = ["torchaudio"]
+    dependencies = ["librosa", "torchaudio"]
 
     def __init__(self, src_fps: float | None = None) -> None:
         super().__init__()

fastvideo/eval/metrics/videoscore2/metric.py

@@ -158,7 +158,13 @@ class VideoScore2Metric(BaseMetric):
         infer_fps: float = 2.0,
         max_tokens: int = 1024,
         temperature: float = 0.7,
-        do_sample: bool = True,
+        # Greedy by default — eval metrics should be deterministic on a
+        # given input.  Upstream's vs2_inference.py defaults to
+        # do_sample=True with temperature=0.7, but stochastic decoding
+        # makes the score irreproducible run-to-run AND can land the
+        # chain-of-thought outside the regex anchor (yielding score=0).
+        # Pass do_sample=True explicitly if you want BoN/ensembling.
+        do_sample: bool = False,
     ) -> None:
         super().__init__()
         self._model_name = model_name
@@ -214,11 +220,23 @@ class VideoScore2Metric(BaseMetric):
 
     def _subsample_frames(self,
                           pil_frames: list[Image.Image],
+                          source_fps: float | None,
                           max_frames: int = 64,
                           max_resolution: int = 960) -> list[Image.Image]:
-        """Subsample to ~infer_fps worth of frames (max 64), resize if too large."""
+        """Subsample to ``infer_fps`` worth of frames, capped at ``max_frames``.
+
+        Mirrors upstream's path-based qwen_vl_utils flow, which samples at
+        ``infer_fps``.  Without ``source_fps`` we can't convert frame count
+        to duration, so we fall back to ``max_frames`` evenly spaced — the
+        pre-fps-derivation behavior.  Pass ``sample["fps"]`` from the
+        caller to get the upstream-aligned sampling rate.
+        """
         n = len(pil_frames)
-        target = min(n, max_frames)
+        if source_fps is not None and source_fps > 0:
+            duration = n / source_fps
+            target = max(1, min(max_frames, int(round(duration * self.infer_fps))))
+        else:
+            target = min(n, max_frames)
         if target < n:
             indices = np.linspace(0, n - 1, target, dtype=int)
             pil_frames = [pil_frames[i] for i in indices]
@@ -244,7 +262,11 @@ class VideoScore2Metric(BaseMetric):
             text = text[0] if text else ""
 
         pil_frames = self._tensor_to_pil_list(video)
-        pil_frames = self._subsample_frames(pil_frames)
+        source_fps = sample.get("fps")
+        pil_frames = self._subsample_frames(
+            pil_frames,
+            source_fps=float(source_fps) if source_fps is not None else None,
+        )
         user_prompt = VS2_QUERY_TEMPLATE.substitute(t2v_prompt=text)
 
         messages = [{

fastvideo/fastvideo_args.py

@@ -273,11 +273,17 @@ class FastVideoArgs:
         dit_config = getattr(self.pipeline_config, "dit_config", None)
         if dit_config is None:
             return
+        # Resolve a registry name (e.g. "nvfp4_qat_train" from the CLI) to a
+        # QuantizationConfig instance; a bare string has no get_quant_method.
+        tq = self.transformer_quant
+        if isinstance(tq, str):
+            from fastvideo.layers.quantization import get_quantization_config
+            tq = get_quantization_config(tq)()
         # Don't overwrite if the caller already set it explicitly on
         # dit_config (e.g. via ``pipeline_config.dit_config.quant_config = NVFP4Config()``);
         # the explicit setter wins.
         if getattr(dit_config, "quant_config", None) is None:
-            dit_config.quant_config = self.transformer_quant
+            dit_config.quant_config = tq
 
     def _resolve_refine_args(self) -> None:
         """Map generic refine_* args to LTX-2-specific refine fields."""
@@ -1018,6 +1024,12 @@ class TrainingArgs(FastVideoArgs):
     @staticmethod
     def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parser.add_argument("--data-path", type=str, required=True, help="Path to parquet files")
+        parser.add_argument("--transformer-quant",
+                            type=str,
+                            default=None,
+                            help="Quantization config name for the DiT (e.g. nvfp4_qat_train for "
+                            "QAT-finetune FP4 linear with a straight-through estimator). "
+                            "Resolved to a QuantizationConfig and pinned on dit_config.quant_config.")
         parser.add_argument("--dataloader-num-workers",
                             type=int,
                             required=True,

fastvideo/layers/fp8linear.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FP8 quantization-aware training for linear layers.
+
+Mirror of ``fp4linear.py`` but for FP8 (e4m3). The forward pass quantizes both
+activations and weights to FP8 and runs ``torch._scaled_mm``; the backward pass
+is a bf16 straight-through estimator so the high-precision master weights stay
+trainable. Falls back to a bf16 fake-quant forward on GPUs older than sm89.
+"""
+import torch
+
+FP8_DTYPE = torch.float8_e4m3fn
+FP8_MAX = float(torch.finfo(FP8_DTYPE).max)  # 448.0
+FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0)
+
+
+def _supports_fp8_compute() -> bool:
+    """Whether the active device supports FP8 ``_scaled_mm`` (sm89+)."""
+    if not torch.cuda.is_available():
+        return False
+    cap = torch.cuda.get_device_capability()
+    return cap[0] > 8 or (cap[0] == 8 and cap[1] >= 9)
+
+
+def _quantize_tensorwise(x_2d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]:
+    """Returns ``(x_fp8 [M, K], x_scale [1] float32)``."""
+    x_absmax = x_2d.abs().amax().float()
+    x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
+    x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+    return x_fp8, x_scale.view(1)
+
+
+def _quantize_rowwise(x_2d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]:
+    """Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``."""
+    x_absmax = x_2d.abs().amax(dim=-1, keepdim=True).float()
+    x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
+    x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+    return x_fp8, x_scale
+
+
+def _fake_quant(x_2d: torch.Tensor, granularity: str) -> torch.Tensor:
+    """bf16 fake-quant (quantize then dequantize) for pre-sm89 fallback."""
+    if granularity == "channel":
+        x_fp8, x_scale = _quantize_rowwise(x_2d)
+    else:
+        x_fp8, x_scale = _quantize_tensorwise(x_2d)
+    return x_fp8.to(x_2d.dtype) * x_scale.to(x_2d.dtype)
+
+
+class _LinearFWD8BWD16Fn(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, weight, bias, granularity="tensor"):
+        # assert/normalize activation dtype
+        if x.dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(dtype=torch.bfloat16)
+
+        # cast params (can be fp32) to activation dtype for quantization
+        weight_cast = weight.to(dtype=x.dtype)
+        bias_cast = bias.to(dtype=x.dtype) if bias is not None else None
+
+        orig_shape = x.shape
+        k = weight_cast.shape[1]
+        n = weight_cast.shape[0]
+        x2d = x.reshape(-1, k).contiguous()
+
+        if not _supports_fp8_compute():
+            # bf16 fake-quant fallback: simulate the FP8 rounding error but
+            # compute the matmul in bf16.
+            x_fq = _fake_quant(x2d, granularity)
+            w_fq = _fake_quant(weight_cast, granularity)
+            out2d = x_fq.matmul(w_fq.t())
+            if bias_cast is not None:
+                out2d = out2d + bias_cast
+            ctx.save_for_backward(x2d, weight, bias)
+            ctx.n = n
+            ctx.orig_shape = orig_shape
+            return out2d.reshape(*orig_shape[:-1], n)
+
+        if granularity == "channel":
+            x_fp8, x_scale = _quantize_rowwise(x2d)
+            w_fp8, w_scale = _quantize_rowwise(weight_cast)
+            scale_b = w_scale.view(1, -1)
+        else:
+            x_fp8, x_scale = _quantize_tensorwise(x2d)
+            w_fp8, w_scale = _quantize_tensorwise(weight_cast)
+            scale_b = w_scale
+
+        out2d = torch._scaled_mm(
+            x_fp8,
+            w_fp8.t(),
+            scale_a=x_scale,
+            scale_b=scale_b,
+            out_dtype=x.dtype,
+        )
+        if isinstance(out2d, tuple):
+            out2d = out2d[0]
+
+        if bias_cast is not None:
+            out2d = out2d + bias_cast
+
+        # save tensors for backward (keep original dtypes)
+        ctx.save_for_backward(x2d, weight, bias)
+        ctx.n = n
+        ctx.orig_shape = orig_shape
+        return out2d.reshape(*orig_shape[:-1], n)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        x2d, weight, bias = ctx.saved_tensors
+        M = x2d.shape[0]
+        n = ctx.n
+
+        grad_out_2d = grad_out.reshape(M, n).contiguous()
+
+        # bf16 straight-through estimator: gradients flow through the
+        # full-precision master weights, not the FP8 quantized values.
+        weight_cast = weight.to(dtype=grad_out.dtype)
+        x_cast = x2d.to(dtype=grad_out.dtype)
+
+        grad_x = grad_out_2d.matmul(weight_cast).reshape(*ctx.orig_shape)
+        grad_w = grad_out_2d.t().matmul(x_cast)
+        grad_b = grad_out_2d.sum(dim=0) if bias is not None else None
+
+        # None for the extra forward arg (granularity)
+        return grad_x, grad_w, grad_b, None
+
+
+def fp8_linear_forward(self, x: torch.Tensor) -> torch.Tensor:
+    # pass config **positionally**; autograd.Function.apply ignores kwargs
+    return _LinearFWD8BWD16Fn.apply(x, self.weight, self.bias, "tensor"), None

fastvideo/layers/quantization/__init__.py

@@ -2,7 +2,7 @@ from typing import Literal, get_args
 
 from fastvideo.layers.quantization.base_config import QuantizationConfig
 
-QuantizationMethods = Literal[None, "AbsMaxFP8", "NVFP4", "nvfp4_qat"]
+QuantizationMethods = Literal[None, "AbsMaxFP8", "FP8", "NVFP4", "nvfp4_qat", "nvfp4_qat_train", "fp8_qat_train"]
 
 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
 
@@ -51,13 +51,19 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
 
     # lazy import to avoid triggering `torch.compile` too early
     from .absmax_fp8 import AbsMaxFP8Config
+    from .fp8_config import FP8Config
     from .nvfp4_config import NVFP4Config
     from .nvfp4_qat_config import NVFP4QATConfig
+    from .nvfp4_qat_train_config import NVFP4QATTrainConfig
+    from .fp8_qat_train_config import FP8QATTrainConfig
 
     method_to_config: dict[str, type[QuantizationConfig]] = {
         "AbsMaxFP8": AbsMaxFP8Config,
+        "FP8": FP8Config,
         "NVFP4": NVFP4Config,
         "nvfp4_qat": NVFP4QATConfig,
+        "nvfp4_qat_train": NVFP4QATTrainConfig,
+        "fp8_qat_train": FP8QATTrainConfig,
     }
     # Update the `method_to_config` with customized quantization methods.
     method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)

fastvideo/layers/quantization/fp8_config.py

@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generic FP8 quantization backed by ``torch._scaled_mm``.
+
+Matches linear layers by suffix (``to_q/k/v/to_out``, ``ffn.fc_in/fc_out``).
+Supports per-tensor (default, fast) and per-channel (higher accuracy) granularity.
+Falls back to bf16 dequant on GPUs older than sm89.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter
+
+from fastvideo.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from fastvideo.models.utils import set_weight_attrs
+
+logger = logging.getLogger(__name__)
+
+FP8_DTYPE = torch.float8_e4m3fn
+FP8_MAX = float(torch.finfo(FP8_DTYPE).max)  # 448.0
+FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0)
+
+_FP8_SUFFIXES = (
+    "ffn.fc_in",
+    "ffn.fc_out",
+    "to_q",
+    "to_k",
+    "to_v",
+    "to_out",
+)
+
+
+def _supports_fp8_compute() -> bool:
+    """Whether the active device supports FP8 ``_scaled_mm`` (sm89+)."""
+    if not torch.cuda.is_available():
+        return False
+    cap = torch.cuda.get_device_capability()
+    return cap[0] > 8 or (cap[0] == 8 and cap[1] >= 9)
+
+
+def _quantize_tensorwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Returns ``(x_fp8 [M, K], x_scale [1] float32)``."""
+    x_absmax = x_2d.abs().amax().float()
+    x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
+    x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+    return x_fp8, x_scale.view(1)
+
+
+def _quantize_rowwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``."""
+    x_absmax = x_2d.abs().amax(dim=-1, keepdim=True).float()
+    x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
+    x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+    return x_fp8, x_scale
+
+
+class FP8QuantizeMethod(QuantizeMethodBase):
+    """FP8 linear method.
+
+    ``granularity='tensor'`` (default): per-tensor weight + per-tensor
+    dynamic activation scales — the fast tensorwise ``_scaled_mm`` path.
+    ``granularity='channel'``: per-output-channel weight + per-token
+    activation scales (rowwise) — higher accuracy but slower ``_scaled_mm``.
+    """
+
+    def __init__(self, granularity: str = "tensor"):
+        super().__init__()
+        self.granularity = granularity
+
+    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,
+    ):
+        weight = Parameter(
+            torch.empty(
+                sum(output_partition_sizes),
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+        layer.register_parameter("weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+    def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]:
+        """Pre-quantize an activation for reuse across q/k/v projections."""
+        assert x.dtype in (torch.bfloat16, torch.float16), (f"only allow bf16/fp16 inputs to fp8 linear, got {x.dtype}")
+        x_2d = x.view(-1, x.shape[-1])
+        if self.granularity == "channel":
+            x_fp8, x_scale = _quantize_rowwise(x_2d)
+        else:
+            x_fp8, x_scale = _quantize_tensorwise(x_2d)
+        return x_fp8, x_scale, None
+
+    def wants_prequantized_input(self) -> bool:
+        return True
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        pre_quantized: tuple[torch.Tensor, torch.Tensor, Any] | None = None,
+    ) -> torch.Tensor:
+        out_dim = layer._fp8_weight.shape[0]
+        original_shape = x.shape
+
+        if not _supports_fp8_compute():
+            return self._apply_dequant(layer, x, bias)
+
+        if pre_quantized is not None:
+            x_fp8, x_scale, _ = pre_quantized
+            if x_fp8.dim() > 2:
+                x_fp8 = x_fp8.reshape(-1, x_fp8.shape[-1])
+            if x_scale.dim() > 2:
+                x_scale = x_scale.reshape(-1, x_scale.shape[-1])
+        elif self.granularity == "channel":
+            x_fp8, x_scale = _quantize_rowwise(x.reshape(-1, x.shape[-1]))
+        else:
+            x_fp8, x_scale = _quantize_tensorwise(x.reshape(-1, x.shape[-1]))
+
+        w_fp8 = layer._fp8_weight
+        w_scale = layer._fp8_weight_scale
+        scale_b = w_scale.view(1, -1) if self.granularity == "channel" else w_scale
+
+        out = torch._scaled_mm(
+            x_fp8,
+            w_fp8.t(),
+            scale_a=x_scale,
+            scale_b=scale_b,
+            out_dtype=torch.bfloat16,
+        )
+        if isinstance(out, tuple):
+            out = out[0]
+        if bias is not None:
+            out = out + bias
+        return out.view(*original_shape[:-1], out_dim)
+
+    def _apply_dequant(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """bf16 fallback for pre-sm89 GPUs."""
+        out_dim = layer._fp8_weight.shape[0]
+        original_shape = x.shape
+        w_fp8 = layer._fp8_weight
+        w_scale = layer._fp8_weight_scale.to(x.dtype)
+        weight = w_fp8.to(x.dtype) * w_scale.unsqueeze(1)
+        out = F.linear(x, weight, bias)
+        return out.view(*original_shape[:-1], out_dim)
+
+
+class FP8Config(QuantizationConfig):
+    """FP8 (e4m3) quantization via suffix matching on standard linear layer names."""
+
+    def __init__(self, granularity: str = "tensor"):
+        super().__init__()
+        if granularity not in ("tensor", "channel"):
+            raise ValueError(f"granularity must be 'tensor' or 'channel', got {granularity!r}")
+        self.granularity = granularity
+
+    def get_name(self) -> str:
+        return "FP8"
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 89
+
+    @staticmethod
+    def get_config_filenames() -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> FP8Config:
+        return cls(granularity=config.get("granularity", "tensor"))
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        from fastvideo.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase) and any(s in prefix for s in _FP8_SUFFIXES):
+            return FP8QuantizeMethod(granularity=self.granularity)
+        return None
+
+
+def convert_model_to_fp8(model: torch.nn.Module) -> None:
+    """Quantize all FP8-tagged linear layers in-place after weights are loaded."""
+    import gc
+    from torch.distributed.tensor import DTensor  # type: ignore
+
+    with torch.no_grad():
+        for mod in model.modules():
+            qm = getattr(mod, "quant_method", None)
+            if not isinstance(qm, FP8QuantizeMethod):
+                continue
+            weight = getattr(mod, "weight", None)
+            if weight is None:
+                continue
+            weight_local = weight.to_local() if isinstance(weight, DTensor) else weight  # type: ignore[arg-type]
+            if getattr(qm, "granularity", "tensor") == "channel":
+                w_absmax = weight_local.detach().abs().amax(dim=1).nan_to_num().float()
+                w_scale = (w_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
+                w_fp8 = (weight_local / w_scale.to(weight_local.dtype).unsqueeze(1)).clamp(-FP8_MAX,
+                                                                                           FP8_MAX).to(FP8_DTYPE)
+            else:
+                w_absmax = weight_local.detach().abs().amax().nan_to_num().to(torch.float32)
+                w_scale = (w_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE).view(1)
+                w_fp8 = (weight_local / w_scale.to(weight_local.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+            mod.register_buffer("_fp8_weight", w_fp8.contiguous(), persistent=False)
+            mod.register_buffer("_fp8_weight_scale", w_scale.to(torch.float32), persistent=False)
+            removed_weight = mod._parameters.pop("weight", None)
+            if removed_weight is not None:
+                removed_weight.grad = None
+            del removed_weight, weight, weight_local, w_absmax, w_scale, w_fp8
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+__all__ = [
+    "FP8Config",
+    "FP8QuantizeMethod",
+    "convert_model_to_fp8",
+]

fastvideo/layers/quantization/fp8_qat_train_config.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FP8 (e4m3) quantization-aware *training* linear method (straight-through estimator).
+
+Mirror of ``nvfp4_qat_train_config.py`` but for FP8. The weight stays a trainable
+bf16/fp32 master that is fake-quantized to FP8 on every forward, with a
+full-precision backward (STE), so the model learns to absorb FP8 linear error.
+
+The STE lives in ``fastvideo.layers.fp8linear._LinearFWD8BWD16Fn`` (FP8 forward
+via ``torch._scaled_mm`` on sm89+, with a bf16 fake-quant fallback on older GPUs;
+full-precision backward). This method bridges it into the standard
+``quant_config`` path, so it activates via ``transformer_quant="fp8_qat_train"``
+on the same Wan-2.1 layers as the FP4 path (to_q/k/v/out + ffn). No conversion is
+needed: the weight is kept in full precision and quantized on the fly each step.
+
+Unlike the FP4 path this needs no flashinfer and runs on any sm89+ GPU (and even
+older ones via the bf16 fallback), not just Blackwell.
+"""
+import logging
+from typing import Any
+
+import torch
+from torch.nn.parameter import Parameter
+
+from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase
+from fastvideo.models.utils import set_weight_attrs
+
+logger = logging.getLogger(__name__)
+
+
+class FP8QATTrainQuantizeMethod(QuantizeMethodBase):
+
+    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):
+        # Trainable master weight, fake-quantized to FP8 on each forward.
+        weight = Parameter(torch.empty(
+            sum(output_partition_sizes),
+            input_size_per_partition,
+            dtype=params_dtype,
+        ),
+                           requires_grad=True)
+        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+        layer.register_parameter("weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+    def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+        # FP8 forward + full-precision backward (STE).
+        from fastvideo.layers.fp8linear import _LinearFWD8BWD16Fn
+        return _LinearFWD8BWD16Fn.apply(x, layer.weight, bias, "tensor")
+
+
+class FP8QATTrainConfig(QuantizationConfig):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def get_name(self):
+        return "fp8_qat_train"
+
+    def get_supported_act_dtypes(self):
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls):
+        return 89
+
+    @staticmethod
+    def get_config_filenames():
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "FP8QATTrainConfig":
+        return cls()
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        from fastvideo.layers.linear import LinearBase
+        fp8_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 fp8_layers):
+            return FP8QATTrainQuantizeMethod()
+        return None

fastvideo/layers/quantization/nvfp4_qat_config.py

@@ -1,38 +1,90 @@
 # SPDX-License-Identifier: Apache-2.0
+"""NVFP4 quantization-aware (QAD) linear method, inference path.
+
+Quantizes every targeted linear's weight to NVFP4 once at load time and
+runs each forward as a registered flashinfer-backed FP4 matmul. The
+original fp16/bf16 weight is *popped* immediately after quantization so
+the half-precision copy does not keep occupying GPU memory — that's
+what lets a Wan-2.1 pipeline stay fully resident on a single GPU
+without any CPU offloading.
+
+The quantize / matmul custom ops are owned by
+:mod:`fastvideo.layers.quantization.nvfp4_config` and registered under
+the ``fastvideo_fp4::`` namespace. We reuse them here for two reasons:
+
+1. Re-registering the same op name in a second module would raise.
+2. The registered ops have ``register_fake`` shape/dtype kernels, which
+   is what makes the inference pipeline's per-block ``torch.compile``
+   trace through without graph breaks. Calling raw flashinfer functions
+   (the old behavior of this file, plus a ``@torch.compile`` on
+   ``apply``) graph-breaks at every quantize and every matmul.
+
+For QAT *training*, see ``nvfp4_qat_train_config`` which keeps the
+weight trainable and fake-quantizes on the fly via a straight-through
+estimator.
+"""
+from __future__ import annotations
+
+import gc
 import logging
 from typing import Any
 
 import torch
 from torch.nn.parameter import Parameter
 
-from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase
+from fastvideo.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from fastvideo.layers.quantization.nvfp4_config import (
+    _mm_fp4,
+    _nvfp4_quantize,
+    _require_flashinfer,
+)
 from fastvideo.models.utils import set_weight_attrs
 
-try:
-    import flashinfer
-except ImportError:
-    flashinfer = None
-
 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").
+DEFAULT_FP4_LAYERS = (
+    "ffn.fc_in",
+    "ffn.fc_out",
+    "to_q",
+    "to_k",
+    "to_v",
+    "to_out",
+)
 
-def _require_flashinfer() -> Any:
-    if flashinfer is None:
-        raise ImportError("flashinfer is required for NVFP4 QAT quantization. "
-                          "Please install flashinfer to use the nvfp4_qat quantization backend.")
-    return flashinfer
+
+def _layout_128x4() -> Any:
+    SfLayout, _, _ = _require_flashinfer()
+    return SfLayout.layout_128x4
 
 
 class NVFP4QATQuantizeMethod(QuantizeMethodBase):
+    """Inference-only NVFP4 linear method with weight popping.
+
+    The dense ``weight`` parameter is materialized at load time only so
+    that :func:`convert_model_to_fp4` can read it once; the loader then
+    removes it via ``mod._parameters.pop('weight')``. From that point
+    forward, ``apply`` reads only ``_fp4_weight`` / ``_fp4_weight_scale``
+    / ``_weight_global_sf``.
+    """
 
     def __init__(self) -> None:
         super().__init__()
         self.weight_fp4 = None
         self.weight_scale = None
+        # Static input global scale factor. Matches the FastVideo-Quantization
+        # production path; recomputing it per-call via a ``.max()`` reduction
+        # (the previous behavior) adds a sync point, costs a kernel launch,
+        # and produces a data-dependent value that prevents CUDA-graph
+        # capture under ``torch.compile(mode='reduce-overhead')``.
+        self.x_global_sf = torch.tensor(1.0, device="cuda", dtype=torch.float32)
 
     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):
-        """Create weights for a linear layer. Note the corrected signature to match LinearMethodBase."""
+                       input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs) -> None:
         weight = Parameter(torch.empty(
             sum(output_partition_sizes),
             input_size_per_partition,
@@ -43,28 +95,27 @@ class NVFP4QATQuantizeMethod(QuantizeMethodBase):
         layer.register_parameter("weight", weight)
         set_weight_attrs(weight, extra_weight_attrs)
 
-    @torch.compile
     def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
-        """Apply NVFP4 QAT quantized computation."""
-        flashinfer_mod = _require_flashinfer()
-        out_dim = layer.weight.shape[0]
+        # ``_fp4_weight`` carries the (out, in/2) packed fp4 weight, so its
+        # row count is the output dim even after the dense weight is popped.
+        out_dim = layer._fp4_weight.shape[0]
         original_shape = x.shape
-        assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}"
 
+        assert x.dtype in (torch.bfloat16, torch.float16), (f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
         x = x.view(-1, x.shape[-1])
-
-        x_global_sf = (448 * 6) / x.float().abs().nan_to_num().max()
-        x_fp4, x_scale = flashinfer_mod.nvfp4_quantize(
+        x_global_sf = self.x_global_sf
+        x_fp4, x_scale = _nvfp4_quantize(
             x,
             x_global_sf,
-            sfLayout=flashinfer_mod.SfLayout.layout_128x4,
+            sfLayout=_layout_128x4(),
             do_shuffle=False,
         )
+
         weight_fp4 = layer._fp4_weight
         weight_scale = layer._fp4_weight_scale
         weight_global_sf = layer._weight_global_sf
 
-        out = flashinfer_mod.mm_fp4(
+        out = _mm_fp4(
             x_fp4,
             weight_fp4.T,
             x_scale,
@@ -76,67 +127,106 @@ class NVFP4QATQuantizeMethod(QuantizeMethodBase):
         )
 
         if bias is not None:
-            if bias.device != out.device or bias.dtype != out.dtype:
-                bias = bias.to(device=out.device, dtype=out.dtype)
             out = out + bias
-
-        if len(original_shape) == 3:
-            out = out.view(original_shape[0], original_shape[1], out_dim)
-
+        out = out.view(*original_shape[:-1], out_dim)
         return out
 
 
 class NVFP4QATConfig(QuantizationConfig):
+    """NVFP4 (Wan-style) linear quantization, inference.
 
-    def __init__(self) -> None:
+    Args:
+        target_layers: Substrings matched against each linear layer's
+            prefix. A layer is quantized if any substring is contained in
+            its prefix. Defaults to the standard Wan attention + FFN
+            projections (:data:`DEFAULT_FP4_LAYERS`).
+    """
+
+    def __init__(self, target_layers: tuple[str, ...] | None = None) -> None:
         super().__init__()
+        self.target_layers = (tuple(target_layers) if target_layers else DEFAULT_FP4_LAYERS)
 
-    def get_name(self):
+    def get_name(self) -> str:
         return "nvfp4_qat"
 
-    def get_supported_act_dtypes(self):
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
         return [torch.bfloat16, torch.float16]
 
     @classmethod
-    def get_min_capability(cls):
+    def get_min_capability(cls) -> int:
         return 100
 
     @staticmethod
-    def get_config_filenames():
+    def get_config_filenames() -> list[str]:
         return []
 
     @classmethod
-    def from_config(cls, config: dict[str, Any]) -> "NVFP4QATConfig":
-        return cls()
+    def from_config(cls, config: dict[str, Any]) -> NVFP4QATConfig:
+        target_layers = config.get("target_layers")
+        if target_layers is not None:
+            target_layers = tuple(target_layers)
+        return cls(target_layers=target_layers)
 
     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 isinstance(layer, LinearBase) and any(name in prefix for name in self.target_layers):
             return NVFP4QATQuantizeMethod()
         return None
 
 
-@torch.compile
-def convert_model_to_fp4(model: torch.nn.Module):
-    flashinfer_mod = _require_flashinfer()
+def convert_model_to_fp4(model: torch.nn.Module) -> None:
+    """Prequantize every FP4-tagged linear and drop its dense weight.
+
+    Walks the module tree, and for each layer whose ``quant_method`` is
+    an :class:`NVFP4QATQuantizeMethod`, computes the NVFP4 packed weight
+    / scale / global-scale buffers, then pops the original fp16/bf16
+    ``weight`` parameter so it no longer occupies GPU memory.
+    """
+    SfLayout, _, _ = _require_flashinfer()
     from torch.distributed.tensor import DTensor  # type: ignore
-    for mod in model.modules():
-        qm = getattr(mod, "quant_method", None)
-        if isinstance(qm, NVFP4QATQuantizeMethod):
+
+    with torch.no_grad():
+        for mod in model.modules():
+            qm = getattr(mod, "quant_method", None)
+            if not isinstance(qm, NVFP4QATQuantizeMethod):
+                continue
+
             weight = getattr(mod, "weight", None)
             if weight is None:
                 continue
+
             weight_local = weight.to_local() if isinstance(weight, DTensor) else weight  # type: ignore[arg-type]
-            weight_global_sf = (448 * 6) / weight_local.float().abs().nan_to_num().max()
-            fp4_w, fp4_s = flashinfer_mod.nvfp4_quantize(
+
+            # Only the reduced scalar needs fp32; avoid a full fp32 copy.
+            weight_absmax = (weight_local.detach().abs().nan_to_num().amax().to(dtype=torch.float32))
+            weight_global_sf = (448 * 6) / weight_absmax
+            fp4_w, fp4_s = _nvfp4_quantize(
                 weight_local,
                 weight_global_sf,
-                sfLayout=flashinfer_mod.SfLayout.layout_128x4,
+                sfLayout=SfLayout.layout_128x4,
                 do_shuffle=False,
             )
             mod.register_buffer("_fp4_weight", fp4_w, persistent=False)
             mod.register_buffer("_fp4_weight_scale", fp4_s, persistent=False)
-            mod.register_buffer("_weight_global_sf",
-                                torch.tensor(weight_global_sf, dtype=torch.bfloat16),
-                                persistent=False)
+            mod.register_buffer(
+                "_weight_global_sf",
+                weight_global_sf.to(dtype=torch.bfloat16),
+                persistent=False,
+            )
+
+            # Drop the dense weight as soon as the fp4 buffers are installed
+            # so it cannot keep occupying GPU memory.
+            removed_weight = mod._parameters.pop("weight", None)
+            if removed_weight is not None:
+                removed_weight.grad = None
+            del removed_weight, weight, weight_local, weight_absmax
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+__all__ = [
+    "NVFP4QATConfig",
+    "NVFP4QATQuantizeMethod",
+    "convert_model_to_fp4",
+    "DEFAULT_FP4_LAYERS",
+]

fastvideo/layers/quantization/nvfp4_qat_train_config.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""NVFP4 quantization-aware *training* linear method (straight-through estimator).
+
+The inference ``nvfp4_qat`` config quantizes each weight to FP4 once at load time
+(``convert_model_to_fp4``) and has no gradient path — it is inference only. For
+QAT *finetuning* the weight must stay a trainable bf16/fp32 master that is
+fake-quantized to FP4 on every forward, with a full-precision backward (a
+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.
+"""
+import logging
+from typing import Any
+
+import torch
+from torch.nn.parameter import Parameter
+
+from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase
+from fastvideo.models.utils import set_weight_attrs
+
+logger = logging.getLogger(__name__)
+
+
+class NVFP4QATTrainQuantizeMethod(QuantizeMethodBase):
+
+    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):
+        # Trainable master weight, fake-quantized to FP4 on each forward.
+        weight = Parameter(torch.empty(
+            sum(output_partition_sizes),
+            input_size_per_partition,
+            dtype=params_dtype,
+        ),
+                           requires_grad=True)
+        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+        layer.register_parameter("weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+    def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+        # FP4 forward + full-precision backward (STE).
+        from fastvideo.layers.fp4linear import _LinearFWD4BWD16Fn
+        return _LinearFWD4BWD16Fn.apply(x, layer.weight, bias, "cutlass", 16, True)
+
+
+class NVFP4QATTrainConfig(QuantizationConfig):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def get_name(self):
+        return "nvfp4_qat_train"
+
+    def get_supported_act_dtypes(self):
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls):
+        return 100
+
+    @staticmethod
+    def get_config_filenames():
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "NVFP4QATTrainConfig":
+        return cls()
+
+    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):
+            return NVFP4QATTrainQuantizeMethod()
+        return None

fastvideo/models/loader/fsdp_load.py

@@ -28,31 +28,29 @@ from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available
 logger = init_logger(__name__)
 
 
-def _maybe_convert_model_to_nvfp4(model: nn.Module) -> None:
-    """Quantize NVFP4-tagged linear layers in-place after weights are loaded.
+def _maybe_quantize_model(model: nn.Module) -> None:
+    """Quantize NVFP4- or FP8-tagged linear layers in-place after weights are loaded.
 
     Walks the module tree once, looking for layers whose ``quant_method``
-    is an :class:`NVFP4QuantizeMethod` (attached at construction time by
-    :meth:`NVFP4Config.get_quant_method`). When at least one such layer
-    exists, calls :func:`convert_model_to_nvfp4` to register the
-    ``_nvfp4_weight*`` / ``_nvfp4_alpha`` / ``_weight_global_sf`` buffers
-    on each targeted layer.
+    is an :class:`NVFP4QuantizeMethod` or :class:`FP8QuantizeMethod` (attached
+    at construction time by the respective ``get_quant_method``). When at least
+    one such layer exists, calls the matching conversion function to register
+    quantized weight buffers on each targeted layer.
 
-    The walk returns on the first NVFP4 layer found so non-NVFP4 callers
-    pay only an ``isinstance`` check per module. flashinfer is imported
-    lazily inside :func:`convert_model_to_nvfp4` so this helper is a
-    no-op on hosts without the NVFP4 backend.
+    The walk returns on the first quantized layer found so unquantized callers
+    pay only an ``isinstance`` check per module. Both imports are deferred so
+    this is a no-op on hosts without the relevant backends.
     """
-    # Defer the import: nvfp4_config imports heavy diffusers /
-    # torch.distributed symbols at module-load time, and unconditional
-    # import would penalize every loader call regardless of whether
-    # NVFP4 is wired.
+    # Defer imports: these modules pull in heavy symbols at module-load time.
     from fastvideo.layers.quantization.nvfp4_config import (
         NVFP4QuantizeMethod, convert_model_to_nvfp4,
     )
     from fastvideo.layers.quantization.nvfp4_qat_config import (
         NVFP4QATQuantizeMethod, convert_model_to_fp4,
     )
+    from fastvideo.layers.quantization.fp8_config import (
+        FP8QuantizeMethod, convert_model_to_fp8,
+    )
 
     for mod in model.modules():
         qm = getattr(mod, "quant_method", None)
@@ -64,6 +62,10 @@ def _maybe_convert_model_to_nvfp4(model: nn.Module) -> None:
             logger.info("Converting loaded model weights for NVFP4-QAT linear layers")
             convert_model_to_fp4(model)
             return
+        if isinstance(qm, FP8QuantizeMethod):
+            logger.info("Converting loaded model weights for FP8 linear layers")
+            convert_model_to_fp8(model)
+            return
 
 
 # TODO(PY): move this to utils elsewhere
@@ -196,14 +198,13 @@ def maybe_load_fsdp_model(
         if isinstance(p, torch.nn.Parameter):
             p.requires_grad = False
 
-    # NVFP4 weight prequantization. We detect by the registered
-    # ``quant_method`` on linear layers rather than by a separate flag —
-    # construction-time ``NVFP4Config.get_quant_method`` already attached
-    # ``NVFP4QuantizeMethod`` to every targeted layer, so the loader's
-    # responsibility is just to materialize the per-layer nvfp4 weight /
-    # scale buffers from the freshly-loaded bf16 weights. No-op when
-    # ``flashinfer`` is not installed (lazy import inside the helper).
-    _maybe_convert_model_to_nvfp4(model)
+    # Post-load weight quantization. We detect the active scheme by the
+    # ``quant_method`` attached to each linear layer at construction time
+    # (via ``QuantizationConfig.get_quant_method``). The loader's
+    # responsibility is just to materialize the quantized weight buffers
+    # from the freshly-loaded bf16 weights. No-op when no quantized layers
+    # are present (lazy imports inside the helper).
+    _maybe_quantize_model(model)
 
     compile_in_loader = enable_torch_compile and training_mode
     if compile_in_loader:

fastvideo/performance_dashboard/api.py

@@ -111,10 +111,17 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
         days: int = Query(DEFAULT_DAYS, ge=1, le=3650),
         model_id: str | None = None,
         gpu_type: str | None = None,
+        run_source: str | None = None,
         success: bool | None = None,
     ) -> dict[str, Any]:
         loaded = data_store.load_records(days=days)
-        filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type, success=success)
+        filtered = filter_records(
+            loaded,
+            model_id=model_id,
+            gpu_type=gpu_type,
+            run_source=run_source,
+            success=success,
+        )
         return {
             "records": filtered,
             "count": len(filtered),
@@ -122,6 +129,7 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
                 "days": days,
                 "model_id": model_id,
                 "gpu_type": gpu_type,
+                "run_source": run_source,
                 "success": success,
             },
             "sync": data_store.health(),
@@ -132,6 +140,7 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
         days: int = Query(DEFAULT_DAYS, ge=1, le=3650),
         model_id: str | None = None,
         gpu_type: str | None = None,
+        run_source: str | None = None,
     ) -> dict[str, Any]:
         # Latest status should be stable when users change the trend window.
         # Use all cached records for latest/baseline computation; the ``days``
@@ -139,7 +148,11 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
         # endpoints without affecting the summary semantics.
         loaded = data_store.load_records(days=None)
         filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type)
-        rows = build_latest_summary(filtered, max_regression=float(os.environ.get("PERF_MAX_REGRESSION", "0.05")))
+        rows = build_latest_summary(
+            filtered,
+            max_regression=float(os.environ.get("PERF_MAX_REGRESSION", "0.05")),
+            run_source=run_source,
+        )
         return {
             "rows": rows,
             "count": len(rows),
@@ -152,6 +165,7 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
                 "trend_window_days": days,
                 "model_id": model_id,
                 "gpu_type": gpu_type,
+                "run_source": run_source,
             },
             "sync": data_store.health(),
         }
@@ -161,9 +175,10 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
         days: int = Query(DEFAULT_DAYS, ge=1, le=3650),
         model_id: str | None = None,
         gpu_type: str | None = None,
+        run_source: str | None = None,
     ) -> dict[str, Any]:
         loaded = data_store.load_records(days=days)
-        filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type)
+        filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type, run_source=run_source)
         groups = build_trends(filtered)
         return {
             "groups": groups,
@@ -172,6 +187,7 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
                 "days": days,
                 "model_id": model_id,
                 "gpu_type": gpu_type,
+                "run_source": run_source,
             },
             "sync": data_store.health(),
         }

fastvideo/performance_dashboard/service.py

@@ -13,7 +13,7 @@ from collections import defaultdict
 from datetime import datetime, timezone
 from typing import Any
 
-from fastvideo.tests.performance.hf_store import safe_float
+from fastvideo.tests.performance.hf_store import is_baseline_eligible_record, safe_float
 
 from .metrics import METRICS
 
@@ -45,6 +45,7 @@ def filter_records(
     *,
     model_id: str | None = None,
     gpu_type: str | None = None,
+    run_source: str | None = None,
     success: bool | None = None,
 ) -> list[Record]:
     filtered = records
@@ -52,11 +53,31 @@ def filter_records(
         filtered = [record for record in filtered if record.get("model_id") == model_id]
     if gpu_type:
         filtered = [record for record in filtered if record.get("gpu_type") == gpu_type]
+    if run_source:
+        filtered = [record for record in filtered if record_run_source(record) == run_source]
     if success is not None:
         filtered = [record for record in filtered if bool(record.get("success", True)) == success]
     return sorted(filtered, key=record_sort_key)
 
 
+def record_run_source(record: Record) -> str:
+    value = str(record.get("run_source") or "unknown")
+    return value if value in {"pr", "local", "scheduled_main", "unknown"} else "unknown"
+
+
+def record_metadata(record: Record) -> Record:
+    return {
+        "run_source": record_run_source(record),
+        "baseline_eligible": is_baseline_eligible_record(record),
+        "branch": record.get("branch") or "",
+        "pr_number": record.get("pr_number") or "",
+        "test_scope": record.get("test_scope") or "",
+        "build_url": record.get("build_url") or "",
+        "build_id": record.get("build_id") or "",
+        "job_id": record.get("job_id") or "",
+    }
+
+
 def group_by_model_gpu(records: list[Record]) -> dict[tuple[str, str], list[Record]]:
     groups: dict[tuple[str, str], list[Record]] = defaultdict(list)
     for record in records:
@@ -86,12 +107,22 @@ def regression_percent(metric_key: str, current: float | None, baseline: float |
 def build_latest_summary(records: list[Record],
                          *,
                          baseline_window: int = 5,
-                         max_regression: float = 0.05) -> list[Record]:
+                         max_regression: float = 0.05,
+                         run_source: str | None = None) -> list[Record]:
     rows: list[Record] = []
     for (model_id, gpu_type), group in group_by_model_gpu(records).items():
-        latest = group[-1]
-        earlier_successes = [record for record in group[:-1] if record.get("success", True)]
-        baseline_records = earlier_successes[-baseline_window:]
+        latest_candidates = group
+        if run_source:
+            latest_candidates = [record for record in group if record_run_source(record) == run_source]
+        if not latest_candidates:
+            continue
+
+        latest = latest_candidates[-1]
+        baseline_pool = [
+            record for record in group
+            if record is not latest and record.get("success", True) and is_baseline_eligible_record(record)
+        ]
+        baseline_records = baseline_pool[-baseline_window:]
 
         metrics: dict[str, Record] = {}
         regressions: list[float] = []
@@ -123,6 +154,7 @@ def build_latest_summary(records: list[Record],
             latest.get("timestamp"),
             "commit_sha":
             latest.get("commit_sha"),
+            **record_metadata(latest),
             "success":
             success,
             "baseline_n":
@@ -150,6 +182,7 @@ def build_trends(records: list[Record]) -> list[Record]:
             point = {
                 "timestamp": record.get("timestamp"),
                 "commit_sha": record.get("commit_sha"),
+                **record_metadata(record),
                 "success": bool(record.get("success", True)),
                 "metrics": {
                     metric.key: safe_float(record.get(metric.key))

fastvideo/registry.py

@@ -615,6 +615,8 @@ def _register_configs() -> None:
                 "cosmos-predict2.5",
             )) and "14b" not in path.lower(),
         ],
+        model_family="cosmos25",
+        default_preset="cosmos25_predict2_2b",
     )
 
     # Cosmos 2.5 (14B)

fastvideo/tests/api/test_presets.py

@@ -464,6 +464,18 @@ class TestCosmosPresets:
         assert len(presets) == 1
         assert presets[0].name == "cosmos25_predict2_2b"
 
+    def test_cosmos25_2b_path_resolves_to_preset(self) -> None:
+        """Regression: the Cosmos-Predict2.5 2B model path must select the
+        cosmos25 preset (guidance_scale=7.0), not fall back to the generic
+        default (guidance_scale=1.0) which produced blurry output.
+        """
+        import fastvideo.registry  # noqa: F401
+        from fastvideo.api.sampling_param import SamplingParam
+        sp = SamplingParam.from_pretrained(
+            "KyleShao/Cosmos-Predict2.5-2B-Diffusers")
+        assert sp.guidance_scale == 7.0
+        assert sp.num_inference_steps == 35
+
     def test_cosmos_and_cosmos25_different_fps(self) -> None:
         import fastvideo.registry  # noqa: F401
         c = get_preset("cosmos_predict2_2b", "cosmos")

fastvideo/tests/eval/asset/ltx2.metadata.json

@@ -0,0 +1,24 @@
+{
+  "_comment": "Metadata for the eval regression asset video. Used by the upstream reference-implementation runs to produce reference scores, and by the fastvideo eval regression tests to feed every metric the inputs it expects. Keep in sync with examples/inference/eval/all_metrics_demo.py (PROMPT and OUTPUT_PATH).",
+  "video": "fastvideo/tests/eval/asset/ltx2.mp4",
+  "fps": 24.0,
+  "width": 832,
+  "height": 480,
+  "duration_seconds": 5.0,
+  "num_frames": 121,
+  "has_audio": true,
+  "audio_sample_rate": 24000,
+  "audio_channels": 2,
+  "language": "en",
+  "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+  "visual_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions.",
+  "dialogue": [
+    { "speaker": "woman", "text": "That's it... Dad's lost it. And we've lost Dad." },
+    { "speaker": "man", "text": "Stop being so dramatic, Jess." }
+  ],
+  "reference_text": "That's it. Dad's lost it. And we've lost Dad. Stop being so dramatic, Jess.",
+  "generator": {
+    "name": "FastVideo/LTX2-Distilled-Diffusers",
+    "script": "examples/inference/eval/all_metrics_demo.py"
+  }
+}

fastvideo/tests/eval/conftest.py

@@ -0,0 +1,123 @@
+"""Shared fixtures for the eval regression suite.
+
+Two session fixtures keep model loads off the per-test path:
+
+* ``gold_results`` — runs every reference-backed metric once via the
+  top-level :func:`samples_from` + :meth:`Evaluator.evaluate` API and
+  caches the per-sample dict for the asset.
+* ``invariant_results`` — runs every identity-invariant metric on
+  N duplicate copies of the asset (paired-and-corpus shape).
+
+Both use ``skip_missing_deps=True`` so a slim test venv missing one
+metric's deps drops only that metric, never the whole run.
+"""
+from __future__ import annotations
+
+import json
+import shutil
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from fastvideo.eval import EvalResults, create_evaluator, samples_from
+
+ASSET_DIR = Path(__file__).resolve().parent / "asset"
+ASSET_VIDEO = ASSET_DIR / "ltx2.mp4"
+ASSET_AUDIO = ASSET_DIR / "ltx2.wav"
+ASSET_META = ASSET_DIR / "ltx2.metadata.json"
+REF_DIR = Path(__file__).resolve().parent / "reference_scores"
+
+N_DUP = 4  # duplicates for set-metric invariants (fvd/fad need ≥2)
+
+
+def _all_reference_files() -> list[Path]:
+    return sorted(REF_DIR.glob("*.json"))
+
+
+def _reference_metric_names() -> list[str]:
+    return [json.loads(p.read_text())["metric"] for p in _all_reference_files()]
+
+
+# Identity invariants — both fastvideo and the math itself give the
+# perfect-match value on duplicate (gen, ref). No reference JSON needed.
+INVARIANT_SPEC = [
+    # (metric, expected, abs_tolerance, location)
+    ("common.ssim",                  1.0, 1e-3, "per_sample"),
+    ("common.psnr",                 99.0, None, "per_sample"),  # ≥ threshold (identical → ~inf, capped by floats)
+    ("common.lpips",                 0.0, 1e-2, "per_sample"),
+    ("optical_flow.gt_optical_flow", 0.0, 1e-2, "per_sample"),
+    ("common.fvd",                   0.0, 1.0,  "corpus"),
+    ("audio.frechet_distance",       0.0, 1.0,  "corpus"),
+    ("audio.kl_divergence",          0.0, 1e-3, "per_sample"),
+]
+
+
+def load_reference(metric: str) -> dict:
+    return json.loads((REF_DIR / f"{metric}.json").read_text())
+
+
+@pytest.fixture(scope="session")
+def asset_meta() -> dict:
+    return json.loads(ASSET_META.read_text())
+
+
+@pytest.fixture(scope="session")
+def gold_results(asset_meta) -> dict:
+    """Top-level API run on the single regression asset.
+
+    Equivalent to four lines in user code::
+
+        ev = create_evaluator(metrics=[...], skip_missing_deps=True)
+        samples = samples_from(video=ASSET, audio=ASSET_WAV,
+                               text_prompt=..., fps=..., extras={...})
+        return ev.evaluate(samples=samples)[0]
+    """
+    metrics = _reference_metric_names()
+    ev = create_evaluator(metrics=metrics, skip_missing_deps=True)
+    samples = samples_from(
+        video=ASSET_VIDEO,
+        audio=ASSET_AUDIO,
+        text_prompt=asset_meta["text_prompt"],
+        fps=asset_meta["fps"],
+        extras={
+            "reference_text": asset_meta["reference_text"],
+            "video_path": str(ASSET_VIDEO),
+            "language": asset_meta["language"],
+        },
+    )
+    results: EvalResults = ev.evaluate(samples=samples)
+    yield results[0]
+    ev.shutdown()
+
+
+@pytest.fixture(scope="session")
+def invariant_results(asset_meta) -> EvalResults:
+    """Run identity-invariant metrics on N duplicate copies of the asset.
+
+    Same top-level shape — five lines.  ``extract_audio`` pulls a ``.wav``
+    next to each copy so audio set-metrics
+    (``audio.frechet_distance``, ``audio.kl_divergence``) see paired
+    audio without us pre-staging it.
+    """
+    tmp = Path(tempfile.mkdtemp(prefix="eval_regression_inv_"))
+    try:
+        gen = tmp / "gen"; ref = tmp / "ref"
+        gen.mkdir(); ref.mkdir()
+        for i in range(N_DUP):
+            shutil.copyfile(ASSET_VIDEO, gen / f"clip_{i:03d}.mp4")
+            shutil.copyfile(ASSET_VIDEO, ref / f"clip_{i:03d}.mp4")
+        ev = create_evaluator(
+            metrics=[m for m, *_ in INVARIANT_SPEC],
+            skip_missing_deps=True,
+        )
+        samples = samples_from(
+            video=gen, reference=ref,
+            text_prompt=asset_meta["text_prompt"],
+            fps=asset_meta["fps"],
+            extract_audio=tmp / "audio_cache",
+        )
+        yield ev.evaluate(samples=samples)
+        ev.shutdown()
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)

fastvideo/tests/eval/reference_scores/audio.audiobox_aesthetics.json

@@ -0,0 +1,12 @@
+{
+  "metric": "audio.audiobox_aesthetics",
+  "score": 6.900569438934326,
+  "details": {
+    "CE": 5.3528289794921875,
+    "CU": 6.002830982208252,
+    "PC": 2.6171281337738037,
+    "PQ": 6.900569438934326,
+    "upstream": "facebookresearch/audiobox-aesthetics",
+    "asset_audio": "fastvideo/tests/eval/asset/ltx2.wav"
+  }
+}

fastvideo/tests/eval/reference_scores/audio.clap_score.json

@@ -0,0 +1,11 @@
+{
+  "metric": "audio.clap_score",
+  "score": 0.1640956997871399,
+  "details": {
+    "model_name": "laion/clap-htsat-fused",
+    "upstream": "transformers.ClapModel (HF mirror of LAION-AI/CLAP 630k-audioset-fusion-best)",
+    "sample_rate": 48000,
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_audio": "fastvideo/tests/eval/asset/ltx2.wav"
+  }
+}

fastvideo/tests/eval/reference_scores/audio.desync.json

@@ -0,0 +1,12 @@
+{
+  "metric": "audio.desync",
+  "score": 0.0,
+  "details": {
+    "front_desync_s": 0.0,
+    "back_desync_s": 0.0,
+    "num_segments_used": 14,
+    "upstream": "hkchengrex/av-benchmark (Synchformer)",
+    "checkpoint": "synchformer_state_dict.pth (MMAudio v0.1 release)",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/audio.imagebind_score.json

@@ -0,0 +1,11 @@
+{
+  "metric": "audio.imagebind_score",
+  "score": 0.33291810750961304,
+  "details": {
+    "model": "imagebind_huge",
+    "upstream": "facebookresearch/ImageBind",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4",
+    "asset_audio": "fastvideo/tests/eval/asset/ltx2.wav",
+    "decoder": "pyav (aarch64 patch for upstream's decord default)"
+  }
+}

fastvideo/tests/eval/reference_scores/audio.wer.json

@@ -0,0 +1,12 @@
+{
+  "metric": "audio.wer",
+  "score": 0.0,
+  "details": {
+    "transcription": "thats it dads lost it and weve lost dad stop being so dramatic jess",
+    "reference_text": "thats it dads lost it and weve lost dad stop being so dramatic jess",
+    "asr_backend": "whisper",
+    "model_name": "whisper-base",
+    "upstream": "openai/whisper",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.aesthetic_quality.json

@@ -0,0 +1,16 @@
+{
+  "metric": "vbench.aesthetic_quality",
+  "score": 0.5733335614204407,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": 0.5733335614204407
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.appearance_style.json

@@ -0,0 +1,9 @@
+{
+  "metric": "vbench.appearance_style",
+  "score": null,
+  "details": {
+    "skipped": "upstream VBench.check_dimension_requires_extra_info refuses appearance_style in mode='custom_input' — it needs the structured 'style:' auxiliary_info from the curated VBench prompt set.",
+    "upstream": "Vchitect/VBench",
+    "deferred_to_stage": 2
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.background_consistency.json

@@ -0,0 +1,18 @@
+{
+  "metric": "vbench.background_consistency",
+  "score": 0.9454950925682774,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": 0.9454950925682774,
+        "video_sim": 112.513916015625,
+        "cnt_per_video": 119
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.dynamic_degree.json

@@ -0,0 +1,16 @@
+{
+  "metric": "vbench.dynamic_degree",
+  "score": 1.0,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": true
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.human_action.json

@@ -0,0 +1,9 @@
+{
+  "metric": "vbench.human_action",
+  "score": null,
+  "details": {
+    "skipped": "upstream UMT (vbench/third_party/umt/datasets/mae.py) imports decord unconditionally; no aarch64 decord wheel. Also: our CI asset prompt is not a Kinetics-400 action label, so the score would be meaningless.",
+    "upstream": "Vchitect/VBench (UMT ViT-L/16)",
+    "deferred_to_stage": 2
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.imaging_quality.json

@@ -0,0 +1,16 @@
+{
+  "metric": "vbench.imaging_quality",
+  "score": 0.6644590053558349,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": 66.44590053558349
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.motion_smoothness.json

@@ -0,0 +1,16 @@
+{
+  "metric": "vbench.motion_smoothness",
+  "score": 0.9897254431173182,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": 0.9897254431173182
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.overall_consistency.json

@@ -0,0 +1,16 @@
+{
+  "metric": "vbench.overall_consistency",
+  "score": 0.3105851113796234,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": 0.3105851113796234
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.subject_consistency.json

@@ -0,0 +1,16 @@
+{
+  "metric": "vbench.subject_consistency",
+  "score": 0.9549573574747358,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": 0.9549573574747358
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.temporal_flickering.json

@@ -0,0 +1,16 @@
+{
+  "metric": "vbench.temporal_flickering",
+  "score": 0.979451680650898,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": 0.979451680650898
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/vbench.temporal_style.json

@@ -0,0 +1,16 @@
+{
+  "metric": "vbench.temporal_style",
+  "score": 0.3105851113796234,
+  "details": {
+    "video_results": [
+      {
+        "video_path": "/tmp/claude-1039/vbench_gold_kvlismkw/videos/ltx2.mp4",
+        "video_results": 0.3105851113796234
+      }
+    ],
+    "upstream": "Vchitect/VBench",
+    "mode": "custom_input",
+    "text_prompt": "A warm sunny backyard. The camera starts in a tight cinematic close-up of a woman and a man in their 30s, facing each other with serious expressions. The woman, emotional and dramatic, says softly, \"That's it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly annoyed: \"Stop being so dramatic, Jess.\"",
+    "asset_video": "fastvideo/tests/eval/asset/ltx2.mp4"
+  }
+}

fastvideo/tests/eval/reference_scores/videoscore2.json

@@ -0,0 +1,18 @@
+{
+  "metric": "videoscore2",
+  "score": 3.6666333333333334,
+  "details": {
+    "visual_quality": 4.0,
+    "text_alignment": 2.9999,
+    "physical_consistency": 4.0,
+    "visual_quality_hard": 4,
+    "text_alignment_hard": 3,
+    "physical_consistency_hard": 4,
+    "raw_output": "\n<think>\nmulti-dimensional analysis:\nLet me analyze this video based on the provided frames and the text prompt.\n\nText Prompt Analysis:\nThe prompt describes a warm sunny backyard with a tight cinematic close-up of two adults facing each other, where the woman expresses drama about losing someone named Dad, and the man tries to calm her down by telling her to stop being dramatic. In the video, the setting and framing match the description, but the exact dialogue is not clearly conveyed; the emotional tension is present but the specific lines are not verifiable from visuals alone.\n\nVisual Quality Analysis:\nThe video maintains solid visual quality overall. Resolution and detail are good, lighting is warm and consistent with a sunny outdoor scene, and color grading remains cohesive across shots. There are no major distortions or artifacts, and transitions feel smooth. However, there are occasional minor softness and slight noise in some frames, which prevents a perfect top-tier score. Overall, Visual Quality aligns with a strong 4/5.\n\nText-to-Video Alignment Analysis:\nThe video largely reflects the intended scenario: a warm backyard setting, a tight close-up, and two adults facing each other with serious expressions. The emotional dynamics are suggested, but the precise dialogue content cannot be confirmed from visuals alone. Given these constraints, the alignment is moderate rather than perfect, corresponding to a 3/5.\n\nPhysical/Common-sense Consistency:\nInteractions appear natural and coherent within the scene\u2019s context. Body language, spacing, and lighting support a believable outdoor confrontation, with no obvious logical inconsistencies. This dimension is well-supported, warranting a 4/5.\n\nGround-truth scores:\nVisual Quality: 4\nText-to-Video Alignment: 3\nPhysical Consistency: 4\n</think>\n\n(1) visual quality \u2013 clarity, smoothness, artifacts: 4\n(2) text-to-video alignment \u2013 fidelity to the prompt: 3\n(3) physical/common-sense consistency \u2013 naturalness and physics plausibility: 4\n",
+    "upstream": "TIGER-AI-Lab/VideoScore2 (greedy)",
+    "model_name": "TIGER-Lab/VideoScore2",
+    "infer_fps": 2.0,
+    "do_sample": false,
+    "frame_sampler": "fastvideo tensor->PIL, np.linspace at infer_fps*duration (~10 frames for 5s clip); upstream vs2_inference.py uses qwen_vl_utils path-based sampler which picks slightly different frames \u2014 same model, ~1-soft-score difference on the physics dim is the residual gap."
+  }
+}

fastvideo/tests/eval/test_datasets_vbench.py

@@ -1,66 +0,0 @@
-"""Smoke tests for the VBench prompt dataset and dataset registry."""
-from __future__ import annotations
-
-import pytest
-
-from fastvideo.eval.datasets import (PromptDataset, VBenchPromptDataset,
-                                     get_dataset, list_datasets)
-
-
-def test_registry_contains_vbench():
-    assert "vbench" in list_datasets()
-
-
-def test_get_dataset_returns_typed_instance():
-    ds = get_dataset("vbench", dimensions=["color"])
-    assert isinstance(ds, PromptDataset)
-    assert isinstance(ds, VBenchPromptDataset)
-    assert ds.name == "vbench"
-    assert ds.supports_dimensions is True
-
-
-def test_full_corpus_size():
-    ds = VBenchPromptDataset()
-    assert len(ds) == 946
-    assert len(ds.dimensions) == 16
-
-
-def test_rows_are_dicts_with_required_keys():
-    ds = VBenchPromptDataset(dimensions=["color"])
-    sample = ds[0]
-    assert isinstance(sample, dict)
-    assert "prompt" in sample
-    assert "n_samples" in sample
-    assert "dimensions" in sample
-
-
-def test_temporal_flickering_n_samples():
-    ds = VBenchPromptDataset(dimensions=["temporal_flickering"])
-    assert all(s["n_samples"] == 25 for s in ds)
-
-
-def test_default_n_samples():
-    ds = VBenchPromptDataset(dimensions=["subject_consistency"])
-    assert all(s["n_samples"] == 5 for s in ds)
-
-
-def test_color_aux_info_is_flat():
-    ds = VBenchPromptDataset(dimensions=["color"])
-    sample = ds[0]
-    aux = sample["auxiliary_info"]
-    # Flat: {"color": "<color name>"}, not nested under the dimension.
-    assert "color" in aux
-    assert isinstance(aux["color"], str)
-
-
-def test_unknown_dimension_raises():
-    with pytest.raises(ValueError, match="Unknown VBench dimensions"):
-        VBenchPromptDataset(dimensions=["bogus"])
-
-
-def test_by_dimension_groups_correctly():
-    ds = VBenchPromptDataset(dimensions=["subject_consistency", "color"])
-    groups = ds.by_dimension()
-    assert set(groups) == {"subject_consistency", "color"}
-    assert all(isinstance(s["auxiliary_info"].get("color"), str)
-               for s in groups["color"])

fastvideo/tests/eval/test_evaluator_multi_gpu.py

@@ -1,107 +0,0 @@
-"""Multi-replica eval through the public ``Evaluator`` API.
-
-Skipped automatically when fewer than 2 CUDA devices are visible.
-"""
-from __future__ import annotations
-
-import pytest
-import torch
-
-from fastvideo.eval import create_evaluator
-
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.device_count() < 2,
-    reason="multi-GPU evaluator tests require at least 2 visible CUDA devices",
-)
-
-
-_T, _C, _H, _W = 6, 3, 32, 32
-
-
-def _make_samples(n: int) -> list[dict]:
-    torch.manual_seed(7)
-    samples: list[dict] = []
-    for i in range(n):
-        gen = torch.rand(_T, _C, _H, _W)
-        # Slightly perturb each row's reference so scores vary by index.
-        ref = gen + 0.01 * (i + 1) * torch.rand_like(gen)
-        samples.append({"video": gen, "reference": ref})
-    return samples
-
-
-@pytest.fixture
-def baseline_scores():
-    """Reference scores computed on a single-GPU evaluator. The multi-GPU
-    runs must reproduce these exactly when handed the same input list —
-    that's the only way to verify round-robin dispatch isn't dropping or
-    reordering samples."""
-    samples = _make_samples(8)
-    ev = create_evaluator(
-        metrics=["common.psnr", "common.ssim"],
-        device="cuda:0",
-        num_gpus=1,
-    )
-    try:
-        out = ev.evaluate(samples=samples)
-    finally:
-        ev.shutdown()
-    return samples, out
-
-
-def test_multi_gpu_evaluator_reports_two_workers():
-    ev = create_evaluator(metrics=["common.psnr"], num_gpus=2)
-    try:
-        assert ev.num_gpus == 2
-    finally:
-        ev.shutdown()
-
-
-def test_multi_gpu_dispatch_preserves_order_and_scores(baseline_scores):
-    """Same samples, multi-GPU dispatch — results must match the single-GPU
-    baseline element-for-element. This verifies (a) the round-robin doesn't
-    reorder, (b) every sample is scored exactly once, (c) the workers
-    don't share mutable state."""
-    samples, expected = baseline_scores
-
-    ev = create_evaluator(
-        metrics=["common.psnr", "common.ssim"],
-        num_gpus=2,
-    )
-    try:
-        got = ev.evaluate(samples=samples)
-    finally:
-        ev.shutdown()
-
-    assert len(got) == len(expected)
-    for i, (g, e) in enumerate(zip(got, expected)):
-        assert set(g.keys()) == {"common.psnr", "common.ssim"}, f"row {i}"
-        assert g["common.psnr"].score == pytest.approx(e["common.psnr"].score), \
-            f"row {i} psnr drift"
-        assert g["common.ssim"].score == pytest.approx(e["common.ssim"].score), \
-            f"row {i} ssim drift"
-
-
-def test_multi_gpu_evaluator_kwargs_form_runs_on_one_replica():
-    """The kwargs form (single sample) is documented to always hit worker
-    0; this test pins the contract so future refactors don't accidentally
-    fan out a single call."""
-    ev = create_evaluator(metrics=["common.psnr"], num_gpus=2)
-    try:
-        torch.manual_seed(0)
-        gen = torch.rand(_T, _C, _H, _W)
-        out = ev.evaluate(video=gen, reference=gen)
-        assert out["common.psnr"].score > 50.0     # PSNR(x, x) is huge
-    finally:
-        ev.shutdown()
-
-
-def test_multi_gpu_release_cuda_memory_runs_clean():
-    """``release_cuda_memory`` must hit every replica without crashing."""
-    ev = create_evaluator(metrics=["common.psnr"], num_gpus=2)
-    try:
-        samples = _make_samples(2)
-        _ = ev.evaluate(samples=samples)
-        ev.release_cuda_memory()                   # should not raise
-    finally:
-        ev.shutdown()

fastvideo/tests/eval/test_evaluator_paths.py

@@ -1,158 +0,0 @@
-"""Path-input variants of the public Evaluator API.
-
-The worker boundary accepts ``video`` / ``reference`` as either a
-pre-loaded ``(T, C, H, W)`` tensor or a path-like (``str`` / ``Path``).
-These tests pin the path-form so future refactors don't accidentally
-re-require pre-loaded tensors.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-import cv2
-import numpy as np
-import pytest
-import torch
-
-from fastvideo.eval import MetricResult, create_evaluator, evaluate
-
-_T, _C, _H, _W = 6, 3, 32, 32
-
-
-def _write_tensor_as_mp4(tensor: torch.Tensor, path: Path) -> None:
-    """Write a (T, C, H, W) float [0, 1] tensor to *path* as an mp4."""
-    frames = (tensor * 255).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()
-    # cv2 expects BGR; the tests don't care about colour fidelity (they
-    # just need the bytes to round-trip), but flipping keeps the
-    # written file faithful to what an mp4 would carry on disk.
-    frames = frames[..., ::-1]
-    path.parent.mkdir(parents=True, exist_ok=True)
-    h, w = frames.shape[1], frames.shape[2]
-    writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 8, (w, h))
-    if not writer.isOpened():
-        pytest.skip("cv2.VideoWriter could not open the mp4v codec on this host")
-    for f in frames:
-        writer.write(np.ascontiguousarray(f))
-    writer.release()
-
-
-@pytest.fixture
-def video_paths(tmp_path):
-    """Two reproducible mp4s on disk + their pre-loaded tensors for parity."""
-    torch.manual_seed(0)
-    paths: list[Path] = []
-    tensors: list[torch.Tensor] = []
-    for i in range(2):
-        t = torch.rand(_T, _C, _H, _W)
-        p = tmp_path / f"clip_{i}.mp4"
-        _write_tensor_as_mp4(t, p)
-        paths.append(p)
-        tensors.append(t)
-    return paths, tensors
-
-
-@pytest.fixture
-def evaluator():
-    ev = create_evaluator(metrics=["common.psnr", "common.ssim"], device="cpu")
-    yield ev
-    ev.shutdown()
-
-
-def test_kwargs_form_accepts_string_path(evaluator, video_paths):
-    paths, _ = video_paths
-    out = evaluator.evaluate(video=str(paths[0]), reference=str(paths[0]))
-    assert isinstance(out, dict)
-    assert isinstance(out["common.psnr"], MetricResult)
-    # Self-paired video → PSNR is huge, SSIM = 1.
-    assert out["common.psnr"].score > 50.0
-    assert out["common.ssim"].score == pytest.approx(1.0, abs=1e-5)
-
-
-def test_kwargs_form_accepts_pathlib_path(evaluator, video_paths):
-    paths, _ = video_paths
-    out = evaluator.evaluate(video=paths[0], reference=paths[0])
-    assert out["common.psnr"].score > 50.0
-
-
-def test_samples_list_accepts_paths(evaluator, video_paths):
-    paths, _ = video_paths
-    samples = [{"video": str(p), "reference": str(p)} for p in paths]
-    out = evaluator.evaluate(samples=samples)
-    assert isinstance(out, list)
-    assert len(out) == len(paths)
-    for row in out:
-        assert row["common.psnr"].score > 50.0
-
-
-def test_samples_list_can_mix_paths_and_tensors(evaluator, video_paths):
-    """A single ``samples`` call can mix path and tensor entries."""
-    paths, tensors = video_paths
-    samples = [
-        {"video": str(paths[0]), "reference": tensors[0]},   # path + tensor
-        {"video": tensors[1], "reference": str(paths[1])},   # tensor + path
-    ]
-    out = evaluator.evaluate(samples=samples)
-    assert len(out) == 2
-    for row in out:
-        assert row["common.psnr"].score > 0.0
-
-
-def test_path_form_score_matches_tensor_form(evaluator, video_paths):
-    """Loading via path must produce the same score as loading via the
-    public ``load_video`` helper and passing the tensor in directly."""
-    from fastvideo.eval.io import load_video
-
-    paths, _ = video_paths
-    via_path = evaluator.evaluate(video=str(paths[0]), reference=str(paths[0]))
-    tensor = load_video(str(paths[0]))
-    via_tensor = evaluator.evaluate(video=tensor, reference=tensor)
-    assert via_path["common.psnr"].score == pytest.approx(
-        via_tensor["common.psnr"].score, abs=1e-4)
-    assert via_path["common.ssim"].score == pytest.approx(
-        via_tensor["common.ssim"].score, abs=1e-4)
-
-
-def test_one_shot_evaluate_accepts_paths(video_paths):
-    """The top-level ``fastvideo.eval.evaluate`` helper also flows paths."""
-    paths, _ = video_paths
-    out = evaluate(generated=str(paths[0]), reference=str(paths[0]),
-                   metrics=["common.psnr"], device="cpu")
-    assert out["common.psnr"].score > 50.0
-
-
-def test_missing_path_surfaces_as_exception(evaluator, tmp_path):
-    """Decode failures must propagate, not silently produce a None score."""
-    bogus = tmp_path / "does_not_exist.mp4"
-    with pytest.raises(Exception):
-        evaluator.evaluate(video=str(bogus), reference=str(bogus))
-
-
-def test_dispatcher_holds_paths_not_tensors_in_queue(tmp_path):
-    """Memory invariant: when many paths are passed, the queued samples
-    are tiny strings, not full tensors. Verify by checking the length
-    of the per-sample reference set the dispatcher materializes."""
-    torch.manual_seed(1)
-    n = 8
-    paths: list[Path] = []
-    for i in range(n):
-        t = torch.rand(_T, _C, _H, _W)
-        p = tmp_path / f"clip_{i}.mp4"
-        _write_tensor_as_mp4(t, p)
-        paths.append(p)
-
-    samples = [{"video": str(p), "reference": str(p)} for p in paths]
-    # Each sample dict is just two strings — no tensor allocations until
-    # the worker's _resolve_video_input runs.
-    for s in samples:
-        assert isinstance(s["video"], str)
-        assert isinstance(s["reference"], str)
-
-    ev = create_evaluator(metrics=["common.psnr"], device="cpu")
-    try:
-        out = ev.evaluate(samples=samples)
-    finally:
-        ev.shutdown()
-    assert len(out) == n
-    # Self-paired ⇒ all PSNRs should be very high.
-    for row in out:
-        assert row["common.psnr"].score > 50.0

fastvideo/tests/eval/test_evaluator_single.py

@@ -1,136 +0,0 @@
-"""End-to-end tests for single-replica eval through the public API.
-
-Runs the lightweight pixel-space metrics — ``common.psnr`` and
-``common.ssim`` — under both shapes that real callers use:
-
-* one-shot ``evaluate(video=..., reference=...)`` (the helper in
-  ``fastvideo.eval.api``);
-* a long-lived ``Evaluator``, called once per sample;
-* a long-lived ``Evaluator``, called with a list of sample dicts to
-  fan out (``samples=[...]``).
-
-GPU-only metrics live in separate test modules / classes; everything
-here runs on CPU so the suite stays cheap to invoke.
-"""
-from __future__ import annotations
-
-import pytest
-import torch
-
-from fastvideo.eval import MetricResult, create_evaluator, evaluate
-
-
-# Tight resolution + frame count so CPU SSIM stays under a couple of seconds.
-_T, _C, _H, _W = 6, 3, 32, 32
-
-
-@pytest.fixture
-def gen_ref():
-    """Reproducible (gen, ref) pair shaped (T, C, H, W)."""
-    torch.manual_seed(0)
-    gen = torch.rand(_T, _C, _H, _W)
-    ref = torch.rand(_T, _C, _H, _W)
-    return gen, ref
-
-
-@pytest.fixture
-def evaluator():
-    ev = create_evaluator(
-        metrics=["common.psnr", "common.ssim"],
-        device="cpu",
-    )
-    yield ev
-    ev.shutdown()
-
-
-def _assert_well_formed(result: MetricResult, name: str) -> None:
-    assert isinstance(result, MetricResult)
-    assert result.name == name
-    assert result.score is not None
-    assert isinstance(result.score, float)
-    # PSNR / SSIM both populate per-frame details.
-    assert "per_frame" in result.details
-    assert len(result.details["per_frame"]) == _T
-
-
-# ---------------------------------------------------------------------------
-# One-shot helper
-# ---------------------------------------------------------------------------
-
-
-def test_evaluate_one_shot_returns_dict_of_metric_results(gen_ref):
-    gen, ref = gen_ref
-    out = evaluate(generated=gen, reference=ref,
-                   metrics=["common.psnr", "common.ssim"], device="cpu")
-    assert isinstance(out, dict)
-    assert set(out.keys()) == {"common.psnr", "common.ssim"}
-    _assert_well_formed(out["common.psnr"], "common.psnr")
-    _assert_well_formed(out["common.ssim"], "common.ssim")
-
-
-# ---------------------------------------------------------------------------
-# Long-lived Evaluator, single-sample form
-# ---------------------------------------------------------------------------
-
-
-def test_evaluator_single_sample_returns_dict(evaluator, gen_ref):
-    gen, ref = gen_ref
-    out = evaluator.evaluate(video=gen, reference=ref)
-    assert isinstance(out, dict)
-    assert set(out.keys()) == {"common.psnr", "common.ssim"}
-    for name, mr in out.items():
-        _assert_well_formed(mr, name)
-
-
-def test_evaluator_accepts_legacy_5d_input(evaluator, gen_ref):
-    """Callers that still pass ``(1, T, C, H, W)`` should get unwrapped."""
-    gen, ref = gen_ref
-    out = evaluator.evaluate(video=gen.unsqueeze(0), reference=ref.unsqueeze(0))
-    _assert_well_formed(out["common.psnr"], "common.psnr")
-
-
-def test_evaluator_score_is_deterministic(evaluator, gen_ref):
-    gen, ref = gen_ref
-    a = evaluator.evaluate(video=gen, reference=ref)
-    b = evaluator.evaluate(video=gen.clone(), reference=ref.clone())
-    assert a["common.psnr"].score == pytest.approx(b["common.psnr"].score)
-    assert a["common.ssim"].score == pytest.approx(b["common.ssim"].score)
-
-
-def test_evaluator_psnr_identical_videos_is_high(evaluator, gen_ref):
-    """PSNR(x, x) is unbounded above; with our clamp it caps near 100 dB."""
-    gen, _ = gen_ref
-    out = evaluator.evaluate(video=gen, reference=gen)
-    assert out["common.psnr"].score > 50.0
-    assert out["common.ssim"].score == pytest.approx(1.0, abs=1e-5)
-
-
-# ---------------------------------------------------------------------------
-# Long-lived Evaluator, list (fan-out) form
-# ---------------------------------------------------------------------------
-
-
-def test_evaluator_samples_list_preserves_input_order(evaluator):
-    """When ``samples=[...]`` is passed, results must come back per sample."""
-    torch.manual_seed(1)
-    samples = []
-    for i in range(4):
-        # Vary the reference enough that scores differ across rows.
-        gen = torch.rand(_T, _C, _H, _W)
-        ref = gen + 0.01 * (i + 1) * torch.rand_like(gen)
-        samples.append({"video": gen, "reference": ref})
-
-    out = evaluator.evaluate(samples=samples)
-    assert isinstance(out, list)
-    assert len(out) == len(samples)
-    for row in out:
-        assert set(row.keys()) == {"common.psnr", "common.ssim"}
-        _assert_well_formed(row["common.psnr"], "common.psnr")
-
-    # Re-running should give bit-identical scores (no nondeterministic
-    # scheduling effects under single-GPU dispatch).
-    out2 = evaluator.evaluate(samples=samples)
-    for a, b in zip(out, out2):
-        assert a["common.psnr"].score == pytest.approx(b["common.psnr"].score)
-
-

fastvideo/tests/eval/test_evaluator_with_dataset.py

@@ -1,98 +0,0 @@
-"""End-to-end test: prompt dataset → Evaluator.
-
-Mirrors the canonical user flow:
-
-    ds = get_dataset("vbench", dimensions=[...])
-    ev = create_evaluator(metrics=[...], device=...)
-    for row in ds:
-        video = my_generator(row["prompt"])
-        scores = ev.evaluate(video=video, **row)
-
-We don't actually generate videos — that would pull in a diffusion
-model. Instead we synthesize a reproducible random tensor per row, so
-the test exercises the dataset-iteration → evaluator-call wiring
-without depending on any model weights.
-"""
-from __future__ import annotations
-
-import pytest
-import torch
-
-from fastvideo.eval import MetricResult, create_evaluator
-from fastvideo.eval.datasets import get_dataset, list_datasets
-
-
-_T, _C, _H, _W = 6, 3, 32, 32
-
-
-def _synth_video(seed: int) -> torch.Tensor:
-    g = torch.Generator().manual_seed(seed)
-    return torch.rand(_T, _C, _H, _W, generator=g)
-
-
-def test_vbench_dataset_is_registered():
-    assert "vbench" in list_datasets()
-
-
-def test_vbench_dataset_rows_drop_into_evaluator():
-    """Every row from the corpus must be a kwargs-friendly dict for
-    ``Evaluator.evaluate``: extra keys flow through without breaking the
-    metric, and the metric returns a well-formed ``MetricResult``."""
-    ds = get_dataset("vbench", dimensions=["color"])
-    rows = [ds[i] for i in range(3)]
-    assert all("prompt" in r for r in rows)
-    assert all("auxiliary_info" in r for r in rows)
-
-    ev = create_evaluator(metrics=["common.psnr"], device="cpu")
-    try:
-        for i, row in enumerate(rows):
-            video = _synth_video(seed=i)
-            reference = _synth_video(seed=100 + i)
-            # Real-world shape: pass dataset row through verbatim, plus
-            # the generated/reference tensors. Extra dataset fields
-            # ('prompt', 'n_samples', 'dimensions', 'auxiliary_info') are
-            # ignored by common.psnr — that's the contract we want to pin.
-            scores = ev.evaluate(video=video, reference=reference, **row)
-            mr = scores["common.psnr"]
-            assert isinstance(mr, MetricResult)
-            assert mr.name == "common.psnr"
-            assert mr.score is not None
-    finally:
-        ev.shutdown()
-
-
-def test_vbench_dataset_full_corpus_iteration():
-    """Iterating the whole dataset should be cheap (no evaluator calls).
-    This guards against a future refactor that accidentally makes
-    ``__iter__`` do real work."""
-    ds = get_dataset("vbench")
-    rows = list(ds)
-    assert len(rows) == len(ds)
-    assert all(isinstance(r, dict) for r in rows)
-
-
-def test_dataset_samples_form_through_evaluator():
-    """``Evaluator.evaluate(samples=[...])`` is the canonical batched
-    entry point; verify it works when the per-row dicts come from a
-    dataset (kwargs form) rather than being hand-built in the test."""
-    ds = get_dataset("vbench", dimensions=["color"])
-    rows = [ds[i] for i in range(3)]
-
-    samples = []
-    for i, row in enumerate(rows):
-        samples.append({
-            "video": _synth_video(seed=i),
-            "reference": _synth_video(seed=100 + i),
-            **row,
-        })
-
-    ev = create_evaluator(metrics=["common.psnr"], device="cpu")
-    try:
-        out = ev.evaluate(samples=samples)
-    finally:
-        ev.shutdown()
-
-    assert len(out) == len(rows)
-    for row in out:
-        assert "common.psnr" in row
-        assert row["common.psnr"].score is not None

fastvideo/tests/eval/test_metric_identity_invariants.py

@@ -0,0 +1,40 @@
+"""Identity-pair / identity-set invariants.
+
+For paired and set-vs-set metrics, duplicating the asset into ``gen/``
+and ``ref/`` should yield the perfect-match value (SSIM=1, PSNR≳99,
+LPIPS=0, FVD=0, FAD=0, KL=0, gt_optical_flow=0). No reference JSON
+needed — the math itself is the gold.
+"""
+from __future__ import annotations
+
+import pytest
+
+from fastvideo.tests.eval.conftest import INVARIANT_SPEC
+
+REQUIRED_GPUS = 1
+
+
+@pytest.mark.parametrize("metric,expected,tol,where", INVARIANT_SPEC)
+def test_metric_identity_invariant(metric: str, expected: float, tol: float | None,
+                                   where: str, invariant_results) -> None:
+    if where == "corpus":
+        result = invariant_results.corpus.get(metric)
+    else:
+        result = invariant_results[0].get(metric)
+
+    if result is None:
+        pytest.skip(f"{metric} not registered on this Evaluator (missing dep)")
+    if result.score is None:
+        pytest.skip(f"fastvideo skipped {metric}: {result.details.get('skipped', 'no score')}")
+
+    if tol is None:
+        # threshold form (e.g. PSNR ≥ 99); for identical inputs PSNR
+        # approaches infinity but float epsilon caps it at ~100.
+        assert result.score >= expected, (
+            f"{metric}: {result.score:.4f} below threshold {expected}"
+        )
+    else:
+        delta = abs(result.score - expected)
+        assert delta <= tol, (
+            f"{metric}: got {result.score:.6f}, expected ≈ {expected} (|Δ|={delta:.6f} > tol={tol})"
+        )

fastvideo/tests/eval/test_metric_score_regression.py

@@ -0,0 +1,97 @@
+"""Score-regression tests against reference values from each metric's upstream repo.
+
+Each reference JSON under ``reference_scores/`` was produced by running
+the canonical upstream implementation against the asset video under
+``asset/`` (provenance recorded in each JSON). This test asserts that
+FastVideo's in-process eval, driven through the top-level
+``samples_from`` + ``Evaluator.evaluate`` API, lands within tolerance.
+"""
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from fastvideo.tests.eval.conftest import _reference_metric_names, load_reference
+
+REQUIRED_GPUS = 1
+
+
+# Per-metric absolute tolerances. Tight for closed-form math, looser for
+# anything that runs CLIP/DINO/ViCLIP/VLM under bf16 or cuDNN.
+TOLERANCE: dict[str, float] = {
+    # WER tolerates ~1 word/clip drift: reference runs openai-whisper with an
+    # explicit language hint; fastvideo runs transformers WhisperForConditionalGeneration
+    # without one (auto-detect). Same model weights, different decode paths.
+    "audio.wer":                       0.2,
+    "audio.desync":                    1e-3,
+    "audio.clap_score":                1e-2,
+    "audio.audiobox_aesthetics":       1e-2,
+    "audio.imagebind_score":           1e-2,
+    # videoscore2 is handled separately (per-dimension hard-score check) —
+    # see _assert_videoscore2 below.
+    "vbench.aesthetic_quality":        1e-2,
+    "vbench.background_consistency":   1e-2,
+    "vbench.dynamic_degree":           0.0,   # binary 0/1 — must match exactly
+    "vbench.imaging_quality":          1e-2,
+    "vbench.motion_smoothness":        1e-2,
+    "vbench.overall_consistency":      1e-2,
+    "vbench.subject_consistency":      1e-2,
+    "vbench.temporal_flickering":      1e-2,
+    "vbench.temporal_style":           1e-2,
+}
+
+
+def _assert_videoscore2(result, reference: dict) -> None:
+    """Regress VideoScore2 on its per-dimension integer (hard) scores, ±1.
+
+    VideoScore2 is a bf16 VLM judge: its three sub-scores are argmaxed
+    integers in 1-5, and they can shift by ±1 across GPU architectures /
+    transformers versions (the gold is generated on a different GPU than CI
+    runs on). A single ±1 flip moves the combined soft mean by ~0.33, so the
+    aggregate score is too brittle to regress directly. The integer hard
+    scores are far more portable; assert each is within 1 of the reference,
+    which still fails loudly on a parse miss (hard score ``None``) or a ≥2
+    point divergence.
+    """
+    for dim in ("visual_quality_hard", "text_alignment_hard",
+                "physical_consistency_hard"):
+        ref_h = reference["details"].get(dim)
+        got_h = result.details.get(dim)
+        assert ref_h is not None, f"reference missing {dim}"
+        assert got_h is not None, (
+            f"videoscore2 did not parse {dim} "
+            f"(raw_output head: {result.details.get('raw_output', '')[:160]!r})")
+        assert abs(got_h - ref_h) <= 1, (
+            f"videoscore2 {dim}: fastvideo={got_h} vs reference={ref_h} "
+            f"(integer sub-score drifted >1 across environments)")
+
+
+@pytest.mark.parametrize("metric_name", _reference_metric_names())
+def test_metric_score_regression(metric_name: str, gold_results: dict) -> None:
+    reference = load_reference(metric_name)
+    if reference["score"] is None:
+        pytest.skip(reference["details"].get("skipped", "deferred"))
+
+    result = gold_results.get(metric_name)
+    if result is None:
+        pytest.skip(f"{metric_name} not in registered metrics this run "
+                    f"(missing dep or filtered)")
+    if result.score is None:
+        pytest.skip(f"fastvideo skipped {metric_name}: "
+                    f"{result.details.get('skipped', 'no score')}")
+
+    if metric_name == "videoscore2":
+        _assert_videoscore2(result, reference)
+        return
+
+    tol = TOLERANCE.get(metric_name)
+    if tol is None:
+        raise AssertionError(f"no tolerance configured for {metric_name}")
+
+    delta = abs(result.score - reference["score"])
+    assert math.isfinite(result.score), f"{metric_name} returned non-finite score"
+    assert delta <= tol, (
+        f"{metric_name}: fastvideo={result.score:.6f} vs reference={reference['score']:.6f} "
+        f"(|Δ|={delta:.6f} > tol={tol})"
+    )

fastvideo/tests/eval/test_registry.py

@@ -1,76 +0,0 @@
-"""Smoke tests for the metric registry surface.
-
-These exercise the public ``fastvideo.eval`` API only:
-``list_metrics``, ``get_metric``, and the group-resolution logic that
-``create_evaluator(metrics="vbench")`` uses.
-"""
-from __future__ import annotations
-
-import pytest
-
-from fastvideo.eval import create_evaluator, get_metric, list_metrics
-
-
-# Metrics that should ship out of the box. Keep this list short so it's
-# resistant to future renames — we want to catch wholesale registry
-# breakage, not chase every individual sub-metric rename.
-_CORE_METRICS = (
-    "common.psnr",
-    "common.ssim",
-    "common.lpips",
-    "common.fvd",
-    "optical_flow.gt_optical_flow",
-    "optical_flow.synthetic_optical_flow",
-    "physics_iq",
-    "vbench.aesthetic_quality",
-    "vbench.subject_consistency",
-)
-
-
-def test_list_metrics_returns_sorted_unique():
-    names = list_metrics()
-    assert len(names) > 0
-    assert names == sorted(names)
-    assert len(names) == len(set(names))
-
-
-def test_core_metrics_are_registered():
-    names = set(list_metrics())
-    for m in _CORE_METRICS:
-        assert m in names, f"expected core metric {m!r} to be registered"
-
-
-def test_get_metric_unknown_raises_keyerror():
-    with pytest.raises(KeyError, match="Unknown metric"):
-        get_metric("nope.nonsense")
-
-
-def test_create_evaluator_resolves_group_prefix():
-    """``metrics="<group>"`` should expand to every ``<group>.*`` sub-metric.
-
-    Use the ``physics_iq`` group because it has multiple sub-metrics and
-    none of them load model weights — the group-resolution behavior is
-    what we're testing, not metric setup."""
-    ev = create_evaluator(metrics="physics_iq", device="cpu")
-    try:
-        names = ev.metric_names
-        assert all(n.startswith("physics_iq.") for n in names)
-        # Don't pin an exact count — sub-metrics may be added later.
-        assert "physics_iq.spatial_iou" in names
-        assert "physics_iq.spatiotemporal_iou" in names
-        assert "physics_iq.mse" in names
-    finally:
-        ev.shutdown()
-
-
-def test_create_evaluator_with_explicit_list_preserves_order():
-    ev = create_evaluator(metrics=["common.psnr", "common.ssim"], device="cpu")
-    try:
-        assert ev.metric_names == ["common.psnr", "common.ssim"]
-    finally:
-        ev.shutdown()
-
-
-def test_create_evaluator_unknown_metric_raises():
-    with pytest.raises(KeyError, match="Unknown metric"):
-        create_evaluator(metrics=["common.nope"], device="cpu")

fastvideo/tests/modal/pr_test.py

@@ -26,6 +26,12 @@ image = (modal.Image.from_registry(
     os.environ.get("BUILDKITE_PULL_REQUEST", ""),
     "BUILDKITE_BRANCH":
     os.environ.get("BUILDKITE_BRANCH", ""),
+    "BUILDKITE_BUILD_URL":
+    os.environ.get("BUILDKITE_BUILD_URL", ""),
+    "BUILDKITE_BUILD_ID":
+    os.environ.get("BUILDKITE_BUILD_ID", ""),
+    "BUILDKITE_JOB_ID":
+    os.environ.get("BUILDKITE_JOB_ID", ""),
     "TEST_SCOPE":
     os.environ.get("TEST_SCOPE", ""),
     "IMAGE_VERSION":
@@ -315,6 +321,32 @@ def seed_grad_norm_references():
     )
 
 
+@app.function(gpu="L40S:1",
+              image=image,
+              timeout=3600,
+              secrets=[
+                  modal.Secret.from_dict(
+                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
+              ],
+              volumes={"/root/data": model_vol})
+def run_eval_tests():
+    # Eval metric regression: drives the high-level fastvideo.eval API on a
+    # fixed asset and asserts each score matches the upstream reference number
+    # checked into fastvideo/tests/eval/reference_scores/. Pulls several scorer
+    # checkpoints (VideoScore2 VLM, VBench nets, audio models) on first run;
+    # they cache on the hf-model-weights volume thereafter.
+    #
+    # Installs [eval-full] (eval + vbench + audio extras) on top of [test]:
+    # the dev image only ships [dev], and without the extras skip_missing_deps
+    # in conftest would silently drop nearly every metric and the lane would
+    # pass vacuously. detectron2-backed vbench metrics remain skipped by
+    # design (not pip-installable; see fastvideo/eval/README.md).
+    run_test_command(
+        'uv pip install -e ".[test,eval-full]" && '
+        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/eval -vs",
+        build_kernel=True)
+
+
 @app.function(gpu="L40S:1",
               image=image,
               timeout=3600,
@@ -337,18 +369,30 @@ def run_lora_extraction_tests():
               ],
               volumes={"/root/data": model_vol})
 def run_performance_tests():
-    # compare_baseline.py runs only after pytest passes, so normalized_perf_*.json
-    # artifacts are emitted for rolling-baseline failures, not fixed-threshold
-    # pytest failures. dashboard.py still runs on red CI for observability.
+    # PR/direct records are uploaded only on pass; scheduled main uploads pass
+    # and fail so the dashboard records every canonical baseline attempt.
     run_test(
         "export HF_HOME='/root/data/.cache' && "
         "export PERFORMANCE_TRACKING_ROOT='/tmp/perf-tracking' && "
         "hf auth login --token $HF_API_KEY && "
+        "if [ \"${BUILDKITE_BRANCH:-}\" = 'main' ] && [ \"${TEST_SCOPE:-}\" = 'full' ]; then "
+        "export PERF_RUN_SOURCE='scheduled_main'; "
+        "export PERF_UPLOAD_POLICY='always'; "
+        "elif [ -n \"${BUILDKITE_PULL_REQUEST:-}\" ] && [ \"${BUILDKITE_PULL_REQUEST:-false}\" != 'false' ]; then "
+        "export PERF_RUN_SOURCE='pr'; "
+        "export PERF_UPLOAD_POLICY='pass'; "
+        "elif [ \"${TEST_SCOPE:-}\" = 'direct' ]; then "
+        "export PERF_RUN_SOURCE='unknown'; "
+        "export PERF_UPLOAD_POLICY='pass'; "
+        "else "
+        "export PERF_RUN_SOURCE='unknown'; "
+        "export PERF_UPLOAD_POLICY='never'; "
+        "fi; "
         "pytest ./fastvideo/tests/performance -vs; "
         "PYTEST_RC=$?; "
         "PERF_RC=0; "
-        "if [ $PYTEST_RC -eq 0 ]; then "
-        "python ./fastvideo/tests/performance/compare_baseline.py; "
+        "if [ $PYTEST_RC -eq 0 ] || [ \"$PERF_UPLOAD_POLICY\" = 'always' ]; then "
+        "PERF_PYTEST_RC=$PYTEST_RC python ./fastvideo/tests/performance/compare_baseline.py; "
         "PERF_RC=$?; "
         "fi; "
         "python ./fastvideo/tests/performance/dashboard.py || true; "

fastvideo/tests/performance/compare_baseline.py

@@ -4,10 +4,10 @@
 This script:
 1) reads current benchmark results from fastvideo/tests/performance/results,
 2) syncs the canonical baseline from the configured HF dataset repo,
-3) compares each current record against the median of up to 5 prior records
-   (filtered by gpu_type, successful only),
-4) on persist runs (full-suite on main branch), writes the normalized record
-   back to the HF dataset repo,
+3) compares each current record against the median of up to 5 prior
+   baseline-eligible successful records (filtered by gpu_type),
+4) writes normalized records back to the HF dataset repo according to
+   PERF_UPLOAD_POLICY,
 5) exits non-zero if any metric regresses by more than PERF_MAX_REGRESSION
    (default 5%).
 """
@@ -20,13 +20,22 @@ import sys
 from datetime import datetime, timezone
 from typing import Any
 
-from hf_store import (
-    load_records_for_model,
-    safe_float,
-    sanitize,
-    sync_from_hf,
-    upload_record,
-)
+try:
+    from .hf_store import (
+        load_records_for_model,
+        safe_float,
+        sanitize,
+        sync_from_hf,
+        upload_record,
+    )
+except ImportError:
+    from hf_store import (
+        load_records_for_model,
+        safe_float,
+        sanitize,
+        sync_from_hf,
+        upload_record,
+    )
 
 RESULTS_DIR = os.path.join(
     os.path.dirname(os.path.abspath(__file__)),
@@ -38,6 +47,9 @@ TRACKING_ROOT = os.environ.get(
 )
 PERF_REPORTS_DIR = os.environ.get("PERF_REPORTS_DIR", "/root/data/perf_reports")
 MAX_REGRESSION = float(os.environ.get("PERF_MAX_REGRESSION", "0.05"))
+UPLOAD_POLICY = os.environ.get("PERF_UPLOAD_POLICY", "never").strip().lower()
+VALID_UPLOAD_POLICIES = {"never", "pass", "always"}
+VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
 METRICS = (
     ("latency", "Latency", 3),
     ("throughput", "Throughput", 3),
@@ -56,9 +68,74 @@ LOWER_IS_BETTER_METRICS = {
 
 
 def _should_persist_tracking() -> bool:
-    test_scope = os.environ.get("TEST_SCOPE", "")
-    branch = os.environ.get("BUILDKITE_BRANCH", "")
-    return test_scope == "full" and branch == "main"
+    return _normalized_upload_policy() != "never"
+
+
+def _normalized_upload_policy() -> str:
+    if UPLOAD_POLICY in VALID_UPLOAD_POLICIES:
+        return UPLOAD_POLICY
+    print(f"Invalid PERF_UPLOAD_POLICY={UPLOAD_POLICY!r}; using 'never'")
+    return "never"
+
+
+def _truthy_pr_number(value: str | None) -> bool:
+    return bool(value and value not in {"false", "0", "None", "none"})
+
+
+def _detect_run_source() -> str:
+    explicit = os.environ.get("PERF_RUN_SOURCE", "").strip().lower()
+    if explicit in VALID_RUN_SOURCES:
+        return explicit
+    if explicit:
+        print(f"Invalid PERF_RUN_SOURCE={explicit!r}; inferring run source")
+
+    if _truthy_pr_number(os.environ.get("BUILDKITE_PULL_REQUEST")):
+        return "pr"
+    if os.environ.get("BUILDKITE_BRANCH") == "main" and os.environ.get("TEST_SCOPE") == "full":
+        return "scheduled_main"
+    if not os.environ.get("BUILDKITE_COMMIT"):
+        return "local"
+    return "unknown"
+
+
+def _is_baseline_eligible(run_source: str, success: bool) -> bool:
+    return run_source == "scheduled_main" and success
+
+
+def _upload_allowed(record: dict[str, Any]) -> bool:
+    policy = _normalized_upload_policy()
+    if policy == "always":
+        return True
+    if policy == "pass":
+        return bool(record.get("success", True))
+    return False
+
+
+def _result_failed_static_thresholds() -> bool:
+    value = os.environ.get("PERF_PYTEST_RC", "")
+    if not value:
+        return False
+    try:
+        return int(value) != 0
+    except ValueError:
+        return False
+
+
+def _record_metadata(run_source: str, result: dict[str, Any]) -> dict[str, Any]:
+    pr_number = result.get("pr_number") or os.environ.get("BUILDKITE_PULL_REQUEST", "")
+    if not _truthy_pr_number(str(pr_number)):
+        pr_number = ""
+    return {
+        "run_source": run_source,
+        "baseline_eligible": False,
+        "branch": os.environ.get("BUILDKITE_BRANCH", ""),
+        "pr_number": pr_number,
+        "test_scope": os.environ.get("TEST_SCOPE", ""),
+        "build_url": os.environ.get("BUILDKITE_BUILD_URL", ""),
+        "build_id": os.environ.get("BUILDKITE_BUILD_ID", ""),
+        "job_id": os.environ.get("BUILDKITE_JOB_ID", ""),
+    }
+
 
 def _load_current_results() -> list[dict[str, Any]]:
     pattern = os.path.join(RESULTS_DIR, "perf_*.json")
@@ -104,6 +181,7 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
         "dit_time_s": dit_time,
         "vae_decode_time_s": vae_decode_time,
         "success": True,
+        **_record_metadata(_detect_run_source(), result),
     }
 
 
@@ -297,8 +375,11 @@ def _emit_markdown_summary(markdown: str, commit_sha: str) -> None:
 
 def main() -> int:
     persist_tracking = _should_persist_tracking()
+    upload_policy = _normalized_upload_policy()
+    static_threshold_failed = _result_failed_static_thresholds()
 
-    # Strict on persist: a silent sync failure would pollute the baseline.
+    # Strict on upload-enabled runs: silent sync failure would make comparison
+    # and upload state ambiguous.
     sync_from_hf(TRACKING_ROOT, strict=persist_tracking)
 
     current_results = _load_current_results()
@@ -310,10 +391,12 @@ def main() -> int:
     summary_rows: list[dict[str, Any]] = []
 
     if persist_tracking:
-        print("Tracking persistence enabled: full-suite run on main branch")
+        print(f"Tracking persistence enabled: PERF_UPLOAD_POLICY={upload_policy}")
     else:
-        print("Tracking persistence disabled: "
-              "only full-suite runs on main branch are persisted")
+        print("Tracking persistence disabled: PERF_UPLOAD_POLICY=never")
+
+    if static_threshold_failed:
+        print(f"Static-threshold phase failed: PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')}")
 
     for raw in current_results:
         record = _normalize_record(raw)
@@ -324,6 +407,7 @@ def main() -> int:
             record["gpu_type"],
             last_n=5,
             successful_only=True,
+            baseline_eligible_only=True,
         )
 
         if not baseline_records:
@@ -333,15 +417,22 @@ def main() -> int:
             record["success"] = True
         else:
             failures = _check_regressions(record, baseline_records, MAX_REGRESSION)
-            record["success"] = not failures
-            all_failures.extend(failures)
+        if static_threshold_failed:
+            failures.append(f"{record['model_id']} fixed-threshold phase failed "
+                            f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
+
+        record["success"] = not failures
+        record["baseline_eligible"] = _is_baseline_eligible(record["run_source"], record["success"])
+        all_failures.extend(failures)
 
         _write_normalized_artifact(record)
 
-        # Strict upload: a silent failure would freeze the rolling baseline.
-        if persist_tracking:
+        if _upload_allowed(record):
             current_path = _write_tracking_record(record)
             upload_record(current_path, record, strict=True)
+        else:
+            print("Tracking upload skipped for "
+                  f"{record['model_id']} ({record['run_source']}, success={record['success']})")
 
         summary_rows.append(_build_summary_row(record, baseline_records, bool(failures)))
 

fastvideo/tests/performance/hf_store.py

@@ -48,6 +48,20 @@ def safe_float(value: Any) -> float | None:
         return None
 
 
+def is_baseline_eligible_record(record: dict[str, Any]) -> bool:
+    """Return whether *record* may contribute to rolling baselines.
+
+    Legacy records predate ``baseline_eligible`` and ``run_source``. They were
+    uploaded only by the old successful main/full-suite path, so keep them
+    eligible until the HF history naturally rolls forward.
+    """
+    if record.get("baseline_eligible") is True:
+        return True
+    if "baseline_eligible" not in record and "run_source" not in record:
+        return True
+    return False
+
+
 def resolve_hf_token() -> str | None:
     """Return the first configured Hugging Face token env var."""
     for env_var in HF_TOKEN_ENV_VARS:
@@ -197,6 +211,7 @@ def load_records(
     *,
     days: int | None = None,
     successful_only: bool = False,
+    baseline_eligible_only: bool = False,
 ) -> list[dict[str, Any]]:
     """Return raw JSON dicts from *local_dir*.
 
@@ -206,6 +221,9 @@ def load_records(
             many days. Records with a missing/unparsable timestamp are kept.
         successful_only: When True, only records with ``success=True`` are
             returned. Useful when building a regression baseline.
+        baseline_eligible_only: When True, only baseline-eligible records are
+            returned. Legacy records missing both ``baseline_eligible`` and
+            ``run_source`` are treated as eligible.
 
     Returns:
         List of raw dicts sorted by ``timestamp`` ascending (records that could
@@ -227,6 +245,9 @@ def load_records(
         if successful_only and not data.get("success", True):
             continue
 
+        if baseline_eligible_only and not is_baseline_eligible_record(data):
+            continue
+
         if cutoff is not None:
             raw_ts = data.get("timestamp")
             if raw_ts:
@@ -251,6 +272,7 @@ def load_records_for_model(
     *,
     last_n: int | None = None,
     successful_only: bool = True,
+    baseline_eligible_only: bool = False,
 ) -> list[dict[str, Any]]:
     """Return records for a specific *model_id*, optionally filtered by GPU.
 
@@ -261,6 +283,7 @@ def load_records_for_model(
         last_n: When set, return only the most recent *n* records (after all
             other filters). Useful for sliding-window baseline calculations.
         successful_only: Passed through to :func:`load_records`.
+        baseline_eligible_only: Passed through to :func:`load_records`.
 
     Returns:
         List of matching dicts sorted by timestamp ascending.
@@ -269,7 +292,11 @@ def load_records_for_model(
     if not os.path.isdir(model_dir):
         return []
 
-    records = load_records(model_dir, successful_only=successful_only)
+    records = load_records(
+        model_dir,
+        successful_only=successful_only,
+        baseline_eligible_only=baseline_eligible_only,
+    )
 
     if gpu_type is not None:
         records = [r for r in records if r.get("gpu_type") == gpu_type]

fastvideo/tests/performance/test_compare_baseline_policy.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.tests.performance import compare_baseline
+
+
+def _raw_result():
+    return {
+        "benchmark_id": "wan-t2v-1.3b-2gpu",
+        "device": "NVIDIA L40S",
+        "avg_generation_time_s": 10.0,
+        "throughput_fps": 4.5,
+        "max_peak_memory_mb": 10000.0,
+        "commit": "a" * 40,
+        "timestamp": "2026-06-16T00:00:00+00:00",
+        "pr_number": "123",
+    }
+
+
+def test_detect_run_source_prefers_explicit_env(monkeypatch):
+    monkeypatch.setenv("PERF_RUN_SOURCE", "local")
+    monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "123")
+
+    assert compare_baseline._detect_run_source() == "local"
+
+
+def test_detect_run_source_infers_pr(monkeypatch):
+    monkeypatch.delenv("PERF_RUN_SOURCE", raising=False)
+    monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "123")
+
+    assert compare_baseline._detect_run_source() == "pr"
+
+
+def test_detect_run_source_infers_scheduled_main(monkeypatch):
+    monkeypatch.delenv("PERF_RUN_SOURCE", raising=False)
+    monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false")
+    monkeypatch.setenv("BUILDKITE_BRANCH", "main")
+    monkeypatch.setenv("TEST_SCOPE", "full")
+
+    assert compare_baseline._detect_run_source() == "scheduled_main"
+
+
+def test_upload_policy_pass_requires_success(monkeypatch):
+    monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "pass")
+
+    assert compare_baseline._upload_allowed({"success": True}) is True
+    assert compare_baseline._upload_allowed({"success": False}) is False
+
+
+def test_upload_policy_always_uploads_failures(monkeypatch):
+    monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "always")
+
+    assert compare_baseline._upload_allowed({"success": False}) is True
+
+
+def test_normalized_record_includes_source_metadata(monkeypatch):
+    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
+    monkeypatch.setenv("BUILDKITE_BRANCH", "feature/perf")
+    monkeypatch.setenv("TEST_SCOPE", "direct")
+    monkeypatch.setenv("BUILDKITE_BUILD_URL", "https://buildkite.example/build")
+    monkeypatch.setenv("BUILDKITE_BUILD_ID", "build-1")
+    monkeypatch.setenv("BUILDKITE_JOB_ID", "job-1")
+
+    record = compare_baseline.normalize_performance_result(_raw_result())
+
+    assert record["run_source"] == "pr"
+    assert record["baseline_eligible"] is False
+    assert record["branch"] == "feature/perf"
+    assert record["pr_number"] == "123"
+    assert record["test_scope"] == "direct"
+    assert record["build_url"] == "https://buildkite.example/build"
+    assert record["build_id"] == "build-1"
+    assert record["job_id"] == "job-1"
+
+
+def test_baseline_eligibility_only_for_successful_scheduled_main():
+    assert compare_baseline._is_baseline_eligible("scheduled_main", True) is True
+    assert compare_baseline._is_baseline_eligible("scheduled_main", False) is False
+    assert compare_baseline._is_baseline_eligible("pr", True) is False
+    assert compare_baseline._is_baseline_eligible("local", True) is False
+

fastvideo/tests/performance/test_dashboard_api.py

@@ -36,8 +36,8 @@ class FakeStore(PerformanceDataStore):
         return records
 
 
-def _record(model_id, gpu_type, ts, commit, latency, throughput, success=True):
-    return {
+def _record(model_id, gpu_type, ts, commit, latency, throughput, success=True, **metadata):
+    record = {
         "model_id": model_id,
         "gpu_type": gpu_type,
         "timestamp": ts,
@@ -50,6 +50,8 @@ def _record(model_id, gpu_type, ts, commit, latency, throughput, success=True):
         "vae_decode_time_s": 3.0,
         "success": success,
     }
+    record.update(metadata)
+    return record
 
 
 def test_summary_endpoint_returns_latest_group_status():
@@ -87,6 +89,46 @@ def test_summary_status_is_independent_of_days_window():
     assert len(trends["groups"][0]["points"]) == 1
 
 
+def test_dashboard_endpoints_filter_and_return_run_source_metadata():
+    app = create_app(FakeStore([
+        _record(
+            "wan",
+            "NVIDIA L40S",
+            "2026-01-01T00:00:00+00:00",
+            "a" * 40,
+            10.0,
+            10.0,
+            run_source="pr",
+            pr_number="123",
+            branch="feature/perf",
+            baseline_eligible=False,
+        ),
+        _record(
+            "wan",
+            "NVIDIA L40S",
+            "2026-01-02T00:00:00+00:00",
+            "b" * 40,
+            11.0,
+            9.0,
+            run_source="scheduled_main",
+            baseline_eligible=True,
+        ),
+    ]))
+    client = TestClient(app)
+
+    summary = client.get("/api/performance/summary", params={"run_source": "pr"}).json()
+    trends = client.get("/api/performance/trends", params={"run_source": "pr"}).json()
+
+    assert summary["count"] == 1
+    assert summary["rows"][0]["run_source"] == "pr"
+    assert summary["rows"][0]["pr_number"] == "123"
+    assert summary["rows"][0]["baseline_n"] == 1
+    assert summary["rows"][0]["metrics"]["latency"]["baseline"] == 11.0
+    assert summary["filters"]["run_source"] == "pr"
+    assert trends["count"] == 1
+    assert trends["groups"][0]["points"][0]["run_source"] == "pr"
+
+
 def test_records_and_trends_endpoints_filter_by_model_and_gpu():
     app = create_app(FakeStore([
         _record("wan", "NVIDIA L40S", "2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),

fastvideo/tests/performance/test_dashboard_service.py

@@ -3,8 +3,8 @@ from fastvideo.performance_dashboard.service import build_latest_summary, build_
 from fastvideo.tests.performance import hf_store
 
 
-def _record(ts, commit, latency, throughput, success=True):
-    return {
+def _record(ts, commit, latency, throughput, success=True, **metadata):
+    record = {
         "model_id": "wan-t2v-1.3b-2gpu",
         "gpu_type": "NVIDIA L40S",
         "timestamp": ts,
@@ -17,6 +17,8 @@ def _record(ts, commit, latency, throughput, success=True):
         "vae_decode_time_s": 3.0,
         "success": success,
     }
+    record.update(metadata)
+    return record
 
 
 def test_build_latest_summary_uses_previous_successful_records_for_baseline():
@@ -50,6 +52,37 @@ def test_build_latest_summary_status_uses_latest_record_success_field():
     assert rows[0]["success"] is False
 
 
+def test_build_latest_summary_run_source_filter_keeps_canonical_baseline():
+    records = [
+        _record(
+            "2026-01-01T00:00:00+00:00",
+            "a" * 40,
+            10.0,
+            10.0,
+            run_source="scheduled_main",
+            baseline_eligible=True,
+        ),
+        _record(
+            "2026-01-02T00:00:00+00:00",
+            "b" * 40,
+            11.0,
+            9.0,
+            run_source="pr",
+            baseline_eligible=False,
+            pr_number="123",
+        ),
+    ]
+
+    rows = build_latest_summary(records, max_regression=0.05, run_source="pr")
+
+    assert len(rows) == 1
+    assert rows[0]["run_source"] == "pr"
+    assert rows[0]["pr_number"] == "123"
+    assert rows[0]["baseline_n"] == 1
+    assert rows[0]["metrics"]["latency"]["baseline"] == 10.0
+    assert rows[0]["computed_regression_status"] == "fail"
+
+
 def test_filter_records_and_trends_preserve_metric_points():
     records = [
         _record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
@@ -64,6 +97,34 @@ def test_filter_records_and_trends_preserve_metric_points():
     assert trends[0]["points"][1]["metrics"]["latency"] == 12.0
 
 
+def test_trends_include_source_metadata_with_legacy_defaults():
+    records = [
+        _record(
+            "2026-01-01T00:00:00+00:00",
+            "a" * 40,
+            10.0,
+            10.0,
+            run_source="pr",
+            baseline_eligible=False,
+            pr_number="123",
+            branch="feature/dashboard",
+            build_url="https://buildkite.example/build",
+        ),
+        _record("2026-01-02T00:00:00+00:00", "b" * 40, 12.0, 8.0),
+    ]
+
+    filtered = filter_records(records, run_source="pr")
+    trends = build_trends(records)
+
+    assert len(filtered) == 1
+    assert trends[0]["points"][0]["run_source"] == "pr"
+    assert trends[0]["points"][0]["pr_number"] == "123"
+    assert trends[0]["points"][0]["branch"] == "feature/dashboard"
+    assert trends[0]["points"][0]["build_url"] == "https://buildkite.example/build"
+    assert trends[0]["points"][1]["run_source"] == "unknown"
+    assert trends[0]["points"][1]["baseline_eligible"] is True
+
+
 def test_hf_token_resolution_accepts_standard_env_names(monkeypatch):
     for env_var in hf_store.HF_TOKEN_ENV_VARS:
         monkeypatch.delenv(env_var, raising=False)
@@ -71,3 +132,28 @@ def test_hf_token_resolution_accepts_standard_env_names(monkeypatch):
     monkeypatch.setenv("HF_TOKEN", "hf_local")
 
     assert hf_store.resolve_hf_token() == "hf_local"
+
+
+def test_load_records_can_filter_baseline_eligible_records(tmp_path):
+    model_dir = tmp_path / "wan"
+    model_dir.mkdir()
+    (model_dir / "pr.json").write_text(
+        '{"timestamp": "2026-01-01T00:00:00+00:00", "success": true, "baseline_eligible": false}',
+        encoding="utf-8",
+    )
+    (model_dir / "main.json").write_text(
+        '{"timestamp": "2026-01-02T00:00:00+00:00", "success": true, "baseline_eligible": true}',
+        encoding="utf-8",
+    )
+    (model_dir / "legacy.json").write_text(
+        '{"timestamp": "2026-01-03T00:00:00+00:00", "success": true}',
+        encoding="utf-8",
+    )
+
+    records = hf_store.load_records(str(tmp_path), successful_only=True, baseline_eligible_only=True)
+
+    assert len(records) == 2
+    assert {record["timestamp"] for record in records} == {
+        "2026-01-02T00:00:00+00:00",
+        "2026-01-03T00:00:00+00:00",
+    }

fastvideo/tests/performance/test_inference_performance.py

@@ -98,10 +98,15 @@ def _extract_component_times(result: dict) -> dict[str, float | None]:
         return component_times
     logger.info("Discovered pipeline stages: %s", list(stages.keys()))
     for stage_name, stage_data in stages.items():
-        metric_key = STAGE_METRIC_MAP.get(stage_name)
+        if not isinstance(stage_data, Mapping):
+            logger.debug("Skipping malformed stage '%s' data: %r", stage_name, stage_data)
+            continue
+        stage_class = stage_data.get("stage_class", stage_name)
+        metric_key = STAGE_METRIC_MAP.get(stage_class)
         if metric_key is None:
-            logger.debug("Unmapped stage '%s' (%.3fs)",
+            logger.debug("Unmapped stage '%s' class '%s' (%.3fs)",
                          stage_name,
+                         stage_class,
                          stage_data.get("execution_time", 0))
             continue
         elapsed = stage_data.get("execution_time")

fastvideo/tests/performance/test_inference_performance_component_times.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.pipelines.pipeline_batch_info import PipelineLoggingInfo
+from fastvideo.tests.performance.test_inference_performance import _extract_component_times
+
+
+def test_extract_component_times_handles_pipeline_logging_info_object():
+    logging_info = PipelineLoggingInfo()
+    logging_info.add_stage_execution_time("prompt_encoding_stage", 1.25)
+    logging_info.add_stage_metric("prompt_encoding_stage", "stage_class", "TextEncodingStage")
+
+    assert _extract_component_times({"logging_info": logging_info}) == {
+        "text_encoder_time_s": 1.25,
+        "dit_time_s": None,
+        "vae_decode_time_s": None,
+    }
+
+
+def test_extract_component_times_uses_stage_class_for_pipeline_stage_keys():
+    # Regression guard for #1377: pre-fix code looked up the pipeline stage key
+    # and returned all component metrics as None for this shape.
+    result = {
+        "logging_info": {
+            "stages": {
+                "prompt_encoding_stage": {
+                    "execution_time": 1.2,
+                    "stage_class": "TextEncodingStage",
+                },
+                "denoising_stage": {
+                    "execution_time": 3.4,
+                    "stage_class": "DenoisingStage",
+                },
+                "decoding_stage": {
+                    "execution_time": 0.8,
+                    "stage_class": "DecodingStage",
+                },
+            },
+        },
+    }
+
+    assert _extract_component_times(result) == {
+        "text_encoder_time_s": 1.2,
+        "dit_time_s": 3.4,
+        "vae_decode_time_s": 0.8,
+    }
+
+
+def test_extract_component_times_keeps_legacy_class_name_keys():
+    # Backward-compatibility check for logs produced before pipeline-unique
+    # stage keys carried a separate stage_class field.
+    result = {
+        "logging_info": {
+            "stages": {
+                "TextEncodingStage": {"execution_time": 1.0},
+                "DenoisingStage": {"execution_time": 2.0},
+                "DecodingStage": {"execution_time": 3.0},
+            },
+        },
+    }
+
+    assert _extract_component_times(result) == {
+        "text_encoder_time_s": 1.0,
+        "dit_time_s": 2.0,
+        "vae_decode_time_s": 3.0,
+    }
+
+
+def test_extract_component_times_accumulates_duplicate_component_classes():
+    result = {
+        "logging_info": {
+            "stages": {
+                "base_denoising_stage": {
+                    "execution_time": 2.0,
+                    "stage_class": "DenoisingStage",
+                },
+                "refine_denoising_stage": {
+                    "execution_time": 3.5,
+                    "stage_class": "DenoisingStage",
+                },
+            },
+        },
+    }
+
+    assert _extract_component_times(result) == {
+        "text_encoder_time_s": None,
+        "dit_time_s": 5.5,
+        "vae_decode_time_s": None,
+    }
+
+
+def test_extract_component_times_ignores_unmapped_stages():
+    # Generator-side bookkeeping timings are intentionally excluded from the
+    # component gates.
+    result = {
+        "logging_info": {
+            "stages": {
+                "PostDecodeFrameProcessStage": {"execution_time": 0.2},
+                "VideoSaveStage": {"execution_time": 0.4},
+                "AudioMuxStage": {"execution_time": 0.1},
+            },
+        },
+    }
+
+    assert _extract_component_times(result) == {
+        "text_encoder_time_s": None,
+        "dit_time_s": None,
+        "vae_decode_time_s": None,
+    }
+
+
+def test_extract_component_times_skips_malformed_stage_data():
+    result = {
+        "logging_info": {
+            "stages": {
+                "prompt_encoding_stage": None,
+                "denoising_stage": "not-a-stage-metric-dict",
+                "decoding_stage": {
+                    "execution_time": 0.8,
+                    "stage_class": "DecodingStage",
+                },
+            },
+        },
+    }
+
+    assert _extract_component_times(result) == {
+        "text_encoder_time_s": None,
+        "dit_time_s": None,
+        "vae_decode_time_s": 0.8,
+    }

performance_dashboard/README.md

@@ -14,6 +14,12 @@ Defaults:
 - `PERFORMANCE_TRACKING_ROOT=/tmp/fastvideo-perf-dashboard`
 - `PERF_MAX_REGRESSION=0.05`
 
+Records can include source metadata:
+
+- `run_source`: `pr`, `local`, `scheduled_main`, or `unknown`
+- `baseline_eligible`: only successful scheduled-main records should be true
+- Buildkite metadata such as branch, PR number, build URL, build ID, and job ID
+
 Set one of `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, or `HF_TOKEN` if the
 configured dataset repo requires authenticated access:
 
@@ -68,13 +74,31 @@ ngrok http 8000
 The ngrok URL will serve the dashboard UI and all `/api/performance/*`
 endpoints from the same local port.
 
+## Dashboard Behavior
+
+The dashboard supports model, GPU, source, and day-window filters.
+
+Trend charts show metric-specific axes and exact point details on hover/focus:
+
+- metric value and unit
+- timestamp
+- commit SHA
+- run source
+- stored status
+- baseline eligibility
+- PR number, branch, and Buildkite URL when present
+
+The latest status table uses the stored JSON `success` value. Recomputed
+baseline context is shown separately and does not override stored status.
+
 ## API
 
 - `GET /api/performance/health`
 - `POST /api/performance/refresh`
-- `GET /api/performance/summary?days=90`
-- `GET /api/performance/trends?days=90`
-- `GET /api/performance/records?days=90`
+- `GET /api/performance/summary?days=90&run_source=pr`
+- `GET /api/performance/trends?days=90&run_source=scheduled_main`
+- `GET /api/performance/records?days=90&run_source=local`
 
 The current v1 grouping key is `(model_id, gpu_type)`. Baselines are computed
-from the latest five previous successful records in each group.
+from the latest five previous successful records in each group for dashboard
+context. CI gating uses only records marked `baseline_eligible=true`.

performance_dashboard/frontend/src/App.tsx

@@ -1,8 +1,63 @@
 import { useEffect, useMemo, useState } from "react";
 
-import { fetchSummary, fetchTrends, refreshData, SummaryResponse, TrendGroup } from "./api";
+import { fetchSummary, fetchTrends, refreshData, RunSource, SummaryResponse, TrendGroup, TrendPoint } from "./api";
 
 const METRIC_KEYS = ["latency", "throughput", "memory", "text_encoder_time_s", "dit_time_s", "vae_decode_time_s"];
+const RUN_SOURCES: Array<{ value: "" | RunSource; label: string }> = [
+  { value: "", label: "All sources" },
+  { value: "scheduled_main", label: "Scheduled main" },
+  { value: "pr", label: "PR" },
+  { value: "local", label: "Local" },
+  { value: "unknown", label: "Unknown" }
+];
+
+const METRIC_DEFINITIONS: Record<
+  string,
+  {
+    label: string;
+    unit: string;
+    precision: number;
+    tooltipPrecision: number;
+    secondary?: (value: number) => string;
+  }
+> = {
+  latency: {
+    label: "Latency",
+    unit: "s",
+    precision: 2,
+    tooltipPrecision: 3,
+    secondary: (value) => `${formatNumber(value * 1000, 0)} ms`
+  },
+  throughput: { label: "Throughput", unit: "FPS", precision: 2, tooltipPrecision: 3 },
+  memory: {
+    label: "Memory",
+    unit: "MB",
+    precision: 0,
+    tooltipPrecision: 1,
+    secondary: (value) => `${formatNumber(value / 1024, 2)} GB`
+  },
+  text_encoder_time_s: {
+    label: "Text Encoder",
+    unit: "s",
+    precision: 2,
+    tooltipPrecision: 3,
+    secondary: (value) => `${formatNumber(value * 1000, 0)} ms`
+  },
+  dit_time_s: {
+    label: "DiT",
+    unit: "s",
+    precision: 2,
+    tooltipPrecision: 3,
+    secondary: (value) => `${formatNumber(value * 1000, 0)} ms`
+  },
+  vae_decode_time_s: {
+    label: "VAE Decode",
+    unit: "s",
+    precision: 2,
+    tooltipPrecision: 3,
+    secondary: (value) => `${formatNumber(value * 1000, 0)} ms`
+  }
+};
 
 function formatNumber(value: number | null | undefined, precision = 2) {
   if (value === null || value === undefined || Number.isNaN(value)) {
@@ -26,52 +81,195 @@ function formatTime(value: string | null | undefined) {
   return date.toLocaleString();
 }
 
+function formatDate(value: string | null | undefined) {
+  if (!value) {
+    return "unknown";
+  }
+  const date = new Date(value);
+  if (Number.isNaN(date.getTime())) {
+    return value;
+  }
+  return date.toLocaleDateString(undefined, { month: "short", day: "numeric" });
+}
+
+function runSourceLabel(value: string | null | undefined) {
+  if (value === "scheduled_main") {
+    return "Scheduled main";
+  }
+  if (value === "pr") {
+    return "PR";
+  }
+  if (value === "local") {
+    return "Local";
+  }
+  return "Unknown";
+}
+
+function metricLabel(metricKey: string) {
+  return METRIC_DEFINITIONS[metricKey]?.label ?? metricKey;
+}
+
+function formatMetricValue(metricKey: string, value: number | null | undefined, tooltip = false) {
+  const definition = METRIC_DEFINITIONS[metricKey];
+  if (!definition) {
+    return formatNumber(value, tooltip ? 3 : 2);
+  }
+  const formatted = formatNumber(value, tooltip ? definition.tooltipPrecision : definition.precision);
+  return formatted === "n/a" ? formatted : `${formatted} ${definition.unit}`;
+}
+
+type ChartPoint = {
+  plotIndex: number;
+  value: number;
+  point: TrendPoint;
+  x: number;
+  y: number;
+};
+
 function TrendChart({ group, metricKey }: { group: TrendGroup; metricKey: string }) {
+  const [activePoint, setActivePoint] = useState<ChartPoint | null>(null);
   const points = group.points
-    .map((point, index) => ({
-      index,
-      value: point.metrics[metricKey],
-      success: point.success
+    .map((point) => ({
+      point,
+      value: point.metrics[metricKey]
     }))
     .filter((point) => point.value !== null && point.value !== undefined) as Array<{
-      index: number;
+      point: TrendPoint;
       value: number;
-      success: boolean;
     }>;
 
   if (points.length === 0) {
     return <div className="empty-chart">No data</div>;
   }
 
-  const width = 280;
-  const height = 96;
-  const pad = 12;
+  const width = 360;
+  const height = 190;
+  const margin = { top: 16, right: 18, bottom: 34, left: 54 };
+  const plotWidth = width - margin.left - margin.right;
+  const plotHeight = height - margin.top - margin.bottom;
   const min = Math.min(...points.map((point) => point.value));
   const max = Math.max(...points.map((point) => point.value));
   const span = max - min || 1;
-  const maxIndex = Math.max(...points.map((point) => point.index)) || 1;
-  const xy = (point: { index: number; value: number }) => {
-    const x = pad + (point.index / maxIndex) * (width - pad * 2);
-    const y = height - pad - ((point.value - min) / span) * (height - pad * 2);
-    return `${x},${y}`;
-  };
+  const xDenominator = Math.max(points.length - 1, 1);
+  const yTicks = [max, min + span / 2, min];
+  const chartPoints: ChartPoint[] = points.map((point, plotIndex) => {
+    const x = margin.left + (plotIndex / xDenominator) * plotWidth;
+    const y = margin.top + (1 - (point.value - min) / span) * plotHeight;
+    return { ...point, plotIndex, x, y };
+  });
+  const rawXTicks = chartPoints.length === 1
+    ? [chartPoints[0]]
+    : [chartPoints[0], chartPoints[Math.floor((chartPoints.length - 1) / 2)], chartPoints[chartPoints.length - 1]];
+  const xTicks = rawXTicks.filter(
+    (point, index, items) => items.findIndex((candidate) => candidate.plotIndex === point.plotIndex) === index
+  );
+  const metric = METRIC_DEFINITIONS[metricKey];
+  const selectedPoint = activePoint ?? chartPoints[chartPoints.length - 1];
+  const activePointStyle = activePoint
+    ? {
+        left: `${(activePoint.x / width) * 100}%`,
+        top: `${(activePoint.y / height) * 100}%`
+      }
+    : undefined;
+  const ariaLabel = `${metricLabel(metricKey)} trend for ${group.model_id} on ${group.gpu_type}`;
 
   return (
-    <svg className="trend-chart" viewBox={`0 0 ${width} ${height}`} role="img">
-      <polyline points={points.map(xy).join(" ")} fill="none" stroke="currentColor" strokeWidth="2.2" />
-      {points.map((point) => {
-        const [cx, cy] = xy(point).split(",");
-        return (
-          <circle
-            key={`${point.index}-${point.value}`}
-            cx={cx}
-            cy={cy}
-            r="3"
-            className={point.success ? "point-pass" : "point-fail"}
-          />
-        );
-      })}
-    </svg>
+    <div className="chart-shell">
+      <svg className="trend-chart" viewBox={`0 0 ${width} ${height}`} role="img" aria-label={ariaLabel}>
+        <line className="axis-line" x1={margin.left} y1={margin.top} x2={margin.left} y2={height - margin.bottom} />
+        <line
+          className="axis-line"
+          x1={margin.left}
+          y1={height - margin.bottom}
+          x2={width - margin.right}
+          y2={height - margin.bottom}
+        />
+        {yTicks.map((tick) => {
+          const y = margin.top + (1 - (tick - min) / span) * plotHeight;
+          return (
+            <g key={`y-${tick}`}>
+              <line className="grid-line" x1={margin.left} y1={y} x2={width - margin.right} y2={y} />
+              <text className="axis-label" x={margin.left - 8} y={y + 4} textAnchor="end">
+                {formatMetricValue(metricKey, tick)}
+              </text>
+            </g>
+          );
+        })}
+        {xTicks.map((point) => (
+          <text
+            className="axis-label"
+            key={`x-${point.plotIndex}-${point.point.timestamp ?? ""}`}
+            x={point.x}
+            y={height - 10}
+            textAnchor={point.plotIndex === 0 ? "start" : point.plotIndex === chartPoints.length - 1 ? "end" : "middle"}
+          >
+            {formatDate(point.point.timestamp)}
+          </text>
+        ))}
+        <polyline
+          points={chartPoints.map((point) => `${point.x},${point.y}`).join(" ")}
+          fill="none"
+          stroke="currentColor"
+          strokeWidth="2.2"
+        />
+        {chartPoints.map((point) => {
+          const pointLabel = `${metricLabel(metricKey)} ${formatMetricValue(metricKey, point.value, true)} at ${formatTime(
+            point.point.timestamp
+          )}, commit ${shortSha(point.point.commit_sha)}, ${runSourceLabel(point.point.run_source)}`;
+          return (
+            <g
+              key={`${point.plotIndex}-${point.value}-${point.point.commit_sha ?? ""}`}
+              onMouseEnter={() => setActivePoint(point)}
+              onMouseLeave={() => setActivePoint(null)}
+            >
+              <title>{pointLabel}</title>
+              <circle
+                className="point-hit-area"
+                cx={point.x}
+                cy={point.y}
+                r="12"
+                tabIndex={0}
+                aria-label={pointLabel}
+                onBlur={() => setActivePoint(null)}
+                onFocus={() => setActivePoint(point)}
+              />
+              <circle
+                cx={point.x}
+                cy={point.y}
+                r={activePoint?.plotIndex === point.plotIndex ? 5 : 4}
+                className={point.point.success ? "point-pass point-marker" : "point-fail point-marker"}
+              />
+            </g>
+          );
+        })}
+      </svg>
+      {activePoint ? (
+        <div className="hover-tooltip" style={activePointStyle} role="tooltip">
+          <strong>{formatMetricValue(metricKey, activePoint.value, true)}</strong>
+          {metric?.secondary ? <span>{metric.secondary(activePoint.value)}</span> : null}
+          <span>{shortSha(activePoint.point.commit_sha)}</span>
+          <span>{runSourceLabel(activePoint.point.run_source)}</span>
+        </div>
+      ) : null}
+      <div className="point-tooltip" aria-live="polite">
+        <strong>
+          {formatMetricValue(metricKey, selectedPoint.value, true)}
+          {metric?.secondary ? <span> ({metric.secondary(selectedPoint.value)})</span> : null}
+        </strong>
+        <span>{formatTime(selectedPoint.point.timestamp)}</span>
+        <span>Commit {shortSha(selectedPoint.point.commit_sha)}</span>
+        <span>{runSourceLabel(selectedPoint.point.run_source)}</span>
+        <span>{selectedPoint.point.success ? "Stored status: pass" : "Stored status: fail"}</span>
+        <span>{selectedPoint.point.baseline_eligible ? "Baseline eligible" : "Not baseline eligible"}</span>
+        {selectedPoint.point.pr_number ? <span>PR #{selectedPoint.point.pr_number}</span> : null}
+        {selectedPoint.point.branch ? <span>Branch {selectedPoint.point.branch}</span> : null}
+        {selectedPoint.point.build_url ? (
+          <a href={selectedPoint.point.build_url} target="_blank" rel="noreferrer">
+            Buildkite
+          </a>
+        ) : null}
+      </div>
+    </div>
   );
 }
 
@@ -79,6 +277,7 @@ export default function App() {
   const [days, setDays] = useState(90);
   const [modelFilter, setModelFilter] = useState("");
   const [gpuFilter, setGpuFilter] = useState("");
+  const [sourceFilter, setSourceFilter] = useState<"" | RunSource>("");
   const [summary, setSummary] = useState<SummaryResponse | null>(null);
   const [trends, setTrends] = useState<TrendGroup[]>([]);
   const [loading, setLoading] = useState(true);
@@ -90,8 +289,8 @@ export default function App() {
     setError(null);
     try {
       const [summaryData, trendData] = await Promise.all([
-        fetchSummary(days, modelFilter || undefined, gpuFilter || undefined),
-        fetchTrends(days, modelFilter || undefined, gpuFilter || undefined)
+        fetchSummary(days, modelFilter || undefined, gpuFilter || undefined, sourceFilter || undefined),
+        fetchTrends(days, modelFilter || undefined, gpuFilter || undefined, sourceFilter || undefined)
       ]);
       setSummary(summaryData);
       setTrends(trendData.groups);
@@ -119,7 +318,7 @@ export default function App() {
     load();
     const interval = window.setInterval(load, 5 * 60 * 1000);
     return () => window.clearInterval(interval);
-  }, [days, modelFilter, gpuFilter]);
+  }, [days, modelFilter, gpuFilter, sourceFilter]);
 
   const models = useMemo(() => {
     const values = new Set(summary?.rows.map((row) => row.model_id) ?? []);
@@ -182,6 +381,16 @@ export default function App() {
             ))}
           </select>
         </label>
+        <label>
+          Source
+          <select value={sourceFilter} onChange={(event) => setSourceFilter(event.target.value as "" | RunSource)}>
+            {RUN_SOURCES.map((source) => (
+              <option key={source.value || "all"} value={source.value}>
+                {source.label}
+              </option>
+            ))}
+          </select>
+        </label>
       </section>
 
       {error && <div className="notice error">Failed to load dashboard data: {error}</div>}
@@ -224,6 +433,8 @@ export default function App() {
                   <th>Model</th>
                   <th>GPU</th>
                   <th>Commit</th>
+                  <th>Source</th>
+                  <th>Baseline</th>
                   <th>Baseline N</th>
                   <th>Latency</th>
                   <th>Throughput</th>
@@ -245,6 +456,10 @@ export default function App() {
                     <td>{row.model_id}</td>
                     <td>{row.gpu_type}</td>
                     <td>{shortSha(row.commit_sha)}</td>
+                    <td>
+                      <span className={`source-badge source-${row.run_source}`}>{runSourceLabel(row.run_source)}</span>
+                    </td>
+                    <td>{row.baseline_eligible ? "eligible" : "excluded"}</td>
                     <td>{row.baseline_n}</td>
                     <td>{formatNumber(row.metrics.latency?.current, 3)}</td>
                     <td>{formatNumber(row.metrics.throughput?.current, 3)}</td>
@@ -274,7 +489,7 @@ export default function App() {
               METRIC_KEYS.map((metricKey) => (
                 <article className="trend-card" key={`${group.model_id}-${group.gpu_type}-${metricKey}`}>
                   <div>
-                    <h3>{summary?.rows[0]?.metrics[metricKey]?.label ?? metricKey}</h3>
+                    <h3>{metricLabel(metricKey)}</h3>
                     <p>
                       {group.model_id} | {group.gpu_type}
                     </p>

performance_dashboard/frontend/src/api.ts

@@ -18,9 +18,19 @@ export type SummaryRow = {
   regression_threshold_pct: number;
   computed_regression_status: "pass" | "fail";
   status: "pass" | "fail";
+  run_source: RunSource;
+  baseline_eligible: boolean;
+  branch: string;
+  pr_number: string;
+  test_scope: string;
+  build_url: string;
+  build_id: string;
+  job_id: string;
   metrics: Record<string, MetricValue>;
 };
 
+export type RunSource = "pr" | "local" | "scheduled_main" | "unknown";
+
 export type SummaryResponse = {
   rows: SummaryRow[];
   count: number;
@@ -29,9 +39,11 @@ export type SummaryResponse = {
     fail: number;
   };
   filters: {
-    days: number;
+    days: number | null;
+    trend_window_days?: number;
     model_id: string | null;
     gpu_type: string | null;
+    run_source: string | null;
   };
   sync: SyncState;
 };
@@ -40,6 +52,14 @@ export type TrendPoint = {
   timestamp: string | null;
   commit_sha: string | null;
   success: boolean;
+  run_source: RunSource;
+  baseline_eligible: boolean;
+  branch: string;
+  pr_number: string;
+  test_scope: string;
+  build_url: string;
+  build_id: string;
+  job_id: string;
   metrics: Record<string, number | null>;
 };
 
@@ -85,15 +105,15 @@ async function getJson<T>(path: string): Promise<T> {
   return response.json() as Promise<T>;
 }
 
-export async function fetchSummary(days = 90, modelId?: string, gpuType?: string) {
+export async function fetchSummary(days = 90, modelId?: string, gpuType?: string, runSource?: string) {
   return getJson<SummaryResponse>(
-    `/api/performance/summary?${params({ days, model_id: modelId, gpu_type: gpuType })}`
+    `/api/performance/summary?${params({ days, model_id: modelId, gpu_type: gpuType, run_source: runSource })}`
   );
 }
 
-export async function fetchTrends(days = 90, modelId?: string, gpuType?: string) {
+export async function fetchTrends(days = 90, modelId?: string, gpuType?: string, runSource?: string) {
   return getJson<TrendsResponse>(
-    `/api/performance/trends?${params({ days, model_id: modelId, gpu_type: gpuType })}`
+    `/api/performance/trends?${params({ days, model_id: modelId, gpu_type: gpuType, run_source: runSource })}`
   );
 }
 

performance_dashboard/frontend/src/styles.css

@@ -84,7 +84,7 @@ h3 {
 
 .filters {
   display: grid;
-  grid-template-columns: 120px minmax(220px, 1fr) minmax(220px, 1fr);
+  grid-template-columns: 120px minmax(200px, 1fr) minmax(200px, 1fr) minmax(180px, 0.8fr);
   gap: 14px;
   margin-bottom: 18px;
 }
@@ -186,7 +186,7 @@ h3 {
 
 table {
   width: 100%;
-  min-width: 900px;
+  min-width: 1120px;
   border-collapse: collapse;
 }
 
@@ -235,6 +235,33 @@ td {
   opacity: 0.78;
 }
 
+.source-badge {
+  display: inline-flex;
+  align-items: center;
+  border-radius: 999px;
+  padding: 4px 9px;
+  color: #1f2933;
+  background: #e8edf2;
+  font-size: 0.75rem;
+  font-weight: 800;
+  white-space: nowrap;
+}
+
+.source-scheduled_main {
+  color: #065f46;
+  background: #d1fae5;
+}
+
+.source-pr {
+  color: #1d4ed8;
+  background: #dbeafe;
+}
+
+.source-local {
+  color: #7c2d12;
+  background: #ffedd5;
+}
+
 .empty {
   padding: 28px 16px;
   color: #607080;
@@ -246,7 +273,7 @@ td {
 
 .trend-grid {
   display: grid;
-  grid-template-columns: repeat(auto-fill, minmax(280px, 1fr));
+  grid-template-columns: repeat(auto-fill, minmax(340px, 1fr));
   gap: 12px;
   padding: 14px;
 }
@@ -260,16 +287,101 @@ td {
 
 .trend-chart {
   width: 100%;
-  min-height: 96px;
+  min-height: 190px;
   color: #0f6b8f;
+  overflow: visible;
+}
+
+.chart-shell {
+  position: relative;
+  display: grid;
+  gap: 10px;
+  min-width: 0;
+}
+
+.axis-line {
+  stroke: #9aa8b6;
+  stroke-width: 1;
+}
+
+.grid-line {
+  stroke: #e4e9ee;
+  stroke-width: 1;
+}
+
+.axis-label {
+  fill: #667789;
+  font-size: 10px;
 }
 
 .point-pass {
   fill: #0f6b8f;
+  outline: none;
 }
 
 .point-fail {
   fill: #d94f4f;
+  outline: none;
+}
+
+.point-hit-area {
+  fill: transparent;
+  outline: none;
+  cursor: pointer;
+}
+
+.point-marker {
+  pointer-events: none;
+}
+
+.point-hit-area:focus + .point-marker {
+  stroke: #17212b;
+  stroke-width: 2;
+}
+
+.hover-tooltip {
+  position: absolute;
+  z-index: 2;
+  display: grid;
+  gap: 2px;
+  min-width: 132px;
+  max-width: 190px;
+  border: 1px solid #22313f;
+  border-radius: 6px;
+  padding: 8px 10px;
+  color: #ffffff;
+  background: #17212b;
+  font-size: 0.78rem;
+  pointer-events: none;
+  transform: translate(10px, -100%);
+  box-shadow: 0 10px 24px rgb(15 23 42 / 22%);
+}
+
+.hover-tooltip strong {
+  font-size: 0.9rem;
+}
+
+.point-tooltip {
+  display: grid;
+  grid-template-columns: repeat(2, minmax(0, 1fr));
+  gap: 4px 10px;
+  border: 1px solid #d5dde5;
+  border-radius: 6px;
+  padding: 10px;
+  color: #263646;
+  background: #f8fafc;
+  font-size: 0.78rem;
+}
+
+.point-tooltip strong {
+  grid-column: 1 / -1;
+  color: #132232;
+  font-size: 0.9rem;
+}
+
+.point-tooltip a {
+  color: #0f6b8f;
+  font-weight: 700;
 }
 
 .empty-chart {

pyproject.toml

@@ -27,9 +27,10 @@ dependencies = [
     "timm>=1.0.11",
     "peft>=0.15.0",
     "diffusers>=0.38.0",
-    "torch==2.11.0",
+    "torch==2.12.0",
     "torchvision",
     "torchaudio",
+    "flashinfer-python",
 
     # Acceleration & Optimization
     "accelerate==1.0.1",
@@ -67,7 +68,6 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.6",
     "wheel",
 
     # Training Dependencies
@@ -95,15 +95,15 @@ members = ["apps/dreamverse"]
 [tool.uv.sources]
 torch = [
   { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux'" },
 ]
 torchvision = [
   { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux'" },
 ]
 torchaudio = [
   { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux'" },
 ]
 imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
 flash-attn-cute = { git = "https://github.com/XOR-op/flash-attention.git", branch = "fa4-compile", subdirectory = "flash_attn/cute" }
@@ -114,8 +114,8 @@ url = "https://download.pytorch.org/whl/cpu"
 explicit = true
 
 [[tool.uv.index]]
-name = "pytorch-cu128"
-url = "https://download.pytorch.org/whl/cu128"
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
 [project.optional-dependencies]

pyproject_other.toml

@@ -65,7 +65,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.6",
+    "fastvideo-kernel==0.3.0",
     "wheel",
 
     # Training Dependencies