[docs]: note the Hub build's Hopper-only arch subset is deliberate

Main's kernel wheels went multi-arch (0.3.x); the Hub build.toml stays sm_90a because the TK attention kernels it packages require Hopper.
[kernel]: declare mutation + return aliases in Hub torch-op schemas
2026-07-10 19:31:18 -07:00 · 2026-07-10 19:31:10 -07:00 · 2026-07-10 19:31:00 -07:00 · 2026-07-10 19:30:47 -07:00 · 2026-07-10 19:25:34 -07:00 · 2026-07-10 21:39:17 +00:00
96 changed files with 7439 additions and 355 deletions
@@ -80,6 +80,23 @@ MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUI

 POST_RUN_HOOK=""

+is_truthy() {
+    case "${1:-}" in
+        1|true|TRUE|yes|YES|on|ON) return 0 ;;
+        *) return 1 ;;
+    esac
+}
+
+ssim_bootstrap_args() {
+    local title="${PR_TITLE:-}"
+    local message="${BUILDKITE_MESSAGE:-}"
+    if is_truthy "${FASTVIDEO_SSIM_BOOTSTRAP_MODE:-}" \
+        || [[ "$title" == *"[new-model]"* ]] \
+        || [[ "$message" == *"[new-model]"* ]]; then
+        printf ' --bootstrap-mode'
+    fi
+}
+
 upload_performance_artifacts() {
    SHORT_SHA=${BUILDKITE_COMMIT:0:7}
    LOCAL_DIR="downloaded_reports"
@@ -172,7 +189,12 @@ case "$TEST_TYPE" in
        ;;
    "ssim")
        log "Running SSIM tests..."
-        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_SSIM_TEST_FILE::run_ssim_tests"
+        SSIM_BOOTSTRAP_ARGS=$(ssim_bootstrap_args)
+        if [ -n "$SSIM_BOOTSTRAP_ARGS" ]; then
+            log "SSIM bootstrap mode enabled for new-model reference draft generation"
+        fi
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run "
+        MODAL_COMMAND+="$MODAL_SSIM_TEST_FILE::run_ssim_tests$SSIM_BOOTSTRAP_ARGS"
        ;;
    "training")
        log "Running training tests..."
@@ -52,6 +52,7 @@ jobs:
            core.setOutput('pr_sha', pr.head.sha);
            core.setOutput('pr_branch', pr.head.ref);
            core.setOutput('pr_number', String(prNumber));
+            core.setOutput('pr_title', pr.title);

      - name: Trigger Full Suite
        if: steps.perm.outputs.has_write == 'true'
@@ -60,6 +61,7 @@ jobs:
          PR_SHA: ${{ steps.label.outputs.pr_sha }}
          PR_BRANCH: ${{ steps.label.outputs.pr_branch }}
          PR_NUMBER: ${{ steps.label.outputs.pr_number }}
+          PR_TITLE: ${{ steps.label.outputs.pr_title }}
          BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
          BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
        run: |
@@ -71,6 +73,7 @@ jobs:
              --arg commit "$PR_SHA" \
              --arg branch "$PR_BRANCH" \
              --arg message "Full Suite for PR #${PR_NUMBER} (via /merge)" \
+              --arg pr_title "$PR_TITLE" \
              --argjson pr_id "$PR_NUMBER" \
              '{
                commit: $commit,
@@ -80,11 +83,12 @@ jobs:
                pull_request_id: $pr_id,
                pull_request_base_branch: "main",
                env: {
-                  TEST_SCOPE: "full",
-                  FULL_SUITE: "true",
-                  PR_NUMBER: ($pr_id | tostring)
-                }
-              }')"
+                    TEST_SCOPE: "full",
+                    FULL_SUITE: "true",
+                    PR_NUMBER: ($pr_id | tostring),
+                    PR_TITLE: $pr_title
+                  }
+                }')"

  parse-command:
    if: >-
@@ -241,6 +245,7 @@ jobs:
          TEST_SCOPE: ${{ needs.parse-command.outputs.test_scope }}
          FULL_SUITE: ${{ needs.parse-command.outputs.full_suite }}
          TEST_TYPE: ${{ needs.parse-command.outputs.test_type }}
+          PR_TITLE: ${{ github.event.issue.title }}
          BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
          BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
        run: |
@@ -257,6 +262,7 @@ jobs:
              --arg full_suite "$FULL_SUITE" \
              --arg test_type "$TEST_TYPE" \
              --arg pr_number "$PR_NUMBER" \
+              --arg pr_title "$PR_TITLE" \
              '{
                commit: $commit,
                branch: $branch,
@@ -266,8 +272,9 @@ jobs:
                pull_request_base_branch: "main",
                env: {
                  TEST_SCOPE: $test_scope,
-                  FULL_SUITE: $full_suite,
-                  TEST_TYPE: $test_type,
-                  PR_NUMBER: $pr_number
-                }
-              }')"
+                    FULL_SUITE: $full_suite,
+                    TEST_TYPE: $test_type,
+                    PR_NUMBER: $pr_number,
+                    PR_TITLE: $pr_title
+                  }
+                }')"
@@ -56,6 +56,7 @@ jobs:
          PR_SHA: ${{ github.event.pull_request.head.sha }}
          PR_BRANCH: ${{ github.event.pull_request.head.ref }}
          PR_NUMBER: ${{ github.event.pull_request.number }}
+          PR_TITLE: ${{ github.event.pull_request.title }}
          BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
          BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
        run: |
@@ -67,6 +68,7 @@ jobs:
              --arg commit "$PR_SHA" \
              --arg branch "$PR_BRANCH" \
              --arg message "Full Suite for PR #${PR_NUMBER}" \
+              --arg pr_title "$PR_TITLE" \
              --argjson pr_id "$PR_NUMBER" \
              '{
                commit: $commit,
@@ -78,6 +80,7 @@ jobs:
                env: {
                  TEST_SCOPE: "full",
                  FULL_SUITE: "true",
-                  PR_NUMBER: ($pr_id | tostring)
+                  PR_NUMBER: ($pr_id | tostring),
+                  PR_TITLE: $pr_title
                }
              }')"
@@ -0,0 +1,95 @@
+name: Publish FastVideo Kernel to Hugging Face Hub
+
+# Version-gated like the PyPI flow (publish-kernel.yml): a push to main only
+# publishes when the fastvideo-kernel version actually changes. Use
+# workflow_dispatch as the manual override.
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - "fastvideo-kernel/pyproject.toml"
+  workflow_dispatch:
+    inputs:
+      branch:
+        description: "Hub branch to upload the build"
+        default: ""
+        required: false
+
+# A Hub build takes hours; never race two publishes to the same repo.
+concurrency:
+  group: ${{ github.workflow }}
+  cancel-in-progress: false
+
+jobs:
+  check-version-change:
+    runs-on: ubuntu-latest
+    outputs:
+      version-changed: ${{ steps.check-version.outputs.changed }}
+      new-version: ${{ steps.check-version.outputs.new-version }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 2
+
+      - name: Check if version changed
+        id: check-version
+        run: |
+          cd fastvideo-kernel
+          # Get current commit's version from pyproject.toml
+          # Use ^ to match start of line to avoid matching minimum-version
+          NEW_VERSION=$(grep -oP '^version\s*=\s*"\K[^"]+' pyproject.toml)
+          echo "New version: $NEW_VERSION"
+
+          # Get previous version from git history
+          # Note: git show expects path relative to repo root
+          OLD_VERSION=$(git show HEAD~1:fastvideo-kernel/pyproject.toml | grep -oP '^version\s*=\s*"\K[^"]+' || echo "0.0.0")
+          echo "Old version: $OLD_VERSION"
+
+          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
+            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
+            echo "changed=true" >> "$GITHUB_OUTPUT"
+            echo "new-version=$NEW_VERSION" >> "$GITHUB_OUTPUT"
+          else
+            echo "Version did not change"
+            echo "changed=false" >> "$GITHUB_OUTPUT"
+          fi
+
+  build:
+    needs: check-version-change
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build kernel via HF Jobs and upload to hao-ai-lab/fastvideo-kernel
+        # Pinned to a commit SHA: this third-party action receives HF_TOKEN.
+        uses: huggingface/kernel-builder-job@ee590cee1ff01e922310401f18d5554b8da448ac # main, 2026-06-03
+        with:
+          token: ${{ secrets.HF_TOKEN }}
+          namespace: hao-ai-lab
+          flavor: cpu-xl
+          timeout: "21600"
+          script: |
+            set +x
+            export HF_TOKEN="${{ secrets.HF_TOKEN }}"
+            export GIT_LFS_SKIP_SMUDGE=1
+
+            git clone "${{ github.server_url }}/${{ github.repository }}" FastVideo
+            cd FastVideo
+            git checkout "${{ github.sha }}"
+            git submodule update --init --recursive fastvideo-kernel/include/cutlass fastvideo-kernel/include/tk
+            cd fastvideo-kernel
+
+            UPLOAD_BRANCH="${{ github.event.inputs.branch }}"
+            BRANCH_ARGS=()
+            if [ -n "${UPLOAD_BRANCH}" ]; then
+              BRANCH_ARGS=(--branch "${UPLOAD_BRANCH}")
+            fi
+
+            nix run github:huggingface/kernels#kernel-builder -- build-and-upload \
+              --max-jobs 4 \
+              --cores 8 \
+              --repo-id hao-ai-lab/fastvideo-kernel \
+              "${BRANCH_ARGS[@]}"
@@ -0,0 +1,22 @@
+{
+  "alpha_yaw": 0.08734091699186919,
+  "alpha_pitch": 0.08169667696275307,
+  "alpha_turn": 5.724587470723463e-17,
+  "beta_fwd": 0.02842768078408099,
+  "beta_strafe": 0.022531015077067108,
+  "focal_length": 457.0,
+  "frame_shape": [
+    352,
+    640
+  ],
+  "calibrated_from": [
+    "1_wasd_only",
+    "camera",
+    "camera4hold_alpha1",
+    "fully_random",
+    "wasdonly_alpha1",
+    "wasd4holdrandview_simple_1key1mouse1"
+  ],
+  "residual_rms": 15.890399609478676,
+  "n_equations": 4125232
+}
@@ -79,10 +79,11 @@ fastvideo/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, baseline-eligible records for that
-model+GPU. PR and local records are visible in the dashboard but are not
-baseline eligible.
+normalized JSON per run. For v2 records, the rolling baseline is the median of
+the last 5 successful, baseline-eligible records in the same comparison cohort:
+`model_id`, `gpu_type`, `workload_id`, `variant_id`, `benchmark_version`,
+`recipe_fingerprint`, `hardware_profile_id`, and `software_profile_id`. PR and
+local records are visible in the dashboard but are not baseline eligible.

 ## Planned Coverage

@@ -158,13 +159,16 @@ unrealistic memory growth, and optionally large component-specific slowdowns
 even when the rolling baseline is empty. They are hand-set with generous
 headroom and almost never need touching.

-### Rolling baseline (per `(model_id, gpu_type)`)
+### Rolling baseline (per comparison cohort)

 `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 evaluates the current run with the metric's
-rolling regression policy. For latency, memory, and component times, higher
-values are regressions. For throughput, lower values are regressions.
+for the same comparison cohort from the HF dataset, computes the median for
+each available metric, and evaluates the current run with the metric's
+rolling regression policy. For v2 records, that cohort is `model_id`,
+`gpu_type`, `workload_id`, `variant_id`, `benchmark_version`,
+`recipe_fingerprint`, `hardware_profile_id`, and `software_profile_id`. For
+latency, memory, and component times, higher values are regressions. For
+throughput, lower values are regressions.

 A metric exceeds its rolling threshold when both of these are true:

@@ -230,8 +234,10 @@ Recipe fingerprinting, hardware/software profile IDs, exact-identity
 comparison, and dashboard cohort grouping land with this change: v2 records
 compare only within their identity cohort, and a record that opens a NEW
 cohort is marked `baseline_status: "initialized_new_cohort"` (regression
-gating starts once that cohort accumulates history). Legacy v1 configs keep
-the `(model_id, gpu_type)` rolling-baseline comparison. Metric-specific
+gating starts once that cohort accumulates history). Legacy v1 configs still
+run and are normalized for reporting, but their records skip rolling-baseline
+comparison entirely (`baseline_status: "skipped_missing_identity"`, never
+baseline eligible); only static thresholds gate them. Metric-specific
 threshold policies and promoted baselines remain separate follow-ups.

 ### Raw record (`results/perf_*.json`)
@@ -241,7 +247,7 @@ Written by `test_inference_performance.py`. One file per benchmark run.
 ```jsonc
 {
  "benchmark_id": "wan-t2v-1.3b-2gpu",
-  "config_schema_version": 2,
+  "result_schema_version": 2,
  "workload_id": "wan-t2v",
  "variant_id": "1.3b-sp2",
  "benchmark_version": 2,
@@ -270,8 +276,15 @@ Written by `test_inference_performance.py`. One file per benchmark run.
    }
  },
  "commit": "<full sha>",
+  "run_source": "pr",
+  "branch": "feature/perf-change",
  "pr_number": "1234",
+  "test_scope": "direct",
+  "build_url": "https://buildkite.example/build",
+  "build_id": "<buildkite-build-id>",
+  "job_id": "<buildkite-job-id>",
  "timestamp": "2026-05-08T22:00:00+00:00",
+  "quality_metadata": { "quality_status": "canonical" },
  "text_encoder_time_s": 2.141,
  "dit_time_s": 8.437,
  "vae_decode_time_s": 3.208,
@@ -322,6 +335,7 @@ result, used as the rolling-baseline source of truth.
 ```jsonc
 {
  "model_id": "wan-t2v-1.3b-2gpu",
+  "result_schema_version": 2,
  "workload_id": "wan-t2v",
  "variant_id": "1.3b-sp2",
  "benchmark_version": 2,
@@ -345,17 +359,30 @@ result, used as the rolling-baseline source of truth.
  "hardware_profile_id": "hw-<sha256-prefix>",
  "software_profile_id": "sw-<sha256-prefix>",
  "environment_fingerprint": "env-<sha256-prefix>",
+  "run_source": "pr",
+  "branch": "feature/perf-change",
+  "pr_number": "1234",
+  "test_scope": "direct",
+  "build_url": "https://buildkite.example/build",
+  "build_id": "<buildkite-build-id>",
+  "job_id": "<buildkite-job-id>",
+  "quality_metadata": { "quality_status": "canonical" },
  "success": true
 }
 ```

 ### Compatibility with legacy records

-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.
+Older records in the HF dataset may not have `result_schema_version`,
+component timing fields, or v2 identity/profile fields. Records without
+`result_schema_version` are treated as v1. 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.
+Current `perf_*.json` artifacts that lack the v2 comparison identity are
+normalized for reporting but skip rolling-baseline comparison and are not marked
+baseline eligible.

 New records compare only against the same `model_id`, `gpu_type`,
 `workload_id`, `variant_id`, `benchmark_version`, `recipe_fingerprint`,
@@ -377,7 +404,7 @@ FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
 | `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` | `fastvideo/performance/hf_store.py` | HF dataset repo holding rolling-baseline records. |
 | `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `fastvideo/performance/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_RUN_SOURCE` | inferred | `compare_baseline.py`, `test_inference_performance.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`. |
@@ -394,11 +421,14 @@ point is `fastvideo/tests/modal/pr_test.py:run_performance_tests` and the
 Buildkite artifact upload is in
 `.buildkite/scripts/pr_test.sh:upload_performance_artifacts`.

-Each performance build runs pytest first. If that fixed-threshold phase fails,
-PR/direct runs skip `compare_baseline.py` because they only upload passing
-records. Scheduled-main runs still execute `compare_baseline.py` with
-`PERF_PYTEST_RC` set so the failed canonical attempt is visible in normalized
-JSON and dashboard history. The dashboard runs best-effort for observability.
+Each performance build runs pytest first. PR and direct runs only continue to
+`compare_baseline.py` when that fixed-threshold phase passes; if pytest fails,
+Markdown summaries and normalized JSON artifacts are not emitted. Scheduled
+main runs set `PERF_UPLOAD_POLICY=always`, so they still run
+`compare_baseline.py` (with `PERF_PYTEST_RC` set) after a fixed-threshold
+failure. Those failed scheduled main runs emit summaries and normalized
+records, upload records with `success=false`, and are excluded from future
+rolling baselines. The dashboard still runs best-effort for observability.
 When the rolling-baseline phase runs, it emits:

 * **Markdown summary** — appended to `$GITHUB_STEP_SUMMARY` when that variable
@@ -406,7 +436,7 @@ When the rolling-baseline phase runs, it emits:
  per-benchmark row with current vs. baseline values for latency, throughput,
  memory, text encoder time, DiT time, and VAE decode time.
 * **Plotly dashboard** — `dashboard_<sha>_<ts>.html` showing time-series for
-  each metric grouped by `(model_id, gpu_type)`.
+  each metric grouped by comparison cohort.
 * **Normalized records** — `normalized_perf_*.json`, one per benchmark.
  Useful as input to the
  [`reseed-performance-baseline`](https://github.com/hao-ai-lab/FastVideo/blob/main/.agents/skills/reseed-performance-baseline/SKILL.md)
@@ -480,8 +510,8 @@ When the rolling-baseline phase runs, it emits:

 ## Troubleshooting

-**"No baseline for ... Initializing"** — first run for this `(model_id,
-gpu_type)`. Run will pass and (if persisting) seed the first record.
+**"No baseline for ... Initializing"** — first run for this comparison cohort.
+Run will pass and (if persisting) seed the first record.

 **Persistent failure right after a torch / kernel / image upgrade** —
 genuine regression *or* baseline drift. Compare the failing normalized record
@@ -180,6 +180,30 @@ python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
  --device-folder L40S_reference_videos
 ```

+### SSIM Bootstrap Mode
+
+Normal SSIM runs are strict: if a reference video or latent is missing, the
+test fails. For new-model PRs, CI can run SSIM in bootstrap mode so missing
+references are uploaded as draft artifacts for review instead of immediately
+blocking on a missing canonical reference.
+
+Buildkite enables SSIM bootstrap mode when either condition is true:
+
+- the PR title or Buildkite message contains `[new-model]`;
+- `FASTVIDEO_SSIM_BOOTSTRAP_MODE=1` is set for the Buildkite job.
+
+Bootstrap mode passes `--ssim-bootstrap-mode` to pytest. When a generated
+artifact is available, the test uploads it under the `drafts/...` namespace in
+the SSIM reference repo and marks that case as expected-failed. After reviewing
+the draft, promote it into the canonical reference layout:
+
+```bash
+python fastvideo/tests/ssim/reference_videos_cli.py promote-draft \
+  --quality-tier default \
+  --device-folder L40S_reference_videos \
+  --model-id <model_id>
+```
+
 ## CI Integration

 FastVideo CI tests are orchestrated by Buildkite and run on Modal GPU
@@ -161,6 +161,18 @@ training:
    decay_interval_steps: 0
 ```

+`training.data.data_path` can also mix multiple preprocessed datasets by using a mapping from dataset path to repeat count:
+
+```yaml
+training:
+  data:
+    data_path:
+      data/zeldam2-clean: 1
+      data/multi3d_games: 2
+```
+
+The repeat count duplicates that dataset's parquet file list before shuffling/sampling, so the example above trains with roughly twice as much `multi3d_games` exposure as `zeldam2-clean`. Paths are just suggested locations; use any local path that contains a FastVideo preprocessed parquet dataset.
+
 See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab,
 including SwanLab installation and authentication.

@@ -326,6 +338,40 @@ Self-Forcing inherits all DMD2 parameters, plus:
 | `enable_gradient_in_rollout` | `true` | Enable backprop through rollout |
 | `start_gradient_frame` | `0` | Frame index where gradients begin |

+### Streaming Long Tuning
+
+`StreamingLongTuningMethod` extends Self-Forcing for LongLive-style rollouts. It
+keeps a streaming state, generates overlapping chunks, and trains only the new
+frames while preserving context from earlier chunks.
+
+For the MatrixGame2/Zelda world-model example, self-forcing and long tuning are
+separate runs: first train or load the 1k-step self-forcing checkpoint using
+`examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml`,
+then run
+`examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml`
+from that checkpoint for the 3k-step streaming long-tuning stage.
+
+```yaml
+method:
+  _target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
+  streaming_chunk_size: 9
+  streaming_max_length: 39
+  streaming_fixed_overlap_latents: 3
+  streaming_reencode_overlap_anchor: true
+  streaming_anchor_inject_k: 1
+  streaming_require_full_blocks: true
+  multi_phased_distill_schedule:
+    - stage: streaming_long
+      start_step: 0
+      end_step: 3000
+      num_latent_t: 39
+      streaming_training: true
+```
+
+See
+`examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml`
+for a complete MatrixGame2/Zelda configuration.
+
 ---

 ## Callbacks
@@ -0,0 +1,69 @@
+# LTX-2.3 distilled inference configs
+
+Ready-to-run `fastvideo generate` run configs for the LTX-2.3
+distilled model (`FastVideo/LTX-2.3-Distilled-Diffusers`), covering both
+workloads (t2v / i2v), both two-stage step schedules (`5+2`, `8+3` = denoise
+ refine), and four resolutions.
+
+```bash
+fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml
+```
+
+Each config is self-contained (no preset registry needed): the two-stage
+refine is wired via `generator.pipeline.preset_overrides.refine`, and the
+base sampling knobs live under `request.sampling`. The refine upsampler
+auto-resolves from the model's `spatial_upscaler`.
+
+## Configs
+
+| workload | schedule | resolution (HxW) | file |
+|---|---|---|---|
+| t2v | 5+2 | 1280x832 | `t2v_5s2_1280x832.yaml` |
+| t2v | 5+2 | 1024x1536 | `t2v_5s2_1024x1536.yaml` |
+| t2v | 5+2 | 768x1280 | `t2v_5s2_768x1280.yaml` |
+| t2v | 5+2 | 512x768 | `t2v_5s2_512x768.yaml` |
+| t2v | 8+3 | 1280x832 | `t2v_8s3_1280x832.yaml` |
+| t2v | 8+3 | 1024x1536 | `t2v_8s3_1024x1536.yaml` |
+| t2v | 8+3 | 768x1280 | `t2v_8s3_768x1280.yaml` |
+| t2v | 8+3 | 512x768 | `t2v_8s3_512x768.yaml` |
+| i2v | 5+2 | 1280x832 | `i2v_5s2_1280x832.yaml` |
+| i2v | 5+2 | 1024x1536 | `i2v_5s2_1024x1536.yaml` |
+| i2v | 5+2 | 768x1280 | `i2v_5s2_768x1280.yaml` |
+| i2v | 5+2 | 512x768 | `i2v_5s2_512x768.yaml` |
+| i2v | 8+3 | 1280x832 | `i2v_8s3_1280x832.yaml` |
+| i2v | 8+3 | 1024x1536 | `i2v_8s3_1024x1536.yaml` |
+| i2v | 8+3 | 768x1280 | `i2v_8s3_768x1280.yaml` |
+| i2v | 8+3 | 512x768 | `i2v_8s3_512x768.yaml` |
+
+## Overriding without editing a file
+
+Dotted overrides (prefixes `generator.` / `request.`) let you tweak any field:
+
+```bash
+# swap prompt
+fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml \
+  --request.prompt "a red fox running through fresh snow"
+
+# change output path / gpu count
+fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_512x768.yaml \
+  --request.output.output_path outputs/preview.mp4 \
+  --generator.engine.num_gpus 4
+```
+
+## i2v
+
+The `i2v_*` configs take a first-frame image via
+`request.extensions.ltx2_images` (`[[path, frame_offset, weight]]`). Edit the
+path in the file, or override it:
+
+```bash
+fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1280x832.yaml \
+  --request.extensions.ltx2_images '[["/data/portrait.jpg", 0, 1.0]]'
+```
+
+## Schedules
+
+`5+2` is the fast preview schedule; `8+3` is the higher-quality distilled
+recipe. Refine (`preset_overrides.refine.num_inference_steps`) only accepts 2
+or 3 steps.
+
@@ -0,0 +1,38 @@
+# LTX-2.3 distilled i2v — 5+2 two-stage at 1024x1536.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_1024x1536.yaml
+#
+# Stage 1 denoises for 5 steps at half resolution; the latents are then
+# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: i2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 2
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
+  sampling:
+    height: 1024
+    width: 1536
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 5
+  # i2v conditioning — replace the path with your own first-frame image.
+  extensions:
+    ltx2_images:
+      - ["/path/to/your/first_frame.jpg", 0, 1.0]
+    ltx2_image_crf: 0.0
+  output:
+    output_path: outputs/ltx2_3_i2v_5s2_1024x1536.mp4
+    save_video: true
@@ -0,0 +1,38 @@
+# LTX-2.3 distilled i2v — 5+2 two-stage at 1280x832.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_1280x832.yaml
+#
+# Stage 1 denoises for 5 steps at half resolution; the latents are then
+# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: i2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 2
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
+  sampling:
+    height: 1280
+    width: 832
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 5
+  # i2v conditioning — replace the path with your own first-frame image.
+  extensions:
+    ltx2_images:
+      - ["/path/to/your/first_frame.jpg", 0, 1.0]
+    ltx2_image_crf: 0.0
+  output:
+    output_path: outputs/ltx2_3_i2v_5s2_1280x832.mp4
+    save_video: true
@@ -0,0 +1,38 @@
+# LTX-2.3 distilled i2v — 5+2 two-stage at 512x768.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_512x768.yaml
+#
+# Stage 1 denoises for 5 steps at half resolution; the latents are then
+# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: i2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 2
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
+  sampling:
+    height: 512
+    width: 768
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 5
+  # i2v conditioning — replace the path with your own first-frame image.
+  extensions:
+    ltx2_images:
+      - ["/path/to/your/first_frame.jpg", 0, 1.0]
+    ltx2_image_crf: 0.0
+  output:
+    output_path: outputs/ltx2_3_i2v_5s2_512x768.mp4
+    save_video: true
@@ -0,0 +1,38 @@
+# LTX-2.3 distilled i2v — 5+2 two-stage at 768x1280.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/i2v_5s2_768x1280.yaml
+#
+# Stage 1 denoises for 5 steps at half resolution; the latents are then
+# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: i2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 2
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
+  sampling:
+    height: 768
+    width: 1280
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 5
+  # i2v conditioning — replace the path with your own first-frame image.
+  extensions:
+    ltx2_images:
+      - ["/path/to/your/first_frame.jpg", 0, 1.0]
+    ltx2_image_crf: 0.0
+  output:
+    output_path: outputs/ltx2_3_i2v_5s2_768x1280.mp4
+    save_video: true
@@ -0,0 +1,38 @@
+# LTX-2.3 distilled i2v — 8+3 two-stage at 1024x1536.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1024x1536.yaml
+#
+# Stage 1 denoises for 8 steps at half resolution; the latents are then
+# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: i2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 3
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
+  sampling:
+    height: 1024
+    width: 1536
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 8
+  # i2v conditioning — replace the path with your own first-frame image.
+  extensions:
+    ltx2_images:
+      - ["/path/to/your/first_frame.jpg", 0, 1.0]
+    ltx2_image_crf: 0.0
+  output:
+    output_path: outputs/ltx2_3_i2v_8s3_1024x1536.mp4
+    save_video: true
@@ -0,0 +1,38 @@
+# LTX-2.3 distilled i2v — 8+3 two-stage at 1280x832.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_1280x832.yaml
+#
+# Stage 1 denoises for 8 steps at half resolution; the latents are then
+# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: i2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 3
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
+  sampling:
+    height: 1280
+    width: 832
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 8
+  # i2v conditioning — replace the path with your own first-frame image.
+  extensions:
+    ltx2_images:
+      - ["/path/to/your/first_frame.jpg", 0, 1.0]
+    ltx2_image_crf: 0.0
+  output:
+    output_path: outputs/ltx2_3_i2v_8s3_1280x832.mp4
+    save_video: true
@@ -0,0 +1,38 @@
+# LTX-2.3 distilled i2v — 8+3 two-stage at 512x768.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_512x768.yaml
+#
+# Stage 1 denoises for 8 steps at half resolution; the latents are then
+# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: i2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 3
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
+  sampling:
+    height: 512
+    width: 768
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 8
+  # i2v conditioning — replace the path with your own first-frame image.
+  extensions:
+    ltx2_images:
+      - ["/path/to/your/first_frame.jpg", 0, 1.0]
+    ltx2_image_crf: 0.0
+  output:
+    output_path: outputs/ltx2_3_i2v_8s3_512x768.mp4
+    save_video: true
@@ -0,0 +1,38 @@
+# LTX-2.3 distilled i2v — 8+3 two-stage at 768x1280.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/i2v_8s3_768x1280.yaml
+#
+# Stage 1 denoises for 8 steps at half resolution; the latents are then
+# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: i2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 3
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    The subject slowly turns toward the camera with a soft, natural expression, hair and clothing swaying gently, shallow depth of field, subtle cinematic motion.
+  sampling:
+    height: 768
+    width: 1280
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 8
+  # i2v conditioning — replace the path with your own first-frame image.
+  extensions:
+    ltx2_images:
+      - ["/path/to/your/first_frame.jpg", 0, 1.0]
+    ltx2_image_crf: 0.0
+  output:
+    output_path: outputs/ltx2_3_i2v_8s3_768x1280.mp4
+    save_video: true
@@ -0,0 +1,33 @@
+# LTX-2.3 distilled t2v — 5+2 two-stage at 1024x1536.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_1024x1536.yaml
+#
+# Stage 1 denoises for 5 steps at half resolution; the latents are then
+# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: t2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 2
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
+  sampling:
+    height: 1024
+    width: 1536
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 5
+  output:
+    output_path: outputs/ltx2_3_t2v_5s2_1024x1536.mp4
+    save_video: true
@@ -0,0 +1,33 @@
+# LTX-2.3 distilled t2v — 5+2 two-stage at 1280x832.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_1280x832.yaml
+#
+# Stage 1 denoises for 5 steps at half resolution; the latents are then
+# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: t2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 2
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
+  sampling:
+    height: 1280
+    width: 832
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 5
+  output:
+    output_path: outputs/ltx2_3_t2v_5s2_1280x832.mp4
+    save_video: true
@@ -0,0 +1,33 @@
+# LTX-2.3 distilled t2v — 5+2 two-stage at 512x768.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_512x768.yaml
+#
+# Stage 1 denoises for 5 steps at half resolution; the latents are then
+# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: t2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 2
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
+  sampling:
+    height: 512
+    width: 768
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 5
+  output:
+    output_path: outputs/ltx2_3_t2v_5s2_512x768.mp4
+    save_video: true
@@ -0,0 +1,33 @@
+# LTX-2.3 distilled t2v — 5+2 two-stage at 768x1280.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/t2v_5s2_768x1280.yaml
+#
+# Stage 1 denoises for 5 steps at half resolution; the latents are then
+# spatially upsampled and refined for 2 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: t2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 2
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
+  sampling:
+    height: 768
+    width: 1280
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 5
+  output:
+    output_path: outputs/ltx2_3_t2v_5s2_768x1280.mp4
+    save_video: true
@@ -0,0 +1,33 @@
+# LTX-2.3 distilled t2v — 8+3 two-stage at 1024x1536.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1024x1536.yaml
+#
+# Stage 1 denoises for 8 steps at half resolution; the latents are then
+# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: t2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 3
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
+  sampling:
+    height: 1024
+    width: 1536
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 8
+  output:
+    output_path: outputs/ltx2_3_t2v_8s3_1024x1536.mp4
+    save_video: true
@@ -0,0 +1,33 @@
+# LTX-2.3 distilled t2v — 8+3 two-stage at 1280x832.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_1280x832.yaml
+#
+# Stage 1 denoises for 8 steps at half resolution; the latents are then
+# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: t2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 3
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
+  sampling:
+    height: 1280
+    width: 832
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 8
+  output:
+    output_path: outputs/ltx2_3_t2v_8s3_1280x832.mp4
+    save_video: true
@@ -0,0 +1,33 @@
+# LTX-2.3 distilled t2v — 8+3 two-stage at 512x768.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_512x768.yaml
+#
+# Stage 1 denoises for 8 steps at half resolution; the latents are then
+# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: t2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 3
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
+  sampling:
+    height: 512
+    width: 768
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 8
+  output:
+    output_path: outputs/ltx2_3_t2v_8s3_512x768.mp4
+    save_video: true
@@ -0,0 +1,33 @@
+# LTX-2.3 distilled t2v — 8+3 two-stage at 768x1280.
+#
+# Run:
+#   fastvideo generate --config examples/inference/ltx2_3/t2v_8s3_768x1280.yaml
+#
+# Stage 1 denoises for 8 steps at half resolution; the latents are then
+# spatially upsampled and refined for 3 steps (refine only supports 2 or 3).
+# The refine upsampler auto-resolves from the model's `spatial_upscaler`.
+generator:
+  model_path: FastVideo/LTX-2.3-Distilled-Diffusers
+  engine:
+    num_gpus: 1
+  pipeline:
+    workload_type: t2v
+    preset_overrides:
+      refine:
+        enabled: true
+        num_inference_steps: 3
+        guidance_scale: 1.0
+        add_noise: true
+request:
+  prompt: >-
+    A cinematic drone shot flying over dramatic coastal cliffs at sunrise, golden light spilling across the water, gentle waves breaking on the rocks below, ultra-detailed, smooth camera motion.
+  sampling:
+    height: 768
+    width: 1280
+    num_frames: 121
+    fps: 24
+    guidance_scale: 1.0
+    num_inference_steps: 8
+  output:
+    output_path: outputs/ltx2_3_t2v_8s3_768x1280.mp4
+    save_video: true
@@ -100,9 +100,6 @@ callbacks:
    sampling_steps: [4]
    sampling_timesteps: [1000, 750, 500, 250]
    num_frames: 81
-    # Validation/inference uses standard CFG in both clean and Self-Forcing,
-    # so this directly matches Self-Forcing guidance_scale=3.0.
-    guidance_scale: 3.0

 pipeline:
  flow_shift: 5
@@ -1,32 +1,119 @@
 # World-Model: Matrix-Game 2.0 I2V

-Three training scenarios for the Matrix-Game 2.0 I2V world model on the
-new YAML-driven trainer (`fastvideo/train/entrypoint/train.py`).
+Training scenarios for the Matrix-Game 2.0 I2V world model on Solaris (Minecraft)
+data and Zelda data, using the new YAML-driven trainer
+(`fastvideo/train/entrypoint/train.py`).
+
+## Solaris Configs

 | Config | Method | Student | Notes |
 |---|---|---|---|
-| `finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Multi-step SFT from `mg_bidirectional_Solaris`. |
-| `dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Diffusion-Forcing SFT with chunkwise timesteps. |
-| `self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | DMD/Self-Forcing distillation; teacher = bidirectional, critic = bidirectional. |
+| `solaris/finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Multi-step SFT from `mg_bidirectional_Solaris`. |
+| `solaris/dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Diffusion-Forcing SFT with chunkwise timesteps. |
+| `solaris/self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | Matrix-Game 2.0 DMD/Self-Forcing distillation; teacher = bidirectional, critic = bidirectional. |
+
+## Zelda Configs
+
+| Config | Method | Student | Notes |
+|---|---|---|---|
+| `zelda/finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Zelda bidirectional I2V finetuning from `FastVideo/Matrix-Game-2.0-Base-Diffusers`. Uses 33-frame clips and Zelda validation with action overlays. |
+| `zelda/dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Zelda causal Diffusion-Forcing SFT from `mignonjia/mg_bidirectional_zelda`. Uses the same Zelda data, resolution, optimizer, and validation defaults as the Zelda finetune config. |
+| `zelda/self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | Zelda DMD/Self-Forcing distillation; student init = `mignonjia/mg_causal_zelda`, teacher = bidirectional (`mignonjia/mg_bidirectional_zelda`), critic = bidirectional. |
+| `zelda/streaming_long_tuning_causal_i2v.yaml` | `StreamingLongTuningMethod` | `MatrixGame2CausalModel` | LongLive-style streaming long tuning from the 1k-step Zelda self-forcing checkpoint. |
+
+Zelda world-model distillation is a two-run workflow: first run
+`zelda/self_forcing_causal_i2v.yaml` to train or load the 1k-step
+self-forcing checkpoint (`mignonjia/mg_sf_distilled_zelda_1k_steps`), then run
+`zelda/streaming_long_tuning_causal_i2v.yaml` for the 3k-step streaming
+long-tuning stage. The long-tuning YAML starts from that 1k-step checkpoint; it
+does not run the short self-forcing stage inside the same config.
+
+## Zelda Training Data
+
+The Zelda training configs use `data/zeldam2-clean` as a suggested local path.
+Download the dataset from Hugging Face before running those configs:
+
+```bash
+python scripts/huggingface/download_hf.py \
+    --repo_id mignonjia/zeldam2-clean \
+    --local_dir data/zeldam2-clean \
+    --repo_type dataset
+```
+
+You can store the dataset elsewhere; update `training.data.data_path` in the
+YAML to point at that location.
+
+## Multi3D Training Data
+
+`zelda/finetune_i2v.yaml` includes an optional, commented-out Multi3D entry.
+Enable it only when you want to mix Zelda with multi-game data from
+`data/multi3d_games`. You can store this dataset anywhere; before enabling it,
+update the matching commented `training.data.data_path` key in the YAML to the
+correct location.
+
+To mix datasets in a training YAML, set `training.data.data_path` to a
+path-to-repeat-count mapping. For example, `zelda/finetune_i2v.yaml` can use
+`data/zeldam2-clean: 1` and `# data/multi3d_games: 10`; uncommenting the
+Multi3D entry repeats the multi-game parquet list ten times before training
+samples are shuffled.
+
+## World Model Validation Data
+
+The Zelda validation configs expect a small public validation bundle under
+`data/zelda_validation_data`.
+
+Download it from Hugging Face before running the Zelda scenarios:
+
+```bash
+python scripts/huggingface/download_hf.py \
+    --repo_id mignonjia/zelda_validation_data \
+    --local_dir data/zelda_validation_data \
+    --repo_type dataset
+```
+
+The bundle contains `validation_zelda.json`, `images/`, and `actions/`.
+The Zelda configs point
+`callbacks.validation.dataset_file` at
+`data/zelda_validation_data/validation_zelda.json`.

 ## Usage

+### Solaris
+
 ```bash
 bash examples/train/run.sh \
-    examples/train/scenario/worldmodel/finetune_i2v.yaml
+    examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml

 bash examples/train/run.sh \
-    examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml
+    examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml

 bash examples/train/run.sh \
-    examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml
+    examples/train/scenario/worldmodel/solaris/self_forcing_causal_i2v.yaml
+```
+
+### Zelda
+
+```bash
+# Finetuning / DFSFT
+bash examples/train/run.sh \
+    examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml
+
+bash examples/train/run.sh \
+    examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml
+
+# Distillation / long tuning
+bash examples/train/run.sh \
+    examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml
+
+bash examples/train/run.sh \
+    examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml
 ```

 Override any field on the command line:

 ```bash
 bash examples/train/run.sh \
-    examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml \
+    examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml \
    --training.distributed.num_gpus 8 \
    --training.optimizer.learning_rate 1e-5
 ```
@@ -97,4 +97,4 @@ callbacks:
    guidance_scale: 6.0

 pipeline:
-  flow_shift: 5
+  flow_shift: 5
@@ -0,0 +1,94 @@
+# Diffusion-Forcing SFT: Zelda world model I2V Causal
+
+models:
+  student:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
+  chunk_size: 3
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path:
+      data/zeldam2-clean: 1
+    dataloader_num_workers: 1
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 9
+    num_height: 480
+    num_width: 832
+    num_frames: 33
+
+  optimizer:
+    learning_rate: 2.0e-5
+    betas: [0.9, 0.95]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 60000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/matrixgame_finetune/checkpoints/zelda_causal_dfsft
+    training_state_checkpointing_steps: 5000
+    checkpoints_total_limit: 3
+
+  tracker:
+    entity: hapo-exp
+    project_name: mg_1.3b_zelda
+    run_name: zelda_causal_dfsft
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
+    dataset_file: data/zelda_validation_data/validation_zelda.json
+    every_steps: 200
+    sampling_steps: [40]
+    sampling_timesteps: [1000, 975, 950, 925, 900, 875, 850, 825, 800, 775,
+                         750, 725, 700, 675, 650, 625, 600, 575, 550, 525,
+                         500, 475, 450, 425, 400, 375, 350, 325, 300, 275,
+                         250, 225, 200, 175, 150, 125, 100, 75, 50, 25]
+    num_frames: 33
+    overlay_actions: true
+    guidance_scale: 6.0
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - vbench.temporal_flickering
+        - vbench.motion_smoothness
+        - vbench.subject_consistency
+        - vbench.background_consistency
+        - vbench.dynamic_degree
+        - optical_flow.synthetic_optical_flow
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+pipeline:
+  flow_shift: 5
+  dit_config:
+    local_attn_size: 6
+    sink_size: 1
@@ -0,0 +1,88 @@
+# Matrix-Game 2.0 Zelda + multi-game I2V finetune.
+
+models:
+  student:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: FastVideo/Matrix-Game-2.0-Base-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path:
+      data/zeldam2-clean: 1
+      # data/multi3d_games: 10
+    dataloader_num_workers: 1
+    train_batch_size: 1
+    training_cfg_rate: 0.0  # unused for MatrixGame2 I2V; no text_embedding CFG dropout
+    seed: 42
+    num_latent_t: 9
+    num_height: 480
+    num_width: 832
+    num_frames: 33
+
+  optimizer:
+    learning_rate: 2.0e-5
+    betas: [0.9, 0.95]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 60000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/matrixgame_finetune/checkpoints/zelda_with_mg_init
+    training_state_checkpointing_steps: 5000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: mg_1.3b_zelda
+    run_name: zelda_with_mg_init
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_i2v_pipeline.MatrixGame2I2VPipeline
+    dataset_file: data/zelda_validation_data/validation_zelda.json
+    every_steps: 200
+    sampling_steps: [40]
+    num_frames: 33
+    overlay_actions: true
+    guidance_scale: 6.0
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - vbench.temporal_flickering
+        - vbench.motion_smoothness
+        - vbench.subject_consistency
+        - vbench.background_consistency
+        - vbench.dynamic_degree
+        - optical_flow.synthetic_optical_flow
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+pipeline:
+  flow_shift: 5
@@ -0,0 +1,120 @@
+# Self-forcing distillation: Zelda world model I2V Causal
+
+models:
+  student:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
+    init_from: mignonjia/mg_causal_zelda
+    trainable: true
+  teacher:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+  warp_denoising_step: true
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+
+  chunk_size: 3
+  student_sample_type: sde
+  same_step_across_blocks: true
+  last_step_only: false
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  # Critic optimizer
+  fake_score_learning_rate: 3.0e-7
+  fake_score_betas: [0.9, 0.95]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/zeldam2-clean
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1001
+    num_latent_t: 9
+    num_height: 480
+    num_width: 832
+    num_frames: 33
+
+  optimizer:
+    learning_rate: 3.0e-6
+    betas: [0.9, 0.95]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/zelda_causal_self_forcing
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 1
+
+  tracker:
+    project_name: wangame_sf
+    run_name: mg2_self_forcing_9_latents
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+callbacks:
+  ema:
+    decay: 0.99
+    start_iter: 200
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
+    dataset_file: data/zelda_validation_data/validation_zelda.json
+    every_steps: 100
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    num_frames: 153
+    overlay_actions: true
+    keyboard_value_scale: 1.0
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - vbench.temporal_flickering
+        - vbench.motion_smoothness
+        - vbench.subject_consistency
+        - vbench.background_consistency
+        - vbench.dynamic_degree
+        - optical_flow.synthetic_optical_flow
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+pipeline:
+  flow_shift: 5
+  dit_config:
+    local_attn_size: 6
+    sink_size: 1
@@ -0,0 +1,140 @@
+# MatrixGame2 I2V LongLive-style streaming distillation: Zelda world model I2V Causal
+# Student init from self forcing after 1k steps
+
+models:
+  student:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
+    init_from: mignonjia/mg_sf_distilled_zelda_1k_steps 
+    trainable: true
+    # transformer_override_safetensor: outputs/matrixgame_dmd/checkpoints/mg_zelda_sf_m2/checkpoint-1000_weight_only/ema/generator_ema.safetensors
+  teacher:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+  warp_denoising_step: true
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+
+  chunk_size: 3
+  student_sample_type: sde
+  same_step_across_blocks: true
+  last_step_only: false
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  streaming_training: true
+  streaming_chunk_size: 9
+  streaming_max_length: 39
+  streaming_fixed_overlap_latents: 3
+  streaming_reencode_overlap_anchor: true
+  streaming_anchor_inject_k: 1
+  streaming_require_full_blocks: true
+
+  multi_phased_distill_schedule:
+    - stage: streaming_long
+      start_step: 0
+      end_step: 3000
+      num_latent_t: 39
+      streaming_training: true
+      streaming_chunk_size: 9
+      streaming_max_length: 39
+      streaming_fixed_overlap_latents: 3
+
+  # Critic optimizer
+  fake_score_learning_rate: 3.0e-7
+  fake_score_betas: [0.9, 0.95]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/zeldam2-clean
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1001
+    num_latent_t: 39
+    num_height: 480
+    num_width: 832
+    num_frames: 153
+
+  optimizer:
+    learning_rate: 3.0e-6
+    betas: [0.9, 0.95]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 3000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/zelda_causal_long_tuning
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 2
+
+  tracker:
+    project_name: wangame_sf
+    run_name: mg2_39only_streaming_long
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+callbacks:
+  ema:
+    decay: 0.99
+    start_iter: 200
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
+    dataset_file: data/zelda_validation_data/validation_zelda.json
+    every_steps: 100
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    num_frames: 153
+    overlay_actions: true
+    keyboard_value_scale: 1.0
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - vbench.temporal_flickering
+        - vbench.motion_smoothness
+        - vbench.subject_consistency
+        - vbench.background_consistency
+        - vbench.dynamic_degree
+        - optical_flow.synthetic_optical_flow
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+pipeline:
+  flow_shift: 5
+  dit_config:
+    local_attn_size: 6
+    sink_size: 1
@@ -0,0 +1,37 @@
+---
+license: apache-2.0
+---
+
+# FastVideo Kernel
+
+Hub-compatible FastVideo CUDA kernels packaged with Hugging Face `kernel-builder`.
+
+## Usage
+
+```python
+from kernels import get_kernel
+
+fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1)
+```
+
+The module exposes:
+
+- `sta_fwd`
+- `block_sparse_fwd`
+- `block_sparse_bwd`
+- `rms_norm`
+- `layer_norm`
+- `int8_quant`
+- `int8_gemm`
+
+These APIs are thin wrappers around the native FastVideo kernels. Higher-level
+fallback logic remains in the regular `fastvideo-kernel` Python package.
+
+## Requirements
+
+This first Hub build targets CUDA Hopper (`sm_90a`) because the native attention
+kernels rely on Hopper/ThunderKittens features.
+
+## Source
+
+Source repository: https://github.com/hao-ai-lab/FastVideo/tree/main/fastvideo-kernel
@@ -40,6 +40,25 @@ Notes:

 ## Installation

+### Hugging Face Kernel Hub
+
+FastVideo kernels are built and published with the Hugging Face
+[`kernels`](https://huggingface.co/docs/kernels/index) project. Downstream code
+loads the precompiled kernels from
+[`hao-ai-lab/fastvideo-kernel`](https://huggingface.co/hao-ai-lab/fastvideo-kernel)
+without building locally:
+
+```python
+from kernels import get_kernel
+
+fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1)
+```
+
+The Hub package exposes the native CUDA entrypoints (`sta_fwd`,
+`block_sparse_fwd`, `block_sparse_bwd`, `rms_norm`, `layer_norm`, `int8_quant`,
+and `int8_gemm`). The full FastVideo Python package still provides the
+high-level fallback logic used during local development.
+
 ### Standard Installation (Local Development)
 This will automatically detect your GPU architecture. If an NVIDIA Hopper (H100/sm_90a) GPU is detected, ThunderKittens kernels will be enabled. Otherwise, they will be skipped, and the package will use Triton fallbacks at runtime.

@@ -0,0 +1,52 @@
+[general]
+name = "fastvideo_kernel"
+version = 1
+license = "Apache-2.0"
+backends = ["cuda"]
+upstream = "https://github.com/hao-ai-lab/FastVideo.git"
+source = "https://github.com/hao-ai-lab/FastVideo.git"
+
+[general.hub]
+repo-id = "hao-ai-lab/fastvideo-kernel"
+
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+
+[kernel.fastvideo_kernel]
+backend = "cuda"
+depends = ["torch"]
+include = [
+  "csrc",
+  "csrc/turbodiffusion",
+  "include/cutlass/include",
+  "include/tk/include",
+  "include/tk/prototype",
+]
+src = [
+  "csrc/attention/st_attn_h100.cu",
+  "csrc/attention/block_sparse_h100.cu",
+  "csrc/turbodiffusion/gemm/gemm.cu",
+  "csrc/turbodiffusion/norm/rmsnorm.cu",
+  "csrc/turbodiffusion/norm/layernorm.cu",
+  "csrc/turbodiffusion/quant/quant.cu",
+]
+# Deliberately Hopper-only (sm_90a): the TK attention kernels here require
+# Hopper, unlike the multi-arch PyPI wheels (see .github/workflows/publish-kernel.yml).
+cuda-capabilities = ["9.0a"]
+cuda-flags = [
+  "-DNDEBUG",
+  "-O3",
+  "-std=c++20",
+  "--use_fast_math",
+  "--expt-extended-lambda",
+  "--expt-relaxed-constexpr",
+  "-Xcompiler=-fno-strict-aliasing",
+  "-Xcompiler=-fPIC",
+  "-DTORCH_COMPILE",
+  "-DKITTENS_HOPPER",
+  "-DFASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION",
+  "-Xptxas=--warn-on-spills",
+]
@@ -71,11 +71,12 @@ void int8_gemm(

 }

+#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION
 void register_gemm(pybind11::module_ &m) {
    m.def("gemm_cuda", &int8_gemm);
 }
+#endif


  

-
@@ -9,10 +9,10 @@
 #include "norm/layernorm.hpp"

 auto layer_norm(
-  at::Tensor const Input, 
+  at::Tensor const Input,
  float eps,
-  std::optional<at::Tensor const> W,
-  std::optional<at::Tensor const> const B,
+  std::optional<at::Tensor> W,
+  std::optional<at::Tensor> const B,
  std::optional<at::Tensor> Output
 ) {
  int64_t const m = Input.size(0);
@@ -93,7 +93,8 @@ auto layer_norm(
  return Output;
 }

+#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION
 void register_layer_norm(pybind11::module_ &m) {
    m.def("layer_norm_cuda", &layer_norm);
 }
-
+#endif
@@ -75,6 +75,8 @@ auto rms_norm(
  return Output;
 }

+#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION
 void register_rms_norm(pybind11::module_ &m) {
    m.def("rms_norm_cuda", &rms_norm);
 }
+#endif
@@ -70,6 +70,8 @@ auto quant(
  return std::make_tuple(Output, Output_S);
 }

+#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION
 void register_quant(pybind11::module_ &m) {
    m.def("quant_cuda", &quant);
 }
+#endif
@@ -0,0 +1,11 @@
+import torch
+from kernels import get_kernel
+
+
+fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1)
+
+x = torch.randn((16, 1024), device="cuda", dtype=torch.float16)
+weight = torch.ones((1024,), device="cuda", dtype=torch.float16)
+y = fastvideo_kernel.rms_norm(x, 1e-6, weight)
+
+print(y.shape)
@@ -0,0 +1,18 @@
+{
+  description = "Flake for FastVideo Kernel Hub builds";
+
+  inputs = {
+    self.submodules = true;
+    kernel-builder.url = "github:huggingface/kernels";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}
@@ -0,0 +1,122 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+
+def sta_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    kernel_t_size: int,
+    kernel_h_size: int,
+    kernel_w_size: int,
+    text_length: int,
+    process_text: bool,
+    has_text: bool,
+    kernel_aspect_ratio_flag: int,
+) -> torch.Tensor:
+    return ops.sta_fwd(
+        q,
+        k,
+        v,
+        out,
+        kernel_t_size,
+        kernel_h_size,
+        kernel_w_size,
+        text_length,
+        process_text,
+        has_text,
+        kernel_aspect_ratio_flag,
+    )
+
+
+def block_sparse_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_block_sparse_index: torch.Tensor,
+    q2k_block_sparse_num: torch.Tensor,
+    kv_block_size: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    out, lse = ops.block_sparse_fwd(
+        q,
+        k,
+        v,
+        q2k_block_sparse_index,
+        q2k_block_sparse_num,
+        kv_block_size,
+    )
+    return out, lse
+
+
+def block_sparse_bwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    l_vec: torch.Tensor,
+    out_grad: torch.Tensor,
+    k2q_block_sparse_index: torch.Tensor,
+    k2q_block_sparse_num: torch.Tensor,
+    kv_block_size: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    q_grad, k_grad, v_grad = ops.block_sparse_bwd(
+        q,
+        k,
+        v,
+        out,
+        l_vec,
+        out_grad,
+        k2q_block_sparse_index,
+        k2q_block_sparse_num,
+        kv_block_size,
+    )
+    return q_grad, k_grad, v_grad
+
+
+def rms_norm(
+    input: torch.Tensor,
+    eps: float,
+    weight: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    return ops.rms_norm_cuda(input, eps, weight, output)
+
+
+def layer_norm(
+    input: torch.Tensor,
+    eps: float,
+    weight: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    return ops.layer_norm_cuda(input, eps, weight, bias, output)
+
+
+def int8_quant(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    quantized, scale = ops.quant_cuda(input)
+    return quantized, scale
+
+
+def int8_gemm(
+    a: torch.Tensor,
+    a_scale: torch.Tensor,
+    b: torch.Tensor,
+    b_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    return ops.gemm_cuda(a, a_scale, b, b_scale, out)
+
+
+__all__ = [
+    "sta_fwd",
+    "block_sparse_fwd",
+    "block_sparse_bwd",
+    "rms_norm",
+    "layer_norm",
+    "int8_quant",
+    "int8_gemm",
+]
@@ -0,0 +1,99 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+torch::Tensor rms_norm_cuda(
+    torch::Tensor const& input,
+    double eps,
+    std::optional<torch::Tensor> const& weight,
+    std::optional<torch::Tensor> const& output) {
+  std::optional<at::Tensor> mutable_output = output;
+  auto result = rms_norm(input, static_cast<float>(eps), weight, mutable_output);
+  TORCH_CHECK(result.has_value(), "rms_norm_cuda did not produce an output tensor");
+  return result.value();
+}
+
+torch::Tensor layer_norm_cuda(
+    torch::Tensor const& input,
+    double eps,
+    std::optional<torch::Tensor> const& weight,
+    std::optional<torch::Tensor> const& bias,
+    std::optional<torch::Tensor> const& output) {
+  auto result = layer_norm(input, static_cast<float>(eps), weight, bias, output);
+  TORCH_CHECK(result.has_value(), "layer_norm_cuda did not produce an output tensor");
+  return result.value();
+}
+
+std::tuple<torch::Tensor, torch::Tensor> quant_cuda(torch::Tensor const& input) {
+  std::optional<torch::Tensor> output;
+  std::optional<torch::Tensor> output_scale;
+  auto result = quant(input, output, output_scale);
+  TORCH_CHECK(std::get<0>(result).has_value(), "quant_cuda did not produce an int8 tensor");
+  TORCH_CHECK(std::get<1>(result).has_value(), "quant_cuda did not produce a scale tensor");
+  return std::make_tuple(std::get<0>(result).value(), std::get<1>(result).value());
+}
+
+torch::Tensor gemm_cuda(
+    torch::Tensor const& a,
+    torch::Tensor const& a_scale,
+    torch::Tensor const& b,
+    torch::Tensor const& b_scale,
+    torch::Tensor c) {
+  int8_gemm(a, a_scale, b, b_scale, c);
+  return c;
+}
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def(
+      // `out` is written in place and returned, so the return must be declared
+      // as the same alias set — otherwise torch.compile can reorder/eliminate
+      // the mutation and silently produce wrong results.
+      "sta_fwd(Tensor q, Tensor k, Tensor v, Tensor(a!) out, int kernel_t_size, int kernel_h_size, "
+      "int kernel_w_size, int text_length, bool process_text, bool has_text, "
+      "int kernel_aspect_ratio_flag) -> Tensor(a!)");
+#if defined(CUDA_KERNEL)
+  ops.impl("sta_fwd", torch::kCUDA, &sta_forward);
+#endif
+
+  ops.def(
+      "block_sparse_fwd(Tensor q, Tensor k, Tensor v, Tensor q2k_block_sparse_index, "
+      "Tensor q2k_block_sparse_num, Tensor kv_block_size) -> Tensor[]");
+#if defined(CUDA_KERNEL)
+  ops.impl("block_sparse_fwd", torch::kCUDA, &block_sparse_attention_forward);
+#endif
+
+  ops.def(
+      "block_sparse_bwd(Tensor q, Tensor k, Tensor v, Tensor out, Tensor l_vec, Tensor out_grad, "
+      "Tensor k2q_block_sparse_index, Tensor k2q_block_sparse_num, Tensor kv_block_size) -> Tensor[]");
+#if defined(CUDA_KERNEL)
+  ops.impl("block_sparse_bwd", torch::kCUDA, &block_sparse_attention_backward);
+#endif
+
+  // When `output` is provided it is written in place and aliased by the
+  // return value; the schema must say so for torch.compile correctness.
+  ops.def("rms_norm_cuda(Tensor input, float eps, Tensor? weight=None, Tensor(a!)? output=None) -> Tensor(a!)");
+#if defined(CUDA_KERNEL)
+  ops.impl("rms_norm_cuda", torch::kCUDA, &rms_norm_cuda);
+#endif
+
+  ops.def(
+      "layer_norm_cuda(Tensor input, float eps, Tensor? weight=None, Tensor? bias=None, "
+      "Tensor(a!)? output=None) -> Tensor(a!)");
+#if defined(CUDA_KERNEL)
+  ops.impl("layer_norm_cuda", torch::kCUDA, &layer_norm_cuda);
+#endif
+
+  ops.def("quant_cuda(Tensor input) -> (Tensor, Tensor)");
+#if defined(CUDA_KERNEL)
+  ops.impl("quant_cuda", torch::kCUDA, &quant_cuda);
+#endif
+
+  // `c` is the accumulator written in place and returned (same alias set).
+  ops.def("gemm_cuda(Tensor a, Tensor a_scale, Tensor b, Tensor b_scale, Tensor(a!) c) -> Tensor(a!)");
+#if defined(CUDA_KERNEL)
+  ops.impl("gemm_cuda", torch::kCUDA, &gemm_cuda);
+#endif
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
@@ -0,0 +1,86 @@
+#pragma once
+
+#include <optional>
+#include <tuple>
+#include <vector>
+
+#include <torch/torch.h>
+
+torch::Tensor sta_forward(
+    torch::Tensor q,
+    torch::Tensor k,
+    torch::Tensor v,
+    torch::Tensor o,
+    int kernel_t_size,
+    int kernel_h_size,
+    int kernel_w_size,
+    int text_length,
+    bool process_text,
+    bool has_text,
+    int kernel_aspect_ratio_flag);
+
+std::vector<torch::Tensor> block_sparse_attention_forward(
+    torch::Tensor q,
+    torch::Tensor k,
+    torch::Tensor v,
+    torch::Tensor q2k_block_sparse_index,
+    torch::Tensor q2k_block_sparse_num,
+    torch::Tensor kv_block_size);
+
+std::vector<torch::Tensor> block_sparse_attention_backward(
+    torch::Tensor q,
+    torch::Tensor k,
+    torch::Tensor v,
+    torch::Tensor o,
+    torch::Tensor l_vec,
+    torch::Tensor og,
+    torch::Tensor k2q_block_sparse_index,
+    torch::Tensor k2q_block_sparse_num,
+    torch::Tensor kv_block_size);
+
+std::optional<at::Tensor> rms_norm(
+    at::Tensor const& input,
+    float eps,
+    std::optional<at::Tensor> const& weight,
+    std::optional<at::Tensor>& output);
+
+std::optional<at::Tensor> layer_norm(
+    at::Tensor const input,
+    float eps,
+    std::optional<at::Tensor> weight,
+    std::optional<at::Tensor> const bias,
+    std::optional<at::Tensor> output);
+
+std::tuple<std::optional<torch::Tensor>, std::optional<torch::Tensor>> quant(
+    torch::Tensor const& input,
+    std::optional<torch::Tensor>& output,
+    std::optional<torch::Tensor>& output_scale);
+
+void int8_gemm(
+    at::Tensor const& a,
+    at::Tensor const& a_scale,
+    at::Tensor const& b,
+    at::Tensor const& b_scale,
+    torch::Tensor& c);
+
+torch::Tensor rms_norm_cuda(
+    torch::Tensor const& input,
+    double eps,
+    std::optional<torch::Tensor> const& weight,
+    std::optional<torch::Tensor> const& output);
+
+torch::Tensor layer_norm_cuda(
+    torch::Tensor const& input,
+    double eps,
+    std::optional<torch::Tensor> const& weight,
+    std::optional<torch::Tensor> const& bias,
+    std::optional<torch::Tensor> const& output);
+
+std::tuple<torch::Tensor, torch::Tensor> quant_cuda(torch::Tensor const& input);
+
+torch::Tensor gemm_cuda(
+    torch::Tensor const& a,
+    torch::Tensor const& a_scale,
+    torch::Tensor const& b,
+    torch::Tensor const& b_scale,
+    torch::Tensor c);
@@ -2,6 +2,7 @@
 import os
 import pickle
 import random
+from collections.abc import Sequence
 from typing import Any

 import pyarrow as pa
@@ -94,8 +95,56 @@ class DP_SP_BatchSampler(Sampler[list[int]]):
        return len(self.sp_group_local_indices) // self.batch_size


-def get_parquet_files_and_length(path: str):
-    dataset_root = os.path.realpath(os.path.expanduser(path))
+def _parse_data_path_specs(path: str | Sequence[str] | dict[str, int]) -> list[tuple[str, int]]:
+    """Parse one or more dataset roots with old-framework repeat counts."""
+    if isinstance(path, dict):
+        return [
+            (str(root), int(repeat))
+            for root, repeat in path.items()
+            if int(repeat) > 0
+        ]
+    if isinstance(path, Sequence) and not isinstance(path, str):
+        return [(str(root), 1) for root in path]
+
+    specs: list[tuple[str, int]] = []
+    for part in str(path).split(","):
+        part = part.strip()
+        if not part:
+            continue
+        if ":" in part:
+            dir_path, count_str = part.rsplit(":", 1)
+            count = int(count_str)
+        else:
+            dir_path, count = part, 1
+        if count > 0:
+            specs.append((dir_path.strip(), count))
+    return specs
+
+
+def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int]):
+    specs = _parse_data_path_specs(path)
+    if len(specs) != 1 or specs[0][1] != 1:
+        all_file_names: list[str] = []
+        all_lengths: list[int] = []
+        for root, repeat in specs:
+            file_names, lengths = get_parquet_files_and_length(str(root))
+            for _ in range(repeat):
+                all_file_names.extend(file_names)
+                all_lengths.extend(lengths)
+        if not all_file_names:
+            raise FileNotFoundError(
+                "No parquet files found under dataset paths: "
+                f"{path}. "
+                "Please verify these paths point to preprocessed parquet data."
+            )
+        file_lengths = sorted(
+            zip(all_file_names, all_lengths, strict=True),
+            key=lambda x: x[0],
+        )
+        file_names_sorted, lengths_sorted = zip(*file_lengths, strict=True)
+        return file_names_sorted, lengths_sorted
+
+    dataset_root = os.path.realpath(os.path.expanduser(specs[0][0]))
    # Check if cached info exists
    cache_dir = os.path.join(dataset_root, "map_style_cache")
    cache_file = os.path.join(cache_dir, "file_info.pkl")
@@ -268,7 +317,7 @@ class LatentsParquetMapStyleDataset(Dataset):

    def __init__(
        self,
-        path: str,
+        path: str | Sequence[str] | dict[str, int],
        batch_size: int,
        parquet_schema: pa.Schema,
        cfg_rate: float = 0.0,
@@ -358,7 +407,6 @@ class LatentsParquetMapStyleDataset(Dataset):
    def __len__(self):
        return sum(self.lengths)

-
 # ────────────────────────────────────────────────────────────────────────────
 # 3.  Loader helper – everything else stays just like your original trainer
 # ────────────────────────────────────────────────────────────────────────────
@@ -122,6 +122,44 @@ to run a subset of the Evaluator's registered metrics on this batch
 (useful for scoring different corpora with different metric subsets in
 successive calls without burning model loads).

+### Training-time validation metrics
+
+The modular trainer can score validation videos during training through
+`callbacks.validation.metrics`. Each distributed rank that writes local
+validation videos also runs its local metrics on `cuda:<local_rank>`;
+rank 0 only merges scalar summaries and logs artifacts.
+
+```yaml
+callbacks:
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
+    dataset_file: examples/distill/MatrixGame2.0/validation.json
+    sampling_steps: [4]
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - optical_flow.synthetic_optical_flow
+        - common.fvd
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+      # `synthetic_optical_flow` also needs actions in the validation
+      # manifest (`action_path`) plus a repo-local calibration file.
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+```
+
+Metric summaries are written under
+`<output_dir>/eval/step_<step>/inference_steps_<n>_rank_<rank>.json` and
+scalar means are logged with the `metrics/validation/...` prefix.
+Install `fastvideo[eval]` for optical-flow dependencies such as
+`ptlflow`; VBench metrics use the pinned submodule under
+`fastvideo/third_party/eval/vbench`. FVD uses `ref_video` entries from
+the validation manifest when present, or the standard
+`FASTVIDEO_FVD_REF_FEATURES` / eval-cache reference feature path.
+
 ### CLI

 ```bash
@@ -350,5 +388,3 @@ sweep several baselines into a table, see

 - **MIND** metrics. Depend on a separate `vipe` upstream submodule.
 - **VBench-2.0**. Sibling vbench2 package; needs its own port.
- **Training-time eval callback** (`EvalCallback`) and the
-  `RolloutEvaluator` helper.
@@ -63,7 +63,7 @@ class ThirdPersonCalibration:
    beta_fwd: float
    beta_strafe: float
    focal_length: float
-    r_z: float
+    r_z: float = 0.0
    r_y: float = 0.0
    init_pitch: float = 0.0
    notes: str = ""
@@ -261,6 +261,22 @@ def _mm_fp4(
    )


+def _coerce_fp4_input_dtype(x: torch.Tensor) -> torch.Tensor:
+    """Coerce an activation to a dtype the FP4 linear accepts.
+
+    The pre-attention norm can emit fp32 (e.g. in eager mode, without the
+    torch.compile fusion that keeps it bf16). The FP4 linear emits bf16
+    regardless (see _mm_fp4 out dtype), so cast fp32 -> bf16 rather than
+    failing, matching the sibling fastvideo/layers/fp4linear.py. Non-floating
+    inputs (e.g. int/bool) are a genuine error and are rejected fast.
+    """
+    if not x.is_floating_point():
+        raise TypeError(f"fp4 linear expects floating-point inputs, got {x.dtype}")
+    if x.dtype not in (torch.bfloat16, torch.float16):
+        x = x.to(torch.bfloat16)
+    return x
+
+
 class NVFP4QuantizeMethod(QuantizeMethodBase):

    def __init__(self, layer_prefix: str = ""):
@@ -285,8 +301,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):

    def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        SfLayout, _, _ = _require_flashinfer()
-        assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
-            f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
+        x = _coerce_fp4_input_dtype(x)
        x_2d = x.view(-1, x.shape[-1])
        x_fp4, x_scale = _nvfp4_quantize(
            x_2d,
@@ -332,8 +347,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
            if x_scale.dim() > 2:
                x_scale = x_scale.view(-1, x_scale.shape[-1])
        else:
-            assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
-                f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
+            x = _coerce_fp4_input_dtype(x)
            x = x.view(-1, x.shape[-1])
            x_global_sf = self.x_global_sf
            x_fp4, x_scale = _nvfp4_quantize(
@@ -665,7 +665,8 @@ class CausalWanTransformer3DModel(BaseDiT):
            else:
                causal_kwargs = {
                    "kv_cache": kv_cache[block_index],
-                    "crossattn_cache": crossattn_cache[block_index],
+                    "crossattn_cache": (crossattn_cache[block_index]
+                                        if crossattn_cache is not None else None),
                    "current_start": current_start,
                    "cache_start": cache_start,
                    "block_mask": self.block_mask,
@@ -16,6 +16,8 @@ from fastvideo.layers.rotary_embedding import (
 )
 from fastvideo.platforms import AttentionBackendEnum

+from .utils import retain_kv_with_sink
+

 DISABLE_COMPILE = False
 flex_attention = torch.compile(
@@ -100,6 +102,7 @@ def _update_kv_cache_and_attend(
    start_frame: int,
    num_frame_per_block: int,
    local_attn_size: int,
+    sink_size: int,
    *,
    use_k_for_num_tokens: bool = False,
    store_first_only: bool = False,
@@ -133,7 +136,6 @@ def _update_kv_cache_and_attend(
        k.shape[1] if use_k_for_num_tokens else q.shape[1]
    ) == num_frame_per_block

-    sink_size = 0
    max_attention_size = local_attn_size
    sink_tokens = sink_size * 1
    kv_cache_size = kv_cache["k"].shape[1]
@@ -206,16 +208,22 @@ def _update_kv_cache_and_attend(

    # Store new k, v in cache
    if store_first_only:
-        kv_cache["k"][:, local_start_index:local_end_index] = k[:1]
-        kv_cache["v"][:, local_start_index:local_end_index] = v[:1]
+        kv_cache["k"][:, local_start_index:local_end_index] = k[
+            : kv_cache["k"].shape[0]
+        ]
+        kv_cache["v"][:, local_start_index:local_end_index] = v[
+            : kv_cache["v"].shape[0]
+        ]
    else:
        kv_cache["k"][:, local_start_index:local_end_index] = k
        kv_cache["v"][:, local_start_index:local_end_index] = v

-    # Retrieve from cache and perform attention
-    cache_start = max(0, local_end_index - max_attention_size)
-    cached_k = kv_cache["k"][:, cache_start:local_end_index]
-    cached_v = kv_cache["v"][:, cache_start:local_end_index]
+    cached_k, cached_v = retain_kv_with_sink(
+        kv_cache["k"][:, :local_end_index],
+        kv_cache["v"][:, :local_end_index],
+        min(local_end_index, max_attention_size),
+        sink_tokens,
+    )

    if repeat_factor is not None:
        cached_k = cached_k.repeat(repeat_factor, 1, 1, 1)
@@ -261,6 +269,7 @@ class ActionModule(nn.Module):
        enable_mouse=True,
        enable_keyboard=True,
        local_attn_size=6,
+        sink_size=0,
        blocks: list | None = None,
    ):
        super().__init__()
@@ -274,6 +283,7 @@ class ActionModule(nn.Module):
        )
        blocks = blocks if blocks is not None else []
        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
        self.enable_mouse = enable_mouse
        self.enable_keyboard = enable_keyboard

@@ -565,6 +575,7 @@ class ActionModule(nn.Module):
                    start_frame,
                    num_frame_per_block,
                    self.local_attn_size,
+                    self.sink_size,
                    use_k_for_num_tokens=False,
                )
        else:
@@ -689,6 +700,7 @@ class ActionModule(nn.Module):
                        start_frame,
                        num_frame_per_block,
                        self.local_attn_size,
+                        self.sink_size,
                        use_k_for_num_tokens=True,
                        store_first_only=True,
                        repeat_factor=S,
@@ -719,6 +731,7 @@ class ActionModule(nn.Module):
                        start_frame,
                        num_frame_per_block,
                        self.local_attn_size,
+                        self.sink_size,
                        use_k_for_num_tokens=True,
                    )
            else:
@@ -46,6 +46,7 @@ from fastvideo.platforms import AttentionBackendEnum, current_platform

 from .action_module import ActionModule
 from .model import MatrixGame2CrossAttention
+from .utils import retain_kv_with_sink

 logger = init_logger(__name__)

@@ -345,9 +346,12 @@ class CausalMatrixGame2SelfAttention(nn.Module):
                kv_cache["k"][:, local_start_index:local_end_index] = stored_key
                kv_cache["v"][:, local_start_index:local_end_index] = v

-            kv_start = max(0, local_end_index - max_attention_size)
-            k_for_attn = kv_cache["k"][:, kv_start:local_end_index]
-            v_for_attn = kv_cache["v"][:, kv_start:local_end_index]
+            k_for_attn, v_for_attn = retain_kv_with_sink(
+                kv_cache["k"][:, :local_end_index],
+                kv_cache["v"][:, :local_end_index],
+                min(local_end_index, max_attention_size),
+                sink_tokens,
+            )

            if relativistic:
                window_len, query_lo, _ = relativistic_window_offsets(
@@ -418,6 +422,7 @@ class CausalMatrixGame2TransformerBlock(nn.Module):
        self.hidden_dim = dim
        self.num_attention_heads = num_heads
        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
        dim_head = dim // num_heads
        if qk_norm == "rms_norm":
            self.norm_q = RMSNorm(dim_head, eps=eps)
@@ -474,6 +479,7 @@ class CausalMatrixGame2TransformerBlock(nn.Module):
                ),
                patch_size=action_config["patch_size"],
                local_attn_size=local_attn_size,
+                sink_size=sink_size,
                qk_norm=action_config["qk_norm"],
                qkv_bias=action_config["qkv_bias"],
                vae_time_compression_ratio=action_config[
@@ -647,6 +653,8 @@ class CausalMatrixGame2WanModel(BaseDiT):
        self.num_channels_latents = config.num_channels_latents
        self.patch_size = config.patch_size

+        # Long tuning controls the causal window from YAML; consume it here
+        # like Wan so train wrappers do not need runtime patching.
        arch_cfg = getattr(config, "arch_config", None)
        self.local_attn_size = (
            getattr(
@@ -662,6 +670,12 @@ class CausalMatrixGame2WanModel(BaseDiT):
            if arch_cfg
            else getattr(config, "sink_size", 0)
        )
+        if self.sink_size < 0:
+            raise ValueError("sink_size must be non-negative")
+        if self.local_attn_size != -1 and self.sink_size >= self.local_attn_size:
+            raise ValueError(
+                "sink_size must be smaller than local_attn_size for "
+                "MatrixGame2 causal attention")
        self.rope_cache_policy = (
            getattr(arch_cfg, "rope_cache_policy",
                    getattr(config, "rope_cache_policy", "absolute"))
@@ -749,6 +763,52 @@ class CausalMatrixGame2WanModel(BaseDiT):

        self.__post_init__()

+    def set_causal_attention_window(
+        self,
+        *,
+        local_attn_size: int | None = None,
+        sink_size: int | None = None,
+    ) -> None:
+        local_attn_size = (
+            self.local_attn_size if local_attn_size is None else int(local_attn_size)
+        )
+        sink_size = self.sink_size if sink_size is None else int(sink_size)
+        if sink_size < 0:
+            raise ValueError("sink_size must be non-negative")
+        if local_attn_size != -1 and sink_size >= local_attn_size:
+            raise ValueError(
+                "sink_size must be smaller than local_attn_size for "
+                "MatrixGame2 causal attention")
+
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.block_mask = None
+        self.block_mask_keyboard = None
+        self.block_mask_mouse = None
+
+        for block in self.blocks:
+            if hasattr(block, "local_attn_size"):
+                block.local_attn_size = local_attn_size
+            if hasattr(block, "sink_size"):
+                block.sink_size = sink_size
+
+            attn1 = getattr(block, "attn1", None)
+            if attn1 is not None:
+                if hasattr(attn1, "local_attn_size"):
+                    attn1.local_attn_size = local_attn_size
+                if hasattr(attn1, "sink_size"):
+                    attn1.sink_size = sink_size
+
+            action_model = getattr(block, "action_model", None)
+            if action_model is not None:
+                if hasattr(action_model, "local_attn_size"):
+                    action_model.local_attn_size = local_attn_size
+                if hasattr(action_model, "sink_size"):
+                    action_model.sink_size = sink_size
+
+    def set_sink_size(self, sink_size: int) -> None:
+        self.set_causal_attention_window(sink_size=sink_size)
+
    @staticmethod
    def _prepare_blockwise_causal_attn_mask(
        device: torch.device | str,
@@ -756,6 +816,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
        frame_seqlen: int = 880,
        num_frame_per_block: int = 1,
        local_attn_size: int = -1,
+        sink_size: int = 0,
    ) -> BlockMask:
        total_length = num_frames * frame_seqlen
        padded_length = math.ceil(total_length / 128) * 128 - total_length
@@ -782,7 +843,10 @@ class CausalMatrixGame2WanModel(BaseDiT):
            else:
                return (
                    (kv_idx < ends[q_idx])
-                    & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
+                    & (
+                        (kv_idx < sink_size * frame_seqlen)
+                        | (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
+                    )
                ) | (q_idx == kv_idx)

        block_mask = create_block_mask(
@@ -796,8 +860,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
        )

        if not dist.is_initialized() or dist.get_rank() == 0:
-            print(
-                f" cache a block wise causal mask with block size of {num_frame_per_block} frames"
+            logger.info(
+                "cache a block wise causal mask with block size of %s frames",
+                num_frame_per_block,
            )

        return block_mask
@@ -809,6 +874,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
        frame_seqlen: int = 880,
        num_frame_per_block: int = 1,
        local_attn_size: int = -1,
+        sink_size: int = 0,
    ) -> BlockMask:
        total_length2 = num_frames * frame_seqlen
        padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
@@ -834,7 +900,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
            else:
                return (
                    (kv_idx < ends2[q_idx])
-                    & (kv_idx >= (ends2[q_idx] - local_attn_size))
+                    & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))
                ) | (q_idx == kv_idx)

        block_mask2 = create_block_mask(
@@ -848,8 +914,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
        )

        if not dist.is_initialized() or dist.get_rank() == 0:
-            print(
-                f" cache a block wise causal mask for keyboard with block size of {num_frame_per_block} frames"
+            logger.info(
+                "cache a block wise causal mask for keyboard with block size of %s frames",
+                num_frame_per_block,
            )

        return block_mask2
@@ -861,6 +928,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
        frame_seqlen: int = 1,
        num_frame_per_block: int = 1,
        local_attn_size: int = -1,
+        sink_size: int = 0,
    ) -> BlockMask:
        total_length2 = num_frames * frame_seqlen
        padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
@@ -886,7 +954,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
            else:
                return (
                    (kv_idx < ends2[q_idx])
-                    & (kv_idx >= (ends2[q_idx] - local_attn_size))
+                    & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))
                ) | (q_idx == kv_idx)

        block_mask2 = create_block_mask(
@@ -900,8 +968,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
        )

        if not dist.is_initialized() or dist.get_rank() == 0:
-            print(
-                f" cache a block wise causal mask for action with block size of {num_frame_per_block} frames"
+            logger.info(
+                "cache a block wise causal mask for action with block size of %s frames",
+                num_frame_per_block,
            )

        return block_mask2
@@ -1063,6 +1132,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
            frame_seqlen=post_patch_height * post_patch_width,
            num_frame_per_block=self.num_frame_per_block,
            local_attn_size=self.local_attn_size,
+            sink_size=self.sink_size,
        )
        if self.use_rope_keyboard:
            block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action(
@@ -1071,6 +1141,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                frame_seqlen=1,
                num_frame_per_block=self.num_frame_per_block,
                local_attn_size=self.local_attn_size,
+                sink_size=self.sink_size,
            )
        else:
            block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
@@ -1079,6 +1150,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                frame_seqlen=post_patch_height * post_patch_width,
                num_frame_per_block=self.num_frame_per_block,
                local_attn_size=self.local_attn_size,
+                sink_size=self.sink_size,
            )
        block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
            device=hidden_states.device,
@@ -1086,6 +1158,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
            frame_seqlen=1,
            num_frame_per_block=self.num_frame_per_block,
            local_attn_size=self.local_attn_size,
+            sink_size=self.sink_size,
        )
        if kv_cache is None:
            kv_cache = [None] * len(self.blocks)
@@ -1239,6 +1312,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                frame_seqlen=post_patch_height * post_patch_width,
                num_frame_per_block=self.num_frame_per_block,
                local_attn_size=self.local_attn_size,
+                sink_size=self.sink_size,
            )
        if self.block_mask_keyboard is None:
            if self.use_rope_keyboard:
@@ -1248,6 +1322,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                    frame_seqlen=1,
                    num_frame_per_block=self.num_frame_per_block,
                    local_attn_size=self.local_attn_size,
+                    sink_size=self.sink_size,
                )
            else:
                self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
@@ -1256,6 +1331,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                    frame_seqlen=post_patch_height * post_patch_width,
                    num_frame_per_block=self.num_frame_per_block,
                    local_attn_size=self.local_attn_size,
+                    sink_size=self.sink_size,
                )
        if self.block_mask_mouse is None:
            self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
@@ -1264,6 +1340,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                frame_seqlen=1,
                num_frame_per_block=self.num_frame_per_block,
                local_attn_size=self.local_attn_size,
+                sink_size=self.sink_size,
            )

        hidden_states = self.patch_embedding(hidden_states)
@@ -452,7 +452,7 @@ class MatrixGame2WanModel(BaseDiT):
            and len(encoder_hidden_states_image) > 0
        ):
            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
+        elif not isinstance(encoder_hidden_states_image, torch.Tensor):
            encoder_hidden_states_image = None

        batch_size, num_channels, num_frames, height, width = (
@@ -4,12 +4,18 @@ import asyncio
 import os
 import random

+import numpy as np
 import torch
 from PIL import Image

 from fastvideo.distributed.parallel_state import get_local_torch_device
 from fastvideo.utils import logger

+try:
+    import cv2
+except ImportError:
+    cv2 = None
+

 CAM_VALUE = 0.1
 CAMERA_MAP = {
@@ -34,6 +40,256 @@ KEYBOARD_MAP_7 = {  # templerun_distilled_model: still/w/s/left/right/a/d
    "d": [0, 0, 0, 0, 0, 0, 1],  # d
 }
 KEYBOARD_MAP = KEYBOARD_MAP_4  # Default for backward compatibility
+SOLARIS_MOVEMENT_KEY_INDICES = {
+    "W": 11,
+    "S": 12,
+    "A": 13,
+    "D": 14,
+}
+
+
+def retain_kv_with_sink(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    target_len: int,
+    sink_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if k.shape[1] <= target_len or sink_size <= 0:
+        return k[:, -target_len:], v[:, -target_len:]
+
+    sink_len = min(int(sink_size), target_len, k.shape[1])
+    tail_len = target_len - sink_len
+    sink_k = k[:, :sink_len]
+    sink_v = v[:, :sink_len]
+    if tail_len <= 0:
+        return sink_k, sink_v
+
+    tail_k = k[:, sink_len:][:, -tail_len:]
+    tail_v = v[:, sink_len:][:, -tail_len:]
+    return torch.cat([sink_k, tail_k], dim=1), torch.cat([sink_v, tail_v], dim=1)
+
+
+def _require_cv2() -> bool:
+    if cv2 is None:
+        logger.warning(
+            "OpenCV is not available; skipping MatrixGame2 validation overlay."
+        )
+        return False
+    return True
+
+
+def _matrixgame2_overlay_keys_from_keyboard(
+    keyboard_frame: np.ndarray | torch.Tensor,
+) -> dict[str, bool]:
+    vec = np.asarray(keyboard_frame, dtype=np.float32).reshape(-1)
+    dim = vec.shape[0]
+
+    if dim >= 23:
+        return {
+            key: bool(vec[idx] > 0.5)
+            for key, idx in SOLARIS_MOVEMENT_KEY_INDICES.items()
+        }
+    if dim == 7:
+        return {
+            "W": bool(vec[1] > 0.5),
+            "S": bool(vec[2] > 0.5),
+            "A": bool(vec[5] > 0.5),
+            "D": bool(vec[6] > 0.5),
+        }
+    if dim >= 6:
+        return {
+            "W": bool(vec[0] > 0.5),
+            "S": bool(vec[1] > 0.5),
+            "A": bool(vec[2] > 0.5),
+            "D": bool(vec[3] > 0.5),
+        }
+    if dim == 4:
+        return {
+            "W": bool(vec[0] > 0.5),
+            "S": bool(vec[1] > 0.5),
+            "A": bool(vec[2] > 0.5),
+            "D": bool(vec[3] > 0.5),
+        }
+    if dim == 2:
+        return {
+            "W": bool(vec[0] > 0.5),
+            "S": bool(vec[1] > 0.5),
+            "A": False,
+            "D": False,
+        }
+    return {"W": False, "S": False, "A": False, "D": False}
+
+
+def draw_rounded_rectangle(
+    image: np.ndarray,
+    top_left: tuple[int, int],
+    bottom_right: tuple[int, int],
+    color: tuple[int, int, int],
+    radius: int = 10,
+    alpha: float = 0.5,
+) -> None:
+    if not _require_cv2():
+        return
+
+    overlay = image.copy()
+    x1, y1 = top_left
+    x2, y2 = bottom_right
+
+    cv2.rectangle(overlay, (x1 + radius, y1), (x2 - radius, y2), color, -1)
+    cv2.rectangle(overlay, (x1, y1 + radius), (x2, y2 - radius), color, -1)
+    cv2.ellipse(overlay, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, -1)
+    cv2.ellipse(overlay, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, -1)
+    cv2.ellipse(overlay, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, -1)
+    cv2.ellipse(overlay, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, -1)
+    cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
+
+
+def draw_keys_on_frame(
+    frame: np.ndarray,
+    keys: dict[str, bool],
+    key_size: tuple[int, int] = (30, 30),
+    top_margin: int = 15,
+) -> None:
+    if not _require_cv2():
+        return
+
+    left_margin = 15
+    gap = 3
+    key_positions = {
+        "W": (left_margin + key_size[0] + gap, top_margin),
+        "A": (left_margin, top_margin + key_size[1] + gap),
+        "S": (left_margin + key_size[0] + gap, top_margin + key_size[1] + gap),
+        "D": (
+            left_margin + (key_size[0] + gap) * 2,
+            top_margin + key_size[1] + gap,
+        ),
+    }
+
+    for key, (x, y) in key_positions.items():
+        pressed = keys.get(key, False)
+        color = (0, 255, 0) if pressed else (200, 200, 200)
+        alpha = 0.8 if pressed else 0.5
+        draw_rounded_rectangle(
+            frame,
+            (x, y),
+            (x + key_size[0], y + key_size[1]),
+            color,
+            radius=5,
+            alpha=alpha,
+        )
+        text_size = cv2.getTextSize(
+            key, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
+        )[0]
+        text_x = x + (key_size[0] - text_size[0]) // 2
+        text_y = y + (key_size[1] + text_size[1]) // 2
+        cv2.putText(
+            frame,
+            key,
+            (text_x, text_y),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.5,
+            (0, 0, 0),
+            1,
+        )
+
+
+def draw_mouse_on_frame(
+    frame: np.ndarray,
+    yaw: float,
+    pitch: float,
+    top_margin: int = 15,
+) -> None:
+    if not _require_cv2():
+        return
+
+    _, width, _ = frame.shape
+    right_margin = 15
+    crosshair_radius = 25
+    crosshair_x = width - right_margin - crosshair_radius
+    crosshair_y = top_margin + crosshair_radius
+
+    dx = int(yaw * crosshair_radius * 8)
+    dy = int(-pitch * crosshair_radius * 8)
+    max_arrow = crosshair_radius - 5
+    dx = max(-max_arrow, min(max_arrow, dx))
+    dy = max(-max_arrow, min(max_arrow, dy))
+
+    cv2.circle(frame, (crosshair_x, crosshair_y), crosshair_radius, (50, 50, 50), -1)
+    cv2.circle(frame, (crosshair_x, crosshair_y), crosshair_radius, (200, 200, 200), 1)
+    cv2.line(
+        frame,
+        (crosshair_x - crosshair_radius + 5, crosshair_y),
+        (crosshair_x + crosshair_radius - 5, crosshair_y),
+        (100, 100, 100),
+        1,
+    )
+    cv2.line(
+        frame,
+        (crosshair_x, crosshair_y - crosshair_radius + 5),
+        (crosshair_x, crosshair_y + crosshair_radius - 5),
+        (100, 100, 100),
+        1,
+    )
+
+    if abs(dx) > 1 or abs(dy) > 1:
+        cv2.arrowedLine(
+            frame,
+            (crosshair_x, crosshair_y),
+            (crosshair_x + dx, crosshair_y + dy),
+            (0, 255, 0),
+            2,
+            tipLength=0.3,
+        )
+
+
+def overlay_validation_actions_on_frames(
+    frames: list[np.ndarray],
+    keyboard_cond: np.ndarray | torch.Tensor | None = None,
+    mouse_cond: np.ndarray | torch.Tensor | None = None,
+) -> list[np.ndarray]:
+    if (keyboard_cond is None and mouse_cond is None) or not _require_cv2():
+        return frames
+
+    if keyboard_cond is not None and torch.is_tensor(keyboard_cond):
+        keyboard_cond = keyboard_cond.detach().cpu().float().numpy()
+    if mouse_cond is not None and torch.is_tensor(mouse_cond):
+        mouse_cond = mouse_cond.detach().cpu().float().numpy()
+
+    if keyboard_cond is not None:
+        keyboard_cond = np.asarray(keyboard_cond, dtype=np.float32)
+        if keyboard_cond.ndim == 3 and keyboard_cond.shape[0] == 1:
+            keyboard_cond = keyboard_cond[0]
+    if mouse_cond is not None:
+        mouse_cond = np.asarray(mouse_cond, dtype=np.float32)
+        if mouse_cond.ndim == 3 and mouse_cond.shape[0] == 1:
+            mouse_cond = mouse_cond[0]
+
+    processed_frames: list[np.ndarray] = []
+    for frame_idx, frame in enumerate(frames):
+        frame = np.ascontiguousarray(frame.copy())
+        if keyboard_cond is not None and frame_idx < len(keyboard_cond):
+            draw_keys_on_frame(
+                frame,
+                _matrixgame2_overlay_keys_from_keyboard(keyboard_cond[frame_idx]),
+            )
+        if mouse_cond is not None and frame_idx < len(mouse_cond):
+            # MG action convention: mouse[0]=pitch, mouse[1]=yaw.
+            pitch = float(mouse_cond[frame_idx, 0])
+            yaw = float(mouse_cond[frame_idx, 1])
+            draw_mouse_on_frame(frame, yaw=yaw, pitch=pitch)
+        processed_frames.append(frame)
+
+    return processed_frames
+
+
+def scale_keyboard_condition(
+    keyboard_condition: torch.Tensor,
+    scale: float,
+) -> torch.Tensor:
+    if float(scale) == 1.0:
+        return keyboard_condition
+    return keyboard_condition * float(scale)
+

 def expand_action_to_frames(action: dict, num_frames: int) -> tuple[torch.Tensor, torch.Tensor]:
    result = {}
@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0

+import io
 import os
 import tempfile
+import time
 from collections.abc import Callable
 from typing import Any
 from urllib.parse import unquote, urlparse
@@ -16,6 +18,10 @@ import torch.nn.functional as F
 import torchvision.transforms.functional as TF
 from packaging import version

+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
 if version.parse(version.parse(
        PIL.__version__).base_version) >= version.parse("9.1.0"):
    PIL_INTERPOLATION = {
@@ -89,6 +95,39 @@ def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
    return 2.0 * images - 1.0


+def _fetch_image_bytes(url: str,
+                       attempts: int = 3,
+                       timeout: float = 30.0) -> bytes:
+    """Download ``url`` fully, retrying transient network errors.
+
+    Reads the whole body via ``response.content`` instead of handing PIL a
+    ``.raw`` stream: a connection dropped mid-body (e.g. an HF-CDN
+    ``IncompleteRead``) then fails here, where it can be retried, instead of
+    surfacing as a truncated image inside PIL.
+    """
+    for attempt in range(attempts):
+        try:
+            response = requests.get(url, timeout=timeout)
+            response.raise_for_status()
+            return response.content
+        except requests.RequestException as e:
+            status = getattr(getattr(e, "response", None), "status_code",
+                             None)
+            # Permanent client errors (4xx except 408/429) won't heal on
+            # retry — fail fast.
+            if (status is not None and 400 <= status < 500
+                    and status not in (408, 429)):
+                raise
+            if attempt == attempts - 1:
+                raise
+            backoff = 2**attempt
+            logger.warning(
+                "Failed to download image %s (attempt %d/%d): %s. "
+                "Retrying in %ds.", url, attempt + 1, attempts, e, backoff)
+            time.sleep(backoff)
+    raise AssertionError("unreachable")
+
+
 # adapted from diffusers.utils import load_image
 def load_image(
    image: str | PIL.Image.Image,
@@ -110,7 +149,7 @@ def load_image(
    """
    if isinstance(image, str):
        if image.startswith("http://") or image.startswith("https://"):
-            image = PIL.Image.open(requests.get(image, stream=True).raw)
+            image = PIL.Image.open(io.BytesIO(_fetch_image_bytes(image)))
        elif os.path.isfile(image):
            image = PIL.Image.open(image)
        else:
@@ -86,20 +86,32 @@ def record_metadata(record: Record) -> Record:
    }


+def _cohort_metadata_value(value: Any) -> Any:
+    if value is None:
+        return ""
+    if isinstance(value, str) and not value.strip():
+        return ""
+    return value
+
+
+def _cohort_key_value(value: Any) -> str:
+    return str(_cohort_metadata_value(value))
+
+
 def record_comparison_metadata(record: Record) -> Record:
-    return {key: record.get(key) or "" for key in COMPARISON_COHORT_KEYS}
+    return {key: _cohort_metadata_value(record.get(key)) for key in COMPARISON_COHORT_KEYS}


 def comparison_cohort_key(record: Record) -> CohortKey:
    return (
        str(record.get("model_id") or "unknown"),
        str(record.get("gpu_type") or "unknown"),
-        str(record.get("workload_id") or ""),
-        str(record.get("variant_id") or ""),
-        str(record.get("benchmark_version") or ""),
-        str(record.get("recipe_fingerprint") or ""),
-        str(record.get("hardware_profile_id") or ""),
-        str(record.get("software_profile_id") or ""),
+        _cohort_key_value(record.get("workload_id")),
+        _cohort_key_value(record.get("variant_id")),
+        _cohort_key_value(record.get("benchmark_version")),
+        _cohort_key_value(record.get("recipe_fingerprint")),
+        _cohort_key_value(record.get("hardware_profile_id")),
+        _cohort_key_value(record.get("software_profile_id")),
    )


@@ -146,9 +158,10 @@ def build_latest_summary(records: list[Record],
            continue

        latest = latest_candidates[-1]
+        latest_index = next(index for index, record in enumerate(group) if record is latest)
        baseline_pool = [
-            record for record in group
-            if record is not latest and record.get("success", True) and is_baseline_eligible_record(record)
+            record for record in group[:latest_index]
+            if record.get("success", True) and is_baseline_eligible_record(record)
        ]
        baseline_records = baseline_pool[-baseline_window:]
        metric_policies = resolve_metric_policies(latest.get("regression_thresholds"))
@@ -19,6 +19,39 @@ from fastvideo.pipelines.stages import (DecodingStage, InputValidationStage, LTX
 logger = init_logger(__name__)


+def _resolve_refine_path(model_path: str, value: str | None) -> str | None:
+    """Resolve a refine component path, preferring paths inside the local
+    model snapshot when the value is relative."""
+    if value is None:
+        return None
+    if os.path.isabs(value):
+        return value
+    candidate = os.path.join(model_path, value)
+    if os.path.exists(candidate):
+        return candidate
+    return value
+
+
+def _resolve_refine_upsampler_path(model_path: str, model_index: dict[str, Any]) -> str | None:
+    """Resolve the refine upsampler directory for a local model snapshot.
+
+    model_index.json keys ("fastvideo_refine_upsampler_path", then
+    "spatial_upsampler") are the documented override and win when present.
+    Otherwise probe the snapshot root: published LTX-2 distilled checkpoints
+    bundle the upsampler weights on disk (spelled "spatial_upscaler")
+    without declaring them in model_index.json.
+    """
+    path = _resolve_refine_path(model_path, model_index.get("fastvideo_refine_upsampler_path"))
+    if path is None and "spatial_upsampler" in model_index:
+        path = _resolve_refine_path(model_path, "spatial_upsampler")
+    if path is None:
+        for dirname in ("spatial_upscaler", "spatial_upsampler"):
+            candidate = os.path.join(model_path, dirname)
+            if os.path.isdir(candidate):
+                return candidate
+    return path
+
+
 class LTX2Pipeline(LoRAPipeline):

    _required_config_modules = [
@@ -189,28 +222,16 @@ class LTX2Pipeline(LoRAPipeline):
        # Apply optional FastVideo-specific refine defaults embedded in
        # model_index.json. These are bundled with distilled checkpoints
        # so the pipeline can self-configure without explicit user kwargs.
-        def _resolve_refine_path(value: str | None) -> str | None:
-            if value is None:
-                return None
-            if os.path.isabs(value):
-                return value
-            candidate = os.path.join(self.model_path, value)
-            if os.path.exists(candidate):
-                return candidate
-            return value
-
        if (model_index.get("fastvideo_refine_enabled") is True and fastvideo_args.refine_enabled is None):
            fastvideo_args.ltx2_refine_enabled = True
        if (fastvideo_args.refine_upsampler_path is None and fastvideo_args.ltx2_refine_upsampler_path is None):
-            fastvideo_args.ltx2_refine_upsampler_path = _resolve_refine_path(
-                model_index.get("fastvideo_refine_upsampler_path"))
-            if (fastvideo_args.ltx2_refine_upsampler_path is None and "spatial_upsampler" in model_index):
-                fastvideo_args.ltx2_refine_upsampler_path = (_resolve_refine_path("spatial_upsampler"))
+            fastvideo_args.ltx2_refine_upsampler_path = (_resolve_refine_upsampler_path(self.model_path, model_index))
        if (fastvideo_args.refine_transformer_path is None and fastvideo_args.ltx2_refine_transformer_path is None):
            fastvideo_args.ltx2_refine_transformer_path = (_resolve_refine_path(
-                model_index.get("fastvideo_refine_transformer_path")))
+                self.model_path, model_index.get("fastvideo_refine_transformer_path")))
        if (fastvideo_args.refine_lora_path is None and fastvideo_args.ltx2_refine_lora_path is None):
-            fastvideo_args.ltx2_refine_lora_path = _resolve_refine_path(model_index.get("fastvideo_refine_lora_path"))
+            fastvideo_args.ltx2_refine_lora_path = _resolve_refine_path(self.model_path,
+                                                                        model_index.get("fastvideo_refine_lora_path"))
        if (fastvideo_args.refine_num_inference_steps is None
                and fastvideo_args.ltx2_refine_num_inference_steps == FastVideoArgs.ltx2_refine_num_inference_steps
                and model_index.get("fastvideo_refine_num_inference_steps") is not None):
@@ -226,10 +247,11 @@ class LTX2Pipeline(LoRAPipeline):
        if (fastvideo_args.refine_add_noise is None and model_index.get("fastvideo_refine_add_noise") is not None):
            fastvideo_args.ltx2_refine_add_noise = bool(model_index["fastvideo_refine_add_noise"])
        if (fastvideo_args.refine_noise_path is None and fastvideo_args.ltx2_refine_noise_path is None):
-            fastvideo_args.ltx2_refine_noise_path = _resolve_refine_path(model_index.get("fastvideo_refine_noise_path"))
+            fastvideo_args.ltx2_refine_noise_path = _resolve_refine_path(self.model_path,
+                                                                         model_index.get("fastvideo_refine_noise_path"))
        if (fastvideo_args.refine_audio_noise_path is None and fastvideo_args.ltx2_refine_audio_noise_path is None):
            fastvideo_args.ltx2_refine_audio_noise_path = (_resolve_refine_path(
-                model_index.get("fastvideo_refine_audio_noise_path")))
+                self.model_path, model_index.get("fastvideo_refine_audio_noise_path")))

        model_index.pop("_class_name")
        model_index.pop("_diffusers_version")
@@ -718,6 +718,12 @@ def _register_configs() -> None:
            "FastVideo/Matrix-Game-2.0-Base-Diffusers",
            "FastVideo/Matrix-Game-2.0-GTA-Diffusers",
            "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers",
+            # Zelda World Model
+            "mignonjia/mg_longtuning_distilled_zelda",  # distilled using streaming long tuning which init from mg_sf_distilled_zelda_1k_steps
+            "mignonjia/mg_sf_distilled_zelda_1k_steps",  # distilled using self forcing for 1k steps
+            "mignonjia/mg_sf_distilled_zelda",  # distilled using self forcing for 3k steps
+            "mignonjia/mg_causal_zelda",
+            "mignonjia/mg_bidirectional_zelda",
        ],
        model_detectors=[
            lambda path: any(token in path.lower() for token in (
@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only tests for load_image URL retry behavior.
+
+A HF-CDN IncompleteRead mid-body previously killed SSIM runs (build 4202,
+test_wan_i2v) because load_image handed PIL a partial ``.raw`` stream with no
+retry. These tests mock requests.get: a transient failure must be retried and
+a full-body success must yield a usable PIL image.
+"""
+import io
+from unittest import mock
+
+import PIL.Image
+import pytest
+import requests
+
+from fastvideo.models import vision_utils
+
+URL = "https://example.com/image.png"
+
+
+def _png_bytes() -> bytes:
+    buf = io.BytesIO()
+    PIL.Image.new("RGB", (4, 4), (255, 0, 0)).save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def _good_response() -> mock.Mock:
+    response = mock.Mock()
+    response.content = _png_bytes()
+    response.raise_for_status = mock.Mock()
+    return response
+
+
+def test_load_image_retries_transient_error_then_succeeds(monkeypatch):
+    monkeypatch.setattr(vision_utils.time, "sleep", lambda _s: None)
+    get = mock.Mock(side_effect=[
+        # The shape that broke build 4202: connection dropped mid-body.
+        requests.exceptions.ChunkedEncodingError(
+            "IncompleteRead(1024 bytes read, 2048 more expected)"),
+        _good_response(),
+    ])
+    monkeypatch.setattr(vision_utils.requests, "get", get)
+
+    image = vision_utils.load_image(URL)
+
+    assert get.call_count == 2
+    assert get.call_args.kwargs["timeout"] == 30.0
+    assert isinstance(image, PIL.Image.Image)
+    assert image.size == (4, 4)
+    assert image.mode == "RGB"
+
+
+def test_load_image_raises_after_exhausting_retries(monkeypatch):
+    sleeps: list[float] = []
+    monkeypatch.setattr(vision_utils.time, "sleep", sleeps.append)
+    get = mock.Mock(
+        side_effect=requests.exceptions.ConnectionError("connection reset"))
+    monkeypatch.setattr(vision_utils.requests, "get", get)
+
+    with pytest.raises(requests.exceptions.ConnectionError):
+        vision_utils.load_image(URL)
+
+    assert get.call_count == 3
+    assert sleeps == [1, 2]
+
+
+def test_load_image_does_not_retry_on_client_error(monkeypatch):
+    sleeps: list[float] = []
+    monkeypatch.setattr(vision_utils.time, "sleep", sleeps.append)
+    response = mock.Mock(status_code=404)
+    error = requests.HTTPError("404 Client Error", response=response)
+    response.raise_for_status.side_effect = error
+    get = mock.Mock(return_value=response)
+    monkeypatch.setattr(vision_utils.requests, "get", get)
+
+    with pytest.raises(requests.HTTPError):
+        vision_utils.load_image("https://example.com/missing.png")
+
+    assert get.call_count == 1
+    assert sleeps == []
@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for parquet map-style dataset path parsing."""
+
+from __future__ import annotations
+
+import pickle
+
+import pytest
+
+from fastvideo.dataset import parquet_dataset_map_style as parquet_dataset
+from fastvideo.dataset.parquet_dataset_map_style import _parse_data_path_specs
+
+
+def test_parse_data_path_specs_accepts_old_repeat_string() -> None:
+    # Dataset parsing keeps compatibility with the old "path:repeat" string
+    # form used by existing training configs.
+    assert _parse_data_path_specs("data/path1:2,data/path2:1") == [
+        ("data/path1", 2),
+        ("data/path2", 1),
+    ]
+
+
+def test_parse_data_path_specs_accepts_yaml_mapping() -> None:
+    # New YAML mapping form should reach the dataset layer as path -> repeat
+    # and parse to the same internal spec representation.
+    assert _parse_data_path_specs({
+        "data/path1": 1,
+        "data/path2": 2,
+    }) == [
+        ("data/path1", 1),
+        ("data/path2", 2),
+    ]
+
+
+def test_parse_data_path_specs_accepts_path_list() -> None:
+    assert _parse_data_path_specs(["data/a", "data/b"]) == [
+        ("data/a", 1),
+        ("data/b", 1),
+    ]
+
+
+def test_parse_data_path_specs_drops_non_positive_repeats() -> None:
+    assert _parse_data_path_specs({"data/a": 0, "data/b": -1, "data/c": 2}) == [
+        ("data/c", 2),
+    ]
+    assert _parse_data_path_specs("data/a:0,data/b:2") == [("data/b", 2)]
+
+
+def test_parse_data_path_specs_rejects_malformed_repeats() -> None:
+    # Repeat counts must parse as ints; anything else is an explicit error
+    # rather than a silently mis-weighted dataset.
+    with pytest.raises(ValueError):
+        _parse_data_path_specs("data/a:abc")
+    with pytest.raises(ValueError):
+        _parse_data_path_specs("data/a:1.5")
+    with pytest.raises(ValueError):
+        _parse_data_path_specs({"data/a": "abc"})  # type: ignore[dict-item]
+
+
+class _DummyWorldGroup:
+
+    def barrier(self) -> None:
+        return None
+
+
+def _write_root_cache(dataset_root, filename: str, length: int) -> str:
+    cache_dir = dataset_root / "map_style_cache"
+    cache_dir.mkdir(parents=True)
+    parquet_file = dataset_root / filename
+    parquet_file.touch()
+    with (cache_dir / "file_info.pkl").open("wb") as f:
+        pickle.dump(((str(parquet_file),), (length,)), f)
+    return str(parquet_file)
+
+
+def test_get_parquet_files_and_length_repeats_single_path(tmp_path, monkeypatch) -> None:
+    # get_parquet_files_and_length applies repeat counts after reading the
+    # per-root parquet cache, so a repeated root duplicates both names and rows.
+    dataset_root = tmp_path / "dataset"
+    parquet_file = _write_root_cache(dataset_root, "sample.parquet", 7)
+
+    monkeypatch.setattr(parquet_dataset, "get_world_rank", lambda: 0)
+    monkeypatch.setattr(parquet_dataset, "get_world_group", _DummyWorldGroup)
+
+    file_names, lengths = parquet_dataset.get_parquet_files_and_length({
+        str(dataset_root): 2,
+    })
+
+    assert file_names == (parquet_file, parquet_file)
+    assert lengths == (7, 7)
+
+
+def test_get_parquet_files_and_length_mixes_roots_and_resorts(tmp_path, monkeypatch) -> None:
+    # Multiple roots are expanded per repeat count and then globally re-sorted
+    # by filename, so the mix order is independent of the mapping order.
+    root_a = tmp_path / "dataset_a"
+    root_b = tmp_path / "dataset_b"
+    file_a = _write_root_cache(root_a, "a.parquet", 5)
+    file_b = _write_root_cache(root_b, "b.parquet", 9)
+
+    monkeypatch.setattr(parquet_dataset, "get_world_rank", lambda: 0)
+    monkeypatch.setattr(parquet_dataset, "get_world_group", _DummyWorldGroup)
+
+    file_names, lengths = parquet_dataset.get_parquet_files_and_length({
+        str(root_b): 1,
+        str(root_a): 2,
+    })
+
+    assert file_names == (file_a, file_a, file_b)
+    assert lengths == (5, 5, 9)
+
+
+def test_get_parquet_files_and_length_raises_when_all_repeats_dropped() -> None:
+    # Zero/negative repeats are dropped at parse time; if that leaves nothing
+    # to read, the mix branch fails loudly instead of yielding an empty dataset.
+    with pytest.raises(FileNotFoundError):
+        parquet_dataset.get_parquet_files_and_length({"data/a": 0})
@@ -30,6 +30,15 @@ print(f"Using image: {image_ref}")
 # the baked setting, and a caller override always wins.
 uv_torch_backend_override = resolve_uv_torch_backend(image_tag)

+# INVARIANT: this image definition must be byte-identical for every CI job at
+# a given base image digest -- one build, shared cache across all concurrent
+# lanes. Never put a per-job/per-commit value (BUILDKITE_*, TEST_SCOPE,
+# env-derived overrides) into the image via .env()/run_commands: it becomes an
+# image layer, so whenever the base digest changes every concurrent job
+# rebuilds its own image variant (~15-20 min each), blowing the Buildkite job
+# budget. `image_ref` is the only env-derived input allowed here, because it
+# *selects* the base digest. Per-job values arrive at runtime via
+# `ci_env_secret` below.
 image = (modal.Image.from_registry(
    image_ref, add_python="3.12"
 ).run_commands("rm -rf /FastVideo").apt_install(
@@ -37,41 +46,43 @@ image = (modal.Image.from_registry(
 ).run_commands(
    "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable"
 ).run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc").env({
-    "PATH":
-    "/root/.cargo/bin:$PATH",
-    "BUILDKITE_REPO":
-    os.environ.get("BUILDKITE_REPO", ""),
-    "BUILDKITE_COMMIT":
-    os.environ.get("BUILDKITE_COMMIT", ""),
-    "BUILDKITE_PULL_REQUEST":
-    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":
-    image_version,
-    **({
-        "UV_TORCH_BACKEND": uv_torch_backend_override
-    } if uv_torch_backend_override else {}),
-    # FA4 is opt-in (FASTVIDEO_FA4); CI lanes keep it enabled to match the
-    # SSIM/perf baselines. Caller override wins.
-    "FASTVIDEO_FA4":
-    os.environ.get("FASTVIDEO_FA4", "1"),
-    "HF_REPO_ID":
-    "FastVideo/performance-tracking",
+    "PATH": "/root/.cargo/bin:$PATH",
+    "HF_REPO_ID": "FastVideo/performance-tracking",
 }))

 dreamverse_image = (image.run_commands(
    "curl -fsSL https://deb.nodesource.com/setup_22.x | bash -"
 ).apt_install("nodejs").run_commands("node --version && npm --version"))

+# Per-job/per-invocation values are injected into the container environment at
+# RUNTIME via this secret (attached to every function below), so the image
+# stays identical across jobs. Consumers (run_test_command's checkout,
+# fastvideo/tests/performance/{compare_baseline,identity}.py, the FA4
+# resolver, `uv pip install`) all read os.environ at runtime, so nothing else
+# changes.
+ci_env_secret = modal.Secret.from_dict({
+    "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
+    "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
+    "BUILDKITE_PULL_REQUEST": 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": image_version,
+    **({
+        "UV_TORCH_BACKEND": uv_torch_backend_override
+    } if uv_torch_backend_override else {}),
+    # FA4 is opt-in (FASTVIDEO_FA4); CI lanes keep it enabled to match the
+    # SSIM/perf baselines. Caller override wins.
+    "FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"),
+})
+
+hf_secret = modal.Secret.from_dict(
+    {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
+wandb_secret = modal.Secret.from_dict(
+    {"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
+

 def run_test(pytest_command: str):
    """Helper function to run a test suite with custom pytest command"""
@@ -147,10 +158,7 @@ def run_test_command(test_command: str,
@app.function(gpu="H100:1",
              image=image,
              timeout=1200,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
-              ],
+              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_encoder_tests():
    run_test(
@@ -161,10 +169,7 @@ def run_encoder_tests():
@app.function(gpu="L40S:1",
              image=image,
              timeout=1200,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
-              ],
+              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_vae_tests():
    run_test(
@@ -175,10 +180,7 @@ def run_vae_tests():
@app.function(gpu="L40S:1",
              image=image,
              timeout=900,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
-              ],
+              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_transformer_tests():
    run_test(
@@ -191,10 +193,7 @@ def run_transformer_tests():
              memory=32768,
              image=image,
              timeout=900,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
-              ],
+              secrets=[wandb_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_training_tests():
    run_test(
@@ -207,10 +206,7 @@ def run_training_tests():
              memory=32768,
              image=image,
              timeout=900,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
-              ],
+              secrets=[wandb_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_training_lora_tests():
    run_test(
@@ -221,44 +217,41 @@ def run_training_lora_tests():
@app.function(gpu="H100:2",
              image=image,
              timeout=900,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
-              ])
+              secrets=[wandb_secret, ci_env_secret])
 def run_training_tests_VSA():
    run_test(
        "wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/VSA -srP"
    )


-@app.function(gpu="H100:1", image=image, timeout=900)
+@app.function(gpu="H100:1", image=image, timeout=900, secrets=[ci_env_secret])
 def run_kernel_tests():
    run_test("pytest fastvideo-kernel/tests/ -vs")


-# @app.function(gpu="H100:1", image=image, timeout=900)
+# @app.function(gpu="H100:1", image=image, timeout=900, secrets=[ci_env_secret])
 # def run_precision_tests_VSA():
 #     # VSA correctness is covered by the same file now
 #     run_test("pytest fastvideo-kernel/tests/test_correctness.py")

-# @app.function(gpu="L40S:1", image=image, timeout=900)
+# @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
 # def run_precision_tests_vmoba():
 #     run_test("pytest fastvideo-kernel/tests/test_vmoba_correctness.py")


-@app.function(gpu="L40S:1", image=image, timeout=900)
+@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
 def run_inference_tests_vmoba():
    run_test('python fastvideo/tests/inference/vmoba/test_vmoba_inference.py')


-@app.function(gpu="L40S:1", image=image, timeout=1200)
+@app.function(gpu="L40S:1", image=image, timeout=1200, secrets=[ci_env_secret])
 def run_inference_lora_tests():
    run_test(
        "pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs"
    )


-@app.function(gpu="L40S:2", image=image, timeout=900)
+@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[ci_env_secret])
 def run_distill_dmd_tests():
    run_test(
        "pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
@@ -267,17 +260,14 @@ def run_distill_dmd_tests():
@app.function(gpu="L40S:2",
              image=image,
              timeout=900,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})
-              ])
+              secrets=[wandb_secret, ci_env_secret])
 def run_self_forcing_tests():
    run_test(
        "wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs"
    )


-@app.function(gpu="L40S:1", image=image, timeout=900)
+@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
 def run_unit_test():
    run_test(
        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
@@ -285,7 +275,10 @@ def run_unit_test():


 # TODO: David: GPU only used to resolve import time requirement (not needed for this test). Maybe make those imports lazy?
-@app.function(gpu="L40S:1", image=dreamverse_image, timeout=1800)
+@app.function(gpu="L40S:1",
+              image=dreamverse_image,
+              timeout=1800,
+              secrets=[ci_env_secret])
 def run_dreamverse_app_tests():
    run_test_command(
        install_command="",
@@ -326,10 +319,7 @@ def run_dreamverse_app_tests():
              memory=32768,
              image=image,
              timeout=1800,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
-              ],
+              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_train_framework_tests():
    run_test(
@@ -340,10 +330,7 @@ def run_train_framework_tests():
@app.function(gpu="L40S:1",
              image=image,
              timeout=1800,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
-              ],
+              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def seed_grad_norm_references():
    """Record the per-method grad-norm reference for the **CI GPU (L40S only)**.
@@ -369,10 +356,7 @@ 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", "")})
-              ],
+              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_eval_tests():
    # Eval metric regression: drives the high-level fastvideo.eval API on a
@@ -395,10 +379,7 @@ def run_eval_tests():
@app.function(gpu="L40S:1",
              image=image,
              timeout=3600,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
-              ])
+              secrets=[hf_secret, ci_env_secret])
 def run_lora_extraction_tests():
    run_test(
        "hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/lora_extraction/test_lora_extraction.py"
@@ -410,10 +391,7 @@ def run_lora_extraction_tests():
              memory=32768,
              image=image,
              timeout=1800,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
-              ],
+              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_performance_tests():
    # PR/direct records are uploaded only on pass; scheduled main uploads pass
@@ -455,10 +433,7 @@ def run_performance_tests():
@app.function(gpu="L40S:1",
              image=image,
              timeout=1800,
-              secrets=[
-                  modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
-              ],
+              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
 def run_api_server_tests():
    run_test(
@@ -57,10 +57,12 @@ image = (
    .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
    .env(
        {
+            # INVARIANT: keep this image identical for every CI job at a given
+            # base digest -- per-job/per-commit values (BUILDKITE_*) must never
+            # be baked here as image layers or every job rebuilds its own image
+            # variant. repo/commit/PR already reach the container as
+            # run_ssim_partition() arguments.
            "PATH": "/root/.cargo/bin:$PATH",
-            "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
-            "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
-            "BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
            "IMAGE_VERSION": image_version,
            **({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}),
            # FA4 is opt-in (FASTVIDEO_FA4); the SSIM references were seeded
@@ -548,6 +550,7 @@ def _spawn_ssim_task(
    log_dir: str,
    task_index: int,
    pytest_extra_args: list[str],
+    hf_api_key: str,
 ) -> _RunningTask:
    import shlex

@@ -557,6 +560,7 @@ def _spawn_ssim_task(
    command = f"set -euo pipefail && source $HOME/.local/bin/env && source /opt/venv/bin/activate && {pytest_command}"
    env = os.environ.copy()
    env["HF_HOME"] = "/root/data/.cache"
+    env["HF_API_KEY"] = hf_api_key
    # MultiprocExecutor returns CUDA tensors through mp pipes (CUDA IPC).
    # On kernels without pidfd_open support, PyTorch fails when
    # expandable_segments=True. Force False for CI compatibility.
@@ -673,6 +677,7 @@ def _build_pytest_extra_args(
    ssim_full_quality: bool,
    ssim_reference_repo: str,
    skip_ssim_reference_download: bool,
+    ssim_bootstrap_mode: bool,
    pytest_k: str,
 ) -> list[str]:
    args = []
@@ -682,6 +687,8 @@ def _build_pytest_extra_args(
        args.extend(["--ssim-reference-repo", ssim_reference_repo.strip()])
    if skip_ssim_reference_download:
        args.append("--skip-ssim-reference-download")
+    if ssim_bootstrap_mode:
+        args.append("--ssim-bootstrap-mode")
    if pytest_k.strip():
        args.extend(["-k", pytest_k.strip()])
    return args
@@ -691,6 +698,7 @@ def _schedule_ssim_tasks(
    repo_root: str,
    tasks: list[SSIMTask],
    pytest_extra_args: list[str],
+    hf_api_key: str,
    fail_fast: bool = True,
 ) -> dict[int, _TaskResult]:
    import tempfile
@@ -733,6 +741,7 @@ def _schedule_ssim_tasks(
                log_dir=log_dir,
                task_index=task.task_id,
                pytest_extra_args=pytest_extra_args,
+                hf_api_key=hf_api_key,
            )
            print(f"Started {task.test_name} on GPUs {','.join(assigned_gpu_ids)}")
            running_tasks.append(running_task)
@@ -883,6 +892,7 @@ def run_ssim_partition(
    ssim_full_quality: bool = False,
    ssim_reference_repo: str = "",
    skip_ssim_reference_download: bool = False,
+    ssim_bootstrap_mode: bool = False,
    pytest_k: str = "",
    sync_generated_to_volume: bool = False,
    generated_volume_subdir: str = "",
@@ -911,12 +921,14 @@ def run_ssim_partition(
        ssim_full_quality=ssim_full_quality,
        ssim_reference_repo=ssim_reference_repo,
        skip_ssim_reference_download=skip_ssim_reference_download,
+        ssim_bootstrap_mode=ssim_bootstrap_mode,
        pytest_k=pytest_k,
    )
    results = _schedule_ssim_tasks(
        repo_root,
        partition,
        pytest_extra_args=pytest_extra_args,
+        hf_api_key=hf_api_key,
        fail_fast=fail_fast,
    )
    summaries = _collect_task_summaries(partition, results)
@@ -953,6 +965,7 @@ def run_ssim_tests(
    full_quality: bool = False,
    reference_repo: str = "",
    skip_reference_download: bool = False,
+    bootstrap_mode: bool = False,
    pytest_k: str = "",
    sync_generated_to_volume: bool = False,
    generated_volume_subdir: str = "",
@@ -973,15 +986,24 @@ def run_ssim_tests(
        print(f"Selected model ids: {model_ids}")
    if pytest_k.strip():
        print(f"Using pytest -k filter: {pytest_k}")
+    if bootstrap_mode:
+        print(
+            "SSIM bootstrap mode enabled: missing references will upload "
+            "draft artifacts and xfail."
+        )
    quality_tier = _resolve_output_quality_tier(full_quality)
    if sync_generated_to_volume:
        resolved_subdir = _resolve_generated_volume_subdir(
            generated_volume_subdir,
            resolved_git_commit,
        )
+        generated_volume_path = _build_generated_volume_relative_path(
+            generated_volume_subdir=resolved_subdir,
+            quality_tier=quality_tier,
+        )
        print(
            "Raw generated videos will be saved to Modal volume path: "
-            f"{_build_generated_volume_relative_path(generated_volume_subdir=resolved_subdir, quality_tier=quality_tier)}"
+            f"{generated_volume_path}"
        )
    else:
        resolved_subdir = ""
@@ -996,6 +1018,7 @@ def run_ssim_tests(
        ssim_full_quality=full_quality,
        ssim_reference_repo=reference_repo,
        skip_ssim_reference_download=skip_reference_download,
+        ssim_bootstrap_mode=bootstrap_mode,
        pytest_k=pytest_k,
        sync_generated_to_volume=sync_generated_to_volume,
        generated_volume_subdir=resolved_subdir,
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+
+class _FakeImage:
+
+    @classmethod
+    def from_registry(cls, *_args, **_kwargs):
+        return cls()
+
+    def apt_install(self, *_args, **_kwargs):
+        return self
+
+    def run_commands(self, *_args, **_kwargs):
+        return self
+
+    def env(self, *_args, **_kwargs):
+        return self
+
+
+class _FakeVolume:
+
+    @classmethod
+    def from_name(cls, *_args, **_kwargs):
+        return cls()
+
+
+class _FakeApp:
+
+    def function(self, *_args, **_kwargs):
+        def decorator(func):
+            return func
+
+        return decorator
+
+    def local_entrypoint(self, *_args, **_kwargs):
+        def decorator(func):
+            return func
+
+        return decorator
+
+
+def _load_ssim_test_module(monkeypatch):
+    fake_modal = types.SimpleNamespace(
+        App=lambda: _FakeApp(),
+        Image=_FakeImage,
+        Volume=_FakeVolume,
+    )
+    monkeypatch.setitem(sys.modules, "modal", fake_modal)
+    module_path = Path(__file__).with_name("ssim_test.py")
+    spec = importlib.util.spec_from_file_location("modal_ssim_test_under_test", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_spawn_ssim_task_passes_hf_api_key_to_pytest_env(monkeypatch, tmp_path):
+    module = _load_ssim_test_module(monkeypatch)
+    captured = {}
+
+    class FakeProcess:
+        pass
+
+    def fake_popen(*args, **kwargs):
+        captured["args"] = args
+        captured["kwargs"] = kwargs
+        return FakeProcess()
+
+    monkeypatch.setattr(module.subprocess, "Popen", fake_popen)
+
+    task = module.SSIMTask(
+        task_id=0,
+        test_file="fastvideo/tests/ssim/test_example.py",
+        required_gpus=1,
+        model_id="model-a",
+    )
+
+    running_task = module._spawn_ssim_task(
+        task=task,
+        repo_root=str(tmp_path),
+        assigned_gpu_ids=["2"],
+        log_dir=str(tmp_path),
+        task_index=0,
+        pytest_extra_args=["--ssim-bootstrap-mode"],
+        hf_api_key="hf_test_token",
+    )
+
+    env = captured["kwargs"]["env"]
+    assert env["HF_API_KEY"] == "hf_test_token"
+    assert env["HF_HOME"] == "/root/data/.cache"
+    assert env["CUDA_VISIBLE_DEVICES"] == "2"
+    assert env["FASTVIDEO_SSIM_MODEL_ID"] == "model-a"
+    assert running_task.process.__class__ is FakeProcess
+    running_task.log_handle.close()
@@ -63,3 +63,30 @@ def _raise_module_on_import(name: str) -> types.ModuleType:
            raise ImportError(f"No module named '{name}.{item}'")

    return _RaisingModule(name)
+
+
+def test_coerce_fp4_input_dtype_casts_and_rejects():
+    """The FP4 input-dtype coercion (CPU-only, no flashinfer/CUDA): bf16/fp16
+    pass through, other floats (the fp32 pre-attention norm in eager mode) are
+    cast to bf16, and non-floating inputs are rejected fast."""
+    import torch
+
+    from fastvideo.layers.quantization.nvfp4_config import (
+        _coerce_fp4_input_dtype)
+
+    # bf16 / fp16 pass through untouched.
+    bf16 = torch.zeros(4, 8, dtype=torch.bfloat16)
+    assert _coerce_fp4_input_dtype(bf16) is bf16
+    fp16 = torch.zeros(4, 8, dtype=torch.float16)
+    assert _coerce_fp4_input_dtype(fp16) is fp16
+
+    # Other floating dtypes (fp32 from an unfused norm, fp64) -> bf16.
+    assert _coerce_fp4_input_dtype(
+        torch.zeros(4, 8, dtype=torch.float32)).dtype is torch.bfloat16
+    assert _coerce_fp4_input_dtype(
+        torch.zeros(4, 8, dtype=torch.float64)).dtype is torch.bfloat16
+
+    # Non-floating inputs are a real error, not silently cast.
+    for bad in (torch.int32, torch.int64, torch.bool):
+        with pytest.raises(TypeError, match="floating-point"):
+            _coerce_fp4_input_dtype(torch.zeros(4, 8, dtype=bad))
@@ -66,6 +66,7 @@ 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"}
 IDENTITY_KEYS = (
+    "result_schema_version",
    "workload_id",
    "variant_id",
    "benchmark_version",
@@ -77,6 +78,8 @@ IDENTITY_KEYS = (
    "software_profile_id",
    "environment_metadata",
    "environment_fingerprint",
+    "quality_metadata",
+    "variant_metadata",
 )
 COMPARISON_IDENTITY_KEYS = (
    "workload_id",
@@ -143,18 +146,22 @@ def _result_failed_static_thresholds() -> bool:


 def _record_metadata(run_source: str, result: dict[str, Any]) -> dict[str, Any]:
+    raw_run_source = str(result.get("run_source") or "").strip().lower()
+    if raw_run_source in VALID_RUN_SOURCES:
+        run_source = raw_run_source
+
    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", ""),
+        "branch": result.get("branch") or 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", ""),
+        "test_scope": result.get("test_scope") or os.environ.get("TEST_SCOPE", ""),
+        "build_url": result.get("build_url") or os.environ.get("BUILDKITE_BUILD_URL", ""),
+        "build_id": result.get("build_id") or os.environ.get("BUILDKITE_BUILD_ID", ""),
+        "job_id": result.get("job_id") or os.environ.get("BUILDKITE_JOB_ID", ""),
    }


@@ -178,11 +185,6 @@ def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]:
        key for key in COMPARISON_IDENTITY_KEYS
        if key not in record or record[key] is None or (isinstance(record[key], str) and not record[key].strip())
    ]
-    if len(missing) == len(COMPARISON_IDENTITY_KEYS):
-        # Legacy v1 record: no identity fields at all. Keep the pre-v2
-        # (model_id, gpu_type) rolling-baseline comparison instead of
-        # crashing — v1 configs are documented as loadable.
-        return {}
    if missing:
        raise ValueError("Performance record missing required comparison identity fields: " + ", ".join(missing))
    return {
@@ -191,6 +193,18 @@ def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]:
    }


+def _comparison_identity_filters_or_none(record: dict[str, Any]) -> dict[str, str] | None:
+    try:
+        return _comparison_identity_filters(record)
+    except ValueError as exc:
+        print("Skipping rolling baseline comparison for "
+              f"{record.get('model_id', 'unknown')} on "
+              f"{record.get('gpu_type', 'unknown')}: {exc}. "
+              "The normalized artifact will still be written, but this record "
+              "will not be baseline eligible.")
+        return None
+
+
 def _format_identity_filters(filters: dict[str, str]) -> str:
    if not filters:
        return ""
@@ -481,43 +495,53 @@ def main() -> int:
    for raw in current_results:
        record = _normalize_record(raw)
        metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
+        identity_filters = _comparison_identity_filters_or_none(record)

-        baseline_records = load_records_for_model(
-            TRACKING_ROOT,
-            record["model_id"],
-            record["gpu_type"],
-            **_comparison_identity_filters(record),
-            last_n=5,
-            successful_only=True,
-            baseline_eligible_only=True,
-        )
-
-        if not baseline_records:
-            # A brand-new cohort has NO regression gating until history
-            # accumulates — make that loud and machine-readable instead of
-            # an indistinguishable pass, so a cohort shift (intended or
-            # accidental, e.g. an identity-field change) never silently
-            # blinds the comparison.
-            record["baseline_status"] = "initialized_new_cohort"
-            print("=" * 72)
-            print(f"WARNING: NO BASELINE — initializing a NEW cohort for "
-                  f"{record['model_id']} on {record['gpu_type']}"
-                  f"{_format_identity_filters(_comparison_identity_filters(record))}")
-            print("Regression gating is INACTIVE for this cohort until "
-                  "baseline history accumulates. If this cohort shift is "
-                  "unexpected, check the identity fields above.")
-            print("=" * 72)
+        baseline_records: list[dict[str, Any]] = []
+        if identity_filters is None:
+            # Records without the full v2 identity block skip rolling-baseline
+            # comparison entirely: only the static thresholds gate them and
+            # they never become baseline eligible.
+            record["baseline_status"] = "skipped_missing_identity"
            failures: list[str] = []
-            record["success"] = True
        else:
-            record["baseline_status"] = "compared"
-            failures = _check_regressions(record, baseline_records, metric_policies)
+            baseline_records = load_records_for_model(
+                TRACKING_ROOT,
+                record["model_id"],
+                record["gpu_type"],
+                **identity_filters,
+                last_n=5,
+                successful_only=True,
+                baseline_eligible_only=True,
+            )
+
+            if not baseline_records:
+                # A brand-new cohort has NO regression gating until history
+                # accumulates — make that loud and machine-readable instead of
+                # an indistinguishable pass, so a cohort shift (intended or
+                # accidental, e.g. an identity-field change) never silently
+                # blinds the comparison.
+                record["baseline_status"] = "initialized_new_cohort"
+                print("=" * 72)
+                print(f"WARNING: NO BASELINE — initializing a NEW cohort for "
+                      f"{record['model_id']} on {record['gpu_type']}"
+                      f"{_format_identity_filters(identity_filters)}")
+                print("Regression gating is INACTIVE for this cohort until "
+                      "baseline history accumulates. If this cohort shift is "
+                      "unexpected, check the identity fields above.")
+                print("=" * 72)
+                failures = []
+            else:
+                record["baseline_status"] = "compared"
+                failures = _check_regressions(record, baseline_records, metric_policies)
        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"])
+        record["baseline_eligible"] = (
+            identity_filters is not None and _is_baseline_eligible(record["run_source"], record["success"])
+        )
        all_failures.extend(failures)

        _write_normalized_artifact(record)
@@ -14,9 +14,9 @@ def _v2_config():
    return {
        "benchmark_id": "wan-t2v-1.3b-2gpu",
        "config_schema_version": 2,
-        "workload_id": "wan-t2v-1.3b",
-        "variant_id": "canonical",
-        "benchmark_version": 1,
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
    }


@@ -41,9 +41,9 @@ def test_v2_benchmark_config_identity_validates_and_is_preserved():
    assert _is_v2_config(cfg) is True
    assert _config_identity_metadata(cfg) == {
        "config_schema_version": 2,
-        "workload_id": "wan-t2v-1.3b",
-        "variant_id": "canonical",
-        "benchmark_version": 1,
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
        "quality_metadata": {"some": "data"},
    }

@@ -91,9 +91,9 @@ def test_v2_benchmark_config_rejects_invalid_benchmark_version_values(value):
 def test_partial_v2_identity_requires_schema_version():
    cfg = {
        "benchmark_id": "wan-t2v-1.3b-2gpu",
-        "workload_id": "wan-t2v-1.3b",
-        "variant_id": "canonical",
-        "benchmark_version": 1,
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
    }

    expected = "wan.json: v2 benchmark identity fields require config_schema_version=2"
@@ -105,9 +105,9 @@ def test_optional_v2_metadata_fields_must_be_objects():
    cfg = {
        "benchmark_id": "wan-t2v-1.3b-2gpu",
        "config_schema_version": 2,
-        "workload_id": "wan-t2v-1.3b",
-        "variant_id": "canonical",
-        "benchmark_version": 1,
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
        "quality_metadata": ["not", "an", "object"],
    }

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0

+import json
+
 import pytest

 from fastvideo.tests.performance import compare_baseline
@@ -125,6 +127,7 @@ def test_normalized_record_preserves_identity_metadata(monkeypatch):
    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
    raw = _raw_result()
    raw.update({
+        "result_schema_version": 2,
        "workload_id": "wan-t2v",
        "variant_id": "1.3b-sp2",
        "benchmark_version": 2,
@@ -146,10 +149,14 @@ def test_normalized_record_preserves_identity_metadata(monkeypatch):
            },
        },
        "environment_fingerprint": "env-1",
+        "quality_metadata": {
+            "quality_status": "canonical",
+        },
    })

    record = compare_baseline.normalize_performance_result(raw)

+    assert record["result_schema_version"] == 2
    assert record["workload_id"] == "wan-t2v"
    assert record["variant_id"] == "1.3b-sp2"
    assert record["benchmark_version"] == 2
@@ -161,6 +168,87 @@ def test_normalized_record_preserves_identity_metadata(monkeypatch):
    assert record["software_profile_id"] == "sw-1"
    assert record["environment_metadata"] == {"env": {"IMAGE_VERSION": "latest"}}
    assert record["environment_fingerprint"] == "env-1"
+    assert record["quality_metadata"] == {"quality_status": "canonical"}
+
+
+def test_normalized_record_prefers_raw_v2_provenance(monkeypatch):
+    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
+    monkeypatch.setenv("BUILDKITE_BRANCH", "feature/from-env")
+    monkeypatch.setenv("TEST_SCOPE", "direct")
+    monkeypatch.setenv("BUILDKITE_BUILD_URL", "https://buildkite.example/env")
+    monkeypatch.setenv("BUILDKITE_BUILD_ID", "env-build")
+    monkeypatch.setenv("BUILDKITE_JOB_ID", "env-job")
+    raw = _raw_result()
+    raw.update({
+        "result_schema_version": 2,
+        "run_source": "scheduled_main",
+        "branch": "main",
+        "test_scope": "full",
+        "build_url": "https://buildkite.example/raw",
+        "build_id": "raw-build",
+        "job_id": "raw-job",
+        "pr_number": "false",
+    })
+
+    record = compare_baseline.normalize_performance_result(raw)
+
+    assert record["result_schema_version"] == 2
+    assert record["run_source"] == "scheduled_main"
+    assert record["branch"] == "main"
+    assert record["test_scope"] == "full"
+    assert record["build_url"] == "https://buildkite.example/raw"
+    assert record["build_id"] == "raw-build"
+    assert record["job_id"] == "raw-job"
+    assert record["pr_number"] == ""
+
+
+def test_v1_normalized_record_has_no_result_schema_version(monkeypatch):
+    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
+
+    record = compare_baseline.normalize_performance_result(_raw_result())
+
+    assert "result_schema_version" not in record
+
+
+def test_main_writes_v1_current_artifact_without_comparison_identity(monkeypatch, tmp_path, capsys):
+    results_dir = tmp_path / "results"
+    reports_dir = tmp_path / "reports"
+    tracking_root = tmp_path / "tracking"
+    results_dir.mkdir()
+    (results_dir / "perf_legacy.json").write_text(json.dumps(_raw_result()), encoding="utf-8")
+    uploaded_records = []
+
+    monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main")
+    monkeypatch.delenv("PERF_PYTEST_RC", raising=False)
+    monkeypatch.setattr(compare_baseline, "RESULTS_DIR", str(results_dir))
+    monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir))
+    monkeypatch.setattr(compare_baseline, "TRACKING_ROOT", str(tracking_root))
+    monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "always")
+    monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda local_dir, strict=False: local_dir)
+
+    def fail_load_records_for_model(*_args, **_kwargs):
+        raise AssertionError("legacy current records should skip v2 baseline lookup")
+
+    def fake_upload_record(_path, record, *, strict=False):
+        uploaded_records.append(record.copy())
+
+    monkeypatch.setattr(compare_baseline, "load_records_for_model", fail_load_records_for_model)
+    monkeypatch.setattr(compare_baseline, "upload_record", fake_upload_record)
+
+    assert compare_baseline.main() == 0
+
+    output = capsys.readouterr().out
+    normalized_files = list((reports_dir / "results").glob("normalized_perf_*.json"))
+    assert "Skipping rolling baseline comparison" in output
+    assert "workload_id" in output
+    assert len(normalized_files) == 1
+
+    normalized = json.loads(normalized_files[0].read_text(encoding="utf-8"))
+    assert "result_schema_version" not in normalized
+    assert normalized["success"] is True
+    assert normalized["baseline_eligible"] is False
+    assert len(uploaded_records) == 1
+    assert uploaded_records[0]["baseline_eligible"] is False


 def test_normalized_record_reads_identity_labels_from_recipe(monkeypatch):
@@ -154,8 +154,8 @@ def test_dashboard_endpoints_filter_and_return_run_source_metadata():
    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["rows"][0]["baseline_n"] == 0
+    assert summary["rows"][0]["metrics"]["latency"]["baseline"] is None
    assert summary["filters"]["run_source"] == "pr"
    assert trends["count"] == 1
    assert trends["groups"][0]["points"][0]["run_source"] == "pr"
@@ -142,6 +142,47 @@ def test_build_latest_summary_separates_informational_threshold_crossing():
    assert rows[0]["threshold_exceeded_metrics"] == ["latency"]
    assert rows[0]["failing_metrics"] == []
    assert rows[0]["computed_regression_status"] == "pass"
+
+
+def test_build_latest_summary_run_source_filter_excludes_future_baselines():
+    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",
+        ),
+        _record(
+            "2026-01-03T00:00:00+00:00",
+            "c" * 40,
+            30.0,
+            3.0,
+            run_source="scheduled_main",
+            baseline_eligible=True,
+        ),
+    ]
+
+    rows = build_latest_summary(records, run_source="pr")
+
+    assert len(rows) == 1
+    assert rows[0]["run_source"] == "pr"
+    assert rows[0]["baseline_n"] == 1
+    assert rows[0]["metrics"]["latency"]["baseline"] == 10.0
+    assert rows[0]["metrics"]["throughput"]["baseline"] == 10.0
+
+
+
 def test_build_latest_summary_keeps_identity_cohorts_separate():
    records = [
        _record(
@@ -249,6 +290,58 @@ def test_build_latest_summary_keeps_variant_versions_separate():
    assert version_2_row["metrics"]["latency"]["baseline"] == 20.0


+def test_dashboard_identity_preserves_zero_benchmark_version():
+    records = [
+        _record(
+            "2026-01-01T00:00:00+00:00",
+            "a" * 40,
+            10.0,
+            10.0,
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=0,
+            recipe_fingerprint="recipe-a",
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+            run_source="scheduled_main",
+            baseline_eligible=True,
+        ),
+        _record(
+            "2026-01-02T00:00:00+00:00",
+            "b" * 40,
+            11.0,
+            9.0,
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=0,
+            recipe_fingerprint="recipe-a",
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+        ),
+        _record(
+            "2026-01-03T00:00:00+00:00",
+            "c" * 40,
+            20.0,
+            5.0,
+        ),
+    ]
+
+    rows = build_latest_summary(records)
+    version_zero_row = next(row for row in rows if row["benchmark_version"] == 0)
+    legacy_row = next(row for row in rows if row["benchmark_version"] == "")
+    trends = build_trends(records)
+    version_zero_trend = next(trend for trend in trends if trend["benchmark_version"] == 0)
+    legacy_trend = next(trend for trend in trends if trend["benchmark_version"] == "")
+
+    assert len(rows) == 2
+    assert version_zero_row["baseline_n"] == 1
+    assert version_zero_row["metrics"]["latency"]["baseline"] == 10.0
+    assert legacy_row["baseline_n"] == 0
+    assert len(trends) == 2
+    assert version_zero_trend["points"][0]["benchmark_version"] == 0
+    assert legacy_trend["points"][0]["benchmark_version"] == ""
+
+
 def test_filter_records_and_trends_preserve_metric_points():
    records = [
        _record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
@@ -55,6 +55,9 @@ V2_OPTIONAL_METADATA_FIELDS = (
    "metric_threshold_policy",
    "quality_metadata",
 )
+RESULT_SCHEMA_VERSION = 2
+VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
+OPTIONAL_RESULT_METADATA_FIELDS = ("quality_metadata", "variant_metadata")

 # -- Config discovery -------------------------------------------------------

@@ -231,6 +234,17 @@ def _run_generation(generator, prompt, generation_kwargs):
    component_times = _extract_component_times(result)
    return elapsed, peak_memory_mb, component_times

+
+def _validate_run_counts(run_config: Mapping[str, Any], benchmark_id: str) -> tuple[int, int]:
+    num_warmup = run_config.get("num_warmup_runs", 1)
+    num_measure = run_config.get("num_measurement_runs", 3)
+    if isinstance(num_warmup, bool) or not isinstance(num_warmup, int) or num_warmup < 0:
+        raise ValueError(f"{benchmark_id}: run_config.num_warmup_runs must be a non-negative integer")
+    if isinstance(num_measure, bool) or not isinstance(num_measure, int) or num_measure < 1:
+        raise ValueError(f"{benchmark_id}: run_config.num_measurement_runs must be a positive integer")
+    return num_warmup, num_measure
+
+
 def _write_results(results):
    """Write JSON results to the results directory."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
@@ -351,6 +365,101 @@ def _build_identity_fields(cfg, init_kwargs, prompt, runtime_identity):
    }


+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 _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 _ci_provenance_fields() -> dict[str, str]:
+    pr_number = os.environ.get("BUILDKITE_PULL_REQUEST", "")
+    if not _truthy_pr_number(pr_number):
+        pr_number = ""
+    return {
+        "run_source": _detect_run_source(),
+        "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 _configured_result_metadata(cfg: Mapping[str, Any]) -> dict[str, Any]:
+    return {
+        field: cfg[field]
+        for field in OPTIONAL_RESULT_METADATA_FIELDS
+        if field in cfg and cfg[field] is not None
+    }
+
+
+def _build_result_record(
+    *,
+    cfg: Mapping[str, Any],
+    model_info: Mapping[str, Any],
+    init_kwargs: Mapping[str, Any],
+    gen_kwargs: Mapping[str, Any],
+    num_warmup: int,
+    num_measure: int,
+    thresholds: Mapping[str, Any],
+    times: list[float],
+    peak_memories: list[float],
+    all_component_times: list[dict],
+    prompt: str,
+    runtime_identity: Mapping[str, Any],
+    device_name: str,
+    timestamp: str | None = None,
+) -> dict[str, Any]:
+    if not times or not peak_memories:
+        raise ValueError("Cannot build a performance result record without measurement runs")
+    avg_time = sum(times) / len(times)
+    max_peak_memory = max(peak_memories)
+    num_frames = gen_kwargs.get("num_frames")
+    throughput_fps = (1.0 / avg_time) if avg_time > 0 else None
+    if isinstance(num_frames, (int, float)) and avg_time > 0:
+        throughput_fps = num_frames / avg_time
+
+    return {
+        "benchmark_id": cfg["benchmark_id"],
+        "result_schema_version": RESULT_SCHEMA_VERSION,
+        "model_short_name": model_info.get("model_short_name", ""),
+        "device": device_name,
+        "num_gpus": init_kwargs.get("num_gpus", 1),
+        "num_warmup_runs": num_warmup,
+        "num_measurement_runs": num_measure,
+        "avg_generation_time_s": round(avg_time, 3),
+        "individual_times_s": [round(t, 3) for t in times],
+        "throughput_fps": round(throughput_fps, 3)
+        if throughput_fps is not None else None,
+        "max_peak_memory_mb": round(max_peak_memory, 1),
+        "individual_peak_memories_mb": [round(m, 1) for m in peak_memories],
+        "thresholds": dict(thresholds),
+        "regression_thresholds": cfg.get("regression_thresholds", {}),
+        "commit": os.environ.get("BUILDKITE_COMMIT", ""),
+        **_ci_provenance_fields(),
+        "timestamp": timestamp or datetime.now(timezone.utc).isoformat(),
+        **_configured_result_metadata(cfg),
+        "text_encoder_time_s": _avg_component(all_component_times,
+                                              "text_encoder_time_s"),
+        "dit_time_s": _avg_component(all_component_times, "dit_time_s"),
+        "vae_decode_time_s": _avg_component(all_component_times,
+                                            "vae_decode_time_s"),
+        **_build_identity_fields(cfg, init_kwargs, prompt, runtime_identity),
+    }
+
+
 # -- Test -------------------------------------------------------------------

 def _run_benchmark(cfg):
@@ -366,8 +475,7 @@ def _run_benchmark(cfg):
    prompts = cfg.get("test_prompts", ["A cinematic video."])
    prompt = prompts[0]

-    num_warmup = run_config.get("num_warmup_runs", 1)
-    num_measure = run_config.get("num_measurement_runs", 3)
+    num_warmup, num_measure = _validate_run_counts(run_config, cfg["benchmark_id"])
    thresholds = _get_thresholds(cfg)

    # Remap JSON keys to VideoGenerator kwargs
@@ -414,37 +522,21 @@ def _run_benchmark(cfg):
    avg_time = sum(times) / len(times)
    max_peak_memory = max(peak_memories)
    device_name = torch.cuda.get_device_name()
-    num_frames = gen_kwargs.get("num_frames")
-    throughput_fps = (1.0 / avg_time) if avg_time > 0 else None
-    if isinstance(num_frames, (int, float)) and avg_time > 0:
-        throughput_fps = num_frames / avg_time
-
-    results = {
-        "benchmark_id": cfg["benchmark_id"],
-        **_config_identity_metadata(cfg),
-        "model_short_name": model_info.get("model_short_name", ""),
-        "device": device_name,
-        "num_gpus": init_kwargs.get("num_gpus", 1),
-        "num_warmup_runs": num_warmup,
-        "num_measurement_runs": num_measure,
-        "avg_generation_time_s": round(avg_time, 3),
-        "individual_times_s": [round(t, 3) for t in times],
-        "throughput_fps": round(throughput_fps, 3)
-        if throughput_fps is not None else None,
-        "max_peak_memory_mb": round(max_peak_memory, 1),
-        "individual_peak_memories_mb": [round(m, 1) for m in peak_memories],
-        "thresholds": thresholds,
-        "regression_thresholds": cfg.get("regression_thresholds", {}),
-        "commit": os.environ.get("BUILDKITE_COMMIT", ""),
-        "pr_number": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-        "text_encoder_time_s": _avg_component(all_component_times,
-                                              "text_encoder_time_s"),
-        "dit_time_s": _avg_component(all_component_times, "dit_time_s"),
-        "vae_decode_time_s": _avg_component(all_component_times,
-                                            "vae_decode_time_s"),
-        **_build_identity_fields(cfg, init_kwargs, prompt, runtime_identity),
-    }
+    results = _build_result_record(
+        cfg=cfg,
+        model_info=model_info,
+        init_kwargs=init_kwargs,
+        gen_kwargs=gen_kwargs,
+        num_warmup=num_warmup,
+        num_measure=num_measure,
+        thresholds=thresholds,
+        times=times,
+        peak_memories=peak_memories,
+        all_component_times=all_component_times,
+        prompt=prompt,
+        runtime_identity=runtime_identity,
+        device_name=device_name,
+    )

    logger.info(
        "Performance results: avg_time=%.2fs, "
@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from fastvideo.tests.performance import test_inference_performance as perf_test
+
+
+def _benchmark_config():
+    return {
+        "benchmark_id": "wan-t2v-1.3b-2gpu",
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "quality_metadata": {
+            "quality_status": "canonical",
+        },
+        "model": {
+            "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+            "model_short_name": "Wan2.1-T2V-1.3B",
+        },
+        "generation_kwargs": {
+            "num_frames": 45,
+        },
+    }
+
+
+def _identity_fields():
+    return {
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "recipe": {
+            "recipe_schema_version": 1,
+            "benchmark": {
+                "benchmark_id": "wan-t2v-1.3b-2gpu",
+                "workload_id": "wan-t2v",
+                "variant_id": "1.3b-sp2",
+                "benchmark_version": 2,
+            },
+        },
+        "recipe_fingerprint": "recipe-1",
+        "hardware_profile": {
+            "device_type": "cuda",
+            "gpu_count": 2,
+        },
+        "hardware_profile_id": "hw-l40s",
+        "software_profile": {
+            "python": "3.12",
+            "pytorch": "2.12",
+            "cuda": "13.0",
+        },
+        "software_profile_id": "sw-cu130",
+        "environment_metadata": {
+            "env": {
+                "IMAGE_VERSION": "py3.12-cuda13.0.0",
+            },
+        },
+        "environment_fingerprint": "env-ci",
+    }
+
+
+def test_build_result_record_emits_v2_wan_shape(monkeypatch):
+    monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main")
+    monkeypatch.setenv("BUILDKITE_COMMIT", "a" * 40)
+    monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false")
+    monkeypatch.setenv("BUILDKITE_BRANCH", "main")
+    monkeypatch.setenv("TEST_SCOPE", "full")
+    monkeypatch.setenv("BUILDKITE_BUILD_URL", "https://buildkite.example/build")
+    monkeypatch.setenv("BUILDKITE_BUILD_ID", "build-1")
+    monkeypatch.setenv("BUILDKITE_JOB_ID", "job-1")
+    monkeypatch.setattr(perf_test, "_build_identity_fields", lambda *_args: _identity_fields())
+
+    record = perf_test._build_result_record(
+        cfg=_benchmark_config(),
+        model_info={"model_short_name": "Wan2.1-T2V-1.3B"},
+        init_kwargs={"num_gpus": 2},
+        gen_kwargs={"num_frames": 45},
+        num_warmup=2,
+        num_measure=3,
+        thresholds={
+            "max_generation_time_s": 34.0,
+            "max_peak_memory_mb": 11000.0,
+        },
+        times=[30.0, 31.0, 32.0],
+        peak_memories=[10000.0, 10100.0, 10050.0],
+        all_component_times=[
+            {
+                "text_encoder_time_s": 1.0,
+                "dit_time_s": 8.0,
+                "vae_decode_time_s": 3.0,
+            },
+            {
+                "text_encoder_time_s": 1.2,
+                "dit_time_s": 8.2,
+                "vae_decode_time_s": 3.2,
+            },
+            {
+                "text_encoder_time_s": None,
+                "dit_time_s": 8.4,
+                "vae_decode_time_s": 3.4,
+            },
+        ],
+        prompt="A cinematic video.",
+        runtime_identity={
+            "resolved_attention_backend": "FLASH_ATTN",
+        },
+        device_name="NVIDIA L40S",
+        timestamp="2026-07-05T00:00:00+00:00",
+    )
+
+    assert record["result_schema_version"] == perf_test.RESULT_SCHEMA_VERSION
+    assert record["benchmark_id"] == "wan-t2v-1.3b-2gpu"
+    assert record["workload_id"] == "wan-t2v"
+    assert record["variant_id"] == "1.3b-sp2"
+    assert record["benchmark_version"] == 2
+    assert record["recipe_fingerprint"] == "recipe-1"
+    assert record["hardware_profile_id"] == "hw-l40s"
+    assert record["software_profile_id"] == "sw-cu130"
+    assert record["environment_fingerprint"] == "env-ci"
+    assert record["environment_metadata"]["env"]["IMAGE_VERSION"] == "py3.12-cuda13.0.0"
+    assert record["quality_metadata"] == {"quality_status": "canonical"}
+    assert record["run_source"] == "scheduled_main"
+    assert record["branch"] == "main"
+    assert record["test_scope"] == "full"
+    assert record["build_url"] == "https://buildkite.example/build"
+    assert record["build_id"] == "build-1"
+    assert record["job_id"] == "job-1"
+    assert record["pr_number"] == ""
+    assert record["commit"] == "a" * 40
+    assert record["avg_generation_time_s"] == 31.0
+    assert record["throughput_fps"] == 1.452
+    assert record["max_peak_memory_mb"] == 10100.0
+    assert record["text_encoder_time_s"] == 1.1
+    assert record["dit_time_s"] == 8.2
+    assert record["vae_decode_time_s"] == 3.2
+
+
+def test_validate_run_counts_rejects_zero_measurement_runs():
+    with pytest.raises(ValueError, match="num_measurement_runs"):
+        perf_test._validate_run_counts({
+            "num_warmup_runs": 1,
+            "num_measurement_runs": 0,
+        }, "wan-t2v-1.3b-2gpu")
+
+
+def test_validate_run_counts_rejects_negative_warmup_runs():
+    with pytest.raises(ValueError, match="num_warmup_runs"):
+        perf_test._validate_run_counts({
+            "num_warmup_runs": -1,
+            "num_measurement_runs": 3,
+        }, "wan-t2v-1.3b-2gpu")
+
+
+def test_validate_run_counts_returns_defaults():
+    assert perf_test._validate_run_counts({}, "wan-t2v-1.3b-2gpu") == (1, 3)
+
+
+def test_build_result_record_rejects_empty_measurements(monkeypatch):
+    monkeypatch.setattr(perf_test, "_build_identity_fields", lambda *_args: _identity_fields())
+
+    with pytest.raises(ValueError, match="measurement runs"):
+        perf_test._build_result_record(
+            cfg=_benchmark_config(),
+            model_info={},
+            init_kwargs={},
+            gen_kwargs={},
+            num_warmup=0,
+            num_measure=0,
+            thresholds={},
+            times=[],
+            peak_memories=[],
+            all_component_times=[],
+            prompt="A cinematic video.",
+            runtime_identity={},
+            device_name="NVIDIA L40S",
+        )
@@ -40,6 +40,9 @@ to override the HF reference repo at test time:
 to skip auto-download of missing refs:
 `pytest fastvideo/tests/ssim/ -vs --skip-ssim-reference-download`

+to bootstrap draft refs for a new model:
+`pytest fastvideo/tests/ssim/ -vs --ssim-bootstrap-mode`
+
 generated videos are written under:
 - default params: `generated_videos/default/<GPU>_reference_videos/<model_id>/<ATTENTION_BACKEND>/`
 - full-quality params: `generated_videos/full_quality/<GPU>_reference_videos/<model_id>/<ATTENTION_BACKEND>/`
@@ -48,6 +51,10 @@ HF repo layout mirrors quality + GPU split:
 - `reference_videos/default/<GPU>_reference_videos/...`
 - `reference_videos/full_quality/<GPU>_reference_videos/...`

+Draft bootstrap refs are uploaded under:
+- `drafts/default/<GPU>_reference_videos/<model_id>/<ATTENTION_BACKEND>/...`
+- `drafts/full_quality/<GPU>_reference_videos/<model_id>/<ATTENTION_BACKEND>/...`
+
 reference videos were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
 causal videos were generated on commit b318063c0a4618f1d5d99ea82ca67a06aad0d19d

@@ -73,9 +80,35 @@ If `REQUIRED_GPUS` is omitted, the test defaults to 1 GPU.
 For files that define model maps (`*_MODEL_TO_PARAMS`), CI splits execution
 into one subprocess per model id by setting `FASTVIDEO_SSIM_MODEL_ID`.

+## Bootstrapping New References
+
+Normal SSIM runs are strict: missing references fail the test. For a new model
+PR, add `[new-model]` to the PR title or set `FASTVIDEO_SSIM_BOOTSTRAP_MODE=1`
+for the Buildkite job. CI then passes `--ssim-bootstrap-mode` to pytest.
+
+In bootstrap mode, a missing pixel `.mp4` or latent `.pt` reference is treated
+as an expected draft-reference case after the generated artifact has been
+written. The test uploads that generated artifact to the HF reference repo under
+`drafts/...` and marks the case `xfail`. Bootstrap mode does not weaken normal
+CI because it is opt-in only.
+
+After reviewing a draft, promote it into the canonical reference layout:
+
+```bash
+python fastvideo/tests/ssim/reference_videos_cli.py promote-draft \
+  --quality-tier default \
+  --device-folder L40S_reference_videos \
+  --model-id <model_id>
+```
+
+Use `--attention-backend <backend>` to promote one backend folder. Promotion
+refuses to overwrite existing canonical refs by default; pass `--force` only
+when intentionally replacing reviewed references.
+
 ## Generation Details

-2 x NVIDIA L40S GPUs
+SSIM CI currently schedules work across the GPUs configured in
+`fastvideo/tests/modal/ssim_test.py`.

 ## Generation Parameters

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+
+from fastvideo.tests.ssim.reference_utils import get_output_quality_tier
+from fastvideo.tests.ssim.reference_videos_cli import (
+    BOOTSTRAP_ENV_KEY,
+    DEFAULT_REPO_ID,
+    DEFAULT_REPO_TYPE,
+    HF_REPO_ENV_KEY,
+    HF_REPO_TYPE_ENV_KEY,
+    upload_draft_reference_artifact,
+)
+
+TRUE_VALUES = {"1", "true", "yes", "on"}
+
+
+def bootstrap_mode_enabled() -> bool:
+    return os.environ.get(BOOTSTRAP_ENV_KEY, "").strip().lower() in TRUE_VALUES
+
+
+def xfail_missing_reference_in_bootstrap_mode(
+    *,
+    generated_artifact_path: str,
+    reference_folder: str,
+    artifact_kind: str,
+) -> None:
+    if not bootstrap_mode_enabled():
+        return
+
+    generated_path = Path(generated_artifact_path)
+    if not generated_path.exists():
+        raise FileNotFoundError(
+            f"SSIM bootstrap mode is enabled, but generated {artifact_kind} artifact is missing: {generated_path}"
+        )
+
+    repo_id = os.environ.get(HF_REPO_ENV_KEY, DEFAULT_REPO_ID)
+    repo_type = os.environ.get(HF_REPO_TYPE_ENV_KEY, DEFAULT_REPO_TYPE)
+    draft_path = upload_draft_reference_artifact(
+        repo_id=repo_id,
+        repo_type=repo_type,
+        generated_artifact_path=generated_path,
+        reference_folder=Path(reference_folder),
+    )
+    pytest.xfail(
+        "SSIM bootstrap mode generated a draft "
+        f"{artifact_kind} reference at {repo_id}/{draft_path}. "
+        "Review it, then promote with "
+        "`python fastvideo/tests/ssim/reference_videos_cli.py promote-draft "
+        f"--quality-tier {get_output_quality_tier()} --model-id <model_id>`."
+    )
@@ -5,6 +5,7 @@ from fastvideo.tests.ssim.reference_utils import (
    get_output_quality_tier,
 )
 from fastvideo.tests.ssim.reference_videos_cli import (
+    BOOTSTRAP_ENV_KEY,
    HF_REPO_ENV_KEY,
    ensure_reference_videos_available,
 )
@@ -34,6 +35,12 @@ def pytest_addoption(parser):
        default=False,
        help="Skip auto-download of missing SSIM reference videos from HF.",
    )
+    parser.addoption(
+        "--ssim-bootstrap-mode",
+        action="store_true",
+        default=False,
+        help="Treat missing SSIM references as draft-reference bootstrap cases.",
+    )


 def pytest_configure(config):
@@ -44,6 +51,9 @@ def pytest_configure(config):
    if repo_id:
        os.environ[HF_REPO_ENV_KEY] = repo_id

+    if config.getoption("--ssim-bootstrap-mode"):
+        os.environ[BOOTSTRAP_ENV_KEY] = "1"
+
    skip_download = config.getoption("--skip-ssim-reference-download")
    skip_download = skip_download or os.environ.get(
        "FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD", ""
@@ -8,6 +8,9 @@ from contextlib import contextmanager
 from logging import Logger

 from fastvideo import VideoGenerator
+from fastvideo.tests.ssim.bootstrap_references import (
+    xfail_missing_reference_in_bootstrap_mode,
+)
 from fastvideo.tests.ssim.reference_utils import (
    build_generated_output_dir,
    build_reference_folder_path,
@@ -85,8 +88,14 @@ def _assert_similarity(
    model_id: str,
    attention_backend_name: str,
 ) -> None:
+    generated_video_path = os.path.join(output_dir, output_video_name)
    if not os.path.exists(reference_folder):
        logger.error("Reference folder missing: %s", reference_folder)
+        xfail_missing_reference_in_bootstrap_mode(
+            generated_artifact_path=generated_video_path,
+            reference_folder=reference_folder,
+            artifact_kind="video",
+        )
        error_msg = (
            f"Reference video folder does not exist: {reference_folder}\n"
            f"To download reference videos, run:\n"
@@ -102,10 +111,13 @@ def _assert_similarity(
            prompt,
            attention_backend_name,
        )
+        xfail_missing_reference_in_bootstrap_mode(
+            generated_artifact_path=generated_video_path,
+            reference_folder=reference_folder,
+            artifact_kind="video",
+        )
        raise error

-    generated_video_path = os.path.join(output_dir, output_video_name)
-
    logger.info(
        "Computing SSIM between %s and %s",
        reference_video_path,
@@ -63,6 +63,9 @@ import torch
 from torch.nn.functional import cosine_similarity

 from fastvideo import VideoGenerator
+from fastvideo.tests.ssim.bootstrap_references import (
+    xfail_missing_reference_in_bootstrap_mode,
+)
 from fastvideo.tests.ssim.inference_similarity_utils import (
    attention_backend,
    build_generation_kwargs,
@@ -465,6 +468,11 @@ def run_text_to_latent_similarity_test(
        attention_backend_name,
    )
    if not os.path.exists(reference_folder):
+        xfail_missing_reference_in_bootstrap_mode(
+            generated_artifact_path=generated_latent_path,
+            reference_folder=reference_folder,
+            artifact_kind="latent",
+        )
        raise FileNotFoundError(
            f"Reference folder does not exist: {reference_folder}\n"
            f"To download references, run:\n"
@@ -475,6 +483,11 @@ def run_text_to_latent_similarity_test(
        output_latent_name,
    )
    if not os.path.exists(reference_latent_path):
+        xfail_missing_reference_in_bootstrap_mode(
+            generated_artifact_path=generated_latent_path,
+            reference_folder=reference_folder,
+            artifact_kind="latent",
+        )
        raise FileNotFoundError(
            "Reference latent missing for prompt/backend: "
            f"{reference_latent_path}")
@@ -22,6 +22,7 @@ REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS
 HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN")
 HF_REPO_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO"
 HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE"
+BOOTSTRAP_ENV_KEY = "FASTVIDEO_SSIM_BOOTSTRAP_MODE"

 DEFAULT_REPO_ID = "FastVideo/ssim-reference-videos"
 DEFAULT_REPO_TYPE = "dataset"
@@ -29,6 +30,7 @@ DEFAULT_OUTPUT_QUALITY_TIER = "default"
 FULL_OUTPUT_QUALITY_TIER = "full_quality"
 ALL_OUTPUT_QUALITY_TIERS = "all"
 REFERENCE_VIDEOS_DIRNAME = "reference_videos"
+DRAFTS_DIRNAME = "drafts"
 DEFAULT_DEVICE_REFERENCE_FOLDER = "L40S_reference_videos"
 QUALITY_TIERS = (
    DEFAULT_OUTPUT_QUALITY_TIER,
@@ -317,6 +319,159 @@ def upload_reference_videos(
        )


+def _reference_folder_repo_relative(reference_folder: Path, base_dir: Path | None = None) -> Path:
+    resolved_reference_folder = reference_folder.resolve()
+    resolved_base_dir = (base_dir or _ssim_dir()).resolve()
+    tiered_root = resolved_base_dir / REFERENCE_VIDEOS_DIRNAME
+    try:
+        relative = resolved_reference_folder.relative_to(tiered_root)
+    except ValueError:
+        try:
+            legacy_relative = resolved_reference_folder.relative_to(resolved_base_dir)
+        except ValueError as exc:
+            raise RuntimeError(
+                f"Reference folder {reference_folder} is not under SSIM directory {resolved_base_dir}"
+            ) from exc
+        if not legacy_relative.parts or not legacy_relative.parts[0].endswith("_reference_videos"):
+            raise RuntimeError(
+                f"Reference folder {reference_folder} does not match reference_videos/<tier>/<device>/<model>/<backend>"
+            )
+        relative = Path(DEFAULT_OUTPUT_QUALITY_TIER, *legacy_relative.parts)
+
+    if len(relative.parts) < 4:
+        raise RuntimeError(
+            f"Reference folder {reference_folder} does not include <quality>/<device>/<model>/<backend>"
+        )
+    return relative
+
+
+def upload_draft_reference_artifact(
+    *,
+    repo_id: str,
+    repo_type: str,
+    generated_artifact_path: Path,
+    reference_folder: Path,
+    token: str | None = None,
+    base_dir: Path | None = None,
+    private: bool = False,
+) -> str:
+    resolved_token = token or _get_hf_token()
+    if resolved_token is None:
+        raise RuntimeError(
+            "Hugging Face API key is required for SSIM bootstrap draft upload. "
+            "Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN."
+        )
+    if not generated_artifact_path.exists():
+        raise FileNotFoundError(f"Generated artifact not found: {generated_artifact_path}")
+
+    relative_reference_folder = _reference_folder_repo_relative(
+        reference_folder=reference_folder,
+        base_dir=base_dir,
+    )
+    path_in_repo = f"{DRAFTS_DIRNAME}/{relative_reference_folder.as_posix()}/{generated_artifact_path.name}"
+
+    HfApi, _ = _load_hf_sdk()
+    api = HfApi(token=resolved_token)
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type=repo_type,
+        private=private,
+        exist_ok=True,
+    )
+    print(f"Uploading SSIM draft reference to {repo_id}/{path_in_repo} ...")
+    api.upload_file(
+        repo_id=repo_id,
+        repo_type=repo_type,
+        path_or_fileobj=str(generated_artifact_path),
+        path_in_repo=path_in_repo,
+        token=resolved_token,
+    )
+    print("Draft reference uploaded. Review before promotion.")
+    return path_in_repo
+
+
+def promote_draft_references(
+    *,
+    repo_id: str,
+    repo_type: str,
+    quality_tier: str,
+    device_folder: str,
+    model_id: str,
+    token: str,
+    attention_backend: str | None = None,
+    private: bool = False,
+    force: bool = False,
+) -> None:
+    HfApi, snapshot_download = _load_hf_sdk()
+    api = HfApi(token=token)
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type=repo_type,
+        private=private,
+        exist_ok=True,
+    )
+    draft_prefix_parts = [
+        DRAFTS_DIRNAME,
+        quality_tier,
+        device_folder,
+        model_id,
+    ]
+    reference_prefix_parts = [
+        REFERENCE_VIDEOS_DIRNAME,
+        quality_tier,
+        device_folder,
+        model_id,
+    ]
+    if attention_backend:
+        draft_prefix_parts.append(attention_backend)
+        reference_prefix_parts.append(attention_backend)
+    draft_prefix = "/".join(draft_prefix_parts)
+    reference_prefix = "/".join(reference_prefix_parts)
+
+    try:
+        existing_repo_files = set(api.list_repo_files(repo_id=repo_id, repo_type=repo_type))
+    except Exception:
+        # Fresh repo or list failure — treat as empty so promotion can report
+        # a clear "no drafts found" error instead of surfacing HF internals.
+        existing_repo_files = set()
+    draft_files = sorted(f for f in existing_repo_files if f.startswith(f"{draft_prefix}/"))
+    if not draft_files:
+        raise RuntimeError(f"No draft references found under {repo_id}/{draft_prefix}")
+
+    conflicts = sorted(
+        f for f in existing_repo_files
+        if f.startswith(f"{reference_prefix}/") or f == reference_prefix)
+    if conflicts and not force:
+        preview = "\n".join(f"  - {c}" for c in conflicts[:10])
+        more = f"\n  ... and {len(conflicts) - 10} more" if len(conflicts) > 10 else ""
+        raise RuntimeError(
+            f"Refusing to overwrite existing HF files under {reference_prefix} "
+            f"({len(conflicts)} file(s) already present):\n{preview}{more}\n"
+            f"Re-run with --force to overwrite.")
+
+    with tempfile.TemporaryDirectory(prefix="fv2-ssim-draft-") as tmp_cache_dir:
+        snapshot_path = Path(
+            snapshot_download(
+                repo_id=repo_id,
+                repo_type=repo_type,
+                cache_dir=tmp_cache_dir,
+                token=token,
+                allow_patterns=[f"{draft_prefix}/**"],
+            )
+        )
+        draft_root = snapshot_path / draft_prefix
+        if not draft_root.exists():
+            raise RuntimeError(f"Downloaded draft path is missing: {draft_root}")
+        print(f"Promoting {len(draft_files)} draft reference file(s) to {repo_id}/{reference_prefix} ...")
+        api.upload_folder(
+            repo_id=repo_id,
+            repo_type=repo_type,
+            folder_path=str(draft_root),
+            path_in_repo=reference_prefix,
+            token=token,
+        )
+
+
 def _resolve_upload_reference_dirs(
    *,
    base_dir: Path,
@@ -559,6 +714,55 @@ def _build_parser() -> argparse.ArgumentParser:
            "clobber existing references."),
    )

+    promote_parser = subparsers.add_parser(
+        "promote-draft",
+        help="Promote reviewed draft references into the canonical HF reference layout.",
+    )
+    promote_parser.add_argument(
+        "--repo-id",
+        default=_default_repo_id(),
+        help=f"HF repo id. Default: env {HF_REPO_ENV_KEY} or {DEFAULT_REPO_ID}",
+    )
+    promote_parser.add_argument(
+        "--repo-type",
+        default=_default_repo_type(),
+        help=f"HF repo type. Default: env {HF_REPO_TYPE_ENV_KEY} or {DEFAULT_REPO_TYPE}",
+    )
+    promote_parser.add_argument(
+        "--quality-tier",
+        choices=(
+            DEFAULT_OUTPUT_QUALITY_TIER,
+            FULL_OUTPUT_QUALITY_TIER,
+        ),
+        default=DEFAULT_OUTPUT_QUALITY_TIER,
+        help="Quality tier to promote.",
+    )
+    promote_parser.add_argument(
+        "--device-folder",
+        default=DEFAULT_DEVICE_REFERENCE_FOLDER,
+        help="GPU reference folder to promote, e.g. L40S_reference_videos.",
+    )
+    promote_parser.add_argument(
+        "--model-id",
+        required=True,
+        help="Model subfolder to promote from drafts.",
+    )
+    promote_parser.add_argument(
+        "--attention-backend",
+        default=None,
+        help="Optional backend subfolder to promote. Default promotes all backends under the model.",
+    )
+    promote_parser.add_argument(
+        "--private",
+        action="store_true",
+        help="Create/use a private repo instead of public.",
+    )
+    promote_parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Allow overwriting existing canonical references. Off by default.",
+    )
+
    ensure_parser = subparsers.add_parser(
        "ensure",
        help="Download refs only when the selected local quality tier is missing.",
@@ -671,6 +875,27 @@ def main(argv: Sequence[str] | None = None) -> int:
        print("Upload complete.")
        return 0

+    if args.command == "promote-draft":
+        token = _get_hf_token()
+        if token is None:
+            raise RuntimeError(
+                "Hugging Face API key is required for promote-draft. "
+                "Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN."
+            )
+        promote_draft_references(
+            repo_id=args.repo_id,
+            repo_type=args.repo_type,
+            quality_tier=args.quality_tier,
+            device_folder=args.device_folder,
+            model_id=args.model_id,
+            attention_backend=args.attention_backend,
+            token=token,
+            private=args.private,
+            force=args.force,
+        )
+        print("Promotion complete.")
+        return 0
+
    if args.command == "ensure":
        downloaded = ensure_reference_videos_available(
            local_dir=args.local_dir,
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+
+import pytest
+
+from fastvideo.tests.ssim import reference_videos_cli
+
+
+class _FakeHfApi:
+    instances = []
+    repo_files = []
+
+    def __init__(self, token):
+        self.token = token
+        self.created_repos = []
+        self.uploaded_files = []
+        self.uploaded_folders = []
+        _FakeHfApi.instances.append(self)
+
+    def create_repo(self, **kwargs):
+        self.created_repos.append(kwargs)
+
+    def upload_file(self, **kwargs):
+        self.uploaded_files.append(kwargs)
+
+    def upload_folder(self, **kwargs):
+        self.uploaded_folders.append(kwargs)
+
+    def list_repo_files(self, **kwargs):
+        return list(self.repo_files)
+
+
+def _fake_snapshot_download(*, allow_patterns, cache_dir, **_kwargs):
+    assert len(allow_patterns) == 1
+    draft_prefix = allow_patterns[0].removesuffix("/**")
+    snapshot_root = Path(cache_dir) / "snapshot"
+    draft_root = snapshot_root / draft_prefix
+    draft_root.mkdir(parents=True)
+    (draft_root / "sample.mp4").write_text("draft", encoding="utf-8")
+    return str(snapshot_root)
+
+
+@pytest.fixture(autouse=True)
+def _fake_hf(monkeypatch):
+    _FakeHfApi.instances = []
+    _FakeHfApi.repo_files = []
+    monkeypatch.setattr(
+        reference_videos_cli,
+        "_load_hf_sdk",
+        lambda: (_FakeHfApi, _fake_snapshot_download),
+    )
+
+
+def test_upload_draft_reference_artifact_uses_drafts_prefix(tmp_path):
+    ssim_dir = tmp_path / "ssim"
+    generated = (
+        ssim_dir
+        / "generated_videos"
+        / "default"
+        / "L40S_reference_videos"
+        / "model-a"
+        / "flash"
+        / "sample.mp4"
+    )
+    generated.parent.mkdir(parents=True)
+    generated.write_text("video", encoding="utf-8")
+    reference_folder = (
+        ssim_dir
+        / "reference_videos"
+        / "default"
+        / "L40S_reference_videos"
+        / "model-a"
+        / "flash"
+    )
+
+    draft_path = reference_videos_cli.upload_draft_reference_artifact(
+        repo_id="FastVideo/ssim-reference-videos",
+        repo_type="dataset",
+        generated_artifact_path=generated,
+        reference_folder=reference_folder,
+        token="hf_token",
+        base_dir=ssim_dir,
+    )
+
+    assert draft_path == "drafts/default/L40S_reference_videos/model-a/flash/sample.mp4"
+    fake_api = _FakeHfApi.instances[-1]
+    assert fake_api.uploaded_files[0]["path_in_repo"] == draft_path
+
+
+def test_promote_draft_refuses_existing_canonical_reference():
+    _FakeHfApi.repo_files = [
+        "drafts/default/L40S_reference_videos/model-a/flash/sample.mp4",
+        "reference_videos/default/L40S_reference_videos/model-a/flash/sample.mp4",
+    ]
+
+    with pytest.raises(RuntimeError, match="Refusing to overwrite existing HF files"):
+        reference_videos_cli.promote_draft_references(
+            repo_id="FastVideo/ssim-reference-videos",
+            repo_type="dataset",
+            quality_tier="default",
+            device_folder="L40S_reference_videos",
+            model_id="model-a",
+            attention_backend="flash",
+            token="hf_token",
+        )
+
+
+def test_promote_draft_uploads_downloaded_draft_folder(tmp_path, monkeypatch):
+    _FakeHfApi.repo_files = [
+        "drafts/default/L40S_reference_videos/model-a/flash/sample.mp4",
+    ]
+
+    def fake_snapshot_download(*, allow_patterns, cache_dir, **_kwargs):
+        snapshot_root = tmp_path / "snapshot"
+        draft_root = snapshot_root / allow_patterns[0].removesuffix("/**")
+        draft_root.mkdir(parents=True)
+        (draft_root / "sample.mp4").write_text("draft", encoding="utf-8")
+        return str(snapshot_root)
+
+    monkeypatch.setattr(
+        reference_videos_cli,
+        "_load_hf_sdk",
+        lambda: (_FakeHfApi, fake_snapshot_download),
+    )
+
+    reference_videos_cli.promote_draft_references(
+        repo_id="FastVideo/ssim-reference-videos",
+        repo_type="dataset",
+        quality_tier="default",
+        device_folder="L40S_reference_videos",
+        model_id="model-a",
+        attention_backend="flash",
+        token="hf_token",
+    )
+
+    fake_api = _FakeHfApi.instances[-1]
+    upload = fake_api.uploaded_folders[0]
+    assert upload["path_in_repo"] == "reference_videos/default/L40S_reference_videos/model-a/flash"
+    assert Path(upload["folder_path"]).name == "flash"
+    shutil.rmtree(tmp_path / "snapshot", ignore_errors=True)
@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for LTX-2 refine upsampler path resolution.
+
+Pure tmp_path tests: no GPU, no model downloads.
+"""
+
+from fastvideo.pipelines.basic.ltx2.ltx2_pipeline import _resolve_refine_upsampler_path
+
+
+def test_model_index_key_wins_over_directory_probe(tmp_path):
+    (tmp_path / "custom_upsampler").mkdir()
+    (tmp_path / "spatial_upscaler").mkdir()
+    model_index = {"fastvideo_refine_upsampler_path": "custom_upsampler"}
+    assert _resolve_refine_upsampler_path(str(tmp_path), model_index) == str(tmp_path / "custom_upsampler")
+
+
+def test_model_index_spatial_upsampler_module_entry(tmp_path):
+    (tmp_path / "spatial_upsampler").mkdir()
+    model_index = {"spatial_upsampler": ["diffusers", "LTX2SpatialUpsampler"]}
+    assert _resolve_refine_upsampler_path(str(tmp_path), model_index) == str(tmp_path / "spatial_upsampler")
+
+
+def test_directory_fallback_finds_spatial_upscaler(tmp_path):
+    # The published LTX-2 distilled checkpoints ship the weights on disk as
+    # "spatial_upscaler" without any model_index.json declaration.
+    (tmp_path / "spatial_upscaler").mkdir()
+    assert _resolve_refine_upsampler_path(str(tmp_path), {}) == str(tmp_path / "spatial_upscaler")
+
+
+def test_directory_fallback_finds_spatial_upsampler(tmp_path):
+    (tmp_path / "spatial_upsampler").mkdir()
+    assert _resolve_refine_upsampler_path(str(tmp_path), {}) == str(tmp_path / "spatial_upsampler")
+
+
+def test_returns_none_when_nothing_resolves(tmp_path):
+    assert _resolve_refine_upsampler_path(str(tmp_path), {}) is None
@@ -14,11 +14,21 @@ distributed init and is exercised by Phase 2/3 tests.
 """
 from __future__ import annotations

+from types import SimpleNamespace
+
+import numpy as np
+import pytest
 import torch

 from fastvideo.train.callbacks.callback import CallbackDict
 from fastvideo.train.callbacks.ema import EMACallback
-from fastvideo.train.callbacks.validation import ValidationCallback
+from fastvideo.train.callbacks.validation import (
+    DEFAULT_VALIDATION_VBENCH_METRICS,
+    SYNTHETIC_OPTICAL_FLOW_LOG_KEYS,
+    SYNTHETIC_OPTICAL_FLOW_METRIC,
+    ValidationCallback,
+    _ValidationMetricStats,
+)


 # ---------------------------------------------------------------------------
@@ -37,6 +47,7 @@ def _make_callback(
    num_frames: int | None = None,
    sampling_timesteps: list[int] | None = None,
    output_dir: str | None = None,
+    overlay_actions: bool = False,
 ) -> ValidationCallback:
    return ValidationCallback(
        pipeline_target=_PIPE_TARGET,
@@ -47,6 +58,7 @@ def _make_callback(
        num_frames=num_frames,
        sampling_timesteps=sampling_timesteps,
        output_dir=output_dir,
+        overlay_actions=overlay_actions,
    )


@@ -67,12 +79,14 @@ class TestConstructor:
        assert cb.num_frames is None
        assert cb.sampling_timesteps is None
        assert cb.output_dir is None
+        assert cb.overlay_actions is False
        assert cb.offload_training_state is False
        assert cb.unload_pipeline_after_validation is False
        # Lazy fields not yet populated.
        assert cb._pipeline is None
        assert cb._sampling_param is None
        assert cb.validation_random_generator is None
+        assert cb.metrics_config.enabled is False

    def test_string_inputs_are_coerced(self) -> None:
        # YAML often produces strings for numeric fields; the
@@ -85,6 +99,7 @@ class TestConstructor:
            guidance_scale="4.5",  # type: ignore[arg-type]
            num_frames="77",  # type: ignore[arg-type]
            sampling_timesteps=["1000", "500"],
+            overlay_actions=1,  # type: ignore[arg-type]
            offload_training_state="1",  # type: ignore[arg-type]
            unload_pipeline_after_validation="false",  # type: ignore[arg-type]
        )
@@ -93,6 +108,7 @@ class TestConstructor:
        assert cb.guidance_scale == 4.5
        assert cb.num_frames == 77
        assert cb.sampling_timesteps == [1000, 500]
+        assert cb.overlay_actions is True
        assert cb.offload_training_state is True
        assert cb.unload_pipeline_after_validation is False

@@ -109,6 +125,45 @@ class TestConstructor:
            "another": "value",
        }

+    def test_metrics_true_uses_default_vbench_subset(self) -> None:
+        cb = ValidationCallback(
+            pipeline_target=_PIPE_TARGET,
+            dataset_file="x.json",
+            metrics=True,
+        )
+        assert cb.metrics_config.enabled is True
+        assert cb.metrics_config.names == DEFAULT_VALIDATION_VBENCH_METRICS
+        assert "metrics" not in cb.pipeline_kwargs
+
+    def test_metrics_mapping_is_coerced(self) -> None:
+        cb = ValidationCallback(
+            pipeline_target=_PIPE_TARGET,
+            dataset_file="x.json",
+            metrics={
+                "names": "vbench.aesthetic_quality",
+                "device": "cpu",
+                "calibration_path": "/tmp/calibration.json",
+                "mouse_pitch_sign": "-1",
+                "skip_missing_deps": False,
+                "strict": True,
+                "unload_after_validation": False,
+                "loader_threads": "2",
+                "prefetch_factor": "3",
+                "log_prefix": "custom/validation",
+            },
+        )
+        assert cb.metrics_config.enabled is True
+        assert cb.metrics_config.names == ["vbench.aesthetic_quality"]
+        assert cb.metrics_config.device == "cpu"
+        assert cb.metrics_config.calibration_path == "/tmp/calibration.json"
+        assert cb.metrics_config.mouse_pitch_sign == -1
+        assert cb.metrics_config.skip_missing_deps is False
+        assert cb.metrics_config.strict is True
+        assert cb.metrics_config.unload_after_validation is False
+        assert cb.metrics_config.loader_threads == 2
+        assert cb.metrics_config.prefetch_factor == 3
+        assert cb.metrics_config.log_prefix == "custom/validation"
+

 # ---------------------------------------------------------------------------
 # B. on_validation_begin gating
@@ -238,3 +293,294 @@ class TestStateDict:
            {"validation_rng": torch.tensor([1, 2, 3], dtype=torch.uint8)}
        )
        assert cb.validation_random_generator is None
+
+
+# ---------------------------------------------------------------------------
+# E. metric result aggregation
+# ---------------------------------------------------------------------------
+
+
+class TestMetricAggregation:
+
+    def test_accumulates_scores_and_scalar_details(self) -> None:
+        stats = _ValidationMetricStats()
+        row: dict = {"path": "sample.mp4"}
+        ValidationCallback._accumulate_metric_results(
+            stats,
+            row,
+            {
+                "vbench.aesthetic_quality": SimpleNamespace(
+                    name="vbench.aesthetic_quality",
+                    score=0.5,
+                    details={"ignored": [1, 2, 3]},
+                ),
+                SYNTHETIC_OPTICAL_FLOW_METRIC: SimpleNamespace(
+                    name=SYNTHETIC_OPTICAL_FLOW_METRIC,
+                    score=1.5,
+                    details={
+                        **{key: float(i) for i, key in enumerate(SYNTHETIC_OPTICAL_FLOW_LOG_KEYS)},
+                        "pixel_epe_mean_std": 99.0,
+                        "pixel_epe_mean_max": 100.0,
+                        "pixel_epe_mean_auc": 101.0,
+                    },
+                ),
+            },
+        )
+
+        assert stats.sums["vbench.aesthetic_quality"] == 0.5
+        assert stats.counts["vbench.aesthetic_quality"] == 1.0
+        for i, key in enumerate(SYNTHETIC_OPTICAL_FLOW_LOG_KEYS):
+            metric_key = f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.{key}"
+            assert stats.sums[metric_key] == float(i)
+            assert stats.counts[metric_key] == 1.0
+        assert SYNTHETIC_OPTICAL_FLOW_METRIC not in stats.sums
+        assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_std" not in stats.sums
+        assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_max" not in stats.sums
+        assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_auc" not in stats.sums
+        assert "vbench.aesthetic_quality.ignored" not in stats.sums
+        assert row["vbench.aesthetic_quality"] == 0.5
+
+    def test_merges_metric_stats(self) -> None:
+        dst = _ValidationMetricStats(
+            sums={"a": 1.0},
+            counts={"a": 1.0},
+            per_video=[{"path": "a.mp4"}],
+            errors=["first"],
+        )
+        src = _ValidationMetricStats(
+            sums={"a": 2.0, "b": 3.0},
+            counts={"a": 2.0, "b": 1.0},
+            per_video=[{"path": "b.mp4"}],
+            errors=["second"],
+        )
+
+        ValidationCallback._merge_metric_stats(dst, src)
+
+        assert dst.sums == {"a": 3.0, "b": 3.0}
+        assert dst.counts == {"a": 3.0, "b": 1.0}
+        assert dst.per_video == [{"path": "a.mp4"}, {"path": "b.mp4"}]
+        assert dst.errors == ["first", "second"]
+
+    def test_metric_log_name_replaces_dots(self) -> None:
+        cb = ValidationCallback(
+            pipeline_target=_PIPE_TARGET,
+            dataset_file="x.json",
+            metrics={"names": ["vbench.aesthetic_quality"]},
+        )
+
+        assert (
+            cb._metric_log_name("vbench.aesthetic_quality")
+            == "metrics/validation/vbench/aesthetic_quality"
+        )
+
+    def test_metric_device_uses_local_rank(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        cb = ValidationCallback(
+            pipeline_target=_PIPE_TARGET,
+            dataset_file="x.json",
+            metrics={"names": ["vbench.aesthetic_quality"]},
+        )
+        cb.world_group = SimpleNamespace(local_rank=5)
+        monkeypatch.setattr(
+            torch.cuda,
+            "is_available",
+            lambda: True,
+        )
+
+        assert cb._metric_device() == "cuda:5"
+
+    def test_metric_extras_include_actions_and_calibration(self) -> None:
+        cb = ValidationCallback(
+            pipeline_target=_PIPE_TARGET,
+            dataset_file="x.json",
+            metrics={
+                "names": ["optical_flow.synthetic_optical_flow"],
+                "calibration_path": "/tmp/calibration.json",
+                "mouse_pitch_sign": -1,
+            },
+        )
+        actions = {
+            "keyboard": np.zeros((3, 6)),
+            "mouse": np.zeros((3, 2)),
+        }
+
+        extras = cb._validation_metric_extras(
+            video_filenames=["a.mp4", "b.mp4"],
+            actions=[actions, None],
+            mouse_pitch_signs=[None, 1],
+        )
+
+        assert extras[0]["actions"] is actions
+        assert extras[0]["calibration"] == "/tmp/calibration.json"
+        assert extras[0]["mouse_pitch_sign"] == -1
+        assert "actions" not in extras[1]
+        assert extras[1]["mouse_pitch_sign"] == 1
+
+    def test_validation_actions_prefers_loaded_conditions(self) -> None:
+        keyboard = np.zeros((4, 6))
+        mouse = np.zeros((4, 2))
+
+        actions = ValidationCallback._validation_actions(
+            {
+                "keyboard_cond": keyboard,
+                "mouse_cond": mouse,
+            }
+        )
+
+        assert actions is not None
+        assert np.array_equal(actions["keyboard"], keyboard)
+        assert np.array_equal(actions["mouse"], mouse)
+
+    def test_validation_actions_loads_action_path(
+        self,
+        tmp_path,
+    ) -> None:
+        path = tmp_path / "actions.npy"
+        np.save(
+            path,
+            {
+                "keyboard": np.ones((4, 6)),
+                "mouse": np.ones((4, 2)),
+            },
+            allow_pickle=True,
+        )
+
+        actions = ValidationCallback._validation_actions({"action_path": str(path)})
+
+        assert actions is not None
+        assert actions["keyboard"].shape == (4, 6)
+        assert actions["mouse"].shape == (4, 2)
+
+    def test_available_paths_requires_all_paths(self) -> None:
+        assert ValidationCallback._available_paths(["a.mp4", "b.mp4"]) == ["a.mp4", "b.mp4"]
+        assert ValidationCallback._available_paths(["a.mp4", None]) is None
+
+
+# ---------------------------------------------------------------------------
+# F. action overlay plumbing
+# ---------------------------------------------------------------------------
+
+
+class TestActionOverlay:
+
+    def test_save_validation_videos_supports_overlay_suffix(
+        self,
+        tmp_path,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        cb = _make_callback()
+        cb.global_rank = 3
+        calls: list[tuple[str, int]] = []
+
+        def fake_mimsave(
+            fname: str,
+            video: list[np.ndarray],
+            *,
+            fps: int,
+        ) -> None:
+            calls.append((fname, fps))
+
+        monkeypatch.setattr(
+            "fastvideo.train.callbacks.validation.imageio.mimsave",
+            fake_mimsave,
+        )
+
+        saved = cb._save_validation_videos(
+            [[np.zeros((2, 2, 3), dtype=np.uint8)]],
+            output_dir=str(tmp_path),
+            step=7,
+            num_inference_steps=4,
+            fps=25,
+            suffix="_overlay",
+        )
+
+        assert saved.filenames == [
+            str(tmp_path / "validation_step_7_inference_steps_4_rank_3_video_0_overlay.mp4")
+        ]
+        assert saved.indices == [0]
+        assert calls == [(saved.filenames[0], 25)]
+
+    def test_save_validation_videos_skips_failed_artifact(
+        self,
+        tmp_path,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        cb = _make_callback()
+        cb.global_rank = 0
+        calls: list[str] = []
+
+        def fake_mimsave(
+            fname: str,
+            video: list[np.ndarray],
+            *,
+            fps: int,
+        ) -> None:
+            del video, fps
+            calls.append(fname)
+            if len(calls) == 1:
+                raise OSError("ffmpeg failed")
+
+        monkeypatch.setattr(
+            "fastvideo.train.callbacks.validation.imageio.mimsave",
+            fake_mimsave,
+        )
+
+        saved = cb._save_validation_videos(
+            [
+                [np.zeros((2, 2, 3), dtype=np.uint8)],
+                [np.ones((2, 2, 3), dtype=np.uint8)],
+            ],
+            output_dir=str(tmp_path),
+            step=8,
+            num_inference_steps=5,
+            fps=24,
+        )
+
+        assert saved.indices == [1]
+        assert saved.filenames == [
+            str(tmp_path / "validation_step_8_inference_steps_5_rank_0_video_1.mp4")
+        ]
+        assert len(calls) == 2
+
+    def test_post_process_validation_frames_uses_student_hook(self) -> None:
+        cb = _make_callback(overlay_actions=True)
+        frames = [np.zeros((2, 2, 3), dtype=np.uint8)]
+        action = {
+            "keyboard": np.ones((1, 4), dtype=np.float32),
+            "mouse": np.zeros((1, 2), dtype=np.float32),
+        }
+
+        class Student:
+
+            def post_process_validation_frames(
+                self,
+                received_frames: list[np.ndarray],
+                *,
+                action: dict,
+            ) -> list[np.ndarray]:
+                assert received_frames is frames
+                assert action["keyboard"].shape == (1, 4)
+                return [np.full((2, 2, 3), 255, dtype=np.uint8)]
+
+        cb.method = SimpleNamespace(student=Student())
+
+        overlay = cb._post_process_validation_frames(
+            frames,
+            action=action,
+        )
+
+        assert overlay is not None
+        assert int(overlay[0][0, 0, 0]) == 255
+
+    def test_post_process_validation_frames_without_hook_is_noop(self) -> None:
+        cb = _make_callback(overlay_actions=True)
+        cb.method = SimpleNamespace(student=object())
+
+        overlay = cb._post_process_validation_frames(
+            [np.zeros((2, 2, 3), dtype=np.uint8)],
+            action=None,
+        )
+
+        assert overlay is None
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# Minimum config to run single training steps of StreamingLongTuningMethod
+# on WanCausalModel for the per-method smoke test. Uses the real Wan 2.1 1.3B
+# checkpoint for the causal student and the bidirectional teacher/critic
+# (mirroring the MatrixGame/LongLive production layout at 1.3B scale), with
+# tiny synthetic latents.
+#
+# Deviations from the production zelda config, chosen to keep the smoke test
+# transformer-only:
+# - cfg_uncond.text: zero — avoids loading the UMT5 encoder for the
+#   negative-prompt embedding on the teacher CFG branch.
+# - streaming_reencode_overlap_anchor: false — avoids loading the VAE for the
+#   overlap-anchor re-encode.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: false
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
+  rollout_mode: simulate
+  generator_update_interval: 1
+  dmd_denoising_steps: [1000, 750, 500, 250]
+  warp_denoising_step: true
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+
+  chunk_size: 3
+  student_sample_type: sde
+  same_step_across_blocks: true
+  last_step_only: false
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  cfg_uncond:
+    text: zero
+
+  streaming_training: true
+  streaming_chunk_size: 3
+  streaming_max_length: 6
+  # With chunk_size 3 and fixed overlap 1, the second streaming chunk carries
+  # 2 new frames + 1 overlap latent, exercising the overlap/chunk-mask path.
+  streaming_fixed_overlap_latents: 1
+  streaming_reencode_overlap_anchor: false
+
+  multi_phased_distill_schedule:
+    - stage: self_forcing
+      start_step: 0
+      end_step: 1
+      num_latent_t: 3
+    - stage: streaming_long
+      start_step: 1
+      num_latent_t: 6
+      streaming_training: true
+      streaming_chunk_size: 3
+      streaming_max_length: 6
+      streaming_fixed_overlap_latents: 1
+
+  # Critic optimizer
+  fake_score_learning_rate: 3.0e-7
+  fake_score_betas: [0.9, 0.95]
+  fake_score_lr_scheduler: constant
+
+training:
+  dit_precision: bf16
+
+  # Required for grad-enabled rollouts: switches the student's streaming KV
+  # cache to the checkpoint-safe variant so cache writes don't invalidate
+  # tensors saved for backward.
+  model:
+    enable_gradient_checkpointing_type: full
+
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+
+  data:
+    seed: 42
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    num_latent_t: 6
+    num_height: 64
+    num_width: 64
+    num_frames: 21
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4
+    gradient_accumulation_steps: 1
+
+pipeline: {}
@@ -6,6 +6,9 @@
  "test_matrixgame2_finetune": {
    "H200": 4.8247
  },
+  "test_streaming_long_tuning": {
+    "H200": 0.0049
+  },
  "test_wan_causal_dfsft": {
    "GB200": 2.9781,
    "L40S": 3.2562
@@ -0,0 +1,362 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-method GPU smoke test: ``WanCausalModel`` + ``StreamingLongTuningMethod``.
+
+Mirrors ``test_wan_causal_cd.py``. LongLive-style streaming long tuning runs a
+short self-forcing stage first, then switches to a persistent streaming
+sequence trained chunk-by-chunk with the causal student KV cache. The GPU test
+walks a single method instance through both stages: the short-stage step, the
+first streaming chunk, an overlapped chunk, the chunk that fills the sequence,
+and the post-reset chunk, asserting finite losses, student/critic gradient
+flow, and the streaming-length bookkeeping at every step.
+
+The pure-tensor helpers (stage selection, schedule validation, first-frame
+padding, latent truncation, chunk sizing) are covered by CPU-only tests at the
+bottom — they need no GPU, no mocks, and run wherever this file is collected.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29523")
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.train.methods.distribution_matching.streaming_long_tuning import (
+    DistillStage,
+    StreamingLongTuningMethod,
+    parse_multi_phased_distill_schedule,
+    select_distill_stage,
+)
+from fastvideo.train.models.wan import WanCausalModel, WanModel
+from fastvideo.train.utils.config import load_run_config
+
+from .grad_norm_regression import check_grad_norm_regression
+
+_FIXTURE = str(
+    Path(__file__).resolve().parent.parent / "fixtures"
+    / "wan_causal_t2v_streaming_long_min.yaml")
+
+
+def _build_synthetic_batch(
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, torch.Tensor]:
+    batch_size = 1
+    return {
+        "text_embedding":
+        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
+        "text_attention_mask":
+        torch.ones(batch_size, 16, device=device),
+        "vae_latent":
+        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
+    }
+
+
+def _assert_layer0_grads(model, *, role: str) -> None:
+    blocks = model.transformer.blocks
+    assert blocks is not None and len(blocks) > 0
+    trainable = [p for p in blocks[0].parameters() if p.requires_grad]
+    assert len(trainable) > 0, f"{role} layer 0 has no trainable parameters"
+    for i, p in enumerate(trainable):
+        assert p.grad is not None, f"{role} layer 0 param[{i}] has None grad"
+        assert torch.isfinite(p.grad).all().item(), (
+            f"{role} layer 0 param[{i}] grad contains NaN/Inf")
+    assert any(
+        p.grad.detach().float().norm().item() > 0.0 for p in trainable), (
+            f"all {role} layer-0 grads are exactly zero; the loss did not "
+            "reach the first transformer block")
+
+
+def _assert_finite_losses(loss_map: dict[str, torch.Tensor]) -> None:
+    for name in ("total_loss", "generator_loss", "fake_score_loss"):
+        loss = loss_map[name]
+        assert torch.is_tensor(loss), f"{name} must be a torch.Tensor"
+        assert torch.isfinite(loss).item(), (
+            f"{name} is not finite: {loss.item()}")
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_streaming_long_tuning_multi_stage_train_steps(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+
+    cfg = load_run_config(_FIXTURE)
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+
+    monkeypatch.setattr(
+        "fastvideo.train.utils.dataloader."
+        "build_parquet_t2v_train_dataloader",
+        lambda *args, **kwargs: None,
+    )
+
+    student = WanCausalModel(
+        init_from=cfg.models["student"]["init_from"],
+        training_config=cfg.training,
+        trainable=True,
+    )
+    student.transformer = student.transformer.to(device=device, dtype=dtype)
+
+    teacher = WanModel(
+        init_from=cfg.models["teacher"]["init_from"],
+        training_config=cfg.training,
+        trainable=False,
+    )
+    teacher.transformer = teacher.transformer.to(device=device, dtype=dtype)
+
+    critic = WanModel(
+        init_from=cfg.models["critic"]["init_from"],
+        training_config=cfg.training,
+        trainable=True,
+    )
+    critic.transformer = critic.transformer.to(device=device, dtype=dtype)
+
+    method = StreamingLongTuningMethod(
+        cfg=cfg,
+        role_models={
+            "student": student,
+            "teacher": teacher,
+            "critic": critic,
+        },
+    )
+    method.on_train_start()
+
+    batch = _build_synthetic_batch(device, dtype)
+
+    # Iteration 0 — short self-forcing stage (stage 0, num_latent_t=3).
+    loss_map, outputs, metrics = method.single_train_step(dict(batch), 0)
+    _assert_finite_losses(loss_map)
+    assert metrics["distill_stage_index"] == 0.0
+    assert metrics["active_num_latent_t"] == 3.0
+    assert metrics["update_student"] == 1.0
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+    _assert_layer0_grads(student, role="student")
+    _assert_layer0_grads(critic, role="critic")
+    assert all(not p.requires_grad
+               for p in teacher.transformer.parameters()), (
+                   "teacher must stay frozen")
+
+    # Iteration 1 — first streaming chunk: 3 fresh latents, no overlap.
+    method.optimizers_zero_grad(1)
+    loss_map, outputs, metrics = method.single_train_step(dict(batch), 1)
+    _assert_finite_losses(loss_map)
+    assert metrics["distill_stage_index"] == 1.0
+    assert metrics["streaming_current_length"] == 3.0
+    assert metrics["streaming_max_length"] == 6.0
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+    _assert_layer0_grads(student, role="student")
+    _assert_layer0_grads(critic, role="critic")
+    vis = outputs["dmd_latent_vis_dict"]
+    assert int(vis["streaming_chunk_mask"].sum().item()) == 3
+
+    # Iteration 2 — overlapped chunk: 2 new latents + 1 overlap anchor.
+    method.optimizers_zero_grad(2)
+    loss_map, outputs, metrics = method.single_train_step(dict(batch), 2)
+    _assert_finite_losses(loss_map)
+    assert metrics["streaming_current_length"] == 5.0
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+    mask = outputs["dmd_latent_vis_dict"]["streaming_chunk_mask"]
+    assert mask.dtype == torch.bool
+    assert int(mask.sum().item()) == 2, (
+        "overlap latents must be masked out of the training chunk")
+
+    # Iteration 3 — final chunk fills the sequence (1 new latent).
+    method.optimizers_zero_grad(3)
+    loss_map, outputs, metrics = method.single_train_step(dict(batch), 3)
+    _assert_finite_losses(loss_map)
+    assert metrics["streaming_current_length"] == 6.0
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+
+    # Iteration 4 — the exhausted sequence resets and a fresh chunk starts.
+    method.optimizers_zero_grad(4)
+    loss_map, outputs, metrics = method.single_train_step(dict(batch), 4)
+    _assert_finite_losses(loss_map)
+    assert metrics["streaming_current_length"] == 3.0, (
+        "streaming state must reset once the sequence reaches max length")
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+    _assert_layer0_grads(student, role="student")
+
+    # Grad-norm regression on the post-reset chunk (grads zeroed at the top
+    # of this iteration, so the reference captures exactly one step). Keep
+    # this last: under FASTVIDEO_GRADNORM_UPDATE=1 it records and skips.
+    check_grad_norm_regression(
+        "test_streaming_long_tuning",
+        student.transformer,
+    )
+
+
+# ---------------------------------------------------------------------------
+# CPU-only helper tests (no GPU, no mocks)
+# ---------------------------------------------------------------------------
+
+
+def _method_without_init() -> StreamingLongTuningMethod:
+    # Bare instance for exercising stateless helpers; no attributes beyond
+    # the ones each test sets are touched.
+    return StreamingLongTuningMethod.__new__(StreamingLongTuningMethod)
+
+
+def _make_stage(**overrides) -> DistillStage:
+    defaults = dict(
+        name="streaming_long",
+        start_step=0,
+        end_step=None,
+        num_latent_t=6,
+        streaming_training=True,
+        streaming_chunk_size=3,
+        streaming_max_length=6,
+    )
+    defaults.update(overrides)
+    return DistillStage(**defaults)
+
+
+class TestStageSelection:
+
+    def test_boundaries_between_stages(self) -> None:
+        short = _make_stage(
+            name="self_forcing",
+            start_step=0,
+            end_step=10,
+            streaming_training=False,
+            streaming_chunk_size=None,
+            streaming_max_length=None,
+        )
+        long = _make_stage(start_step=10)
+        stages = [short, long]
+
+        assert select_distill_stage(stages, 0) is short
+        assert select_distill_stage(stages, 9) is short
+        assert select_distill_stage(stages, 10) is long
+        assert select_distill_stage(stages, 10_000) is long
+
+    def test_past_last_closed_stage_falls_back_to_last(self) -> None:
+        short = _make_stage(
+            name="self_forcing",
+            start_step=0,
+            end_step=10,
+            streaming_training=False,
+            streaming_chunk_size=None,
+            streaming_max_length=None,
+        )
+        assert select_distill_stage([short], 99) is short
+
+
+class TestScheduleParsing:
+
+    def test_rejects_overlapping_stages(self) -> None:
+        with pytest.raises(ValueError, match="ordered and non-overlapping"):
+            parse_multi_phased_distill_schedule(
+                [
+                    {"stage": "self_forcing", "start_step": 0, "end_step": 10},
+                    {"stage": "streaming_long", "start_step": 5},
+                ],
+                default_num_latent_t=6,
+                default_streaming_chunk_size=3,
+            )
+
+    def test_rejects_fixed_overlap_not_smaller_than_chunk(self) -> None:
+        with pytest.raises(ValueError, match="smaller than"):
+            parse_multi_phased_distill_schedule(
+                [{
+                    "stage": "streaming_long",
+                    "streaming_chunk_size": 3,
+                    "streaming_fixed_overlap_latents": 3,
+                }],
+                default_num_latent_t=6,
+            )
+
+
+class TestPadFirstFrameLatent:
+
+    def test_pads_short_condition_with_zeros(self) -> None:
+        method = _method_without_init()
+        first_frame = torch.ones(1, 4, 1, 2, 2)
+
+        batch = method._with_padded_first_frame_latent(
+            {"first_frame_latent": first_frame},
+            target_latent_frames=3,
+        )
+
+        padded = batch["first_frame_latent"]
+        assert padded.shape == (1, 4, 3, 2, 2)
+        assert torch.equal(padded[:, :, :1], first_frame)
+        assert padded[:, :, 1:].abs().sum().item() == 0.0
+
+    def test_passthrough_without_condition_or_when_long_enough(self) -> None:
+        method = _method_without_init()
+        raw: dict = {"other": 1}
+        assert method._with_padded_first_frame_latent(
+            raw, target_latent_frames=3) is raw
+
+        full = {"first_frame_latent": torch.zeros(1, 4, 3, 2, 2)}
+        assert method._with_padded_first_frame_latent(
+            full, target_latent_frames=3) is full
+
+    def test_rejects_non_5d_condition(self) -> None:
+        method = _method_without_init()
+        with pytest.raises(ValueError, match=r"\[B, C, T, H, W\]"):
+            method._with_padded_first_frame_latent(
+                {"first_frame_latent": torch.zeros(4, 1, 2, 2)},
+                target_latent_frames=3,
+            )
+
+
+class TestTruncateBatchLatents:
+
+    def test_slices_every_temporal_field(self) -> None:
+        method = _method_without_init()
+        # No _build_attention_metadata on the student → metadata rebuild is
+        # a no-op, so only the tensor slicing is exercised.
+        method.student = object()
+        batch = SimpleNamespace(
+            latents=torch.randn(1, 6, 2, 4, 4),
+            noise_latents=torch.randn(1, 6, 2, 4, 4),
+            noisy_model_input=torch.randn(1, 6, 2, 4, 4),
+            noise=torch.randn(1, 6, 2, 4, 4),
+            timesteps=torch.zeros(1, 6),
+            sigmas=torch.linspace(0, 1, 6),
+            raw_latent_shape=(1, 6, 2, 4, 4),
+            attn_metadata=None,
+            attn_metadata_vsa=None,
+        )
+
+        method._truncate_batch_latents(batch, 3)
+
+        assert batch.latents.shape == (1, 3, 2, 4, 4)
+        assert batch.noise_latents.shape == (1, 3, 2, 4, 4)
+        assert batch.noisy_model_input.shape == (1, 3, 2, 4, 4)
+        assert batch.noise.shape == (1, 3, 2, 4, 4)
+        assert batch.timesteps.shape == (1, 3)
+        assert batch.sigmas.shape == (3, )
+        assert batch.raw_latent_shape == (1, 2, 3, 4, 4)
+
+
+class TestStageChunkSize:
+
+    def test_prefers_stage_then_config_then_rollout_block(self) -> None:
+        method = _method_without_init()
+        method.method_config = {"streaming_chunk_size": 5}
+        method._chunk_size = 7
+
+        assert method._stage_chunk_size(_make_stage()) == 3
+        assert method._stage_chunk_size(
+            _make_stage(streaming_chunk_size=None)) == 5
+
+        method.method_config = {}
+        assert method._stage_chunk_size(
+            _make_stage(streaming_chunk_size=None)) == 7
+
+    def test_rejects_non_positive_chunk_size(self) -> None:
+        method = _method_without_init()
+        method.method_config = {}
+        method._chunk_size = 0
+        with pytest.raises(ValueError, match="must be positive"):
+            method._stage_chunk_size(_make_stage(streaming_chunk_size=None))
@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only integration tests for Trainer validation hook dispatch."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+
+from fastvideo.train.callbacks.validation import ValidationCallback
+from fastvideo.train.trainer import Trainer
+from fastvideo.train.utils.training_config import TrainingConfig
+
+
+class _RecordingValidationCallback(ValidationCallback):
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.run_calls: list[int] = []
+
+    def _run_validation(self, method: Any, step: int) -> None:  # type: ignore[override]
+        del method
+        self.run_calls.append(step)
+
+
+class _DummyTracker:
+
+    def __init__(self) -> None:
+        self.logs: list[tuple[dict[str, float], int]] = []
+        self.finished = False
+
+    def log(self, metrics: dict[str, float], step: int) -> None:
+        self.logs.append((metrics, step))
+
+    def finish(self) -> None:
+        self.finished = True
+
+
+class _DummyMethod:
+
+    def __init__(self) -> None:
+        self.weight = torch.nn.Parameter(torch.tensor(1.0))
+        self.train_start_calls = 0
+        self.zero_grad_steps: list[int] = []
+        self.optimizer_steps: list[int] = []
+        self.backward_calls = 0
+        self.tracker = None
+
+    def set_tracker(self, tracker: Any) -> None:
+        self.tracker = tracker
+
+    def on_train_start(self) -> None:
+        self.train_start_calls += 1
+
+    def manages_optimization(self) -> bool:
+        return False
+
+    def single_train_step(
+        self,
+        batch: dict[str, Any],
+        iteration: int,
+    ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, float]]:
+        assert batch["sample"] == "x"
+        loss = self.weight * 0.0 + 1.0
+        return {"total_loss": loss}, {}, {"iteration_seen": float(iteration)}
+
+    def backward(
+        self,
+        loss_map: dict[str, torch.Tensor],
+        outputs: dict[str, Any],
+        *,
+        grad_accum_rounds: int,
+    ) -> None:
+        del outputs
+        self.backward_calls += 1
+        (loss_map["total_loss"] / grad_accum_rounds).backward()
+
+    def optimizers_schedulers_step(self, iteration: int) -> None:
+        self.optimizer_steps.append(iteration)
+
+    def optimizers_zero_grad(self, iteration: int) -> None:
+        self.zero_grad_steps.append(iteration)
+        self.weight.grad = None
+
+
+def test_trainer_runs_validation_callback_during_training(
+    monkeypatch,
+) -> None:
+    tracker = _DummyTracker()
+    group = SimpleNamespace(rank=0, local_rank=0, rank_in_group=0, world_size=1)
+
+    monkeypatch.setattr("fastvideo.train.trainer.get_world_group", lambda: group)
+    monkeypatch.setattr("fastvideo.train.trainer.get_sp_group", lambda: group)
+    monkeypatch.setattr(
+        "fastvideo.train.callbacks.validation.get_world_group",
+        lambda: group,
+    )
+    monkeypatch.setattr(
+        "fastvideo.train.callbacks.validation.get_sp_group",
+        lambda: group,
+    )
+    monkeypatch.setattr(
+        "fastvideo.train.trainer.build_tracker",
+        lambda *args, **kwargs: tracker,
+    )
+
+    cfg = TrainingConfig()
+    cfg.tracker.project_name = ""
+    cfg.loop.gradient_accumulation_steps = 1
+    callback_configs = {
+        "validation": {
+            "_target_": f"{__name__}._RecordingValidationCallback",
+            "pipeline_target": "unused.pipeline.Target",
+            "dataset_file": "unused.json",
+            "every_steps": 2,
+        }
+    }
+    trainer = Trainer(
+        cfg,
+        callback_configs=callback_configs,
+    )
+    method = _DummyMethod()
+
+    trainer.run(
+        method,
+        dataloader=[{"sample": "x"}],
+        max_steps=3,
+    )
+
+    validation = trainer.callbacks._callbacks["validation"]
+    assert isinstance(validation, _RecordingValidationCallback)
+    assert validation.run_calls == [0, 2]
+    assert method.train_start_calls == 1
+    assert method.backward_calls == 3
+    assert method.zero_grad_steps == [0, 1, 2, 3]
+    assert method.optimizer_steps == [1, 2, 3]
+    assert [step for _, step in tracker.logs] == [1, 2, 3]
+    assert tracker.finished is True
@@ -155,6 +155,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
    assert t.distributed.pin_cpu_memory is True

    assert t.data.train_batch_size == 2
+    assert t.data.data_path == "/some/path"
    assert t.data.num_frames == 33
    assert t.data.seed == 42

@@ -244,6 +245,48 @@ def test_betas_parses_list_and_string_forms(
    assert cfg.training.optimizer.betas == expected


+def test_data_path_mapping_parses_repeat_counts(tmp_path: Path) -> None:
+    # Config loading should preserve structured multi-dataset paths so the
+    # dataset layer can interpret repeat counts later.
+    data = _minimal_yaml()
+    data["training"] = {
+        "data": {
+            "data_path": {
+                "data/path1": 1,
+                "data/path2": 2,
+            }
+        }
+    }
+
+    cfg = load_run_config(_write_yaml(tmp_path, data))
+
+    assert cfg.training.data.data_path == {
+        "data/path1": 1,
+        "data/path2": 2,
+    }
+
+
+def test_dotted_override_replaces_mapping_data_path(tmp_path: Path) -> None:
+    # A dict-valued data_path is a single leaf for overrides: a scalar
+    # --training.data.data_path replaces the whole mapping.
+    data = _minimal_yaml()
+    data["training"] = {
+        "data": {
+            "data_path": {
+                "data/path1": 1,
+                "data/path2": 2,
+            }
+        }
+    }
+
+    cfg = load_run_config(
+        _write_yaml(tmp_path, data),
+        overrides=["--training.data.data_path=data/only"],
+    )
+
+    assert cfg.training.data.data_path == "data/only"
+
+
 def test_dotted_overrides_apply_with_type_coercion(tmp_path: Path) -> None:
    path = _write_yaml(tmp_path, _minimal_yaml())
    overrides = [
@@ -9,8 +9,11 @@ section.  The pipeline class is resolved from
 from __future__ import annotations

 import contextlib
+import gc
+import json
 import os
-from dataclasses import dataclass
+from copy import deepcopy
+from dataclasses import dataclass, field
 from typing import Any, TYPE_CHECKING

 import imageio
@@ -46,6 +49,64 @@ logger = init_logger(__name__)
 class _ValidationStepResult:
    videos: list[list[np.ndarray]]
    captions: list[str]
+    overlay_videos: list[list[np.ndarray]] = field(default_factory=list)
+    overlay_captions: list[str] = field(default_factory=list)
+    ref_videos: list[str | None] = field(default_factory=list)
+    actions: list[dict[str, Any] | None] = field(default_factory=list)
+    mouse_pitch_signs: list[int | None] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class _ValidationMetricStats:
+    sums: dict[str, float] = field(default_factory=dict)
+    counts: dict[str, float] = field(default_factory=dict)
+    per_video: list[dict[str, Any]] = field(default_factory=list)
+    errors: list[str] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class _SavedValidationVideos:
+    filenames: list[str] = field(default_factory=list)
+    indices: list[int] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class _ValidationMetricsConfig:
+    enabled: bool = False
+    names: list[str] = field(default_factory=list)
+    device: str = "cuda"
+    calibration_path: str | None = None
+    mouse_pitch_sign: int = 1
+    skip_missing_deps: bool = True
+    strict: bool = False
+    unload_after_validation: bool = True
+    loader_threads: int = 1
+    prefetch_factor: int = 2
+    log_prefix: str = "metrics/validation"
+
+
+DEFAULT_VALIDATION_VBENCH_METRICS = [
+    "vbench.imaging_quality",
+    "vbench.aesthetic_quality",
+    "vbench.temporal_flickering",
+    "vbench.motion_smoothness",
+    "vbench.subject_consistency",
+    "vbench.background_consistency",
+    "vbench.dynamic_degree",
+]
+
+SYNTHETIC_OPTICAL_FLOW_METRIC = "optical_flow.synthetic_optical_flow"
+SYNTHETIC_OPTICAL_FLOW_LOG_KEYS = (
+    "mf_epe_mean",
+    "mf_angle_err_mean",
+    "mf_cosine_mean",
+    "mf_mag_ratio_mean",
+    "pixel_epe_mean_mean",
+    "px_angle_rmse_mean",
+    "fl_all_mean",
+    "foe_dist_mean",
+    "flow_kl_2d_mean",
+)


 class ValidationCallback(Callback):
@@ -68,6 +129,8 @@ class ValidationCallback(Callback):
        num_frames: int | None = None,
        output_dir: str | None = None,
        sampling_timesteps: list[int] | None = None,
+        overlay_actions: bool = False,
+        keyboard_value_scale: float = 1.0,
        offload_training_state: bool = False,
        unload_pipeline_after_validation: bool = False,
        **pipeline_kwargs: Any,
@@ -80,6 +143,11 @@ class ValidationCallback(Callback):
        self.num_frames = (int(num_frames) if num_frames is not None else None)
        self.output_dir = (str(output_dir) if output_dir is not None else None)
        self.sampling_timesteps = ([int(s) for s in sampling_timesteps] if sampling_timesteps is not None else None)
+        self.overlay_actions = self._coerce_bool(overlay_actions)
+        # Validation-only action amplification for world model; training keeps raw action values.
+        self.keyboard_value_scale = float(keyboard_value_scale)
+        metrics_config = pipeline_kwargs.pop("metrics", None)
+        self.metrics_config = self._parse_metrics_config(metrics_config)
        self.offload_training_state = self._coerce_bool(offload_training_state)
        self.unload_pipeline_after_validation = self._coerce_bool(unload_pipeline_after_validation)
        self.pipeline_kwargs = dict(pipeline_kwargs)
@@ -88,10 +156,56 @@ class ValidationCallback(Callback):
        self._pipeline: Any | None = None
        self._pipeline_key: tuple[Any, ...] | None = None
        self._sampling_param: SamplingParam | None = None
+        self._metric_evaluator: Any | None = None
        self.tracker: Any = DummyTracker()
        self.validation_random_generator: (torch.Generator | None) = None
        self.seed: int = 0

+    @staticmethod
+    def _parse_metrics_config(config: Any) -> _ValidationMetricsConfig:
+        if config is None or config is False:
+            return _ValidationMetricsConfig(enabled=False)
+        if config is True:
+            return _ValidationMetricsConfig(
+                enabled=True,
+                names=list(DEFAULT_VALIDATION_VBENCH_METRICS),
+            )
+        if isinstance(config, str):
+            return _ValidationMetricsConfig(
+                enabled=True,
+                names=[config],
+            )
+        if isinstance(config, list | tuple):
+            return _ValidationMetricsConfig(
+                enabled=bool(config),
+                names=[str(name) for name in config],
+            )
+        if not isinstance(config, dict):
+            raise TypeError("callbacks.validation.metrics must be a bool, string, list, or mapping")
+
+        enabled = bool(config.get("enabled", True))
+        raw_names = config.get("names", config.get("metrics", None))
+        if raw_names is None:
+            names = list(DEFAULT_VALIDATION_VBENCH_METRICS)
+        elif isinstance(raw_names, str):
+            names = [raw_names]
+        else:
+            names = [str(name) for name in raw_names]
+
+        return _ValidationMetricsConfig(
+            enabled=enabled and bool(names),
+            names=names,
+            device=str(config.get("device", "cuda")),
+            calibration_path=(str(config["calibration_path"]) if config.get("calibration_path") is not None else None),
+            mouse_pitch_sign=int(config.get("mouse_pitch_sign", 1)),
+            skip_missing_deps=bool(config.get("skip_missing_deps", True)),
+            strict=bool(config.get("strict", False)),
+            unload_after_validation=bool(config.get("unload_after_validation", True)),
+            loader_threads=int(config.get("loader_threads", 1)),
+            prefetch_factor=int(config.get("prefetch_factor", 2)),
+            log_prefix=str(config.get("log_prefix", "metrics/validation")),
+        )
+
    @staticmethod
    def _coerce_bool(value: Any) -> bool:
        if isinstance(value, str):
@@ -354,6 +468,7 @@ class ValidationCallback(Callback):
        try:
            transformer.eval()
            num_sp_groups = (self.world_group.world_size // self.sp_group.world_size)
+            sp = self._get_sampling_param()

            for num_inference_steps in self.sampling_steps:
                result = self._run_validation_for_steps(
@@ -364,72 +479,536 @@ class ValidationCallback(Callback):
                if self.rank_in_sp_group != 0:
                    continue

+                os.makedirs(
+                    output_dir,
+                    exist_ok=True,
+                )
+                local_videos = self._save_validation_videos(
+                    result.videos,
+                    output_dir=output_dir,
+                    step=step,
+                    num_inference_steps=num_inference_steps,
+                    fps=sp.fps,
+                )
+                local_overlay_videos = self._save_validation_videos(
+                    result.overlay_videos,
+                    output_dir=output_dir,
+                    step=step,
+                    num_inference_steps=num_inference_steps,
+                    fps=sp.fps,
+                    suffix="_overlay",
+                )
+                local_video_filenames = local_videos.filenames
+                local_overlay_video_filenames = local_overlay_videos.filenames
+                local_captions = self._select_by_indices(
+                    result.captions,
+                    local_videos.indices,
+                )
+                local_overlay_captions = self._select_by_indices(
+                    result.overlay_captions,
+                    local_overlay_videos.indices,
+                )
+                local_ref_videos = self._select_by_indices(
+                    result.ref_videos,
+                    local_videos.indices,
+                )
+                local_actions = self._select_by_indices(
+                    result.actions,
+                    local_videos.indices,
+                )
+                local_mouse_pitch_signs = self._select_by_indices(
+                    result.mouse_pitch_signs,
+                    local_videos.indices,
+                )
+                local_metric_stats = self._evaluate_validation_metrics(
+                    video_filenames=local_video_filenames,
+                    captions=local_captions,
+                    ref_videos=local_ref_videos,
+                    actions=local_actions,
+                    mouse_pitch_signs=local_mouse_pitch_signs,
+                    fps=sp.fps,
+                    output_dir=output_dir,
+                    step=step,
+                    num_inference_steps=num_inference_steps,
+                )
+
                if self.global_rank == 0:
-                    all_videos = list(result.videos)
-                    all_captions = list(result.captions)
+                    all_video_filenames = list(local_video_filenames)
+                    all_overlay_video_filenames = list(local_overlay_video_filenames)
+                    all_captions = list(local_captions)
+                    all_overlay_captions = list(local_overlay_captions)
+                    all_metric_stats = local_metric_stats
                    for sp_idx in range(1, num_sp_groups):
                        src = (sp_idx * self.sp_world_size)
                        recv_v = (self.world_group.recv_object(src=src))
                        recv_c = (self.world_group.recv_object(src=src))
-                        all_videos.extend(recv_v)
+                        recv_ov = (self.world_group.recv_object(src=src))
+                        recv_oc = (self.world_group.recv_object(src=src))
+                        recv_m = (self.world_group.recv_object(src=src))
+                        all_video_filenames.extend(recv_v)
+                        all_overlay_video_filenames.extend(recv_ov)
                        all_captions.extend(recv_c)
+                        all_overlay_captions.extend(recv_oc)
+                        self._merge_metric_stats(
+                            all_metric_stats,
+                            recv_m,
+                        )

-                    os.makedirs(
-                        output_dir,
-                        exist_ok=True,
+                    self._log_validation_metrics(
+                        all_metric_stats,
+                        step=step,
                    )
-                    video_filenames: list[str] = []
-                    sp = self._get_sampling_param()
-                    for i, video in enumerate(all_videos):
-                        fname = os.path.join(
-                            output_dir,
-                            f"validation_step_{step}"
-                            f"_inference_steps_"
-                            f"{num_inference_steps}"
-                            f"_video_{i}.mp4",
-                        )
-                        imageio.mimsave(
-                            fname,
-                            video,
+                    self._log_validation_video_artifacts(
+                        all_video_filenames,
+                        all_captions,
+                        key=f"validation_videos_{num_inference_steps}_steps",
+                        step=step,
+                        fps=sp.fps,
+                    )
+                    if all_overlay_video_filenames:
+                        self._log_validation_video_artifacts(
+                            all_overlay_video_filenames,
+                            all_overlay_captions,
+                            key=(f"validation_videos_{num_inference_steps}"
+                                 f"_steps_overlay"),
+                            step=step,
                            fps=sp.fps,
                        )
-                        video_filenames.append(fname)
-
-                    video_logs = []
-                    for fname, cap in zip(
-                            video_filenames,
-                            all_captions,
-                            strict=True,
-                    ):
-                        art = self.tracker.video(
-                            fname,
-                            caption=cap,
-                            fps=sp.fps,
-                        )
-                        if art is not None:
-                            video_logs.append(art)
-                    if video_logs:
-                        logs = {
-                            f"validation_videos_"
-                            f"{num_inference_steps}"
-                            f"_steps": video_logs
-                        }
-                        self.tracker.log_artifacts(
-                            logs,
-                            step,
-                        )
                else:
                    self.world_group.send_object(
-                        result.videos,
+                        local_video_filenames,
                        dst=0,
                    )
                    self.world_group.send_object(
-                        result.captions,
+                        local_captions,
+                        dst=0,
+                    )
+                    self.world_group.send_object(
+                        local_overlay_video_filenames,
+                        dst=0,
+                    )
+                    self.world_group.send_object(
+                        local_overlay_captions,
+                        dst=0,
+                    )
+                    self.world_group.send_object(
+                        local_metric_stats,
                        dst=0,
                    )
        finally:
            if was_training:
                transformer.train()
+            self._release_metric_evaluator()
+
+    def _save_validation_videos(
+        self,
+        videos: list[list[np.ndarray]],
+        *,
+        output_dir: str,
+        step: int,
+        num_inference_steps: int,
+        fps: int,
+        suffix: str = "",
+    ) -> _SavedValidationVideos:
+        saved = _SavedValidationVideos()
+        for i, video in enumerate(videos):
+            fname = os.path.join(
+                output_dir,
+                f"validation_step_{step}"
+                f"_inference_steps_{num_inference_steps}"
+                f"_rank_{self.global_rank}"
+                f"_video_{i}{suffix}.mp4",
+            )
+            # Validation video encoding can fail due transient ffmpeg or
+            # filesystem errors; skip the artifact so training can continue.
+            try:
+                imageio.mimsave(
+                    fname,
+                    video,
+                    fps=fps,
+                )
+            except Exception as exc:
+                logger.exception(
+                    "Failed to save validation video %s on rank %s; skipping artifact: %s",
+                    fname,
+                    self.global_rank,
+                    exc,
+                )
+                with contextlib.suppress(OSError):
+                    os.remove(fname)
+                continue
+            saved.filenames.append(fname)
+            saved.indices.append(i)
+        return saved
+
+    @staticmethod
+    def _select_by_indices(
+        values: list[Any],
+        indices: list[int],
+    ) -> list[Any]:
+        return [values[i] for i in indices if i < len(values)]
+
+    def _log_validation_video_artifacts(
+        self,
+        video_filenames: list[str],
+        captions: list[str],
+        *,
+        key: str,
+        step: int,
+        fps: int,
+    ) -> None:
+        video_logs = []
+        for fname, cap in zip(
+                video_filenames,
+                captions,
+                strict=True,
+        ):
+            art = self.tracker.video(
+                fname,
+                caption=cap,
+                fps=fps,
+            )
+            if art is not None:
+                video_logs.append(art)
+        if video_logs:
+            self.tracker.log_artifacts(
+                {key: video_logs},
+                step,
+            )
+
+    # ----------------------------------------------------------
+    # Metric evaluation
+    # ----------------------------------------------------------
+
+    def _metric_device(self) -> str:
+        device = self.metrics_config.device
+        if device == "cuda" and torch.cuda.is_available():
+            local_rank = int(getattr(self.world_group, "local_rank", 0))
+            return f"cuda:{local_rank}"
+        return device
+
+    def _get_metric_evaluator(self) -> Any:
+        if self._metric_evaluator is not None:
+            return self._metric_evaluator
+        from fastvideo.eval import Evaluator
+
+        cfg = self.metrics_config
+        self._metric_evaluator = Evaluator(
+            metrics=cfg.names,
+            device=self._metric_device(),
+            num_gpus=1,
+            loader_threads=cfg.loader_threads,
+            prefetch_factor=cfg.prefetch_factor,
+            pre_upload=False,
+            skip_missing_deps=cfg.skip_missing_deps,
+        )
+        logger.info(
+            "Initialized validation metrics: %s",
+            ", ".join(self._metric_evaluator.metric_names),
+        )
+        return self._metric_evaluator
+
+    def _release_metric_evaluator(self) -> None:
+        if self._metric_evaluator is None:
+            return
+        if self.metrics_config.unload_after_validation:
+            self._metric_evaluator.unload()
+            self._metric_evaluator = None
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def _evaluate_validation_metrics(
+        self,
+        *,
+        video_filenames: list[str],
+        captions: list[str],
+        ref_videos: list[str | None],
+        actions: list[dict[str, Any] | None],
+        mouse_pitch_signs: list[int | None],
+        fps: int,
+        output_dir: str,
+        step: int,
+        num_inference_steps: int,
+    ) -> _ValidationMetricStats:
+        stats = _ValidationMetricStats()
+        if not self.metrics_config.enabled or not video_filenames:
+            return stats
+
+        try:
+            evaluator = self._get_metric_evaluator()
+            from fastvideo.eval import samples_from
+
+            samples = samples_from(
+                video=video_filenames,
+                reference=self._available_paths(ref_videos),
+                text_prompts=captions,
+                fps=float(fps),
+                extras=self._validation_metric_extras(
+                    video_filenames=video_filenames,
+                    actions=actions,
+                    mouse_pitch_signs=mouse_pitch_signs,
+                ),
+            )
+            results = evaluator.evaluate(samples=samples)
+            for filename, metric_results in zip(
+                    video_filenames,
+                    results,
+                    strict=True,
+            ):
+                row: dict[str, Any] = {"path": filename}
+                self._accumulate_metric_results(
+                    stats,
+                    row,
+                    metric_results,
+                )
+                stats.per_video.append(row)
+            for metric_name, metric_result in results.corpus.items():
+                row = {"path": "<corpus>"}
+                self._accumulate_metric_results(
+                    stats,
+                    row,
+                    {metric_name: metric_result},
+                )
+                stats.per_video.append(row)
+        except Exception as exc:
+            message = ("Validation metric evaluation failed on rank "
+                       f"{self.global_rank}: {exc}")
+            logger.exception(message)
+            stats.errors.append(message)
+            if self.metrics_config.strict:
+                raise
+        finally:
+            if self._metric_evaluator is not None:
+                self._metric_evaluator.release_cuda_memory()
+
+        self._write_metric_summary(
+            stats,
+            output_dir=output_dir,
+            step=step,
+            num_inference_steps=num_inference_steps,
+        )
+        return stats
+
+    @staticmethod
+    def _accumulate_metric_results(
+        stats: _ValidationMetricStats,
+        row: dict[str, Any],
+        metric_results: dict[str, Any],
+    ) -> None:
+        for metric_name, metric_result in metric_results.items():
+            details = getattr(metric_result, "details", {}) or {}
+            if metric_name == SYNTHETIC_OPTICAL_FLOW_METRIC:
+                if not isinstance(details, dict):
+                    continue
+                for detail_name in SYNTHETIC_OPTICAL_FLOW_LOG_KEYS:
+                    key = f"{metric_name}.{detail_name}"
+                    ValidationCallback._accumulate_scalar(
+                        stats,
+                        row,
+                        key,
+                        details.get(detail_name),
+                    )
+                continue
+
+            score = getattr(metric_result, "score", None)
+            ValidationCallback._accumulate_scalar(
+                stats,
+                row,
+                metric_name,
+                score,
+            )
+            if not isinstance(details, dict):
+                continue
+            for detail_name, value in details.items():
+                key = f"{metric_name}.{detail_name}"
+                ValidationCallback._accumulate_scalar(
+                    stats,
+                    row,
+                    key,
+                    value,
+                )
+
+    @staticmethod
+    def _accumulate_scalar(
+        stats: _ValidationMetricStats,
+        row: dict[str, Any],
+        key: str,
+        value: Any,
+    ) -> None:
+        if not isinstance(value, float | int | np.floating | np.integer):
+            return
+        value_float = float(value)
+        if not np.isfinite(value_float):
+            return
+        row[key] = value_float
+        stats.sums[key] = stats.sums.get(key, 0.0) + value_float
+        stats.counts[key] = stats.counts.get(key, 0.0) + 1.0
+
+    @staticmethod
+    def _merge_metric_stats(
+        dst: _ValidationMetricStats,
+        src: _ValidationMetricStats,
+    ) -> None:
+        for key, value in src.sums.items():
+            dst.sums[key] = dst.sums.get(key, 0.0) + float(value)
+        for key, value in src.counts.items():
+            dst.counts[key] = dst.counts.get(key, 0.0) + float(value)
+        dst.per_video.extend(src.per_video)
+        dst.errors.extend(src.errors)
+
+    def _log_validation_metrics(
+        self,
+        stats: _ValidationMetricStats,
+        *,
+        step: int,
+    ) -> None:
+        if stats.errors:
+            message = "Validation metric evaluation failed:\n" + "\n".join(stats.errors)
+            if self.metrics_config.strict:
+                raise RuntimeError(message)
+            logger.warning(message)
+
+        logs: dict[str, float] = {}
+        for key, metric_sum in stats.sums.items():
+            metric_count = stats.counts.get(key, 0.0)
+            if metric_count <= 0:
+                continue
+            value = float(metric_sum / metric_count)
+            if not np.isfinite(value):
+                continue
+            logs[self._metric_log_name(key)] = value
+        if logs:
+            self.tracker.log(
+                logs,
+                step,
+            )
+
+    def _metric_log_name(
+        self,
+        key: str,
+    ) -> str:
+        return f"{self.metrics_config.log_prefix}/{key.replace('.', '/')}"
+
+    def _write_metric_summary(
+        self,
+        stats: _ValidationMetricStats,
+        *,
+        output_dir: str,
+        step: int,
+        num_inference_steps: int,
+    ) -> None:
+        if not self.metrics_config.enabled:
+            return
+        summary_dir = os.path.join(
+            output_dir,
+            "eval",
+            f"step_{step}",
+        )
+        os.makedirs(
+            summary_dir,
+            exist_ok=True,
+        )
+        metrics = {
+            key: float(stats.sums[key] / stats.counts[key])
+            for key in sorted(stats.sums) if stats.counts.get(key, 0.0) > 0
+        }
+        payload = {
+            "step": int(step),
+            "num_inference_steps": int(num_inference_steps),
+            "rank": int(self.global_rank),
+            "metrics": metrics,
+            "per_video": stats.per_video,
+            "errors": stats.errors,
+        }
+        path = os.path.join(
+            summary_dir,
+            f"inference_steps_{num_inference_steps}_rank_{self.global_rank}.json",
+        )
+        with open(
+                path,
+                "w",
+                encoding="utf-8",
+        ) as f:
+            json.dump(
+                payload,
+                f,
+                indent=2,
+            )
+
+    def _validation_metric_extras(
+        self,
+        *,
+        video_filenames: list[str],
+        actions: list[dict[str, Any] | None],
+        mouse_pitch_signs: list[int | None],
+    ) -> list[dict[str, Any]]:
+        extras: list[dict[str, Any]] = []
+        for i, _filename in enumerate(video_filenames):
+            extra: dict[str, Any] = {}
+            action = actions[i] if i < len(actions) else None
+            if action is not None:
+                extra["actions"] = action
+            if self.metrics_config.calibration_path is not None:
+                extra["calibration"] = self.metrics_config.calibration_path
+            mouse_pitch_sign = mouse_pitch_signs[i] if i < len(mouse_pitch_signs) else None
+            extra["mouse_pitch_sign"] = int(mouse_pitch_sign or self.metrics_config.mouse_pitch_sign)
+            extras.append(extra)
+        return extras
+
+    @staticmethod
+    def _available_paths(paths: list[str | None]) -> list[str] | None:
+        if not paths or any(path is None for path in paths):
+            return None
+        return [str(path) for path in paths]
+
+    @staticmethod
+    def _validation_actions(validation_batch: dict[str, Any]) -> dict[str, Any] | None:
+        keyboard = validation_batch.get("keyboard_cond")
+        mouse = validation_batch.get("mouse_cond")
+        if keyboard is not None and mouse is not None:
+            return {
+                "keyboard": np.asarray(keyboard),
+                "mouse": np.asarray(mouse),
+            }
+
+        action_path = validation_batch.get("action_path")
+        if not isinstance(action_path, str) or not os.path.isfile(action_path):
+            return None
+        try:
+            actions_obj = np.load(action_path, allow_pickle=True)
+            if isinstance(actions_obj, np.ndarray) and actions_obj.dtype == object:
+                actions_obj = actions_obj.item()
+        except Exception as exc:
+            logger.warning(
+                "Failed to load validation action file %s for metrics: %s",
+                action_path,
+                exc,
+            )
+            return None
+        if not isinstance(actions_obj, dict):
+            return None
+        keyboard = actions_obj.get("keyboard")
+        mouse = actions_obj.get("mouse")
+        if keyboard is None or mouse is None:
+            return None
+        return {
+            "keyboard": np.asarray(keyboard),
+            "mouse": np.asarray(mouse),
+        }
+
+    def _validation_mouse_pitch_sign(
+        self,
+        validation_batch: dict[str, Any],
+    ) -> int:
+        if validation_batch.get("mouse_pitch_sign") is not None:
+            return int(validation_batch["mouse_pitch_sign"])
+        flipped = validation_batch.get("mouse_pitch_flipped")
+        if isinstance(flipped, str):
+            flipped = flipped.lower() in {"1", "true", "yes", "y"}
+        if bool(flipped):
+            return -1
+        return int(self.metrics_config.mouse_pitch_sign)

    # ----------------------------------------------------------
    # Pipeline management
@@ -437,9 +1016,51 @@ class ValidationCallback(Callback):

    def _get_sampling_param(self) -> SamplingParam:
        if self._sampling_param is None:
-            self._sampling_param = (SamplingParam.from_pretrained(self.training_config.model_path))
+            self._sampling_param = (SamplingParam.from_pretrained(self._pipeline_model_path()))
        return self._sampling_param

+    def _pipeline_model_path(self) -> str:
+        tc = self.training_config
+        pipeline_config = getattr(tc, "pipeline_config", None)
+        model_path = getattr(pipeline_config, "_fastvideo_train_model_path", None)
+        return str(model_path or tc.model_path)
+
+    def _validation_pipeline_config(self, transformer: torch.nn.Module) -> Any:
+        tc = self.training_config
+        pipeline_config = deepcopy(tc.pipeline_config)
+        self._sync_runtime_dit_arch_config(
+            pipeline_config,
+            transformer,
+        )
+        return pipeline_config
+
+    @staticmethod
+    def _sync_runtime_dit_arch_config(
+        pipeline_config: Any,
+        transformer: torch.nn.Module,
+    ) -> None:
+        dit_config = getattr(
+            pipeline_config,
+            "dit_config",
+            None,
+        )
+        arch_config = getattr(
+            dit_config,
+            "arch_config",
+            None,
+        )
+        if arch_config is None:
+            return
+
+        for name in ("local_attn_size", "sink_size"):
+            if not hasattr(arch_config, name) or not hasattr(transformer, name):
+                continue
+            setattr(
+                arch_config,
+                name,
+                getattr(transformer, name),
+            )
+
    def _get_pipeline(
        self,
        *,
@@ -479,9 +1100,18 @@ class ValidationCallback(Callback):
        kwargs.update(self.pipeline_kwargs)

        self._pipeline = PipelineCls.from_pretrained(
-            tc.model_path,
+            self._pipeline_model_path(),
            **kwargs,
        )
+        if tc.pipeline_config is not None:
+            self._pipeline.fastvideo_args.pipeline_config = self._validation_pipeline_config(transformer)
+            arch_config = self._pipeline.fastvideo_args.pipeline_config.dit_config.arch_config
+            logger.info(
+                "Validation pipeline runtime config: local_attn_size=%s sink_size=%s boundary_ratio=%s",
+                getattr(arch_config, "local_attn_size", None),
+                getattr(arch_config, "sink_size", None),
+                getattr(self._pipeline.fastvideo_args.pipeline_config.dit_config, "boundary_ratio", None),
+            )

        self._pipeline_key = key
        return self._pipeline
@@ -552,8 +1182,47 @@ class ValidationCallback(Callback):
        if ("image" in validation_batch and validation_batch["image"] is not None):
            batch.pil_image = validation_batch["image"]

+        self._attach_action_conditions(
+            batch,
+            validation_batch,
+            sampling_param.num_frames,
+        )
+
        return batch

+    def _attach_action_conditions(
+        self,
+        batch: ForwardBatch,
+        validation_batch: dict[str, Any],
+        num_frames: int,
+    ) -> None:
+        for name in ("keyboard_cond", "mouse_cond"):
+            value = validation_batch.get(name)
+            if value is None:
+                continue
+            array = np.asarray(value)
+            if array.size == 0:
+                continue
+            array = array[:num_frames]
+            tensor = torch.as_tensor(
+                array,
+                dtype=torch.bfloat16,
+            )
+            student = getattr(self.method, "student", None)
+            prepare_action = getattr(student, "prepare_validation_action_condition", None)
+            if prepare_action is not None:
+                tensor = prepare_action(
+                    tensor,
+                    name=name,
+                    keyboard_value_scale=self.keyboard_value_scale,
+                )
+            tensor = tensor.unsqueeze(0)
+            setattr(
+                batch,
+                name,
+                tensor,
+            )
+
    # ----------------------------------------------------------
    # Validation loop
    # ----------------------------------------------------------
@@ -579,6 +1248,10 @@ class ValidationCallback(Callback):
            tc,
            model_path=tc.model_path,
        )
+        self._sync_runtime_dit_arch_config(
+            inference_args.pipeline_config,
+            transformer,
+        )

        # Propagate sampling_timesteps to pipeline_config so
        # causal/DMD denoising stages can read them.
@@ -586,7 +1259,12 @@ class ValidationCallback(Callback):
            inference_args.pipeline_config.dmd_denoising_steps = ([int(s) for s in self.sampling_timesteps])

        videos: list[list[np.ndarray]] = []
+        overlay_videos: list[list[np.ndarray]] = []
        captions: list[str] = []
+        overlay_captions: list[str] = []
+        ref_videos: list[str | None] = []
+        actions: list[dict[str, Any] | None] = []
+        mouse_pitch_signs: list[int | None] = []

        for validation_batch in dataloader:
            batch = self._prepare_validation_batch(
@@ -597,6 +1275,11 @@ class ValidationCallback(Callback):

            assert (batch.prompt is not None and isinstance(batch.prompt, str))
            captions.append(batch.prompt)
+            ref_video = validation_batch.get("ref_video")
+            ref_videos.append(ref_video if isinstance(ref_video, str) else None)
+            action = self._validation_actions(validation_batch)
+            actions.append(action)
+            mouse_pitch_signs.append(self._validation_mouse_pitch_sign(validation_batch))

            with torch.no_grad():
                output_batch = pipeline.forward(
@@ -621,10 +1304,38 @@ class ValidationCallback(Callback):
                x = (x.transpose(0, 1).transpose(1, 2).squeeze(-1))
                frames.append((x * 255).numpy().astype(np.uint8))
            videos.append(frames)
+            if self.overlay_actions:
+                overlay_frames = self._post_process_validation_frames(
+                    frames,
+                    action=action,
+                )
+                if overlay_frames is not None:
+                    overlay_videos.append(overlay_frames)
+                    overlay_captions.append(batch.prompt)

        return _ValidationStepResult(
            videos=videos,
            captions=captions,
+            overlay_videos=overlay_videos,
+            overlay_captions=overlay_captions,
+            ref_videos=ref_videos,
+            actions=actions,
+            mouse_pitch_signs=mouse_pitch_signs,
+        )
+
+    def _post_process_validation_frames(
+        self,
+        frames: list[np.ndarray],
+        *,
+        action: dict[str, Any] | None,
+    ) -> list[np.ndarray] | None:
+        student = getattr(self.method, "student", None)
+        post_process = getattr(student, "post_process_validation_frames", None)
+        if post_process is None:
+            return None
+        return post_process(
+            frames,
+            action=action,
        )

    # ----------------------------------------------------------
@@ -3,8 +3,11 @@
 from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
 from fastvideo.train.methods.distribution_matching.self_forcing import (
    SelfForcingMethod, )
+from fastvideo.train.methods.distribution_matching.streaming_long_tuning import (
+    StreamingLongTuningMethod, )

 __all__ = [
    "DMD2Method",
    "SelfForcingMethod",
+    "StreamingLongTuningMethod",
 ]
@@ -13,6 +13,7 @@ from fastvideo.distributed import (
    get_sp_group,
    get_world_group,
 )
+from fastvideo.models.dits.matrixgame2.utils import scale_keyboard_condition
 from fastvideo.pipelines import TrainingBatch
 from fastvideo.training.training_utils import normalize_dit_input

@@ -84,12 +85,14 @@ class MatrixGame2Model(WanModel):
            raw_batch,
            key="keyboard_cond",
            expected_frames=self._expected_action_frames(tc.data.num_latent_t),
+            expected_dim=self._expected_action_dim("keyboard"),
            dtype=dtype,
        )
        mouse_cond = self._get_optional_action(
            raw_batch,
            key="mouse_cond",
            expected_frames=self._expected_action_frames(tc.data.num_latent_t),
+            expected_dim=self._expected_action_dim("mouse"),
            dtype=dtype,
        )

@@ -172,23 +175,64 @@ class MatrixGame2Model(WanModel):
            raise ValueError("text_dict cannot be None for Matrix-Game 2.0")
        hidden_states = noise_input.permute(0, 2, 1, 3, 4)
        if hidden_states.shape[1] == 16:
-            cond_latents = text_dict.get("image_latents")
+            # Streaming long tuning scores a window at a time; use the
+            # window-specific first-image latent when the method provides it.
+            cond_latents = self._get_score_window_image_latents(
+                text_dict,
+                batch_info=text_dict.get("_batch_info"),
+                num_latents=hidden_states.shape[2],
+            )
            if cond_latents is None:
                raise RuntimeError("Matrix-Game 2.0 requires image_latents in conditional_dict "
                                   "when noise_input has 16 channels")
-            num_t = hidden_states.shape[2]
-            cond_latents = cond_latents[:, :, :num_t]
            hidden_states = torch.cat([hidden_states, cond_latents], dim=1)
+        # For streaming long tuning, align actions to the same score window as
+        # the latent/image condition; otherwise this falls back to the prefix.
+        keyboard_cond, mouse_cond = self._get_score_window_actions(
+            text_dict,
+            batch_info=text_dict.get("_batch_info"),
+            num_latents=hidden_states.shape[2],
+        )
        return {
            "hidden_states": hidden_states,
            "encoder_hidden_states": None,
            "timestep": timestep.to(device=self.device, dtype=torch.bfloat16),
            "encoder_hidden_states_image": text_dict["encoder_hidden_states_image"],
-            "keyboard_cond": text_dict["keyboard_cond"],
-            "mouse_cond": text_dict["mouse_cond"],
+            "keyboard_cond": keyboard_cond,
+            "mouse_cond": mouse_cond,
            "return_dict": False,
        }

+    def predict_noise(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor:
+        # Streaming long tuning re-anchors each score window on its first
+        # image/latent, so pass window metadata through existing dict plumbing.
+        info = getattr(batch, "matrixgame_streaming_chunk_info", None)
+        for text_dict in (batch.conditional_dict, batch.unconditional_dict):
+            if isinstance(text_dict, dict):
+                text_dict["_batch_info"] = info
+        try:
+            return super().predict_noise(
+                noisy_latents,
+                timestep,
+                batch,
+                conditional=conditional,
+                cfg_uncond=cfg_uncond,
+                attn_kind=attn_kind,
+            )
+        finally:
+            for text_dict in (batch.conditional_dict, batch.unconditional_dict):
+                if isinstance(text_dict, dict):
+                    text_dict.pop("_batch_info", None)
+
    def _build_matrixgame_cond_concat(
        self,
        image_latents: torch.Tensor,
@@ -235,21 +279,138 @@ class MatrixGame2Model(WanModel):
        ).transpose(1, 2)
        return torch.cat([mask_lat_size, image_latents], dim=1)

+    def _get_score_window_image_latents(
+        self,
+        text_dict: dict[str, Any],
+        *,
+        batch_info: dict[str, Any] | None,
+        num_latents: int,
+    ) -> torch.Tensor | None:
+        # Prefer a window-specific first-image latent when long tuning has
+        # recomputed one; otherwise keep the original MG first-frame condition.
+        if batch_info is not None:
+            condition_latents = batch_info.get("condition_image_latents")
+            if condition_latents is not None:
+                if condition_latents.shape[2] != num_latents:
+                    raise ValueError("condition_image_latents temporal length must match "
+                                     "noise_input: "
+                                     f"{condition_latents.shape[2]} vs {num_latents}")
+                return condition_latents
+
+        image_latents = text_dict.get("image_latents")
+        if image_latents is None:
+            return None
+        return image_latents[:, :, :num_latents]
+
+    def _get_score_window_actions(
+        self,
+        text_dict: dict[str, Any],
+        *,
+        batch_info: dict[str, Any] | None,
+        num_latents: int,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        frame_start = 0
+        if batch_info is not None:
+            frame_start = int(batch_info.get("frame_start", 0) or 0)
+        frame_end = frame_start + int(num_latents)
+        action_frame_start = frame_start * self._temporal_compression_ratio()
+        action_frame_end = ((frame_end - 1) * self._temporal_compression_ratio()) + 1
+
+        def _slice_action(action: torch.Tensor | None) -> torch.Tensor | None:
+            if action is None:
+                return None
+            if action.shape[1] < action_frame_end:
+                raise ValueError("Matrix-Game 2.0 score action tensor is shorter than "
+                                 "the requested streaming window: "
+                                 f"got={action.shape[1]}, required>={action_frame_end}")
+            return action[:, action_frame_start:action_frame_end]
+
+        return (
+            _slice_action(text_dict.get("keyboard_cond")),
+            _slice_action(text_dict.get("mouse_cond")),
+        )
+
    def _get_optional_action(
        self,
        raw_batch: dict[str, Any],
        *,
        key: str,
        expected_frames: int,
+        expected_dim: int,
        dtype: torch.dtype,
    ) -> torch.Tensor | None:
        value = raw_batch.get(key)
        if value is None or value.numel() == 0:
            return None
+        # Actions must follow the same first-image window as the latent
+        # condition; allow only padded zero channels during format alignment.
+        value = self._align_action_feature_dim(
+            value,
+            name=key,
+            expected_dim=expected_dim,
+        )
        if value.shape[1] < expected_frames:
            raise ValueError(f"{key} has {value.shape[1]} frames but requires at least {expected_frames}")
        return value[:, :expected_frames].to(device=self.device, dtype=dtype)

+    def prepare_validation_action_condition(
+        self,
+        action: torch.Tensor,
+        *,
+        name: str,
+        keyboard_value_scale: float = 1.0,
+    ) -> torch.Tensor:
+        channel = name.removesuffix("_cond")
+        action = self._align_action_feature_dim(
+            action,
+            name=name,
+            expected_dim=self._expected_action_dim(channel),
+        )
+        if name == "keyboard_cond":
+            action = scale_keyboard_condition(action, float(keyboard_value_scale))
+        return action
+
+    def post_process_validation_frames(
+        self,
+        frames: list[Any],
+        *,
+        action: dict[str, Any] | None = None,
+    ) -> list[Any] | None:
+        if action is None:
+            return None
+        from fastvideo.models.dits.matrixgame2.utils import (
+            overlay_validation_actions_on_frames, )
+
+        return overlay_validation_actions_on_frames(
+            frames,
+            keyboard_cond=action.get("keyboard"),
+            mouse_cond=action.get("mouse"),
+        )
+
+    def _expected_action_dim(self, channel: str) -> int:
+        action_config = getattr(self.transformer, "action_config", {}) or {}
+        return int(action_config.get(f"{channel}_dim_in", 0) or 0)
+
+    def _align_action_feature_dim(
+        self,
+        action: torch.Tensor,
+        *,
+        name: str,
+        expected_dim: int,
+    ) -> torch.Tensor:
+        if expected_dim <= 0:
+            return action
+        actual_dim = int(action.shape[-1])
+        if actual_dim == expected_dim:
+            return action
+        if actual_dim > expected_dim:
+            extra = action[..., expected_dim:]
+            if bool(torch.count_nonzero(extra).item() == 0):
+                return action[..., :expected_dim]
+            raise ValueError(f"{name} feature dim mismatch: got={actual_dim}, "
+                             f"expected={expected_dim}, and trailing channels are not all zero.")
+        raise ValueError(f"{name} feature dim mismatch: got={actual_dim}, expected={expected_dim}")
+
    def _expected_action_frames(self, num_latent_t: int) -> int:
        return (num_latent_t - 1) * self._temporal_compression_ratio() + 1

@@ -129,6 +129,13 @@ class WanCausalModel(WanModel, CausalModelBase):

        if (self._should_snapshot_streaming_cache() and torch.is_grad_enabled()):
            kv_cache = self._snapshot_kv_cache_indices(kv_cache)
+            # The cross-attn cache is stateful inside the checkpointed block
+            # forward (is_init flip + k/v writes); a checkpoint recompute that
+            # sees the mutated dict skips the k/v projections and trips
+            # torch's saved-tensor-count check. Grad-enabled forwards
+            # recompute cross k/v instead, matching the legacy
+            # gradient_checkpointing branch of _forward_inference.
+            crossattn_cache = None

        model_kwargs: dict[str, Any] = {
            "kv_cache": kv_cache,
@@ -343,6 +343,15 @@ def _build_training_config(
            if init_from is not None:
                model_path = str(init_from)

+    raw_data_path = da.get("data_path", "") or ""
+    data_path: str | list[str] | dict[str, int]
+    if isinstance(raw_data_path, dict):
+        data_path = {str(path): int(repeat) for path, repeat in raw_data_path.items()}
+    elif isinstance(raw_data_path, list | tuple):
+        data_path = [str(path) for path in raw_data_path]
+    else:
+        data_path = str(raw_data_path)
+
    preprocessed_data_type = str(da.get("preprocessed_data_type", "t2v") or "t2v").strip().lower()
    if preprocessed_data_type not in {"t2v", "text_only"}:
        raise ValueError("training.data.preprocessed_data_type must be one of "
@@ -359,7 +368,7 @@ def _build_training_config(
            pin_cpu_memory=bool(d.get("pin_cpu_memory", False)),
        ),
        data=DataConfig(
-            data_path=str(da.get("data_path", "") or ""),
+            data_path=data_path,
            preprocessed_data_type=preprocessed_data_type,
            train_batch_size=int(da.get("train_batch_size", 1) or 1),
            dataloader_num_workers=int(da.get("dataloader_num_workers", 0) or 0),
@@ -392,6 +401,7 @@ def _build_training_config(
        ),
        tracker=TrackerConfig(
            trackers=list(tr.get("trackers", []) or []),
+            entity=str(tr.get("entity", "") or ""),
            project_name=str(tr.get("project_name", "fastvideo") or "fastvideo"),
            run_name=str(tr.get("run_name", "") or ""),
        ),
@@ -39,6 +39,7 @@ def build_tracker(
        tracker_log_dir = os.path.join(tracker_log_dir, "tracker")

    tracker_config_dict = config if trackers else None
+    tracker_entity = tracker_config.entity or None
    tracker_run_name = tracker_config.run_name or None
    project = (tracker_config.project_name or "fastvideo")

@@ -47,5 +48,6 @@ def build_tracker(
        experiment_name=project,
        config=tracker_config_dict,
        log_dir=tracker_log_dir,
+        entity=tracker_entity,
        run_name=tracker_run_name,
    )
@@ -22,7 +22,7 @@ class DistributedConfig:

@dataclass(slots=True)
 class DataConfig:
-    data_path: str = ""
+    data_path: str | list[str] | dict[str, int] = ""
    preprocessed_data_type: str = "t2v"
    train_batch_size: int = 1
    dataloader_num_workers: int = 0
@@ -63,6 +63,7 @@ class CheckpointConfig:
@dataclass(slots=True)
 class TrackerConfig:
    trackers: list[str] = field(default_factory=list)
+    entity: str = ""
    project_name: str = "fastvideo"
    run_name: str = ""

@@ -305,6 +305,7 @@ class WandbTracker(BaseTracker):
        log_dir: str,
        *,
        config: dict[str, Any] | None = None,
+        entity: str | None = None,
        run_name: str | None = None,
    ) -> None:
        super().__init__()
@@ -315,6 +316,7 @@ class WandbTracker(BaseTracker):

        self._wandb = wandb
        self._run = wandb.init(
+            entity=entity,
            project=experiment_name,
            dir=log_dir,
            config=(_sanitize_wandb_config(config) if config is not None else None),
@@ -515,6 +517,7 @@ def initialize_trackers(
    experiment_name: str,
    config: dict[str, Any] | None,
    log_dir: str,
+    entity: str | None = None,
    run_name: str | None = None,
 ) -> BaseTracker:
    """Create tracker instances based on ``trackers`` configuration."""
@@ -538,6 +541,7 @@ def initialize_trackers(
                    experiment_name,
                    os.path.abspath(log_dir),
                    config=config,
+                    entity=entity,
                    run_name=run_name,
                ))
        elif tracker_name == Trackers.SWANLAB.value:
Author	SHA1	Message	Date
SolitaryThinker	0a231e0a1a	[docs]: note the Hub build's Hopper-only arch subset is deliberate Main's kernel wheels went multi-arch (0.3.x); the Hub build.toml stays sm_90a because the TK attention kernels it packages require Hopper.	2026-07-10 19:31:18 -07:00
SolitaryThinker	e9efd52211	[kernel]: declare mutation + return aliases in Hub torch-op schemas sta_fwd and gemm_cuda return the tensor they write in place, and the optional output args of rms_norm_cuda/layer_norm_cuda are mutated when provided — none of which the schemas declared, which silently breaks torch.compile. Annotate with alias sets (Tensor(a!) ... -> Tensor(a!)). Schemas validated with torch.library on torch 2.x (parse + register).	2026-07-10 19:31:10 -07:00
SolitaryThinker	2b3109b68d	[ci]: SHA-pin kernel-builder-job and serialize Hub publishes The third-party action receives HF_TOKEN, so pin it to a commit SHA instead of @main. Add a concurrency group so two merges can't race multi-hour publish jobs against the same Hub repo.	2026-07-10 19:31:00 -07:00
SolitaryThinker	ebd4d59951	[ci]: version-gate the HF Kernel Hub publish like the PyPI flow Publish to the Hub only when the fastvideo-kernel version bumps (same trigger + version-extraction as publish-kernel.yml), instead of on every main push touching csrc/**. workflow_dispatch stays as the manual override.	2026-07-10 19:30:47 -07:00
Aryan Kumar	a6eb382d34	[feat]: add Hugging Face kernel hub packaging	2026-07-10 19:25:34 -07:00
Raghav KandSolitaryThinker	0c90c8c24d	[bugfix] nvfp4: cast fp32 inputs to bf16 instead of asserting (#1488 ) Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>	2026-07-10 21:39:17 +00:00
Mingjia HuoandClaude Fable 5	4c08ffce49	[feat] World model training using third person games (#1443 ) Co-authored-by: Claude Fable 5 <noreply@anthropic.com>	2026-07-10 05:09:07 +00:00
Atharv Ramesh	af4a77553c	[ci]: add SSIM reference bootstrap flow (#1522 ) (#1547 )	2026-07-10 01:49:38 +00:00
alexzmsandSolitaryThinker	c096fda1eb	[docs] Add LTX-2.3 distilled inference run configs (t2v/i2v × 5+2/8+3 × resolutions) (#1568 ) Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>	2026-07-09 18:53:12 +00:00
William Lin	8f47e85be0	[bugfix]: retry remote image downloads in load_image (#1570 )	2026-07-09 07:33:24 -07:00
William Lin	90d3bd19eb	[infra] Deliver per-job Buildkite env to Modal CI at runtime, not as image layers (#1569 )	2026-07-09 06:33:37 -07:00
Mac LeeandSolitaryThinker	afb4f7d3c5	[ci]: emit v2 performance result schema (#1551 ) Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>	2026-07-09 06:11:15 +00:00