[bugfix]: keep text-encoder CPU offload on by default

Revert the default flips from 40c8b7c1 (OffloadConfig.text_encoder and
FastVideoArgs.text_encoder_cpu_offload False -> True). Flipping the
global default changes the baseline VRAM footprint for every model
family, not just Kandinsky, and breaks the unit lane:
fastvideo/tests/api/test_parser.py::test_load_run_config_supports_yaml_roundtrip
still golden-checks text_encoder offload as True.

Nothing in this PR needs the flip: the Kandinsky examples pass
text_encoder_cpu_offload explicitly, and the offload-aware
TextEncodingStage (kept) makes offload work for Kandinsky's CLIP
either way. Offload stays opt-out rather than opt-in.

.buildkite/performance-benchmarks/tests/wan-t2v-1.3b.json

@@ -1,9 +1,9 @@
 {
   "benchmark_id": "wan-t2v-1.3b-2gpu",
   "config_schema_version": 2,
-  "workload_id": "wan-t2v",
-  "variant_id": "1.3b-sp2",
-  "benchmark_version": 2,
+  "workload_id": "wan-t2v-1.3b",
+  "variant_id": "canonical",
+  "benchmark_version": 1,
   "description": "Wan2.1 T2V 1.3B inference performance",
   "model": {
     "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",

.buildkite/pipeline.yml

@@ -481,7 +481,6 @@ steps:
                 - "fastvideo/layers/**"
                 - "fastvideo/worker/**"
                 - "fastvideo/entrypoints/**"
-                - "fastvideo/performance/**"
                 - "fastvideo/tests/performance/**"
                 - ".buildkite/performance-benchmarks/**"
                 - "pyproject.toml"

.buildkite/scripts/pr_test.sh

@@ -80,23 +80,6 @@ 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"
@@ -189,12 +172,7 @@ case "$TEST_TYPE" in
         ;;
     "ssim")
         log "Running 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"
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_SSIM_TEST_FILE::run_ssim_tests"
         ;;
     "training")
         log "Running training tests..."

.github/scripts/gate_full_suite.sh

@@ -1,133 +0,0 @@
-#!/usr/bin/env bash
-# Gate the expensive Buildkite full suite on the cheap GitHub checks.
-#
-# Polls the workflow runs for the PR head commit and only exits 0 once the
-# watched cheap workflows (pre-commit, docs build) have succeeded, so the
-# 'ready' label cannot burn ~20 GPU lanes on a head that a cheap check has
-# already doomed.
-#
-# Semantics:
-#   - watched run completed with a bad conclusion  -> exit 1 (fail CLOSED:
-#     no full suite; the next push re-arms via the 'synchronize' trigger)
-#   - watched run cancelled                         -> still pending: the docs
-#     workflow's repo-global 'pages' concurrency group cancels runs superseded
-#     by unrelated pushes, so 'cancelled' is not a verdict on this PR
-#   - watched runs pending                          -> poll until done
-#   - docs run absent                               -> not applicable after a
-#     short grace period ('Deploy Documentation' is path-filtered on PRs)
-#   - pre-commit run absent                         -> keep polling: pre-commit
-#     is never path-filtered, so its absence is always anomalous
-#   - 'ready' label removed while waiting           -> exit 1 (fail CLOSED:
-#     un-labeling is a deliberate maintainer action)
-#   - GitHub API unreachable or timeout             -> exit 0 (fail OPEN,
-#     loud warning: never brick CI on a GitHub outage)
-#
-# Required env: PR_SHA (PR head commit), PR_NUMBER, GITHUB_REPOSITORY, GH_TOKEN.
-set -euo pipefail
-
-: "${PR_SHA:?PR_SHA (PR head commit) is required}"
-: "${PR_NUMBER:?PR_NUMBER (pull request number) is required}"
-: "${GITHUB_REPOSITORY:?GITHUB_REPOSITORY is required}"
-
-# Workflow-level `name:` values that must be green before the full suite
-# may start. "Deploy Documentation" is path-filtered on PRs, so its run may
-# legitimately never exist; pre-commit always runs, so it must appear.
-WATCHED_NAMES='["pre-commit", "Deploy Documentation"]'
-WATCHED_REGEX='^(pre-commit|Deploy Documentation)$'
-POLL_SECS="${POLL_SECS:-20}"
-GRACE_SECS="${GRACE_SECS:-60}"
-MAX_WAIT_SECS="${MAX_WAIT_SECS:-1500}"
-
-# Bound each API call so a hung connection hits the 3-strike fail-open path
-# instead of pinning the loop until the job timeout (which would fail closed
-# on exactly the GitHub-outage case this script is meant to survive).
-if command -v timeout >/dev/null 2>&1; then
-  gh_api() { timeout 30 gh api "$@"; }
-else
-  gh_api() { gh api "$@"; }  # macOS dev boxes; CI always has coreutils timeout
-fi
-
-# The workflow checked the label before starting the gate, but the wait can
-# last ~25 min: re-check once before any exit 0 and fail closed if 'ready'
-# was removed in the meantime. An API error here proceeds (the label was
-# present when the gate started; never brick CI on an outage).
-recheck_ready_label() {
-  local pr_json
-  if pr_json=$(gh_api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" 2>/dev/null); then
-    if ! jq -e '[.labels[]?.name] | index("ready")' <<<"$pr_json" >/dev/null 2>&1; then
-      echo "::error::PR #${PR_NUMBER} no longer has the 'ready' label —" \
-        "NOT triggering the Buildkite full suite. Re-add the label to re-arm."
-      exit 1
-    fi
-  else
-    echo "::warning::Could not re-check the 'ready' label on PR #${PR_NUMBER}; proceeding (it was present when the gate started)."
-  fi
-}
-
-start=$(date +%s)
-api_fails=0
-missing=""
-
-while true; do
-  elapsed=$(( $(date +%s) - start ))
-
-  if runs_json=$(gh_api "repos/${GITHUB_REPOSITORY}/actions/runs?head_sha=${PR_SHA}&per_page=100" 2>/dev/null) \
-     && state=$(jq --arg re "$WATCHED_REGEX" '
-          [.workflow_runs[]? | select(.name // "" | test($re))]
-          | group_by(.name) | map(max_by(.id))
-          | map({name, status, conclusion})' <<<"$runs_json" 2>/dev/null); then
-    api_fails=0
-    echo "t+${elapsed}s watched checks: $(jq -c . <<<"$state")"
-
-    failed=$(jq -r '[.[] | select(.status == "completed"
-          and (.conclusion | IN("success", "skipped", "neutral", "cancelled") | not))]
-        | map(.name) | join(", ")' <<<"$state")
-    if [ -n "$failed" ]; then
-      echo "::error::Cheap check(s) failed on ${PR_SHA}: ${failed}." \
-        "NOT triggering the Buildkite full suite. Push a fix (the 'ready'" \
-        "label re-arms on every push), or re-run the failed check and then" \
-        "re-run this workflow."
-      exit 1
-    fi
-
-    # 'cancelled' counts as pending: wait for a re-run to reach a real verdict
-    # (bounded by MAX_WAIT, then the fail-open below).
-    pending=$(jq '[.[] | select(.status != "completed" or .conclusion == "cancelled")] | length' <<<"$state")
-    missing=$(jq -r --argjson watched "$WATCHED_NAMES" '($watched - map(.name)) | join(", ")' <<<"$state")
-    if [ "$pending" -eq 0 ]; then
-      if [ -z "$missing" ]; then
-        recheck_ready_label
-        echo "All watched cheap checks are green — full suite may proceed."
-        exit 0
-      fi
-      case "$missing" in
-        *pre-commit*)
-          echo "pre-commit run not found for ${PR_SHA} yet; waiting (pre-commit is never path-filtered, so its absence is anomalous)."
-          ;;
-        *)
-          if [ "$elapsed" -ge "$GRACE_SECS" ]; then
-            recheck_ready_label
-            echo "::warning::Watched run(s) never appeared for ${PR_SHA}: ${missing} (path-filtered, likely not applicable). Proceeding on the checks that did run."
-            exit 0
-          fi
-          echo "Waiting up to ${GRACE_SECS}s grace for path-filtered run(s) to appear: ${missing}."
-          ;;
-      esac
-    fi
-  else
-    api_fails=$(( api_fails + 1 ))
-    echo "::warning::GitHub API error querying workflow runs for ${PR_SHA} (attempt ${api_fails}/3)."
-    if [ "$api_fails" -ge 3 ]; then
-      recheck_ready_label
-      echo "::warning::FAILING OPEN: cannot query GitHub check status — triggering the full suite WITHOUT the cheap-check gate."
-      exit 0
-    fi
-  fi
-
-  if [ "$elapsed" -ge "$MAX_WAIT_SECS" ]; then
-    recheck_ready_label
-    echo "::warning::FAILING OPEN: watched checks still pending after $(( MAX_WAIT_SECS / 60 )) min${missing:+ (never appeared: ${missing})} — triggering the full suite anyway."
-    exit 0
-  fi
-  sleep "$POLL_SECS"
-done

.github/scripts/test_gate_full_suite.sh

@@ -1,122 +0,0 @@
-#!/usr/bin/env bash
-# Self-test for gate_full_suite.sh using a mocked `gh`. No network, runs on
-# any dev box: bash .github/scripts/test_gate_full_suite.sh
-set -u
-here=$(cd "$(dirname "$0")" && pwd)
-tmp=$(mktemp -d)
-trap 'rm -rf "$tmp"' EXIT
-
-# Mock gh. Asserts the exact endpoint (including head_sha) it is called
-# with — an endpoint typo in the gate script fails the test rather than
-# silently serving canned data. On the runs endpoint it serves
-# $MOCK_DIR/response_<call#>.json, sticking on the highest existing file,
-# and exits 1 if none exist (simulates a GitHub API outage). On the pulls
-# endpoint it serves $MOCK_DIR/pr.json, defaulting to a 'ready'-labeled PR.
-cat > "$tmp/gh" <<'EOF'
-#!/usr/bin/env bash
-if [ "${1:-}" != "api" ]; then
-  echo "unexpected gh invocation: $*" >> "$MOCK_DIR/endpoint_error"
-  exit 2
-fi
-case "${2:-}" in
-  "repos/o/r/actions/runs?head_sha=deadbeef&per_page=100")
-    n=$(( $(cat "$MOCK_DIR/count" 2>/dev/null || echo 0) + 1 ))
-    echo "$n" > "$MOCK_DIR/count"
-    while [ "$n" -gt 0 ]; do
-      if [ -f "$MOCK_DIR/response_$n.json" ]; then
-        cat "$MOCK_DIR/response_$n.json"
-        exit 0
-      fi
-      n=$(( n - 1 ))
-    done
-    echo "api outage" >&2
-    exit 1
-    ;;
-  "repos/o/r/pulls/42")
-    if [ -f "$MOCK_DIR/pr.json" ]; then
-      cat "$MOCK_DIR/pr.json"
-    else
-      echo '{"labels": [{"name": "ready"}]}'
-    fi
-    ;;
-  *)
-    echo "unexpected gh endpoint: $2" >> "$MOCK_DIR/endpoint_error"
-    exit 2
-    ;;
-esac
-EOF
-chmod +x "$tmp/gh"
-
-PC_OK='{"name": "pre-commit", "id": 1, "status": "completed", "conclusion": "success"}'
-PC_BAD='{"name": "pre-commit", "id": 1, "status": "completed", "conclusion": "failure"}'
-PC_PENDING='{"name": "pre-commit", "id": 1, "status": "in_progress", "conclusion": null}'
-DOCS_OK='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "success"}'
-DOCS_BAD='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "failure"}'
-DOCS_CANCELLED='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "cancelled"}'
-OTHER='{"name": "Trigger Full Suite", "id": 3, "status": "in_progress", "conclusion": null}'
-NULL_NAME='{"name": null, "id": 4, "status": "completed", "conclusion": "failure"}'
-PC_OK_RERUN='{"name": "pre-commit", "id": 5, "status": "completed", "conclusion": "success"}'
-
-fails=0
-want_log=""  # optional: expect() also greps out.log for this regex, then resets
-pr_json=""   # optional: served for the pulls (label re-check) endpoint, then resets
-raw_body=""  # optional: serve responses verbatim instead of wrapping in workflow_runs
-expect() { # <name> <expected-exit> <response json>...
-  local name=$1 want=$2 dir i=1
-  shift 2
-  dir=$(mktemp -d "$tmp/test_XXXXXX")
-  for body in "$@"; do
-    if [ -n "$raw_body" ]; then
-      printf '%s' "$body" > "$dir/response_$i.json"
-    else
-      printf '{"workflow_runs": [%s]}' "$body" > "$dir/response_$i.json"
-    fi
-    i=$(( i + 1 ))
-  done
-  [ -n "$pr_json" ] && printf '%s' "$pr_json" > "$dir/pr.json"
-  ( export PATH="$tmp:$PATH" MOCK_DIR="$dir" PR_SHA=deadbeef PR_NUMBER=42 \
-      GITHUB_REPOSITORY=o/r POLL_SECS=0 GRACE_SECS=1 MAX_WAIT_SECS=3
-    bash "$here/gate_full_suite.sh" > "$dir/out.log" 2>&1 )
-  local rc=$?
-  if [ "$rc" -ne "$want" ]; then
-    echo "FAIL: $name (exit $rc, want $want)"
-    cat "$dir/out.log"
-    fails=1
-  elif [ -f "$dir/endpoint_error" ]; then
-    echo "FAIL: $name (mock gh got an unexpected call)"
-    cat "$dir/endpoint_error"
-    fails=1
-  elif [ -n "$want_log" ] && ! grep -Eq "$want_log" "$dir/out.log"; then
-    echo "FAIL: $name (log does not match: $want_log)"
-    cat "$dir/out.log"
-    fails=1
-  else
-    echo "ok: $name"
-  fi
-  want_log="" pr_json="" raw_body=""
-}
-
-expect "both green -> proceed" 0 "$PC_OK, $DOCS_OK, $OTHER, $NULL_NAME"
-expect "docs build failed -> blocked" 1 "$PC_OK, $DOCS_BAD"
-expect "pre-commit failed -> blocked" 1 "$PC_BAD"
-expect "pending then green -> proceed" 0 "$PC_PENDING" "$PC_OK, $DOCS_OK"
-want_log="never appeared.*Deploy Documentation"
-expect "docs run absent (path-filtered) -> proceed after grace" 0 "$PC_OK"
-expect "API outage -> fail open" 0
-want_log="FAILING OPEN"
-expect "pending past MAX_WAIT -> fail open" 0 "$PC_PENDING"
-want_log="FAILING OPEN"
-expect "unrelated runs only -> no grace, fail open at MAX_WAIT" 0 "$OTHER"
-expect "cancelled docs then green -> proceed" 0 \
-  "$PC_OK, $DOCS_CANCELLED" "$PC_OK, $DOCS_OK"
-want_log="FAILING OPEN"
-expect "cancelled docs forever -> fail open at MAX_WAIT" 0 "$PC_OK, $DOCS_CANCELLED"
-want_log="FAILING OPEN"
-expect "pre-commit absent -> no grace, fail open at MAX_WAIT" 0 "$DOCS_OK"
-expect "duplicate run names -> latest wins" 0 "$PC_BAD, $PC_OK_RERUN, $DOCS_OK"
-raw_body=1
-expect "garbage response body -> fail open" 0 "this is not json"
-pr_json='{"labels": [{"name": "other"}]}'
-expect "ready label removed mid-gate -> blocked" 1 "$PC_OK, $DOCS_OK"
-
-exit "$fails"

.github/workflows/ci-precommit.yml

@@ -47,6 +47,3 @@ jobs:
     - uses: pre-commit/action@v3.0.1
       with:
         extra_args: --all-files --hook-stage manual
-    # After pre-commit so a self-test failure cannot mask lint failures.
-    - name: Full-suite gate self-test
-      run: bash .github/scripts/test_gate_full_suite.sh

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

@@ -52,7 +52,6 @@ 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'
@@ -61,7 +60,6 @@ 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: |
@@ -73,7 +71,6 @@ 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,
@@ -83,12 +80,11 @@ jobs:
                 pull_request_id: $pr_id,
                 pull_request_base_branch: "main",
                 env: {
-                    TEST_SCOPE: "full",
-                    FULL_SUITE: "true",
-                    PR_NUMBER: ($pr_id | tostring),
-                    PR_TITLE: $pr_title
-                  }
-                }')"
+                  TEST_SCOPE: "full",
+                  FULL_SUITE: "true",
+                  PR_NUMBER: ($pr_id | tostring)
+                }
+              }')"
 
   parse-command:
     if: >-
@@ -245,7 +241,6 @@ 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: |
@@ -262,7 +257,6 @@ 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,
@@ -272,9 +266,8 @@ jobs:
                 pull_request_base_branch: "main",
                 env: {
                   TEST_SCOPE: $test_scope,
-                    FULL_SUITE: $full_suite,
-                    TEST_TYPE: $test_type,
-                    PR_NUMBER: $pr_number,
-                    PR_TITLE: $pr_title
-                  }
-                }')"
+                  FULL_SUITE: $full_suite,
+                  TEST_TYPE: $test_type,
+                  PR_NUMBER: $pr_number
+                }
+              }')"

.github/workflows/ci-trigger-full-suite.yml

@@ -7,7 +7,6 @@ on:
 permissions:
   contents: read
   pull-requests: read
-  actions: read
 
 concurrency:
   group: full-suite-${{ github.event.pull_request.number }}
@@ -19,8 +18,6 @@ jobs:
       (github.event.action == 'labeled' && github.event.label.name == 'ready')
       || github.event.action == 'synchronize'
     runs-on: ubuntu-latest
-    # Gate below may wait for cheap checks (up to MAX_WAIT_SECS = 25 min).
-    timeout-minutes: 35
     steps:
       - name: Check ready label
         id: check
@@ -52,20 +49,6 @@ jobs:
               "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds/${build_num}/cancel"
           done
 
-      # Checks out the BASE branch (default for pull_request_target), so PR
-      # authors cannot tamper with the gate script.
-      - name: Checkout gate script
-        if: steps.check.outputs.has_ready == 'true'
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-
-      - name: Wait for pre-commit and docs build
-        if: steps.check.outputs.has_ready == 'true'
-        env:
-          GH_TOKEN: ${{ github.token }}
-          PR_SHA: ${{ github.event.pull_request.head.sha }}
-          PR_NUMBER: ${{ github.event.pull_request.number }}
-        run: bash .github/scripts/gate_full_suite.sh
-
       - name: Trigger Buildkite Full Suite
         if: steps.check.outputs.has_ready == 'true'
         env:
@@ -73,7 +56,6 @@ 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: |
@@ -85,7 +67,6 @@ 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,
@@ -97,7 +78,6 @@ jobs:
                 env: {
                   TEST_SCOPE: "full",
                   FULL_SUITE: "true",
-                  PR_NUMBER: ($pr_id | tostring),
-                  PR_TITLE: $pr_title
+                  PR_NUMBER: ($pr_id | tostring)
                 }
               }')"

apps/performance_dashboard/frontend/src/App.tsx

@@ -1,7 +1,6 @@
 import { useEffect, useMemo, useState } from "react";
 
-import { fetchSummary, fetchTrends, refreshData } from "./api";
-import type { CohortValue, RunSource, SummaryResponse, TrendGroup, TrendPoint } from "./api";
+import { fetchSummary, fetchTrends, refreshData, RunSource, SummaryResponse, TrendGroup, TrendPoint } from "./api";
 
 const METRIC_KEYS = ["latency", "throughput", "memory", "text_encoder_time_s", "dit_time_s", "vae_decode_time_s"];
 const RUN_SOURCES: Array<{ value: "" | RunSource; label: string }> = [
@@ -110,61 +109,6 @@ function metricLabel(metricKey: string) {
   return METRIC_DEFINITIONS[metricKey]?.label ?? metricKey;
 }
 
-type CohortFields = {
-  model_id: string;
-  gpu_type: string;
-  workload_id: CohortValue;
-  variant_id: CohortValue;
-  benchmark_version: CohortValue;
-  recipe_fingerprint: CohortValue;
-  hardware_profile_id: CohortValue;
-  software_profile_id: CohortValue;
-};
-
-function cohortValue(value: CohortValue) {
-  if (value === null || value === undefined || value === "") {
-    return "legacy";
-  }
-  return String(value);
-}
-
-function shortCohortValue(value: CohortValue) {
-  const text = cohortValue(value);
-  if (text === "legacy" || text.length <= 14) {
-    return text;
-  }
-  return text.slice(0, 12);
-}
-
-function cohortKey(cohort: CohortFields) {
-  return [
-    cohort.model_id,
-    cohort.gpu_type,
-    cohortValue(cohort.workload_id),
-    cohortValue(cohort.variant_id),
-    cohortValue(cohort.benchmark_version),
-    cohortValue(cohort.recipe_fingerprint),
-    cohortValue(cohort.hardware_profile_id),
-    cohortValue(cohort.software_profile_id)
-  ].join("|");
-}
-
-function cohortTitle(cohort: CohortFields) {
-  const workload = cohortValue(cohort.workload_id);
-  const variant = cohortValue(cohort.variant_id);
-  const version = cohortValue(cohort.benchmark_version);
-  const versionLabel = version === "legacy" ? version : `v${version}`;
-  return `${workload} / ${variant} / ${versionLabel}`;
-}
-
-function cohortDetail(cohort: CohortFields) {
-  return [
-    `recipe ${shortCohortValue(cohort.recipe_fingerprint)}`,
-    shortCohortValue(cohort.hardware_profile_id),
-    shortCohortValue(cohort.software_profile_id)
-  ].join(" | ");
-}
-
 function formatMetricValue(metricKey: string, value: number | null | undefined, tooltip = false) {
   const definition = METRIC_DEFINITIONS[metricKey];
   if (!definition) {
@@ -227,9 +171,7 @@ function TrendChart({ group, metricKey }: { group: TrendGroup; metricKey: string
         top: `${(activePoint.y / height) * 100}%`
       }
     : undefined;
-  const ariaLabel = `${metricLabel(metricKey)} trend for ${group.model_id} on ${group.gpu_type}, ${cohortTitle(
-    group
-  )}`;
+  const ariaLabel = `${metricLabel(metricKey)} trend for ${group.model_id} on ${group.gpu_type}`;
 
   return (
     <div className="chart-shell">
@@ -477,7 +419,7 @@ export default function App() {
       <section className="panel">
         <div className="panel-header">
           <h2>Latest Status</h2>
-          <span>{latestRows.length} comparison cohorts</span>
+          <span>{latestRows.length} model/GPU groups</span>
         </div>
         {latestRows.length === 0 ? (
           <div className="empty">No records match the selected filters.</div>
@@ -490,7 +432,6 @@ export default function App() {
                   <th>Recomputed</th>
                   <th>Model</th>
                   <th>GPU</th>
-                  <th>Cohort</th>
                   <th>Commit</th>
                   <th>Source</th>
                   <th>Baseline</th>
@@ -505,7 +446,7 @@ export default function App() {
               </thead>
               <tbody>
                 {latestRows.map((row) => (
-                  <tr key={cohortKey(row)}>
+                  <tr key={`${row.model_id}-${row.gpu_type}`}>
                     <td>
                       <span className={`badge ${row.status}`}>{row.status}</span>
                     </td>
@@ -516,12 +457,6 @@ export default function App() {
                     </td>
                     <td>{row.model_id}</td>
                     <td>{row.gpu_type}</td>
-                    <td>
-                      <div className="cohort-cell">
-                        <strong>{cohortTitle(row)}</strong>
-                        <span>{cohortDetail(row)}</span>
-                      </div>
-                    </td>
                     <td>{shortSha(row.commit_sha)}</td>
                     <td>
                       <span className={`source-badge source-${row.run_source}`}>{runSourceLabel(row.run_source)}</span>
@@ -560,13 +495,11 @@ export default function App() {
           ) : (
             trends.map((group) =>
               METRIC_KEYS.map((metricKey) => (
-                <article className="trend-card" key={`${cohortKey(group)}-${metricKey}`}>
+                <article className="trend-card" key={`${group.model_id}-${group.gpu_type}-${metricKey}`}>
                   <div>
                     <h3>{metricLabel(metricKey)}</h3>
                     <p>
                       {group.model_id} | {group.gpu_type}
-                      <span>{cohortTitle(group)}</span>
-                      <span>{cohortDetail(group)}</span>
                     </p>
                   </div>
                   <TrendChart group={group} metricKey={metricKey} />

apps/performance_dashboard/frontend/src/api.ts

@@ -13,17 +13,6 @@ export type MetricValue = {
   precision: number;
 };
 
-export type CohortValue = string | number | null;
-
-export type ComparisonCohort = {
-  workload_id: CohortValue;
-  variant_id: CohortValue;
-  benchmark_version: CohortValue;
-  recipe_fingerprint: CohortValue;
-  hardware_profile_id: CohortValue;
-  software_profile_id: CohortValue;
-};
-
 export type SummaryRow = {
   model_id: string;
   gpu_type: string;
@@ -45,7 +34,7 @@ export type SummaryRow = {
   build_id: string;
   job_id: string;
   metrics: Record<string, MetricValue>;
-} & ComparisonCohort;
+};
 
 export type RunSource = "pr" | "local" | "scheduled_main" | "unknown";
 
@@ -79,13 +68,13 @@ export type TrendPoint = {
   build_id: string;
   job_id: string;
   metrics: Record<string, number | null>;
-} & ComparisonCohort;
+};
 
 export type TrendGroup = {
   model_id: string;
   gpu_type: string;
   points: TrendPoint[];
-} & ComparisonCohort;
+};
 
 export type TrendsResponse = {
   groups: TrendGroup[];

apps/performance_dashboard/frontend/src/styles.css

@@ -149,11 +149,6 @@ h3 {
   font-size: 0.82rem;
 }
 
-.trend-card p {
-  display: grid;
-  gap: 2px;
-}
-
 .stat strong {
   display: block;
   margin-top: 8px;
@@ -191,7 +186,7 @@ h3 {
 
 table {
   width: 100%;
-  min-width: 1260px;
+  min-width: 1120px;
   border-collapse: collapse;
 }
 
@@ -214,25 +209,6 @@ td {
   font-size: 0.9rem;
 }
 
-.cohort-cell {
-  display: grid;
-  gap: 2px;
-}
-
-.cohort-cell strong,
-.trend-card p span {
-  color: #1b2836;
-  font-size: 0.78rem;
-  font-weight: 700;
-}
-
-.cohort-cell span,
-.trend-card p span + span {
-  color: #607080;
-  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", monospace;
-  font-size: 0.72rem;
-}
-
 .badge {
   display: inline-flex;
   align-items: center;

assets/eval/worldmodel_synthetic_flow_calibration.json

@@ -1,22 +0,0 @@
-{
-  "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
-}

docker/Dockerfile

@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
+ARG FA4_CUTE_REF=940cd9680f3315f2f06b43ab5bea2c2cf2d96806
 
 # Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
 # select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
@@ -161,7 +161,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         uv pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation; \
     fi
 
-# Overlay the CuTe-DSL-4.6-compatible upstream FA4 cute (FA4_CUTE_REF) over the
+# Overlay the cutlass-4.5-safe upstream FA4 cute (FA4_CUTE_REF) over the
 # wheel/source one so the image runs FA4, not the FA2 fallback. This pulls the FA4
 # runtime stack (cutlass-dsl, quack-kernels, apache-tvm-ffi, torch-c-dlpack-ext) --
 # the same deps the [dreamverse] extra already installs in CI; the installed torch

docs/contributing/ci_architecture.md

@@ -99,13 +99,6 @@ status.
 Full Suite is also path-filtered. It validates broader behavior before Mergify
 can merge a PR.
 
-A `ready`-labeled PR does not hit Buildkite immediately:
-`ci-trigger-full-suite.yml` first runs `.github/scripts/gate_full_suite.sh`,
-which waits for the cheap Tier-1 checks (pre-commit, docs build) on the PR
-head. A red cheap check blocks the suite (fail closed; the next push re-arms
-it), while a GitHub outage or a >25 min wait lets it run anyway (fail open).
-`/test full` bypasses the gate.
-
 | Buildkite label | `TEST_TYPE` | Main watched paths |
 |---|---|---|
 | SSIM Tests | `ssim` | `fastvideo/**/*.py`, `pyproject.toml`, `docker/Dockerfile` |

docs/contributing/performance_benchmarks.md

@@ -79,11 +79,10 @@ fastvideo/performance/
 ```
 
 The HF dataset (`FastVideo/performance-tracking` by default) holds one
-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.
+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.
 
 ## Planned Coverage
 
@@ -159,16 +158,13 @@ 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 comparison cohort)
+### Rolling baseline (per `(model_id, gpu_type)`)
 
 `compare_baseline.py` loads the last 5 successful, baseline-eligible records
-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.
+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.
 
 A metric exceeds its rolling threshold when both of these are true:
 
@@ -205,9 +201,9 @@ configs and remain loadable. New or migrated configs should use
 {
   "benchmark_id": "wan-t2v-1.3b-2gpu",
   "config_schema_version": 2,
-  "workload_id": "wan-t2v",
-  "variant_id": "1.3b-sp2",
-  "benchmark_version": 2
+  "workload_id": "wan-t2v-1.3b",
+  "variant_id": "canonical",
+  "benchmark_version": 1
 }
 ```
 
@@ -218,27 +214,21 @@ metadata that make the measured workload explicit:
 
 | Field | Purpose |
 |---|---|
-| `workload_id` | Stable benchmark family, such as `wan-t2v`. |
-| `variant_id` | Intentional recipe family, including model size and parallelism config, such as `1.3b-sp2`. |
+| `workload_id` | Stable benchmark family, such as `wan-t2v-1.3b`. |
+| `variant_id` | Intentional recipe family, such as `canonical`. |
 | `benchmark_version` | Version of the measurement protocol and comparison policy. |
 
 If a config declares `config_schema_version: 2`, loading fails clearly when any
 required v2 identity field is missing. If v2 identity or metadata fields are
 added without `config_schema_version: 2`, loading also fails so partial
 migrations do not silently run as v1 configs. Optional v2 metadata fields
-reserved for follow-up work, such as `metric_threshold_policy` and
-`quality_metadata`, must be JSON objects when present. (`recipe` is emitted
-by the harness and is not config-declarable.)
+reserved for follow-up work, such as `recipe`, `metric_threshold_policy`, and
+`quality_metadata`, must be JSON objects when present.
 
 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 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.
+comparison, metric-specific threshold policy behavior, promoted baselines, and
+dashboard regrouping are separate follow-up changes. Until those land, rolling
+baseline comparison remains keyed by `(model_id, gpu_type)`.
 
 ### Raw record (`results/perf_*.json`)
 
@@ -247,10 +237,10 @@ Written by `test_inference_performance.py`. One file per benchmark run.
 ```jsonc
 {
   "benchmark_id": "wan-t2v-1.3b-2gpu",
-  "result_schema_version": 2,
-  "workload_id": "wan-t2v",
-  "variant_id": "1.3b-sp2",
-  "benchmark_version": 2,
+  "config_schema_version": 2,
+  "workload_id": "wan-t2v-1.3b",
+  "variant_id": "canonical",
+  "benchmark_version": 1,
   "model_short_name": "Wan2.1-T2V-1.3B-Diffusers",
   "device": "NVIDIA L40S",
   "num_gpus": 2,
@@ -276,54 +266,11 @@ 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,
-  "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
-    },
-    "model": { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" },
-    "init_kwargs": { "num_gpus": 2, "sp_size": 2, "tp_size": 1 },
-    "generation_kwargs": { "height": 480, "width": 832, "num_frames": 45 },
-    "inputs": { "prompt_count": 1, "prompt_sha256": ["<measured-prompt-sha256>"] },
-    "attention": { "requested_backend": "FLASH_ATTN", "resolved_backend": "FLASH_ATTN" }
-  },
-  "recipe_fingerprint": "<sha256>",
-  "hardware_profile": {
-    "device_type": "cuda",
-    "gpu_count": 2,
-    "gpus": [{ "name": "NVIDIA L40S", "memory_gb": 48, "compute_capability": "8.9" }],
-    "interconnect": "none_or_partial"
-  },
-  "hardware_profile_id": "hw-<sha256-prefix>",
-  "software_profile": {
-    "python": "3.12",
-    "pytorch": "2.12",
-    "cuda": "13.0",
-    "packages": {
-      "fastvideo_kernel": "0.3.2",
-      "flashinfer": "0.2.11",
-      "nvidia_cutlass_dsl": "4.5.0",
-      "triton": "3.4.1"
-    }
-  },
-  "software_profile_id": "sw-<sha256-prefix>",
-  "environment_metadata": { "env": { "IMAGE_VERSION": "py3.12-cuda13.0.0" } },
-  "environment_fingerprint": "env-<sha256-prefix>"
+  "vae_decode_time_s": 3.208
 }
 ```
 
@@ -335,10 +282,6 @@ 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,
   "timestamp": "2026-05-08T22:00:00+00:00",
   "commit_sha": "<full sha>",
   "gpu_type": "NVIDIA L40S",
@@ -355,46 +298,17 @@ result, used as the rolling-baseline source of truth.
       "gated": true
     }
   },
-  "recipe_fingerprint": "<sha256>",
-  "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 `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`,
-`hardware_profile_id`, and `software_profile_id` cohort.
-`environment_metadata` and `environment_fingerprint` are audit data and are not
-part of the comparison key.
-The recipe prompt digests describe the prompts actually measured by the
-benchmark run; extra configured prompts are ignored unless the benchmark runner
-executes them.
-Software profile package cohorts keep exact versions for relevant
-attention/kernel packages, including FastVideo kernels, FlashAttention,
-FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
+Older records in the HF dataset may not have component timing fields. The
+comparator ignores missing or `null` metrics when computing a median, and the
+dashboard lists skipped plots for metric series that have no non-null values.
+Records missing both `run_source` and `baseline_eligible` are treated as legacy
+successful main/full-suite uploads and remain eligible for rolling baselines.
 
 ## Environment variable reference
 
@@ -404,7 +318,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`, `test_inference_performance.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
+| `PERF_RUN_SOURCE` | inferred | `compare_baseline.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
 | `PERF_UPLOAD_POLICY` | `never` | `compare_baseline.py` | Upload policy: `never`, `pass`, or `always`. |
 | `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Static-threshold pytest exit code, used so scheduled-main failures can be uploaded with `success=false`. |
 | `TEST_SCOPE` | unset | `compare_baseline.py` | CI context used to infer scheduled-main runs together with `BUILDKITE_BRANCH=main`. |
@@ -421,14 +335,11 @@ 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. 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.
+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.
 When the rolling-baseline phase runs, it emits:
 
 * **Markdown summary** — appended to `$GITHUB_STEP_SUMMARY` when that variable
@@ -436,7 +347,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 comparison cohort.
+  each metric grouped by `(model_id, gpu_type)`.
 * **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)
@@ -453,7 +364,7 @@ When the rolling-baseline phase runs, it emits:
      "benchmark_id": "<unique-id>",
      "config_schema_version": 2,
      "workload_id": "<stable-workload-id>",
-     "variant_id": "<variant, e.g. 1.3b-sp2>",
+     "variant_id": "canonical",
      "benchmark_version": 1,
      "model": { "model_path": "...", "model_short_name": "..." },
      "init_kwargs": { "num_gpus": 1, ... },
@@ -480,9 +391,7 @@ When the rolling-baseline phase runs, it emits:
 
    Legacy v1 configs without `config_schema_version` still load, but should not
    gain v2 identity or metadata fields until they are migrated to
-   `config_schema_version: 2`. For v2 configs, `workload_id`, `variant_id`,
-   and `benchmark_version` are part of the comparison key; benchmark runs
-   fail if any of these identity fields are missing.
+   `config_schema_version: 2`.
 
 2. The pytest test auto-discovers all configs — no test code needed. CI
    picks it up on the next `/test performance` run.
@@ -510,8 +419,8 @@ When the rolling-baseline phase runs, it emits:
 
 ## Troubleshooting
 
-**"No baseline for ... Initializing"** — first run for this comparison cohort.
-Run will pass and (if persisting) seed the first record.
+**"No baseline for ... Initializing"** — first run for this `(model_id,
+gpu_type)`. 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

docs/contributing/testing.md

@@ -180,30 +180,6 @@ 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

docs/design/inference_schema_parity_inventory.yaml

@@ -191,9 +191,6 @@ surfaces:
         sources:
           - fastvideo.configs.pipelines.gen3c.Gen3CConfig
           - fastvideo.configs.pipelines.gen3c.Gen3CInferenceConfig
-      color_correction_strength:
-        sources:
-          - fastvideo.configs.pipelines.dreamx_world.DreamXWorld5BARPipelineConfig
       default_camera_rotation:
         sources:
           - fastvideo.configs.pipelines.gen3c.Gen3CConfig

docs/distillation/anyflow.md

@@ -1,116 +0,0 @@
-# 🌊 AnyFlow Any-Step Video Distillation
-
-**AnyFlow** ([paper](https://arxiv.org/abs/2605.13724), [project page](https://nvlabs.github.io/AnyFlow/), [official code](https://github.com/NVlabs/AnyFlow), [model weights](https://huggingface.co/collections/nvidia/anyflow)) is an any-step video diffusion framework built on flow maps. A single distilled checkpoint can be evaluated at NFE ∈ {1, 2, 4, 8, 16, 32} without retraining, and quality scales **monotonically** with steps — unlike consistency-based distillation, which often degrades as NFE grows.
-
-The student network ``u_θ(x_t, t, r)`` predicts the *average velocity* from time ``t`` back to time ``r``, so one Euler step is
-
-```
-x_r = x_t - ((t - r) / N) · u_θ(x_t, t, r)
-```
-
-for any ``t > r``.
-
-## 📊 Model Overview
-
-NVIDIA publishes four checkpoints under [`nvidia/anyflow`](https://huggingface.co/collections/nvidia/anyflow):
-
-- `nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers` — bidirectional T2V, Wan2.1 1.3B base
-- `nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers` — bidirectional T2V, Wan2.1 14B base
-- `nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers` — frame-autoregressive variant, 1.3B
-- `nvidia/AnyFlow-FAR-Wan2.1-14B-Diffusers` — frame-autoregressive variant, 14B
-
-FastVideo currently supports the bidirectional T2V variants for training; the FAR variants can be loaded for inference through the diffusers integration.
-
-## ⚙️ Inference
-
-For inference, load the published checkpoint directly through diffusers; FastVideo's training-side ``WanModel`` config maps the HF AnyFlow ``delta_embedder`` weights onto its internal layout via ``param_names_mapping`` so the same checkpoint can be used as the ``init_from`` for the on-policy YAML below.
-
-## 🧠 Algorithm
-
-Training runs in two stages. Both use the dual-timestep Wan backbone — enabled by ``pipeline.dit_config.r_embedder: true`` in the YAML, which allocates a sibling ``condition_embedder.delta_embedder`` and fuses its embedding with the standard timestep embedding via either an additive or a gated mixer.
-
-### Stage 1 — Pretrain (flow-map central-difference)
-
-Method: ``AnyFlowPretrainMethod`` (``fastvideo/train/methods/distribution_matching/anyflow_pretrain.py``)
-
-For each batch, sample ``(t, r) ∈ [0, 1]`` as ``(max, min)`` of two uniform draws, then:
-
-- a ``diffusion_ratio`` fraction (default 0.5) gets ``r = t`` — recovers plain flow matching;
-- a ``consistency_ratio`` fraction (default 0.25) gets ``r = 0`` — forces consistency to clean data;
-- the remainder is free.
-
-The student forward at ``(t, r)`` is trained against the central-difference target
-
-```
-target = (eps - x_0) - (t - r) · dF/dt
-```
-
-where ``dF/dt`` is estimated from the student's own forward at ``(t ± δ, r)`` with the sample also moved along the flow trajectory by ``v_pred · (δ / N)``. Per-timestep weighting uses ``beta08`` (``w(t) = t · sqrt(1 - t)``, renormalized). A stop-gradient scale-balance keeps the non-diffusion branches' loss magnitude aligned with the diffusion branch.
-
-### Stage 2 — On-policy DMD
-
-Method: ``AnyFlowMethod`` (``fastvideo/train/methods/distribution_matching/anyflow.py``)
-
-Inherits ``DMD2Method``. The student is rolled out for ``student_sample_steps`` Euler-flow steps from pure noise; one randomly-chosen step is gradient-enabled (broadcast from rank 0 so every worker agrees), the rest run under ``torch.no_grad``. With ``use_mean_velocity: true`` (default) the rollout uses ``r = t_next`` at each step, matching AnyFlow's ``WanAnyFlowPipeline.training_rollout``.
-
-The inherited ``_dmd_loss`` (VSD with fake-score critic) consumes the rollout output and the teacher's CFG prediction. The optional pinned ``t_list_override`` lets configs reproduce the paper's hand-tuned 4-step schedule ``[999, 937, 833, 624, 0]``.
-
-## 🚀 Training Scripts
-
-### Stage 1 — pretrain
-
-```bash
-bash examples/train/run.sh \
-    examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml
-```
-
-**Key configuration** (in ``examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml``):
-
-- Global batch size: 32 (8 GPUs × 4 per-GPU)
-- Learning rate: 5e-5
-- Flow shift: 5.0
-- ``diffusion_ratio`` / ``consistency_ratio``: 0.5 / 0.25
-- ``epsilon`` (finite-difference step): 5 (absolute train-timestep units)
-- ``weight_type``: ``beta08``
-- ``fuse_guidance_scale``: 3.0
-- Training steps: 6000
-
-### Stage 2 — on-policy
-
-```bash
-bash examples/train/run.sh \
-    examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml \
-    --models.student.init_from outputs/wan2.1_anyflow_pretrain/checkpoint-final
-```
-
-(Or point ``models.student.init_from`` directly at ``nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers`` to bootstrap from the paper weights and skip Stage 1.)
-
-**Key configuration**:
-
-- Global batch size: 8 (8 GPUs × 1 per-GPU)
-- Learning rate: 2e-6
-- Flow shift: 5.0
-- ``student_sample_steps``: 4
-- ``t_list_override``: ``[999, 937, 833, 624, 0]``
-- ``use_mean_velocity``: ``true`` (i.e. ``r = t_next`` during rollout)
-- ``real_score_guidance_scale``: 3.0
-- ``generator_update_interval``: 5 (DMD2 alternation)
-- Training steps: 4000
-
-## 🔌 Loading published AnyFlow checkpoints
-
-The HF AnyFlow checkpoints expose ``condition_embedder.delta_embedder.*`` weights that FastVideo internally maps onto its ``condition_embedder.delta_embedder.mlp.*`` layout. This rename happens automatically through the regex in ``WanVideoArchConfig.param_names_mapping`` — no separate adapter is needed. The same regex is a no-op on plain Wan checkpoints (which don't contain any ``delta_embedder`` keys).
-
-Set the YAML's ``pipeline.dit_config.r_embedder: true`` to allocate the ``delta_embedder`` module on the FastVideo side; when initializing from a plain Wan checkpoint the delta weights are deep-copied from ``time_embedder`` (matching AnyFlow's ``setup_flowmap_model()`` behavior).
-
-## 🧭 Note on ``fuse_guidance_scale``
-
-Stage 1 optionally fuses classifier-free guidance into the training target so the resulting checkpoint can be sampled at ``guidance_scale=1.0`` (no extra forward pass at inference time). The transformation is
-
-```
-noise_pred ← (noise_pred - (1 - g) · noise_pred_uncond) / g
-```
-
-with ``g = fuse_guidance_scale``. The negative prompt embedding comes from ``WanModel``'s ``ensure_negative_conditioning()`` — i.e. the dataset's configured ``sampling_param.negative_prompt``. Setting ``fuse_guidance_scale: 1.0`` skips the extra unconditional forward entirely.
-
-The on-policy stage's ``real_score_guidance_scale`` (inherited from DMD2) follows the same parameterization conventions documented in [``dmd.md``](dmd.md#-note-on-real_score_guidance_scale).

docs/inference/support_matrix.md

@@ -58,8 +58,6 @@ pipeline initialization and sampling.
 | FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
-| DreamX-World 5B Cam | `FastVideo/DreamX-World-5B-Cam-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
-| DreamX-World 5B AR | `FastVideo/DreamX-World-5B-Diffusers` | 704px1280p | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |

docs/training/trackers.md

@@ -1,91 +0,0 @@
-# Training Trackers
-
-FastVideo can send training metrics and validation media to Weights & Biases
-or SwanLab. Tracking runs only on global rank 0, and local tracker files are
-stored under `<output_dir>/tracker`.
-
-## Supported Trackers
-
-| Value | Backend | Installation |
-|-------|---------|--------------|
-| `wandb` | Weights & Biases | Included with FastVideo |
-| `swanlab` | SwanLab | Install the optional `swanlab` dependency |
-| `none` | Disable external tracking | No additional package |
-
-You can enable more than one backend, for example `trackers: [wandb, swanlab]`.
-Metrics and validation media are converted to the artifact type required by
-each backend.
-
-## Install SwanLab
-
-For a published FastVideo installation, install the SwanLab extra:
-
-```bash
-uv pip install "fastvideo[swanlab]"
-```
-
-For an editable source checkout, include the same extra during installation:
-
-```bash
-uv pip install -e ".[swanlab]"
-```
-
-If FastVideo is already installed, you can install the compatible SDK directly:
-
-```bash
-uv pip install "swanlab>=0.6.7"
-```
-
-Authenticate once before starting a training run:
-
-```bash
-swanlab login
-```
-
-See the [SwanLab login documentation](https://docs.swanlab.cn/en/api/cli-swanlab-login.html)
-for non-interactive and self-hosted setups.
-
-## Configure Tracking
-
-Select SwanLab in the YAML config used by the modular training framework:
-
-```yaml
-training:
-  checkpoint:
-    output_dir: outputs/my_run
-  tracker:
-    trackers: [swanlab]
-    project_name: my_project
-    run_name: my_run
-```
-
-To log to both supported services:
-
-```yaml
-training:
-  tracker:
-    trackers: [wandb, swanlab]
-    project_name: my_project
-    run_name: my_run
-```
-
-An empty or omitted `trackers` list selects W&B when `project_name` is set.
-Use an explicit `none` entry to disable external tracking:
-
-```yaml
-training:
-  tracker:
-    trackers: [none]
-```
-
-## Validation Videos
-
-SwanLab currently accepts GIF video artifacts. FastVideo converts validation
-MP4 files and in-memory video arrays to GIF automatically before logging them.
-For video files, FastVideo uses the sampling frame rate supplied by the caller,
-or the source file's frame rate when no value is supplied. In-memory arrays use
-the frame rate supplied by the caller. Both forms fall back to 16 FPS when no
-frame rate is available.
-
-For details about configuring validation callbacks, see
-[Training Infrastructure](train_infra.md#callbacks-pluggable-hooks).

docs/training/train_infra.md

@@ -161,21 +161,6 @@ 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.
-
 ### `callbacks` — Pluggable hooks
 
 Callbacks run at specific points in the training loop (before/after optimizer
@@ -338,40 +323,6 @@ 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

examples/inference/basic/basic_dreamx_world.py

@@ -1,64 +0,0 @@
-import os
-
-from fastvideo import VideoGenerator
-
-
-OUTPUT_PATH = os.getenv("DREAMX_WORLD_OUTPUT_PATH", "video_samples_dreamx_world")
-
-
-def _env_int(name: str, default: int) -> int:
-    return int(os.getenv(name, str(default)))
-
-
-def _env_float(name: str, default: float) -> float:
-    return float(os.getenv(name, str(default)))
-
-
-def main():
-    model_name = os.getenv("DREAMX_WORLD_MODEL_DIR", "FastVideo/DreamX-World-5B-Cam-Diffusers")
-    generator = VideoGenerator.from_pretrained(
-        model_name,
-        num_gpus=1,
-        use_fsdp_inference=False,
-        dit_cpu_offload=False,
-        vae_cpu_offload=True,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
-        override_pipeline_cls_name="DreamXWorldPipeline",
-    )
-
-    prompt = os.getenv(
-        "DREAMX_WORLD_PROMPT",
-        "A cinematic first-person drive through a futuristic coastal city at "
-        "sunrise, reflective glass towers, clean streets, soft volumetric light.",
-    )
-    image_path = os.getenv(
-        "DREAMX_WORLD_IMAGE_PATH",
-        "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG",
-    )
-
-    kwargs = {
-        "output_path": OUTPUT_PATH,
-        "save_video": os.getenv("DREAMX_WORLD_SAVE_VIDEO", "1") != "0",
-        "height": _env_int("DREAMX_WORLD_HEIGHT", 480),
-        "width": _env_int("DREAMX_WORLD_WIDTH", 832),
-        "num_frames": _env_int("DREAMX_WORLD_NUM_FRAMES", 161),
-        "num_inference_steps": _env_int("DREAMX_WORLD_STEPS", 30),
-        "guidance_scale": _env_float("DREAMX_WORLD_GUIDANCE", 5.0),
-        "action_list": os.getenv("DREAMX_WORLD_ACTIONS", "w,d,w").split(","),
-        "action_speed_list": [
-            float(value)
-            for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",")
-        ],
-    }
-    if image_path:
-        kwargs["image_path"] = image_path
-
-    try:
-        generator.generate_video(prompt, **kwargs)
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/ltx2_3/README.md

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

examples/inference/ltx2_3/i2v_5s2_1024x1536.yaml

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

examples/inference/ltx2_3/i2v_5s2_1280x832.yaml

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

examples/inference/ltx2_3/i2v_5s2_512x768.yaml

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

examples/inference/ltx2_3/i2v_5s2_768x1280.yaml

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

examples/inference/ltx2_3/i2v_8s3_1024x1536.yaml

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

examples/inference/ltx2_3/i2v_8s3_1280x832.yaml

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

examples/inference/ltx2_3/i2v_8s3_512x768.yaml

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

examples/inference/ltx2_3/i2v_8s3_768x1280.yaml

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

examples/inference/ltx2_3/t2v_5s2_1024x1536.yaml

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

examples/inference/ltx2_3/t2v_5s2_1280x832.yaml

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

examples/inference/ltx2_3/t2v_5s2_512x768.yaml

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

examples/inference/ltx2_3/t2v_5s2_768x1280.yaml

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

examples/inference/ltx2_3/t2v_8s3_1024x1536.yaml

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

examples/inference/ltx2_3/t2v_8s3_1280x832.yaml

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

examples/inference/ltx2_3/t2v_8s3_512x768.yaml

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

examples/inference/ltx2_3/t2v_8s3_768x1280.yaml

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

examples/train/configs/consistency_model/wan/causal_cd_t2v.yaml

@@ -1,78 +0,0 @@
-# Causal Consistency Distillation: Wan 2.1 T2V 1.3B Causal
-#
-# ODE-data-free distillation. A frozen teacher takes a single CFG Euler step;
-# the student matches an EMA copy of itself at the next timestep, all under
-# clean-history teacher forcing.
-#
-# All three roles initialize from the SAME checkpoint (the teacher-forcing
-# AR-diffusion model). Point init_from at that checkpoint for a real run.
-
-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.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-  ema:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-
-method:
-  _target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod
-  discrete_cd_N: 48
-  guidance_scale: 3.0
-  ema_decay: 0.99
-  ema_start_step: 200
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/Wan-Syn_77x448x832_600k
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  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: 3000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_causal_cd
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_cd_shift5
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-
-pipeline:
-  flow_shift: 5

examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml

@@ -1,105 +0,0 @@
-# AnyFlow on-policy DMD — Wan 2.1 T2V 1.3B.
-#
-# Stage 2 of the AnyFlow two-stage recipe. Continues from the pretrain
-# checkpoint; refines the student via DMD2 with a multi-step Euler-flow
-# rollout from pure noise. Teacher provides the real score, critic
-# learns the fake score; both inherited from DMD2Method.
-#
-# Replace <PATH_TO_PRETRAIN_CKPT> with the output of the pretrain stage,
-# or with the NVIDIA-released checkpoint
-# nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers to bootstrap directly from
-# the paper weights (the delta_embedder rename is handled by the
-# param_names_mapping in WanVideoArchConfig).
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanModel
-    init_from: <PATH_TO_PRETRAIN_CKPT>
-    trainable: true
-  teacher:
-    _target_: fastvideo.train.models.wan.WanModel
-    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
-    trainable: false
-    disable_custom_init_weights: true
-  critic:
-    _target_: fastvideo.train.models.wan.WanModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    disable_custom_init_weights: true
-
-method:
-  _target_: fastvideo.train.methods.distribution_matching.anyflow.AnyFlowMethod
-  rollout_mode: simulate
-  generator_update_interval: 5
-  real_score_guidance_scale: 3.0
-  dmd_denoising_steps: [999, 937, 833, 624]
-  warp_denoising_step: false
-
-  # AnyFlow rollout knobs.
-  student_sample_steps: 4
-  use_mean_velocity: true
-  t_list_override: [999.0, 937.0, 833.0, 624.0, 0.0]
-  dmd_score_r_value: 0.0  # DMD scoring conditioning is at r=0 (consistency target).
-
-  # Critic optimizer (DMD2 inherited).
-  fake_score_learning_rate: 8.0e-6
-  fake_score_betas: [0.0, 0.999]
-  fake_score_lr_scheduler: constant
-
-  attn_kind: vsa
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 21
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.0, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 4000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_anyflow_onpolicy
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 3
-    resume_from_checkpoint: latest
-
-  tracker:
-    project_name: anyflow-wan
-    run_name: wan2.1_t2v_anyflow_onpolicy
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-
-pipeline:
-  flow_shift: 5.0
-  dit_config:
-    r_embedder: true
-    r_embedder_fusion: gated
-    r_embedder_gate_value: 0.25
-    r_embedder_deltatime_type: r

examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml

@@ -1,83 +0,0 @@
-# AnyFlow pretrain (flow-map central-difference) — Wan 2.1 T2V 1.3B.
-#
-# Stage 1 of the AnyFlow two-stage recipe. Trains the dual-timestep
-# u_θ(x_t, t, r) on the central-difference target so the same checkpoint
-# can be sampled at arbitrary NFE in the on-policy stage.
-#
-# Initialize from base Wan 2.1 T2V 1.3B. No teacher or critic at this
-# stage; AnyFlowPretrainMethod owns a single student + one optimizer.
-
-models:
-  student:
-    _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.anyflow_pretrain.AnyFlowPretrainMethod
-  diffusion_ratio: 0.5
-  consistency_ratio: 0.25
-  epsilon: 5            # finite-difference step in absolute train-timestep units
-  weight_type: beta08   # per-timestep loss weight = t * sqrt(1 - t), renormalized
-  fuse_guidance_scale: 3.0
-  # shift is taken from pipeline.flow_shift below.
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 4
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 21
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    learning_rate: 5.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 6000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_anyflow_pretrain
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 3
-    resume_from_checkpoint: latest
-
-  tracker:
-    project_name: anyflow-wan
-    run_name: wan2.1_t2v_anyflow_pretrain
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-
-pipeline:
-  flow_shift: 5.0
-  dit_config:
-    # Enable AnyFlow dual-timestep conditioning. The student loads from
-    # base Wan 2.1 — its checkpoint has no delta_embedder weights, so they
-    # get initialized identically to time_embedder via deep-copy in
-    # WanTimeTextImageEmbedding.__init__.
-    r_embedder: true
-    r_embedder_fusion: gated
-    r_embedder_gate_value: 0.25
-    r_embedder_deltatime_type: r

examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v.yaml

@@ -100,6 +100,9 @@ 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

examples/train/configs/fine_tuning/wan/dfsft_causal_t2v_framewise.yaml

@@ -1,73 +0,0 @@
-# DFSFT (Diffusion-Forcing SFT), frame-wise: Wan 2.1 T2V 1.3B Causal
-#
-# - Student: trainable causal Wan model with a block size of 1 frame
-# - Training: each frame gets its own independent noise level (frame-wise
-#   diffusion forcing), versus the chunk-wise variant that shares one noise
-#   level across num_frames_per_block frames.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    num_frames_per_block: 1
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
-  chunk_size: 1
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/Wan-Syn_77x448x832_600k
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  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: 4000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_causal_dfsft_framewise
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_dfsft_framewise_shift5_gauss_weight
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-  validation:
-    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
-    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
-    every_steps: 50
-    sampling_steps: [40]
-    guidance_scale: 6.0
-    num_frames: 69
-
-pipeline:
-  flow_shift: 5

examples/train/configs/fine_tuning/wan/tfsft_causal_t2v.yaml

@@ -1,72 +0,0 @@
-# TFSFT (Teacher-Forcing SFT): Wan 2.1 T2V 1.3B Causal
-#
-# - Student: trainable causal Wan model
-# - Training: inhomogeneous timesteps per chunk, but the causal transformer
-#   denoises the current block while attending to *clean* history (clean_x),
-#   not its own noisy rollout.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod
-  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/Wan-Syn_77x448x832_600k
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  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: 4000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_causal_tfsft
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_tfsft_shift5_gauss_weight
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-  validation:
-    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
-    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
-    every_steps: 50
-    sampling_steps: [40]
-    guidance_scale: 6.0
-    num_frames: 69
-
-pipeline:
-  flow_shift: 5

examples/train/scenario/worldmodel/README.md

@@ -1,119 +1,32 @@
 # World-Model: Matrix-Game 2.0 I2V
 
-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
+Three training scenarios for the Matrix-Game 2.0 I2V world model on the
+new YAML-driven trainer (`fastvideo/train/entrypoint/train.py`).
 
 | Config | Method | Student | Notes |
 |---|---|---|---|
-| `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`.
+| `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. |
 
 ## Usage
 
-### Solaris
-
 ```bash
 bash examples/train/run.sh \
-    examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml
+    examples/train/scenario/worldmodel/finetune_i2v.yaml
 
 bash examples/train/run.sh \
-    examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml
+    examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml
 
 bash examples/train/run.sh \
-    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
+    examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml
 ```
 
 Override any field on the command line:
 
 ```bash
 bash examples/train/run.sh \
-    examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml \
+    examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml \
     --training.distributed.num_gpus 8 \
     --training.optimizer.learning_rate 1e-5
 ```

examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml

@@ -97,4 +97,4 @@ callbacks:
     guidance_scale: 6.0
 
 pipeline:
-  flow_shift: 5
+  flow_shift: 5

examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml

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

examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml

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

examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml

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

examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml

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

fastvideo/attention/backends/flash_attn.py

@@ -1,15 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import importlib.util
 import os
+
 import torch
 import torch.nn.functional as F
 from dataclasses import dataclass
 
-from fastvideo.attention.utils.flash_attn_default import (
-    fa_version,
-    flash_attn_func_compilable,
-)
-
+from fastvideo import envs
 from fastvideo.attention.backends.abstract import (
     AttentionBackend,
     AttentionImpl,
@@ -19,6 +17,119 @@ from fastvideo.attention.backends.abstract import (
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
+
+# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1, mirroring the
+# kernel package's FASTVIDEO_VSA_CUTEDSL: its CuTeDSL kernels JIT-compile per
+# shape family and can fail at runtime on some arch/shape combinations, so it
+# is never auto-selected just because it is installed. Below sm90 a capability
+# gate in flash_attn_cute routes to FA2 the calls FA4 cannot serve there:
+# grad-enabled (its backward asserts sm90+) and GQA (pack_gqa fails CuTeDSL
+# JIT, observed on sm_89).
+if envs.FASTVIDEO_FA4:
+    try:
+        from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
+    except ImportError as e:
+        raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
+                           "fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
+                           "or unset FASTVIDEO_FA4.") from e
+    fa_version = "4"
+else:
+    try:
+        from flash_attn_interface import flash_attn_func as flash_attn_3_func
+
+        # flash_attn 3 no longer have a different API, see following commit:
+        # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
+        flash_attn_func = flash_attn_3_func
+        fa_version = "3"
+    except ImportError:
+        from flash_attn import flash_attn_func as flash_attn_2_func
+        flash_attn_func = flash_attn_2_func
+        fa_version = "2"
+    try:
+        if importlib.util.find_spec("flash_attn.cute") is not None:
+            logger.info("flash_attn.cute (FA4) is installed but not enabled; "
+                        "set FASTVIDEO_FA4=1 to use it for inference.")
+    except ImportError:
+        pass
+
+# torch.compile traceability: the FA4/cute path (fa_version=="4") is
+# already a registered torch.library custom op, so dynamo treats it as a
+# graph node. The external FA2/FA3 `flash_attn_func` is NOT — dynamo
+# breaks the graph at the call site (observed: wanvideo.py self-attn,
+# once per layer every step), which fragments the compiled region and
+# blocks CUDA-graph capture. Wrap the FA2/FA3 default call in a custom
+# op (mirrors the FP4 `flash_attn_cute` template) so it becomes an
+# opaque-but-traceable node. The kernel still runs eager inside the op
+# (correct — flash-attn must run eager); only dynamo's treatment of the
+# boundary changes, so numerics are unchanged (SSIM-gate to confirm).
+if fa_version in ("2", "3"):
+    _fa_default = flash_attn_func
+
+    # Scope: this op covers exactly the q/k/v + softmax_scale + causal
+    # call shape used by FlashAttentionImpl.forward's default branch
+    # (see `flash_attn_func_compilable(...)` call site below). The
+    # masked/no-pad and varlen / cross-attn paths use different
+    # entry points (`flash_attn_no_pad`, `flash_attn_varlen_*`) which
+    # are intentionally out of scope for this PR — wrapping them is a
+    # natural follow-up. The wrapper's signature is the contract: any
+    # extra kwarg (dropout_p, window_size, alibi_slopes, deterministic,
+    # return_attn_probs, ...) raises TypeError at the call site, so
+    # silent loss of kwargs is not a failure mode.
+    @torch.library.custom_op(
+        "fastvideo::_flash_attn_default_forward",
+        mutates_args=(),
+        device_types="cuda",
+    )
+    def _flash_attn_default_forward(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: float | None,
+        causal: bool,
+    ) -> torch.Tensor:
+        return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
+
+    @torch.library.register_fake("fastvideo::_flash_attn_default_forward")
+    def _flash_attn_default_forward_fake(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: float | None,
+        causal: bool,
+    ) -> torch.Tensor:
+        del softmax_scale, causal
+        # FA2/FA3 default path: [batch, seqlen_q, nheads, head_dim_v],
+        # same dtype/device as q (head dim taken from v).
+        return q.new_empty(q.shape[0], q.shape[1], q.shape[2], v.shape[-1])
+
+    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
+        # Autograd carve-out. The custom op above registers a forward + fake
+        # kernel but NO backward (register_autograd), so it is opaque to
+        # autograd. Inference runs under no_grad / inference_mode and routes
+        # through the traceable custom op — that is the torch.compile win, and
+        # the only path this PR claims. Training backprops through attention,
+        # so route grad-enabled calls to the original FA2/FA3 `flash_attn_func`
+        # (itself an autograd.Function, so backward is correct) at the cost of a
+        # dynamo graph break on the training path — i.e. pre-PR behavior, no
+        # regression. Full autograd parity for the custom op (mirroring the FP4
+        # cute template) is a tracked follow-up.
+        if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad):
+            return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
+        return torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
+elif fa_version == "4":
+    # FA4 path: `flash_attn_func` (from `flash_attn_cute`) goes through a
+    # registered torch.library custom op (with an FA4 backward on sm90+;
+    # grad-enabled and GQA calls below sm90 route to FA2), so a passthrough
+    # is enough — no extra registration needed.
+    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
+        return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
+else:
+    # Defensive: the probe above only ever sets fa_version to "2", "3",
+    # or "4"; an unexpected value means an import/probe regression and
+    # we want a loud error at import, not a silent NameError later.
+    raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
+                       f"'2', '3', or '4' from the import probe above.")
+
 logger.info("Using FlashAttention-%s backend", fa_version)
 
 # FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.
@@ -198,17 +309,9 @@ class FlashAttentionImpl(AttentionImpl):
         attn_metadata: FlashAttnMetadata,
     ):
         if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None):
-            # Route through the *_compilable wrappers so dynamo sees one
-            # traceable node for each masked entry point (the unpad/pad
-            # bookkeeping runs eager inside the custom op). On FA2 these
-            # wrappers go through ops with full register_autograd, so
-            # training also backprops through the op (no graph break on
-            # the training path); on FA3/FA4 they carve out to the
-            # autograd.Function for grad-enabled calls — see
-            # fastvideo/attention/utils/flash_attn_no_pad.py.
             from fastvideo.attention.utils.flash_attn_no_pad import (
-                flash_attn_no_pad_compilable as flash_attn_no_pad,
-                flash_attn_varlen_qk_no_pad_compilable as flash_attn_varlen_qk_no_pad,
+                flash_attn_no_pad,
+                flash_attn_varlen_qk_no_pad,
             )
 
             attn_mask = attn_metadata.attn_mask
@@ -235,11 +338,7 @@ class FlashAttentionImpl(AttentionImpl):
                 )
 
             qkv = torch.stack([query, key, value], dim=2)
-            key_padding_mask = _key_padding_mask_from_attn_mask(attn_mask, attn_mask.shape[-1]).to(device=query.device)
-            if key_padding_mask.shape[-1] > qkv.shape[1]:
-                raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
-                                 f"expected at most {qkv.shape[1]}, got {key_padding_mask.shape[-1]}")
-            attn_mask_padded = F.pad(key_padding_mask, (qkv.shape[1] - key_padding_mask.shape[-1], 0), value=True)
+            attn_mask_padded = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True)
             output = flash_attn_no_pad(qkv, attn_mask_padded, causal=False, dropout_p=0, softmax_scale=None)
         elif self.nvfp4_fa4:
             output = self._forward_nvfp4(query, key, value)

fastvideo/attention/backends/sdpa.py

@@ -2,7 +2,6 @@
 
 import torch
 from dataclasses import dataclass
-from torch.nn import functional as F
 from fastvideo.attention.backends.abstract import (  # FlashAttentionMetadata,
     AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder)
 from fastvideo.logger import init_logger
@@ -20,7 +19,7 @@ class SDPABackend(AttentionBackend):
 
     @staticmethod
     def get_name() -> str:
-        return "TORCH_SDPA"
+        return "SDPA"
 
     @staticmethod
     def get_impl_cls() -> type["SDPAImpl"]:
@@ -50,51 +49,9 @@ class SDPAMetadataBuilder(AttentionMetadataBuilder):
             current_timestep: int,
             attn_mask: torch.Tensor,
     ) -> SDPAMetadata:
-        # Store the mask exactly as passed. The metadata is cross-backend:
-        # call sites (HYWorld, HunyuanVideo15) build SDPAMetadata while the
-        # layer's selector may pick FLASH_ATTN, and the shared convention for
-        # padding masks is the tokenizer-style 2D [batch, key_len]. Any
-        # reshaping for torch.sdpa happens inside the SDPA impl
-        # (_normalize_attn_mask_for_sdpa).
         return SDPAMetadata(current_timestep=current_timestep, attn_mask=attn_mask)
 
 
-def _normalize_attn_mask_for_sdpa(
-    attn_mask: torch.Tensor | None,
-    query: torch.Tensor,
-    key: torch.Tensor,
-) -> torch.Tensor | None:
-    if attn_mask is None:
-        return None
-
-    attn_mask = attn_mask.to(device=query.device)
-    # F.scaled_dot_product_attention only accepts bool or float masks;
-    # tokenizers commonly produce int64 0/1 padding masks.
-    if attn_mask.dtype != torch.bool and not attn_mask.dtype.is_floating_point:
-        attn_mask = attn_mask != 0
-
-    key_len = key.shape[-2]
-    if attn_mask.shape[-1] > key_len:
-        raise ValueError("Invalid attention mask length for SDPA: "
-                         f"expected at most {key_len}, got {attn_mask.shape[-1]}")
-    if attn_mask.shape[-1] < key_len:
-        # Front-pad as "attend": double-stream layouts (HYWorld) prepend
-        # non-text tokens the tokenizer mask does not cover.
-        valid_value = True if attn_mask.dtype == torch.bool else 0.0
-        attn_mask = F.pad(attn_mask, (key_len - attn_mask.shape[-1], 0), value=valid_value)
-
-    if attn_mask.dim() == 2:
-        # In-tree producers pass 2D [batch, key_len] padding masks; lift to a
-        # broadcastable [batch, 1, 1, key_len] here so torch.sdpa does not
-        # reinterpret 2D as its documented [query_len, key_len] broadcast.
-        return attn_mask[:, None, None, :]
-    if attn_mask.dim() == 3:
-        return attn_mask[:, None, :, :]
-    if attn_mask.dim() == 4:
-        return attn_mask
-    raise ValueError(f"Unsupported attention mask shape for SDPA: {attn_mask.shape}")
-
-
 class SDPAImpl(AttentionImpl):
 
     def __init__(
@@ -125,7 +82,6 @@ class SDPAImpl(AttentionImpl):
 
         attn_mask = attn_metadata.attn_mask if (attn_metadata is not None
                                                 and hasattr(attn_metadata, "attn_mask")) else None
-        attn_mask = _normalize_attn_mask_for_sdpa(attn_mask, query, key)
         attn_kwargs = {
             "attn_mask": attn_mask,
             "dropout_p": self.dropout,

fastvideo/attention/utils/flash_attn_default.py

@@ -1,251 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""torch.compile-traceable wrapper for the FA2/FA3/FA4 default attention path.
-
-The FA4/cute path (`fa_version == "4"`) is already a registered
-`torch.library.custom_op` in `fastvideo.attention.utils.flash_attn_cute`, so
-dynamo treats it as a graph node. The external FA2/FA3 ``flash_attn_func`` is
-NOT — dynamo breaks the graph at the call site (observed: wanvideo.py
-self-attn, once per layer every step), which fragments the compiled region
-and blocks CUDA-graph capture. Wrap the FA2/FA3 default call in a custom op
-(mirrors the FP4 `flash_attn_cute` template) so it becomes an
-opaque-but-traceable node. The kernel still runs eager inside the op
-(correct — flash-attn must run eager); only dynamo's treatment of the
-boundary changes, so numerics are unchanged (SSIM-gated).
-
-Autograd: FA2 has full ``register_autograd`` parity — the custom op's
-backward calls flash_attn's ``_flash_attn_backward`` directly, so training
-backprops *through* the op (no graph break on the training path either).
-FA3 currently keeps the no-backward + carve-out pattern from PR #1373
-because FA3's private backward signature wants validation on a real Hopper
-box (gated on Kuan-Hao's Modal FA3 setup PR). Once that lands the FA3 path
-can mirror FA2.
-
-Lives in `attention/utils/` (sibling of `flash_attn_cute.py` and
-`flash_attn_no_pad.py`) so it can be imported by any backend that wants the
-traceable FA default call without pulling in backend dispatch logic. The
-backend (`attention/backends/flash_attn.py`) just imports
-`flash_attn_func_compilable` and `fa_version` from here.
-"""
-
-import importlib.util
-
-import torch
-
-from fastvideo import envs
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-# Pick the same backend the rest of FastVideo picked for `flash_attn_func`
-# (FA4/cute → FA3 → FA2). Mirror the precedence used in
-# `attention/utils/flash_attn_no_pad.py` so the two probes always agree.
-#
-# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1: its CuTeDSL
-# kernels JIT-compile per shape family and can fail at runtime on some
-# arch/shape combinations, so it is never auto-selected just because it is
-# installed.
-if envs.FASTVIDEO_FA4:
-    try:
-        from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
-    except ImportError as e:
-        raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
-                           "fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
-                           "or unset FASTVIDEO_FA4.") from e
-    fa_version = "4"
-else:
-    try:
-        from flash_attn_interface import flash_attn_func as flash_attn_3_func
-
-        # flash_attn 3 no longer has a different API, see following commit:
-        # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
-        flash_attn_func = flash_attn_3_func
-        fa_version = "3"
-    except ImportError:
-        from flash_attn import flash_attn_func as flash_attn_2_func
-        flash_attn_func = flash_attn_2_func
-        fa_version = "2"
-    try:
-        if importlib.util.find_spec("flash_attn.cute") is not None:
-            logger.info("flash_attn.cute (FA4) is installed but not enabled; "
-                        "set FASTVIDEO_FA4=1 to use it for inference.")
-    except ImportError:
-        pass
-
-if fa_version == "2":
-    # Scope: this op covers exactly the q/k/v + softmax_scale + causal call
-    # shape used by FlashAttentionImpl.forward's default branch (see
-    # `flash_attn_func_compilable(...)` call site in
-    # `attention/backends/flash_attn.py`). The masked/no-pad and varlen /
-    # cross-attn paths use different entry points
-    # (`flash_attn_no_pad`, `flash_attn_varlen_*`) which live in
-    # `attention/utils/flash_attn_no_pad.py`. The wrapper's signature is the
-    # contract: any extra kwarg (dropout_p, window_size, alibi_slopes,
-    # deterministic, return_attn_probs, ...) raises TypeError at the call
-    # site, so silent loss of kwargs is not a failure mode.
-    from flash_attn.flash_attn_interface import _flash_attn_backward as _fa2_backward
-    _fa_default = flash_attn_func
-
-    @torch.library.custom_op(
-        "fastvideo::_flash_attn_default_forward",
-        mutates_args=(),
-        device_types="cuda",
-    )
-    def _flash_attn_default_forward(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        softmax_scale: float | None,
-        causal: bool,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        # `return_attn_probs=True` asks FA2 to also return softmax_lse +
-        # S_dmask. We need softmax_lse to feed the backward; S_dmask is the
-        # dropout mask (always None here since dropout_p is fixed at 0).
-        out, softmax_lse, _ = _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal, return_attn_probs=True)
-        return out, softmax_lse
-
-    @torch.library.register_fake("fastvideo::_flash_attn_default_forward")
-    def _flash_attn_default_forward_fake(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        softmax_scale: float | None,
-        causal: bool,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        del softmax_scale, causal
-        # FA2 default path: out = [batch, seqlen_q, nheads, head_dim_v],
-        # softmax_lse = [batch, nheads, seqlen_q], fp32 regardless of q dtype.
-        b, sq, hq = q.shape[0], q.shape[1], q.shape[2]
-        out = q.new_empty(b, sq, hq, v.shape[-1])
-        lse = q.new_empty(b, hq, sq, dtype=torch.float32)
-        return out, lse
-
-    def _flash_attn_default_setup_context(ctx, inputs, output):
-        q, k, v, softmax_scale, causal = inputs
-        out, lse = output
-        ctx.save_for_backward(q, k, v, out, lse)
-        # `lse` is an auxiliary output we save to feed FA2's backward; nobody
-        # should differentiate through it. Mark it non-differentiable so
-        # autograd errors loudly if a caller wires it into a loss, rather
-        # than silently producing zero/None grads through the `del grad_lse`
-        # in our backward.
-        ctx.mark_non_differentiable(lse)
-        # FA2's *forward* substitutes `1 / sqrt(head_dim)` for `softmax_scale=None`
-        # internally; FA2's *backward* (`_flash_attn_backward`) demands a concrete
-        # float in its C++ schema and rejects None at the binding boundary. Resolve
-        # the default here so the value saved on ctx (and passed to backward) is
-        # always a real float — matches what FA2's own autograd.Function does.
-        if softmax_scale is None:
-            softmax_scale = q.shape[-1]**-0.5
-        ctx.softmax_scale = softmax_scale
-        ctx.causal = causal
-
-    def _flash_attn_default_backward(ctx, grad_out, grad_lse):
-        # We only differentiate `out`; softmax_lse is saved-for-backward, not
-        # a real differentiable output. (Mirrors the FP4 cute template.)
-        del grad_lse
-        q, k, v, out, lse = ctx.saved_tensors
-        dq = torch.empty_like(q)
-        dk = torch.empty_like(k)
-        dv = torch.empty_like(v)
-        # FA2's `_flash_attn_backward` writes into dq/dk/dv in place. The
-        # extra kwargs (window_size_*, softcap, alibi_slopes, deterministic,
-        # rng_state) are pinned to the same defaults the forward wrapper
-        # uses — flash-attn==2.8.1 (the version FastVideo pins) requires
-        # all of them explicitly. `rng_state=None` is correct for our
-        # `dropout_p=0` configuration.
-        _fa2_backward(
-            grad_out,
-            q,
-            k,
-            v,
-            out,
-            lse,
-            dq,
-            dk,
-            dv,
-            dropout_p=0.0,
-            softmax_scale=ctx.softmax_scale,
-            causal=ctx.causal,
-            window_size_left=-1,
-            window_size_right=-1,
-            softcap=0.0,
-            alibi_slopes=None,
-            deterministic=False,
-            rng_state=None,
-        )
-        return dq, dk, dv, None, None
-
-    torch.library.register_autograd(
-        "fastvideo::_flash_attn_default_forward",
-        _flash_attn_default_backward,
-        setup_context=_flash_attn_default_setup_context,
-    )
-
-    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
-        # Backward is registered: autograd flows through the op (training
-        # path is also traceable; no carve-out needed). Public API matches
-        # `flash_attn_func` — returns just `out`; we drop the saved-for-
-        # backward `lse` here so callers see the original single-tensor
-        # contract.
-        out, _ = torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
-        return out
-elif fa_version == "3":
-    # FA3 path: same forward+fake custom op as the original PR #1373, with
-    # the autograd carve-out kept. The full backward (mirroring the FA2 leg
-    # above) wants a Hopper box for grad-check validation, which we don't
-    # have until Kuan-Hao's Modal FA3 setup PR lands. Until then this keeps
-    # inference traceable + training correct (via the original
-    # autograd.Function path + a pre-PR-style graph break on training).
-    _fa_default = flash_attn_func
-
-    @torch.library.custom_op(
-        "fastvideo::_flash_attn_default_forward",
-        mutates_args=(),
-        device_types="cuda",
-    )
-    def _flash_attn_default_forward(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        softmax_scale: float | None,
-        causal: bool,
-    ) -> torch.Tensor:
-        return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
-
-    @torch.library.register_fake("fastvideo::_flash_attn_default_forward")
-    def _flash_attn_default_forward_fake(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        softmax_scale: float | None,
-        causal: bool,
-    ) -> torch.Tensor:
-        del softmax_scale, causal
-        return q.new_empty(q.shape[0], q.shape[1], q.shape[2], v.shape[-1])
-
-    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
-        # Autograd carve-out. The custom op above registers a forward + fake
-        # kernel but NO backward (register_autograd), so it is opaque to
-        # autograd. Inference runs under no_grad / inference_mode and routes
-        # through the traceable custom op — that is the torch.compile win, and
-        # the only path this PR claims. Training backprops through attention,
-        # so route grad-enabled calls to the original FA2/FA3 `flash_attn_func`
-        # (itself an autograd.Function, so backward is correct) at the cost of a
-        # dynamo graph break on the training path — i.e. pre-PR behavior, no
-        # regression. Full autograd parity for the custom op (mirroring the FP4
-        # cute template) is a tracked follow-up.
-        if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad):
-            return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
-        return torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
-elif fa_version == "4":
-    # FA4 path: `flash_attn_func` is already a torch.library custom op
-    # (registered in `fastvideo.attention.utils.flash_attn_cute`), so a
-    # passthrough is enough — no extra registration needed.
-    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
-        return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
-else:
-    # Defensive: the probe above only ever sets fa_version to "2", "3",
-    # or "4"; an unexpected value means an import/probe regression and
-    # we want a loud error at import, not a silent NameError later.
-    raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
-                       f"'2', '3', or '4' from the import probe above.")

fastvideo/attention/utils/flash_attn_no_pad.py

@@ -24,7 +24,7 @@ from flash_attn.bert_padding import pad_input, unpad_input
 from fastvideo import envs
 
 
-def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
+def _resolve_flash_attn_varlen_func() -> Any:
     if envs.FASTVIDEO_FA4:
         # FA4 cute is explicit opt-in (see fastvideo/attention/backends/
         # flash_attn.py); with FASTVIDEO_FA4=1 an unimportable FA4 build must
@@ -39,28 +39,20 @@ def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
                                "fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
                                "or unset FASTVIDEO_FA4.") from e
 
-        return flash_attn_varlen_func_cute, "4"
+        return flash_attn_varlen_func_cute
     try:
         from flash_attn_interface import (
             flash_attn_varlen_func as flash_attn_varlen_func_interface, )
 
-        return flash_attn_varlen_func_interface, "3"
+        return flash_attn_varlen_func_interface
     except ImportError:
         from flash_attn import (
             flash_attn_varlen_func as flash_attn_varlen_func_flash, )
 
-        return flash_attn_varlen_func_flash, "2"
+        return flash_attn_varlen_func_flash
 
 
-flash_attn_varlen_func_impl, _FA_VARLEN_VERSION = _resolve_flash_attn_varlen_func()
-
-# FA2-only: the private varlen backward we register against the custom ops
-# below. FA3 / FA4 have different private signatures and validation paths
-# (Hopper / Blackwell boxes) — those legs keep the autograd carve-out
-# pattern from PR #1373 until their setup PRs land.
-if _FA_VARLEN_VERSION == "2":
-    from flash_attn.flash_attn_interface import (
-        _flash_attn_varlen_backward as _fa2_varlen_backward, )
+flash_attn_varlen_func_impl = _resolve_flash_attn_varlen_func()
 
 
 def flash_attn_no_pad(
@@ -199,473 +191,3 @@ def flash_attn_varlen_qk_no_pad(
         h=nheads,
     )
     return output
-
-
-# ---------------------------------------------------------------------------
-# torch.compile traceability + register_autograd parity for the masked /
-# varlen attention paths.
-#
-# Wraps the two entry points `FlashAttentionImpl.forward` calls
-# (`flash_attn_no_pad`, `flash_attn_varlen_qk_no_pad`) as
-# `torch.library.custom_op`s so dynamo sees one traceable node — the
-# internal unpad / pad bookkeeping (data-dependent `nnz` shapes) runs
-# eager inside the op, and the op's outputs are the statically-shaped
-# padded tensors. This mirrors the FA2 default-path wrapper in
-# `fastvideo/attention/backends/flash_attn.py`.
-#
-# Autograd: on FA2 we register a real backward (`register_autograd`)
-# that calls FA2's `_flash_attn_varlen_backward` on the unpadded form
-# — re-unpadding the saved padded tensors using the saved mask. The
-# `softmax_lse` from the varlen forward is naturally unpadded
-# (`[nheads, total_q]`); we pad it to `[batch, nheads, seqlen]` on
-# the way out (statically shaped) and re-unpad in backward. So
-# training backprops *through* the op (no graph break on the training
-# path either).
-#
-# FA3 / FA4 keep the autograd carve-out pattern from PR #1373: the
-# custom op has forward + fake only, and `*_compilable` falls back to
-# the original autograd.Function for grad-enabled calls. Those legs
-# are gated on Hopper-class / Blackwell-class boxes for backward
-# validation and ship as separate follow-ups.
-
-if _FA_VARLEN_VERSION == "2":
-    # ---------- masked self-attention: flash_attn_no_pad (FA2) ----------
-
-    @torch.library.custom_op(
-        "fastvideo::_flash_attn_no_pad_forward",
-        mutates_args=(),
-        device_types="cuda",
-    )
-    def _flash_attn_no_pad_forward(
-        qkv: torch.Tensor,
-        key_padding_mask: torch.Tensor,
-        causal: bool,
-        dropout_p: float,
-        softmax_scale: float | None,
-        deterministic: bool,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        b, s, _three, h, d = qkv.shape
-        x = rearrange(qkv, "b s three h d -> b s (three h d)")
-        x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
-        x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
-        out_unpad, lse_unpad, _ = flash_attn_varlen_qkvpacked_func(x_unpad,
-                                                                   cu_seqlens,
-                                                                   max_s,
-                                                                   dropout_p,
-                                                                   softmax_scale=softmax_scale,
-                                                                   causal=causal,
-                                                                   deterministic=deterministic,
-                                                                   return_attn_probs=True)
-        # Pad out: [nnz, h, d] -> [b, s, h, d]
-        out_padded = rearrange(pad_input(rearrange(out_unpad, "nnz h d -> nnz (h d)"), indices, b, s),
-                               "b s (h d) -> b s h d",
-                               h=h)
-        # Pad lse: FA2 varlen returns [nheads, total_q]. Transpose to [total_q,
-        # nheads], pad to [b, s, nheads], permute to [b, nheads, s] — statically
-        # shaped so register_fake matches.
-        lse_padded = pad_input(lse_unpad.t().contiguous(), indices, b, s).permute(0, 2, 1).contiguous()
-        return out_padded, lse_padded
-
-    @torch.library.register_fake("fastvideo::_flash_attn_no_pad_forward")
-    def _flash_attn_no_pad_forward_fake(
-        qkv: torch.Tensor,
-        key_padding_mask: torch.Tensor,
-        causal: bool,
-        dropout_p: float,
-        softmax_scale: float | None,
-        deterministic: bool,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        del key_padding_mask, causal, dropout_p, softmax_scale, deterministic
-        b, s, _three, h, d = qkv.shape
-        out = qkv.new_empty(b, s, h, d)
-        lse = qkv.new_empty(b, h, s, dtype=torch.float32)
-        return out, lse
-
-    def _flash_attn_no_pad_setup_context(ctx, inputs, output):
-        qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic = inputs
-        out, lse = output
-        ctx.save_for_backward(qkv, out, lse, key_padding_mask)
-        # Auxiliary output, not differentiable — see default-path note.
-        ctx.mark_non_differentiable(lse)
-        # FA2's varlen backward requires a concrete float for softmax_scale.
-        if softmax_scale is None:
-            softmax_scale = qkv.shape[-1]**-0.5  # head_dim from qkv's last dim
-        ctx.softmax_scale = softmax_scale
-        ctx.causal = causal
-        ctx.dropout_p = dropout_p
-        ctx.deterministic = deterministic
-
-    def _flash_attn_no_pad_backward(ctx, grad_out, grad_lse):
-        # lse is saved-for-backward, not differentiated.
-        del grad_lse
-        qkv, out_padded, lse_padded, key_padding_mask = ctx.saved_tensors
-        b, s, _three, h, d = qkv.shape
-
-        # One `unpad_input` call (on qkv) gives us indices + cu_seqlens + max_s;
-        # reuse those for out / dout / lse below via direct indexing instead
-        # of redundant `unpad_input` calls (each of which would re-run
-        # `nonzero` + `cumsum` + a `.max().item()` GPU→CPU sync).
-        x = rearrange(qkv, "b s three h d -> b s (three h d)")
-        x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
-        x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
-        q_unpad, k_unpad, v_unpad = (t.contiguous() for t in x_unpad.unbind(dim=1))
-
-        # Direct-index variants reuse `indices` (computed above).
-        out_unpad = out_padded.flatten(0, 1)[indices].view(-1, h, d).contiguous()
-        dout_unpad = grad_out.flatten(0, 1)[indices].view(-1, h, d).contiguous()
-        # lse_padded [b, h, s] -> [b, s, h] -> [nnz, h] -> [h, nnz].
-        lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[indices].t().contiguous()
-
-        dq_unpad = torch.empty_like(q_unpad)
-        dk_unpad = torch.empty_like(k_unpad)
-        dv_unpad = torch.empty_like(v_unpad)
-        _fa2_varlen_backward(
-            dout_unpad,
-            q_unpad,
-            k_unpad,
-            v_unpad,
-            out_unpad,
-            lse_unpad,
-            dq_unpad,
-            dk_unpad,
-            dv_unpad,
-            cu_seqlens_q=cu_seqlens,
-            cu_seqlens_k=cu_seqlens,
-            max_seqlen_q=max_s,
-            max_seqlen_k=max_s,
-            dropout_p=ctx.dropout_p,
-            softmax_scale=ctx.softmax_scale,
-            causal=ctx.causal,
-            window_size_left=-1,
-            window_size_right=-1,
-            softcap=0.0,
-            alibi_slopes=None,
-            deterministic=ctx.deterministic,
-            rng_state=None,
-        )
-
-        # Re-pad each grad and stack into dqkv.
-        def _repad(dt_unpad: torch.Tensor) -> torch.Tensor:
-            padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, b, s)
-            return rearrange(padded, "b s (h d) -> b s h d", h=h)
-
-        dqkv = torch.stack([_repad(dq_unpad), _repad(dk_unpad), _repad(dv_unpad)], dim=2)
-        # 6 inputs total: qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic.
-        return dqkv, None, None, None, None, None
-
-    torch.library.register_autograd(
-        "fastvideo::_flash_attn_no_pad_forward",
-        _flash_attn_no_pad_backward,
-        setup_context=_flash_attn_no_pad_setup_context,
-    )
-
-    # ---------- cross-attention: flash_attn_varlen_qk_no_pad (FA2) ----------
-
-    @torch.library.custom_op(
-        "fastvideo::_flash_attn_varlen_qk_no_pad_forward",
-        mutates_args=(),
-        device_types="cuda",
-    )
-    def _flash_attn_varlen_qk_no_pad_forward(
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        query_padding_mask: torch.Tensor,
-        key_padding_mask: torch.Tensor,
-        causal: bool,
-        dropout_p: float,
-        softmax_scale: float | None,
-        deterministic: bool,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        b, sq, h, d = query.shape
-        q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
-                                                                        query_padding_mask)
-        k_unpad, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"),
-                                                                key_padding_mask)
-        v_unpad, _, _, _, _ = unpad_input(rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
-        q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h)
-        k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h)
-        v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h)
-        out_unpad, lse_unpad, _ = flash_attn_varlen_func_impl(q_unpad,
-                                                              k_unpad,
-                                                              v_unpad,
-                                                              cu_seqlens_q,
-                                                              cu_seqlens_k,
-                                                              max_seqlen_q,
-                                                              max_seqlen_k,
-                                                              dropout_p=dropout_p,
-                                                              softmax_scale=softmax_scale,
-                                                              causal=causal,
-                                                              deterministic=deterministic,
-                                                              return_attn_probs=True)
-        # Pad out: [nnz_q, h, d] -> [b, sq, h, d]
-        out_padded = rearrange(pad_input(rearrange(out_unpad, "nnz h d -> nnz (h d)"), q_indices, b, sq),
-                               "b s (h d) -> b s h d",
-                               h=h)
-        # Pad lse: [h, nnz_q] -> [b, h, sq]
-        lse_padded = pad_input(lse_unpad.t().contiguous(), q_indices, b, sq).permute(0, 2, 1).contiguous()
-        return out_padded, lse_padded
-
-    @torch.library.register_fake("fastvideo::_flash_attn_varlen_qk_no_pad_forward")
-    def _flash_attn_varlen_qk_no_pad_forward_fake(
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        query_padding_mask: torch.Tensor,
-        key_padding_mask: torch.Tensor,
-        causal: bool,
-        dropout_p: float,
-        softmax_scale: float | None,
-        deterministic: bool,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        del key, query_padding_mask, key_padding_mask
-        del causal, dropout_p, softmax_scale, deterministic
-        b, sq, h, _ = query.shape
-        # `out`'s head_dim comes from value (d_v), matching the real forward's
-        # out_padded ([b, sq, h, d_v]); it can differ from query's d_q.
-        out = query.new_empty(b, sq, h, value.shape[-1])
-        lse = query.new_empty(b, h, sq, dtype=torch.float32)
-        return out, lse
-
-    def _flash_attn_varlen_qk_no_pad_setup_context(ctx, inputs, output):
-        (query, key, value, query_padding_mask, key_padding_mask, causal, dropout_p, softmax_scale,
-         deterministic) = inputs
-        out, lse = output
-        ctx.save_for_backward(query, key, value, out, lse, query_padding_mask, key_padding_mask)
-        # Auxiliary output, not differentiable — see default-path note.
-        ctx.mark_non_differentiable(lse)
-        if softmax_scale is None:
-            softmax_scale = query.shape[-1]**-0.5
-        ctx.softmax_scale = softmax_scale
-        ctx.causal = causal
-        ctx.dropout_p = dropout_p
-        ctx.deterministic = deterministic
-
-    def _flash_attn_varlen_qk_no_pad_backward(ctx, grad_out, grad_lse):
-        del grad_lse
-        (query, key, value, out_padded, lse_padded, query_padding_mask, key_padding_mask) = ctx.saved_tensors
-        b, sq, h, d = query.shape
-        sk = key.shape[1]
-
-        # One `unpad_input` call per distinct mask; reuse the returned
-        # indices via direct indexing for everything else that shares
-        # the same mask (v with k_mask; out/dout/lse with q_mask; the
-        # final repad of dk/dv also reuses k_indices). Avoids ~4
-        # redundant `unpad_input` calls + their GPU→CPU `.max().item()`
-        # syncs.
-        q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
-                                                                        query_padding_mask)
-        k_unpad, k_indices, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"),
-                                                                        key_padding_mask)
-        q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
-        k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
-        v_unpad = value.flatten(0, 1)[k_indices].view(-1, h, d).contiguous()
-
-        # out / dout / lse follow q's shape, so index with q_indices.
-        out_unpad = out_padded.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
-        dout_unpad = grad_out.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
-        lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[q_indices].t().contiguous()
-
-        dq_unpad = torch.empty_like(q_unpad)
-        dk_unpad = torch.empty_like(k_unpad)
-        dv_unpad = torch.empty_like(v_unpad)
-        _fa2_varlen_backward(
-            dout_unpad,
-            q_unpad,
-            k_unpad,
-            v_unpad,
-            out_unpad,
-            lse_unpad,
-            dq_unpad,
-            dk_unpad,
-            dv_unpad,
-            cu_seqlens_q=cu_seqlens_q,
-            cu_seqlens_k=cu_seqlens_k,
-            max_seqlen_q=max_seqlen_q,
-            max_seqlen_k=max_seqlen_k,
-            dropout_p=ctx.dropout_p,
-            softmax_scale=ctx.softmax_scale,
-            causal=ctx.causal,
-            window_size_left=-1,
-            window_size_right=-1,
-            softcap=0.0,
-            alibi_slopes=None,
-            deterministic=ctx.deterministic,
-            rng_state=None,
-        )
-
-        # k_indices is already available from the unpad_input above —
-        # no need to recompute it for the dk/dv repad.
-        def _repad(dt_unpad: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
-            padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, batch, seqlen)
-            return rearrange(padded, "b s (h d) -> b s h d", h=h)
-
-        dq_padded = _repad(dq_unpad, q_indices, b, sq)
-        dk_padded = _repad(dk_unpad, k_indices, b, sk)
-        dv_padded = _repad(dv_unpad, k_indices, b, sk)
-        # 9 inputs total: query, key, value, q_mask, k_mask, causal, dropout_p,
-        # softmax_scale, deterministic.
-        return dq_padded, dk_padded, dv_padded, None, None, None, None, None, None
-
-    torch.library.register_autograd(
-        "fastvideo::_flash_attn_varlen_qk_no_pad_forward",
-        _flash_attn_varlen_qk_no_pad_backward,
-        setup_context=_flash_attn_varlen_qk_no_pad_setup_context,
-    )
-
-    # ---------- public dispatchers (FA2: autograd flows through the op) -----
-
-
-    def flash_attn_no_pad_compilable(qkv,
-                                     key_padding_mask,
-                                     causal=False,
-                                     dropout_p=0.0,
-                                     softmax_scale=None,
-                                     deterministic=False):
-        """dynamo-traceable wrapper around ``flash_attn_no_pad`` (registered op,
-        full register_autograd on FA2 — both inference and training go through
-        the op, no graph break on either)."""
-        out, _ = torch.ops.fastvideo._flash_attn_no_pad_forward(qkv, key_padding_mask, causal, dropout_p, softmax_scale,
-                                                                deterministic)
-        return out
-
-    def flash_attn_varlen_qk_no_pad_compilable(query,
-                                               key,
-                                               value,
-                                               query_padding_mask,
-                                               key_padding_mask,
-                                               causal=False,
-                                               dropout_p=0.0,
-                                               softmax_scale=None,
-                                               deterministic=False):
-        """dynamo-traceable wrapper around ``flash_attn_varlen_qk_no_pad`` (registered
-        op, full register_autograd on FA2)."""
-        out, _ = torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward(query, key, value, query_padding_mask,
-                                                                          key_padding_mask, causal, dropout_p,
-                                                                          softmax_scale, deterministic)
-        return out
-
-else:
-    # ---------- FA3 / FA4: carve-out (forward+fake only, no real backward) ---
-    # Same pattern as the parked varlen-extension and the FA3 default leg in
-    # `fastvideo/attention/backends/flash_attn.py`. Real backward for these
-    # versions is a follow-up gated on Hopper / Blackwell box validation.
-
-    @torch.library.custom_op(
-        "fastvideo::_flash_attn_no_pad_forward",
-        mutates_args=(),
-        device_types="cuda",
-    )
-    def _flash_attn_no_pad_forward(
-        qkv: torch.Tensor,
-        key_padding_mask: torch.Tensor,
-        causal: bool,
-        dropout_p: float,
-        softmax_scale: float | None,
-        deterministic: bool,
-    ) -> torch.Tensor:
-        return flash_attn_no_pad(  # type: ignore[no-untyped-call]
-            qkv,
-            key_padding_mask,
-            causal=causal,
-            dropout_p=dropout_p,
-            softmax_scale=softmax_scale,
-            deterministic=deterministic)
-
-    @torch.library.register_fake("fastvideo::_flash_attn_no_pad_forward")
-    def _flash_attn_no_pad_forward_fake(
-        qkv: torch.Tensor,
-        key_padding_mask: torch.Tensor,
-        causal: bool,
-        dropout_p: float,
-        softmax_scale: float | None,
-        deterministic: bool,
-    ) -> torch.Tensor:
-        del key_padding_mask, causal, dropout_p, softmax_scale, deterministic
-        b, s, _three, h, d = qkv.shape
-        return qkv.new_empty(b, s, h, d)
-
-    @torch.library.custom_op(
-        "fastvideo::_flash_attn_varlen_qk_no_pad_forward",
-        mutates_args=(),
-        device_types="cuda",
-    )
-    def _flash_attn_varlen_qk_no_pad_forward(
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        query_padding_mask: torch.Tensor,
-        key_padding_mask: torch.Tensor,
-        causal: bool,
-        dropout_p: float,
-        softmax_scale: float | None,
-        deterministic: bool,
-    ) -> torch.Tensor:
-        return flash_attn_varlen_qk_no_pad(  # type: ignore[no-untyped-call]
-            query,
-            key,
-            value,
-            query_padding_mask=query_padding_mask,
-            key_padding_mask=key_padding_mask,
-            causal=causal,
-            dropout_p=dropout_p,
-            softmax_scale=softmax_scale,
-            deterministic=deterministic)
-
-    @torch.library.register_fake("fastvideo::_flash_attn_varlen_qk_no_pad_forward")
-    def _flash_attn_varlen_qk_no_pad_forward_fake(
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        query_padding_mask: torch.Tensor,
-        key_padding_mask: torch.Tensor,
-        causal: bool,
-        dropout_p: float,
-        softmax_scale: float | None,
-        deterministic: bool,
-    ) -> torch.Tensor:
-        del key, query_padding_mask, key_padding_mask
-        del causal, dropout_p, softmax_scale, deterministic
-        b, sq, h, _ = query.shape
-        # `out`'s head_dim comes from value (d_v), matching the real forward's
-        # output ([b, sq, h, d_v]); it can differ from query's d_q.
-        return query.new_empty(b, sq, h, value.shape[-1])
-
-    def flash_attn_no_pad_compilable(qkv,
-                                     key_padding_mask,
-                                     causal=False,
-                                     dropout_p=0.0,
-                                     softmax_scale=None,
-                                     deterministic=False):
-        if torch.is_grad_enabled() and qkv.requires_grad:
-            return flash_attn_no_pad(qkv,
-                                     key_padding_mask,
-                                     causal=causal,
-                                     dropout_p=dropout_p,
-                                     softmax_scale=softmax_scale,
-                                     deterministic=deterministic)
-        return torch.ops.fastvideo._flash_attn_no_pad_forward(qkv, key_padding_mask, causal, dropout_p, softmax_scale,
-                                                              deterministic)
-
-    def flash_attn_varlen_qk_no_pad_compilable(query,
-                                               key,
-                                               value,
-                                               query_padding_mask,
-                                               key_padding_mask,
-                                               causal=False,
-                                               dropout_p=0.0,
-                                               softmax_scale=None,
-                                               deterministic=False):
-        if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad):
-            return flash_attn_varlen_qk_no_pad(query,
-                                               key,
-                                               value,
-                                               query_padding_mask=query_padding_mask,
-                                               key_padding_mask=key_padding_mask,
-                                               causal=causal,
-                                               dropout_p=dropout_p,
-                                               softmax_scale=softmax_scale,
-                                               deterministic=deterministic)
-        return torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward(query, key, value, query_padding_mask,
-                                                                        key_padding_mask, causal, dropout_p,
-                                                                        softmax_scale, deterministic)

fastvideo/configs/models/dits/__init__.py

@@ -1,6 +1,5 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
-from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig
 from fastvideo.configs.models.dits.flux_2 import Flux2Config
 from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
@@ -14,7 +13,7 @@ from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 
 __all__ = [
-    "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
-    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
-    "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
+    "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
+    "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
+    "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
 ]

fastvideo/configs/models/dits/dreamx_world.py

@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig
-
-
-@dataclass
-class DreamXWorldArchConfig(WanVideoArchConfig):
-    """DreamX-World DiT config with camera PRoPE control fields."""
-
-    add_control_adapter: bool = True
-    cam_method: str | None = "prope"
-    attn_compress: int = 1
-    cam_self_attn_layers: tuple[int, ...] | None = None
-
-
-@dataclass
-class DreamXWorldConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=DreamXWorldArchConfig)
-
-    prefix: str = "Wan"
-
-
-@dataclass
-class DreamXWorldARArchConfig(DreamXWorldArchConfig):
-    """DreamX-World-5B autoregressive causal DiT config."""
-
-    model_type: str = "ti2v"
-    patch_size: tuple[int, int, int] = (1, 2, 2)
-    text_len: int = 512
-    text_dim: int = 4096
-    freq_dim: int = 256
-    attn_compress: int = 4
-    cam_self_attn_layers: tuple[int, ...] | None = tuple(range(30))
-    local_attn_size: int = 12
-    sink_size: int = 3
-    num_frames_per_block: int = 3
-    rope_cache_policy: str = "block_relativistic"
-    # The official AR checkpoint (AMAP-ML/DreamX-World ``model.safetensors``)
-    # already uses FastVideo's native key names and the converter copies the
-    # tensors verbatim, so every rule is an identity. The rules enumerate the
-    # full state-dict surface of ``DreamXWorldARTransformer3DModel`` (norm1 /
-    # norm2 / head.norm are affine-free and have no parameters).
-    param_names_mapping: dict = field(
-        default_factory=lambda: {
-            r"^patch_embedding\.(.*)$": r"patch_embedding.\1",
-            r"^text_embedding\.([02])\.(.*)$": r"text_embedding.\1.\2",
-            r"^time_embedding\.([02])\.(.*)$": r"time_embedding.\1.\2",
-            r"^time_projection\.1\.(.*)$": r"time_projection.1.\1",
-            r"^blocks\.(\d+)\.self_attn\.(q|k|v|o)\.(.*)$": r"blocks.\1.self_attn.\2.\3",
-            r"^blocks\.(\d+)\.self_attn\.norm_(q|k)\.weight$": r"blocks.\1.self_attn.norm_\2.weight",
-            r"^blocks\.(\d+)\.cross_attn\.(q|k|v|o)\.(.*)$": r"blocks.\1.cross_attn.\2.\3",
-            r"^blocks\.(\d+)\.cross_attn\.norm_(q|k)\.weight$": r"blocks.\1.cross_attn.norm_\2.weight",
-            r"^blocks\.(\d+)\.cam_self_attn\.(q_proj|k_proj|v_proj|out_proj)\.(.*)$": r"blocks.\1.cam_self_attn.\2.\3",
-            r"^blocks\.(\d+)\.cam_self_attn\.norm_(q|k)\.weight$": r"blocks.\1.cam_self_attn.norm_\2.weight",
-            r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm3.\2",
-            r"^blocks\.(\d+)\.ffn\.([02])\.(.*)$": r"blocks.\1.ffn.\2.\3",
-            r"^blocks\.(\d+)\.modulation$": r"blocks.\1.modulation",
-            r"^head\.head\.(.*)$": r"head.head.\1",
-            r"^head\.modulation$": r"head.modulation",
-        })
-    reverse_param_names_mapping: dict = field(default_factory=dict)
-    lora_param_names_mapping: dict = field(default_factory=dict)
-
-
-@dataclass
-class DreamXWorldARConfig(DreamXWorldConfig):
-    arch_config: DiTArchConfig = field(default_factory=DreamXWorldARArchConfig)
-
-    prefix: str = "Wan"

fastvideo/configs/models/dits/wanvideo.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Literal
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
 
@@ -20,11 +19,6 @@ class WanVideoArchConfig(DiTArchConfig):
             r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
             r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
             r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
-            # AnyFlow dual-timestep checkpoints expose delta_embedder weights with the
-            # same internal layout as time_embedder. The regex is harmless on plain
-            # Wan checkpoints (no delta_embedder keys to match).
-            r"^condition_embedder\.delta_embedder\.linear_1\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_in.\1",
-            r"^condition_embedder\.delta_embedder\.linear_2\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_out.\1",
             r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
             r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
             r"condition_embedder.image_embedder.ff.fc_in.\1",
@@ -92,14 +86,6 @@ class WanVideoArchConfig(DiTArchConfig):
     # "relativistic" keeps long rollouts in-distribution; a no-op unless sink_size > 0 and local_attn_size > 0.
     rope_cache_policy: str = "absolute"
 
-    # AnyFlow dual-timestep conditioning. Defaults preserve bit-identity with
-    # the legacy single-timestep forward (no delta_embedder allocated, no
-    # extra computation on the embedder forward path).
-    r_embedder: bool = False
-    r_embedder_fusion: Literal["additive", "gated"] = "additive"
-    r_embedder_gate_value: float = 0.25
-    r_embedder_deltatime_type: Literal["r", "t-r"] = "r"
-
     def __post_init__(self):
         super().__post_init__()
         self.out_channels = self.out_channels or self.in_channels

fastvideo/configs/pipelines/__init__.py

@@ -1,7 +1,6 @@
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
@@ -18,6 +17,6 @@ __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
-    "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
-    "Kandinsky5I2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "HYWorldConfig", "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "MatrixGame2I2V480PConfig",
+    "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/dreamx_world.py

@@ -1,127 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World-5B-Cam FastVideo model configuration helpers."""
-
-from collections.abc import Callable
-from dataclasses import dataclass, field
-
-import torch
-
-from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.configs.models.dits.dreamx_world import (DreamXWorldARArchConfig, DreamXWorldARConfig,
-                                                        DreamXWorldArchConfig, DreamXWorldConfig)
-from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
-from fastvideo.configs.models.encoders.t5 import T5ArchConfig
-from fastvideo.configs.models.vaes import WanVAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.configs.pipelines.wan import LucyEditDevConfig, t5_postprocess_text
-
-
-def make_dreamx_world_5b_cam_dit_config() -> DreamXWorldConfig:
-    """Return the DreamX-World DiT config matching DreamX-World-5B-Cam."""
-    return DreamXWorldConfig(arch_config=DreamXWorldArchConfig(
-        num_attention_heads=24,
-        attention_head_dim=128,
-        in_channels=48,
-        out_channels=48,
-        ffn_dim=14336,
-        num_layers=30,
-        cross_attn_norm=True,
-        qk_norm="rms_norm_across_heads",
-        add_control_adapter=True,
-        cam_method="prope",
-        attn_compress=1,
-        cam_self_attn_layers=None,
-    ))
-
-
-def make_dreamx_world_5b_ar_dit_config() -> DreamXWorldARConfig:
-    """Return the DreamX-World-5B autoregressive causal DiT config."""
-    return DreamXWorldARConfig(arch_config=DreamXWorldARArchConfig(
-        model_type="ti2v",
-        num_attention_heads=24,
-        attention_head_dim=128,
-        in_channels=48,
-        out_channels=48,
-        ffn_dim=14336,
-        num_layers=30,
-        cross_attn_norm=True,
-        qk_norm=True,
-        add_control_adapter=True,
-        cam_method="prope",
-        attn_compress=4,
-        cam_self_attn_layers=tuple(range(30)),
-        local_attn_size=12,
-        sink_size=3,
-        num_frames_per_block=3,
-    ))
-
-
-def make_dreamx_world_5b_cam_vae_config() -> WanVAEConfig:
-    """Return the Wan2.2 48-channel VAE config used by DreamX-World-5B-Cam."""
-    return LucyEditDevConfig().vae_config
-
-
-def make_dreamx_world_5b_cam_text_encoder_config() -> T5Config:
-    """Return the UMT5-XXL text encoder config used by DreamX-World-5B-Cam."""
-    return T5Config(
-        arch_config=T5ArchConfig(
-            vocab_size=256384,
-            d_model=4096,
-            d_kv=64,
-            d_ff=10240,
-            num_layers=24,
-            num_decoder_layers=None,
-            num_heads=64,
-            relative_attention_num_buckets=32,
-            dropout_rate=0.0,
-            text_len=512,
-            feed_forward_proj="gelu",
-            is_encoder_decoder=False,
-        ),
-        prefix="umt5",
-    )
-
-
-@dataclass
-class DreamXWorld5BCamPipelineConfig(PipelineConfig):
-    """Pipeline config for the first-scope DreamX-World-5B-Cam mode."""
-
-    dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_cam_dit_config)
-    vae_config: VAEConfig = field(default_factory=make_dreamx_world_5b_cam_vae_config)
-    text_encoder_configs: tuple[EncoderConfig,
-                                ...] = field(default_factory=lambda: (make_dreamx_world_5b_cam_text_encoder_config(), ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda: (t5_postprocess_text, ))
-    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
-    flow_shift: float | None = 3.0
-    ti2v_task: bool = True
-    expand_timesteps: bool = True
-    vae_tiling: bool = False
-    vae_sp: bool = False
-    vae_precision: str = "fp32"
-    vae_decode_precision: str | None = "bf16"
-    dit_precision: str = "bf16"
-
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True
-        self.dit_config.expand_timesteps = self.expand_timesteps
-
-
-@dataclass
-class DreamXWorld5BARPipelineConfig(DreamXWorld5BCamPipelineConfig):
-    """Pipeline config for DreamX-World-5B autoregressive forcing."""
-
-    dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_ar_dit_config)
-    flow_shift: float | None = 5.0
-    ti2v_task: bool = True
-    is_causal: bool = True
-    dmd_denoising_steps: tuple[int, ...] = (1000, 750, 500, 250)
-    warp_denoising_step: bool = True
-    context_noise: float = 0.1
-    num_frames_per_block: int = 3
-    color_correction_strength: float = 1.0
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        self.dit_config.expand_timesteps = True

fastvideo/configs/pipelines/kandinsky5.py

@@ -84,12 +84,6 @@ class Kandinsky5T2VConfig(PipelineConfig):
         if len(self.text_encoder_configs) != 2:
             raise ValueError(f"Kandinsky5 pipeline requires exactly 2 text encoders (qwen and clip), "
                              f"but got {len(self.text_encoder_configs)} encoder(s).")
-        if len(self.text_encoder_precisions) != 2:
-            raise ValueError("Kandinsky5 pipeline requires exactly 2 text encoder precisions, "
-                             f"but got {len(self.text_encoder_precisions)}.")
-        if len(self.text_encoder_max_lengths) != 2:
-            raise ValueError("Kandinsky5 pipeline requires exactly 2 text encoder max lengths, "
-                             f"but got {len(self.text_encoder_max_lengths)}.")
 
         self.vae_config.load_encoder = False
         self.vae_config.load_decoder = True

fastvideo/dataset/parquet_dataset_map_style.py

@@ -2,7 +2,6 @@
 import os
 import pickle
 import random
-from collections.abc import Sequence
 from typing import Any
 
 import pyarrow as pa
@@ -95,56 +94,8 @@ class DP_SP_BatchSampler(Sampler[list[int]]):
         return len(self.sp_group_local_indices) // self.batch_size
 
 
-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]))
+def get_parquet_files_and_length(path: str):
+    dataset_root = os.path.realpath(os.path.expanduser(path))
     # 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")
@@ -317,7 +268,7 @@ class LatentsParquetMapStyleDataset(Dataset):
 
     def __init__(
         self,
-        path: str | Sequence[str] | dict[str, int],
+        path: str,
         batch_size: int,
         parquet_schema: pa.Schema,
         cfg_rate: float = 0.0,
@@ -407,6 +358,7 @@ class LatentsParquetMapStyleDataset(Dataset):
     def __len__(self):
         return sum(self.lengths)
 
+
 # ────────────────────────────────────────────────────────────────────────────
 # 3.  Loader helper – everything else stays just like your original trainer
 # ────────────────────────────────────────────────────────────────────────────

fastvideo/distributed/parallel_state.py

@@ -23,7 +23,6 @@ If you only need to use the distributed environment without model parallelism,
  you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
-import gc
 import os
 import pickle
 import weakref
@@ -1007,13 +1006,6 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     if shutdown_ray:
         import ray  # Lazy import Ray
         ray.shutdown()
-    # Actually free GPU memory, as the name promises. FSDP-wrapped modules sit
-    # in reference cycles (module <-> FSDP state <-> hooks), so dropping the
-    # last user reference does not free their parameters until a gc pass runs.
-    # Without this, back-to-back model-load tests in one pytest session OOM.
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
 
 
 def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0) -> list[int]:

fastvideo/eval/README.md

@@ -122,44 +122,6 @@ 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
@@ -388,3 +350,5 @@ 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.

fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py

@@ -63,7 +63,7 @@ class ThirdPersonCalibration:
     beta_fwd: float
     beta_strafe: float
     focal_length: float
-    r_z: float = 0.0
+    r_z: float
     r_y: float = 0.0
     init_pitch: float = 0.0
     notes: str = ""

fastvideo/layers/quantization/nvfp4_config.py

@@ -261,22 +261,6 @@ 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 = ""):
@@ -301,7 +285,8 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
 
     def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         SfLayout, _, _ = _require_flashinfer()
-        x = _coerce_fp4_input_dtype(x)
+        assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
+            f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
         x_2d = x.view(-1, x.shape[-1])
         x_fp4, x_scale = _nvfp4_quantize(
             x_2d,
@@ -347,7 +332,8 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
             if x_scale.dim() > 2:
                 x_scale = x_scale.view(-1, x_scale.shape[-1])
         else:
-            x = _coerce_fp4_input_dtype(x)
+            assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
+                f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
             x = x.view(-1, x.shape[-1])
             x_global_sf = self.x_global_sf
             x_fp4, x_scale = _nvfp4_quantize(

fastvideo/models/dits/causal_wanvideo.py

@@ -437,7 +437,6 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         # Causal-specific
         self.block_mask = None
-        self.teacher_forcing_block_mask = None
         self.num_frame_per_block = config.arch_config.num_frames_per_block
         assert self.num_frame_per_block <= 3
         self.independent_first_frame = False
@@ -501,70 +500,6 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         return block_mask
 
-    @staticmethod
-    def _prepare_teacher_forcing_mask(
-        device: torch.device | str, num_frames: int = 21,
-        frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1
-    ) -> BlockMask:
-        """Attention mask for the teacher-forcing ``[clean | noisy]`` sequence.
-
-        A noisy token attends to its own block plus the clean context of all
-        strictly previous blocks; clean tokens are block-wise causal.
-        """
-        if local_attn_size != -1:
-            raise NotImplementedError(
-                f"Teacher forcing ignores local_attn_size={local_attn_size}: "
-                "unlike the block-wise causal mask, this mask always attends "
-                "to the full clean context. Windowed teacher forcing is not "
-                "implemented; use local_attn_size=-1 for teacher-forcing "
-                "training.")
-        total_length = num_frames * frame_seqlen * 2
-        padded_length = math.ceil(total_length / 128) * 128 - total_length
-
-        clean_ends = num_frames * frame_seqlen
-        context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-
-        attention_block_size = frame_seqlen * num_frame_per_block
-        frame_indices = torch.arange(
-            start=0, end=num_frames * frame_seqlen,
-            step=attention_block_size, device=device, dtype=torch.long
-        )
-        for start in frame_indices:
-            context_ends[start:start + attention_block_size] = start + attention_block_size
-
-        noisy_image_start_list = torch.arange(
-            num_frames * frame_seqlen, total_length,
-            step=attention_block_size, device=device, dtype=torch.long
-        )
-        noisy_image_end_list = noisy_image_start_list + attention_block_size
-        for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)):
-            noise_noise_starts[start:end] = start
-            noise_noise_ends[start:end] = end
-            noise_context_ends[start:end] = block_index * attention_block_size
-
-        def attention_mask(b, h, q_idx, kv_idx):
-            clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx])
-            c1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx])
-            c2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx])
-            noise_mask = (q_idx >= clean_ends) & (c1 | c2)
-            eye_mask = q_idx == kv_idx
-            return eye_mask | clean_mask | noise_mask
-
-        block_mask = create_block_mask(
-            attention_mask, B=None, H=None,
-            Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length,
-            _compile=False, device=device)
-
-        if not dist.is_initialized() or dist.get_rank() == 0:
-            print(f" cache a teacher-forcing mask with block size of {num_frame_per_block} frames")
-            print(block_mask)
-
-        return block_mask
-
     def _forward_inference(
                 self,
                 hidden_states: torch.Tensor,
@@ -665,8 +600,7 @@ class CausalWanTransformer3DModel(BaseDiT):
             else:
                 causal_kwargs = {
                     "kv_cache": kv_cache[block_index],
-                    "crossattn_cache": (crossattn_cache[block_index]
-                                        if crossattn_cache is not None else None),
+                    "crossattn_cache": crossattn_cache[block_index],
                     "current_start": current_start,
                     "cache_start": cache_start,
                     "block_mask": self.block_mask,
@@ -694,12 +628,9 @@ class CausalWanTransformer3DModel(BaseDiT):
                 encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                 | None = None,
                 start_frame: int = 0,
-                clean_x: torch.Tensor | None = None,
-                aug_t: torch.Tensor | None = None,
                 **kwargs) -> torch.Tensor:
 
         orig_dtype = hidden_states.dtype
-        teacher_forcing = clean_x is not None
         if not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
         if isinstance(encoder_hidden_states_image,
@@ -732,26 +663,15 @@ class CausalWanTransformer3DModel(BaseDiT):
         freqs_cis = (freqs_cos,
                      freqs_sin) if freqs_cos is not None else None
 
-        if teacher_forcing:
-            if self.teacher_forcing_block_mask is None:
-                self.teacher_forcing_block_mask = self._prepare_teacher_forcing_mask(
-                    device=hidden_states.device,
-                    num_frames=num_frames,
-                    frame_seqlen=post_patch_height * post_patch_width,
-                    num_frame_per_block=self.num_frame_per_block,
-                    local_attn_size=self.local_attn_size,
-                )
-            block_mask = self.teacher_forcing_block_mask
-        else:
-            if self.block_mask is None:
-                self.block_mask = self._prepare_blockwise_causal_attn_mask(
-                    device=hidden_states.device,
-                    num_frames=num_frames,
-                    frame_seqlen=post_patch_height * post_patch_width,
-                    num_frame_per_block=self.num_frame_per_block,
-                    local_attn_size=self.local_attn_size
-                )
-            block_mask = self.block_mask
+        # Construct blockwise causal attn mask
+        if self.block_mask is None:
+            self.block_mask = self._prepare_blockwise_causal_attn_mask(
+                device=hidden_states.device,
+                num_frames=num_frames,
+                frame_seqlen=post_patch_height * post_patch_width,
+                num_frame_per_block=self.num_frame_per_block,
+                local_attn_size=self.local_attn_size
+            )
 
         hidden_states = self.patch_embedding(hidden_states)
         grid_sizes = torch.stack(
@@ -759,7 +679,6 @@ class CausalWanTransformer3DModel(BaseDiT):
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
         encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
-        encoder_hidden_states_text = encoder_hidden_states
 
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
                         timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
@@ -775,35 +694,18 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         assert encoder_hidden_states.dtype == orig_dtype
 
-        if teacher_forcing:
-            # Tile RoPE/modulation so clean frame i and noisy frame i share a position.
-            clean_tokens = self.patch_embedding(clean_x).flatten(2).transpose(1, 2)
-            hidden_states = torch.cat([clean_tokens, hidden_states], dim=1)
-            if aug_t is None:
-                aug_t = torch.zeros_like(timestep)
-            _, timestep_proj_clean, _, _ = self.condition_embedder(
-                aug_t.flatten(), encoder_hidden_states_text, None)
-            timestep_proj_clean = timestep_proj_clean.unflatten(
-                1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
-            timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1)
-            freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0),
-                         torch.cat([freqs_sin, freqs_sin], dim=0))
-
         # 4. Transformer blocks
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             for block in self.blocks:
                 hidden_states = self._gradient_checkpointing_func(
                     block, hidden_states, encoder_hidden_states,
                     timestep_proj, freqs_cis,
-                    block_mask=block_mask)
+                    block_mask=self.block_mask)
         else:
             for block in self.blocks:
                 hidden_states = block(hidden_states, encoder_hidden_states,
                                         timestep_proj, freqs_cis,
-                                        block_mask=block_mask)
-
-        if teacher_forcing:
-            hidden_states = hidden_states[:, hidden_states.shape[1] // 2:]
+                                        block_mask=self.block_mask)
 
         # 5. Output norm, projection & unpatchify
         temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)

fastvideo/models/dits/dreamx_world.py

@@ -1,511 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import math
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-from fastvideo.configs.models.dits.dreamx_world import DreamXWorldConfig
-from fastvideo.distributed.communication_op import (
-    sequence_model_parallel_all_gather_with_unpad,
-    sequence_model_parallel_shard)
-from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.layers.layernorm import RMSNorm
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
-from fastvideo.models.dits.wanvideo import (LayerNormScaleShift,
-                                            PatchEmbed,
-                                            WanTimeTextImageEmbedding,
-                                            WanTransformer3DModel,
-                                            WanTransformerBlock)
-from fastvideo.platforms import AttentionBackendEnum, current_platform
-from fastvideo.attention import LocalAttention
-from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.layers.quantization import QuantizationConfig
-from fastvideo.models.dits.base import BaseDiT
-
-
-def _dreamx_invert_se3(transforms: torch.Tensor) -> torch.Tensor:
-    assert transforms.shape[-2:] == (4, 4)
-    rot_inv = transforms[..., :3, :3].transpose(-1, -2)
-    out = torch.zeros_like(transforms)
-    out[..., :3, :3] = rot_inv
-    out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv,
-                                    transforms[..., :3, 3])
-    out[..., 3, 3] = 1.0
-    return out.to(dtype=transforms.dtype)
-
-
-def _dreamx_lift_k(intrinsics: torch.Tensor) -> torch.Tensor:
-    assert intrinsics.shape[-2:] == (3, 3)
-    out = torch.zeros(intrinsics.shape[:-2] + (4, 4),
-                      device=intrinsics.device,
-                      dtype=intrinsics.dtype)
-    out[..., :3, :3] = intrinsics
-    out[..., 3, 3] = 1.0
-    return out
-
-
-def _dreamx_invert_k(intrinsics: torch.Tensor) -> torch.Tensor:
-    assert intrinsics.shape[-2:] == (3, 3)
-    out = torch.zeros_like(intrinsics)
-    out[..., 0, 0] = 1.0 / intrinsics[..., 0, 0]
-    out[..., 1, 1] = 1.0 / intrinsics[..., 1, 1]
-    out[..., 0, 2] = -intrinsics[..., 0, 2] / intrinsics[..., 0, 0]
-    out[..., 1, 2] = -intrinsics[..., 1, 2] / intrinsics[..., 1, 1]
-    out[..., 2, 2] = 1.0
-    return out.to(dtype=intrinsics.dtype)
-
-
-def _dreamx_apply_tiled_projmat(feats: torch.Tensor,
-                                matrix: torch.Tensor) -> torch.Tensor:
-    batch, num_heads, seq_len, feat_dim = feats.shape
-    proj_dim = matrix.shape[-1]
-    assert feat_dim % proj_dim == 0
-
-    if matrix.shape[1] == seq_len:
-        feats = feats.view(batch, num_heads, seq_len, feat_dim // proj_dim,
-                           proj_dim)
-        out = torch.einsum("btij,bntpj->bntpi", matrix, feats)
-        return out.reshape(batch, num_heads, seq_len, feat_dim)
-
-    cameras = matrix.shape[1]
-    assert seq_len > cameras and seq_len % cameras == 0
-    feats = feats.reshape(batch, num_heads, cameras, -1,
-                          feat_dim // proj_dim, proj_dim)
-    out = torch.einsum("bcij,bncpkj->bncpki", matrix, feats)
-    return out.reshape(batch, num_heads, seq_len, feat_dim)
-
-
-def _dreamx_prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
-                      viewmats: torch.Tensor, intrinsics: torch.Tensor):
-    batch, num_heads, seq_len, head_dim = q.shape
-    cameras = viewmats.shape[1]
-    assert q.shape == k.shape == v.shape
-    assert viewmats.shape == (batch, cameras, 4, 4)
-    assert intrinsics.shape == (batch, cameras, 3, 3)
-    assert head_dim % 4 == 0
-
-    intrinsics_norm = torch.zeros_like(intrinsics)
-    intrinsics_norm[..., 0, 0] = intrinsics[..., 0, 0]
-    intrinsics_norm[..., 1, 1] = intrinsics[..., 1, 1]
-    intrinsics_norm[..., 2, 2] = 1.0
-
-    proj = torch.einsum("...ij,...jk->...ik",
-                        _dreamx_lift_k(intrinsics_norm), viewmats)
-    proj_t = proj.transpose(-1, -2).to(dtype=viewmats.dtype)
-    proj_inv = torch.einsum(
-        "...ij,...jk->...ik",
-        _dreamx_invert_se3(viewmats),
-        _dreamx_lift_k(_dreamx_invert_k(intrinsics_norm)),
-    ).to(dtype=viewmats.dtype)
-
-    q = _dreamx_apply_tiled_projmat(q, proj_t)
-    k = _dreamx_apply_tiled_projmat(k, proj_inv)
-    v = _dreamx_apply_tiled_projmat(v, proj_inv)
-    return q, k, v, proj
-
-
-class DreamXPropeSelfAttention(nn.Module):
-    """DreamX-World parallel PRoPE camera self-attention branch."""
-
-    def __init__(self,
-                 dim: int,
-                 attn_dim: int,
-                 num_heads: int,
-                 qk_norm: str | bool = True,
-                 eps: float = 1e-6,
-                 quant_config: QuantizationConfig | None = None,
-                 prefix: str = ""):
-        super().__init__()
-        assert attn_dim % num_heads == 0
-        self.attn_dim = attn_dim
-        self.num_heads = num_heads
-        self.head_dim = attn_dim // num_heads
-        self.qk_norm = qk_norm
-
-        self.q_proj = ReplicatedLinear(dim,
-                                       attn_dim,
-                                       quant_config=quant_config,
-                                       prefix=f"{prefix}.q_proj")
-        self.k_proj = ReplicatedLinear(dim,
-                                       attn_dim,
-                                       quant_config=quant_config,
-                                       prefix=f"{prefix}.k_proj")
-        self.v_proj = ReplicatedLinear(dim,
-                                       attn_dim,
-                                       quant_config=quant_config,
-                                       prefix=f"{prefix}.v_proj")
-        self.out_proj = ReplicatedLinear(attn_dim,
-                                         dim,
-                                         quant_config=quant_config,
-                                         prefix=f"{prefix}.out_proj")
-
-        if qk_norm == "rms_norm":
-            self.norm_q = RMSNorm(self.head_dim, eps=eps)
-            self.norm_k = RMSNorm(self.head_dim, eps=eps)
-        elif qk_norm in (True, "rms_norm_across_heads"):
-            self.norm_q = RMSNorm(attn_dim, eps=eps)
-            self.norm_k = RMSNorm(attn_dim, eps=eps)
-        elif qk_norm is False:
-            self.norm_q = nn.Identity()
-            self.norm_k = nn.Identity()
-        else:
-            raise ValueError(f"Unsupported qk_norm for DreamX PRoPE: {qk_norm}")
-
-        nn.init.zeros_(self.out_proj.weight)
-        if self.out_proj.bias is not None:
-            nn.init.zeros_(self.out_proj.bias)
-
-        self.attn = LocalAttention(
-            num_heads=num_heads,
-            head_size=self.head_dim,
-            dropout_rate=0,
-            softmax_scale=None,
-            causal=False,
-            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
-                                          AttentionBackendEnum.TORCH_SDPA))
-
-    def forward(self, hidden_states: torch.Tensor,
-                y_camera: dict[str, torch.Tensor]) -> torch.Tensor:
-        if get_sp_world_size() > 1:
-            # The transformer shards the sequence before the block loop and
-            # this branch uses LocalAttention (no all-to-all): under
-            # sequence parallelism each rank would attend only within its
-            # own shard — silently wrong output. Fail loudly until this
-            # path is ported to DistributedAttention and validated.
-            raise NotImplementedError(
-                "DreamXPropeSelfAttention does not support sequence "
-                "parallelism yet (LocalAttention on a sharded sequence "
-                "corrupts output). Run with sp_size=1.")
-        batch_size, seq_len, _ = hidden_states.shape
-
-        query, _ = self.q_proj(hidden_states)
-        key, _ = self.k_proj(hidden_states)
-        value, _ = self.v_proj(hidden_states)
-
-        if self.qk_norm == "rms_norm":
-            query = query.view(batch_size, seq_len, self.num_heads,
-                               self.head_dim)
-            key = key.view(batch_size, seq_len, self.num_heads, self.head_dim)
-            query = self.norm_q(query)
-            key = self.norm_k(key)
-        else:
-            query = self.norm_q(query).view(batch_size, seq_len,
-                                            self.num_heads, self.head_dim)
-            key = self.norm_k(key).view(batch_size, seq_len, self.num_heads,
-                                        self.head_dim)
-
-        value = value.view(batch_size, seq_len, self.num_heads, self.head_dim)
-        query = query.transpose(1, 2)
-        key = key.transpose(1, 2)
-        value = value.transpose(1, 2)
-
-        query, key, value, output_projection = _dreamx_prope_qkv(
-            query,
-            key,
-            value,
-            viewmats=y_camera["viewmats"],
-            intrinsics=y_camera["K"],
-        )
-
-        out = self.attn(query.transpose(1, 2), key.transpose(1, 2),
-                        value.transpose(1, 2))
-        out = _dreamx_apply_tiled_projmat(out.transpose(1, 2),
-                                          output_projection).transpose(1, 2)
-        out = out.flatten(2)
-        out, _ = self.out_proj(out)
-        return out
-
-
-class DreamXWorldTransformerBlock(WanTransformerBlock):
-
-    def __init__(self,
-                 dim: int,
-                 ffn_dim: int,
-                 num_heads: int,
-                 qk_norm: str = "rms_norm_across_heads",
-                 cross_attn_norm: bool = False,
-                 eps: float = 1e-6,
-                 added_kv_proj_dim: int | None = None,
-                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                 | None = None,
-                 quant_config: QuantizationConfig | None = None,
-                 prefix: str = "",
-                 add_control_adapter: bool = True,
-                 cam_method: str | None = "prope",
-                 attn_compress: int = 1,
-                 cam_self_attn_layers: tuple[int, ...] | None = None,
-                 layer_idx: int | None = None):
-        super().__init__(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm,
-                         eps, added_kv_proj_dim,
-                         supported_attention_backends, quant_config, prefix)
-        self.cam_self_attn = None
-        add_cam_attn = add_control_adapter and cam_method == "prope"
-        if add_cam_attn and cam_self_attn_layers is not None:
-            add_cam_attn = layer_idx in cam_self_attn_layers
-        if add_cam_attn:
-            if num_heads % attn_compress != 0 or dim % attn_compress != 0:
-                raise ValueError("DreamX attn_compress must divide dim and num_heads")
-            self.cam_self_attn = DreamXPropeSelfAttention(
-                dim,
-                dim // attn_compress,
-                num_heads // attn_compress,
-                qk_norm=qk_norm,
-                eps=eps,
-                quant_config=quant_config,
-                prefix=f"{prefix}.cam_self_attn")
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        freqs_cis: tuple[torch.Tensor, torch.Tensor],
-        original_seq_len: int,
-        y_camera: dict[str, torch.Tensor] | None = None,
-    ) -> torch.Tensor:
-        if hidden_states.dim() == 4:
-            hidden_states = hidden_states.squeeze(1)
-        orig_dtype = hidden_states.dtype
-
-        if temb.dim() == 4:
-            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(
-                    6, dim=2)
-            shift_msa = shift_msa.squeeze(2)
-            scale_msa = scale_msa.squeeze(2)
-            gate_msa = gate_msa.squeeze(2)
-            c_shift_msa = c_shift_msa.squeeze(2)
-            c_scale_msa = c_scale_msa.squeeze(2)
-            c_gate_msa = c_gate_msa.squeeze(2)
-        else:
-            e = self.scale_shift_table + temb.float()
-            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
-                6, dim=1)
-        assert shift_msa.dtype == torch.float32
-
-        norm_hidden_states = (self.norm1(hidden_states.float()) *
-                              (1 + scale_msa) + shift_msa).to(orig_dtype)
-        query, _ = self.to_q(norm_hidden_states)
-        key, _ = self.to_k(norm_hidden_states)
-        value, _ = self.to_v(norm_hidden_states)
-
-        if self.norm_q is not None:
-            query = self.norm_q(query)
-        if self.norm_k is not None:
-            key = self.norm_k(key)
-
-        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
-        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
-        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
-
-        attn_output, _ = self.attn1(
-            query,
-            key,
-            value,
-            original_seq_len,
-            freqs_cis=freqs_cis,
-        )
-        attn_output = attn_output.flatten(2)
-        attn_output, _ = self.to_out(attn_output)
-        attn_output = attn_output.squeeze(1)
-        if self.cam_self_attn is not None and y_camera is not None:
-            attn_output = attn_output + self.cam_self_attn(
-                norm_hidden_states, y_camera)
-
-        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
-        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
-            hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
-
-        attn_output = self.attn2(norm_hidden_states,
-                                 context=encoder_hidden_states,
-                                 context_lens=None)
-        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
-            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
-
-        ff_output = self.ffn(norm_hidden_states)
-        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
-        hidden_states = hidden_states.to(orig_dtype)
-
-        return hidden_states
-
-
-class DreamXWorldTransformer3DModel(WanTransformer3DModel):
-    _fsdp_shard_conditions = DreamXWorldConfig()._fsdp_shard_conditions
-    _compile_conditions = DreamXWorldConfig()._compile_conditions
-    _supported_attention_backends = DreamXWorldConfig(
-    )._supported_attention_backends
-    param_names_mapping = DreamXWorldConfig().param_names_mapping
-    reverse_param_names_mapping = DreamXWorldConfig().reverse_param_names_mapping
-    lora_param_names_mapping = DreamXWorldConfig().lora_param_names_mapping
-
-    def __init__(self, config: DreamXWorldConfig, hf_config: dict[str,
-                                                                  Any]) -> None:
-        BaseDiT.__init__(self, config=config, hf_config=hf_config)
-        self.quant_config = config.quant_config
-
-        inner_dim = config.num_attention_heads * config.attention_head_dim
-        self.hidden_size = config.hidden_size
-        self.num_attention_heads = config.num_attention_heads
-        self.in_channels = config.in_channels
-        self.out_channels = config.out_channels
-        self.num_channels_latents = config.num_channels_latents
-        self.patch_size = config.patch_size
-        self.text_len = config.text_len
-
-        assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"
-
-        self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
-                                          embed_dim=inner_dim,
-                                          patch_size=config.patch_size,
-                                          flatten=False)
-        self.condition_embedder = WanTimeTextImageEmbedding(
-            dim=inner_dim,
-            time_freq_dim=config.freq_dim,
-            text_embed_dim=config.text_dim,
-            image_embed_dim=config.image_dim,
-        )
-        self.blocks = nn.ModuleList([
-            DreamXWorldTransformerBlock(
-                inner_dim,
-                config.ffn_dim,
-                config.num_attention_heads,
-                config.qk_norm,
-                config.cross_attn_norm,
-                config.eps,
-                config.added_kv_proj_dim,
-                self._supported_attention_backends,
-                quant_config=config.quant_config,
-                prefix=f"{config.prefix}.blocks.{i}",
-                add_control_adapter=config.add_control_adapter,
-                cam_method=config.cam_method,
-                attn_compress=config.attn_compress,
-                cam_self_attn_layers=config.cam_self_attn_layers,
-                layer_idx=i)
-            for i in range(config.num_layers)
-        ])
-        self.norm_out = LayerNormScaleShift(inner_dim,
-                                            norm_type="layer",
-                                            eps=config.eps,
-                                            elementwise_affine=False,
-                                            dtype=torch.float32,
-                                            compute_dtype=torch.float32)
-        self.proj_out = nn.Linear(
-            inner_dim, config.out_channels * math.prod(config.patch_size))
-        self.scale_shift_table = nn.Parameter(
-            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
-        self.gradient_checkpointing = False
-        self.__post_init__()
-
-    def forward(self,
-                hidden_states: torch.Tensor,
-                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
-                timestep: torch.LongTensor,
-                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
-                | None = None,
-                guidance=None,
-                y_camera: dict[str, torch.Tensor] | None = None,
-                **kwargs) -> torch.Tensor:
-        orig_dtype = hidden_states.dtype
-        if encoder_hidden_states is not None and not isinstance(
-                encoder_hidden_states, torch.Tensor):
-            encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
-
-        batch_size, _, num_frames, height, width = hidden_states.shape
-        p_t, p_h, p_w = self.patch_size
-        post_patch_num_frames = num_frames // p_t
-        post_patch_height = height // p_h
-        post_patch_width = width // p_w
-
-        d = self.hidden_size // self.num_attention_heads
-        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
-        freqs_cos, freqs_sin = get_rotary_pos_embed(
-            (post_patch_num_frames, post_patch_height, post_patch_width),
-            self.hidden_size,
-            self.num_attention_heads,
-            rope_dim_list,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
-            rope_theta=10000)
-        freqs_cis = (freqs_cos.to(hidden_states.device).float(),
-                     freqs_sin.to(hidden_states.device).float())
-
-        hidden_states = self.patch_embedding(hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)
-        hidden_states, original_seq_len = sequence_model_parallel_shard(
-            hidden_states, dim=1)
-
-        if timestep.dim() == 2:
-            ts_seq_len = timestep.shape[1]
-            timestep = timestep.flatten()
-        else:
-            ts_seq_len = None
-
-        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep,
-            encoder_hidden_states,
-            encoder_hidden_states_image,
-            timestep_seq_len=ts_seq_len)
-        if ts_seq_len is not None:
-            timestep_proj = timestep_proj.unflatten(2, (6, -1))
-        else:
-            timestep_proj = timestep_proj.unflatten(1, (6, -1))
-
-        if encoder_hidden_states_image is not None:
-            if encoder_hidden_states is not None:
-                encoder_hidden_states = torch.concat(
-                    [encoder_hidden_states_image, encoder_hidden_states], dim=1)
-            else:
-                encoder_hidden_states = encoder_hidden_states_image
-
-        if current_platform.is_mps() or current_platform.is_npu():
-            encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
-
-        assert encoder_hidden_states.dtype == orig_dtype
-
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-            for block in self.blocks:
-                hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, encoder_hidden_states,
-                    timestep_proj, freqs_cis, original_seq_len, y_camera)
-        else:
-            for block in self.blocks:
-                hidden_states = block(hidden_states, encoder_hidden_states,
-                                      timestep_proj, freqs_cis, original_seq_len,
-                                      y_camera=y_camera)
-
-        if temb.dim() == 3:
-            shift, scale = (self.scale_shift_table.unsqueeze(0) +
-                            temb.unsqueeze(2)).chunk(2, dim=2)
-            shift = shift.squeeze(2)
-            scale = scale.squeeze(2)
-        else:
-            shift, scale = (self.scale_shift_table +
-                            temb.unsqueeze(1)).chunk(2, dim=1)
-
-        hidden_states = self.norm_out(hidden_states, shift, scale)
-        hidden_states = sequence_model_parallel_all_gather_with_unpad(
-            hidden_states, original_seq_len, dim=1)
-        hidden_states = self.proj_out(hidden_states)
-
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
-        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
-        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
-
-        return output
-
-
-EntryClass = DreamXWorldTransformer3DModel

fastvideo/models/dits/dreamx_world_ar.py

@@ -1,920 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World autoregressive causal DiT.
-
-Adapted from DreamX-World's Apache-2.0
-``wan/modules/causal_camera_model_2_2_prope_infinity.py``. The implementation is
-kept native to FastVideo: no production import from DreamX, Diffusers, or
-Transformers is required.
-"""
-
-from typing import Any
-
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.models.dits.base import BaseDiT
-from fastvideo.models.dits.dreamx_world import (_dreamx_apply_tiled_projmat,
-                                               _dreamx_prope_qkv)
-
-
-def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    # Deliberately raw SDPA rather than fastvideo.attention.LocalAttention:
-    # (1) LocalAttention dispatches through the attention-backend registry, so
-    #     FLASH_ATTN could be selected and its kernel is not bit-identical to
-    #     torch SDPA — the AR KV-cache rollout must stay numerically frozen;
-    # (2) LocalAttention requires an active ForwardContext, which direct
-    #     transformer invocations (parity tests) do not set;
-    # (3) the sibling causal model keeps raw SDPA in the same KV-cache window
-    #     path (matrixgame2/causal_model.py).
-    # Sequence-parallel gap: this model never shards the sequence; run with
-    # sp_size=1 (see fastvideo/layers/AGENTS.md on documenting raw SDPA).
-    q_bhld = q.transpose(1, 2)
-    k_bhld = k.transpose(1, 2)
-    v_bhld = v.transpose(1, 2)
-    out = F.scaled_dot_product_attention(q_bhld, k_bhld, v_bhld, dropout_p=0.0)
-    return out.transpose(1, 2)
-
-
-def prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
-              viewmats: torch.Tensor, Ks: torch.Tensor):
-    q, k, v, output_projection = _dreamx_prope_qkv(q, k, v, viewmats, Ks)
-
-    def apply_fn_o(x: torch.Tensor) -> torch.Tensor:
-        return _dreamx_apply_tiled_projmat(x, output_projection)
-
-    return q, k, v, apply_fn_o
-
-
-def sinusoidal_embedding_1d(dim, position):
-    assert dim % 2 == 0
-    half = dim // 2
-    position = position.type(torch.float64)
-    sinusoid = torch.outer(
-        position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
-    return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
-
-
-def rope_params(max_seq_len, dim, theta=10000):
-    assert dim % 2 == 0
-    freqs = torch.outer(
-        torch.arange(max_seq_len),
-        1.0 / torch.pow(theta,
-                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
-    return torch.polar(torch.ones_like(freqs), freqs)
-
-
-class WanRMSNorm(nn.Module):
-    """Kept private instead of fastvideo.layers.layernorm.RMSNorm.
-
-    The official DreamX-World ``model_2_2.py`` computes the RMS statistics in
-    the *input* dtype — the upstream code has the fp32 upcast explicitly
-    commented out (``# return self._norm(x.float())...``). FastVideo's RMSNorm
-    always normalizes in fp32, which is not bit-identical under bf16, so the
-    verbatim implementation stays.
-    """
-
-    def __init__(self, dim, eps=1e-5):
-        super().__init__()
-        self.dim = dim
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-    def forward(self, x):
-        return self._norm(x).type_as(x) * self.weight
-
-    def _norm(self, x):
-        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
-
-
-class WanLayerNorm(nn.LayerNorm):
-    """Kept private instead of fastvideo.layers.layernorm.FP32LayerNorm.
-
-    The official DreamX-World ``model_2_2.py`` normalizes in the *input* dtype
-    (no ``x.float()`` upcast, unlike Wan2.1). FP32LayerNorm casts input and
-    affine params to fp32, which is not bit-identical under bf16, so the
-    verbatim implementation stays.
-    """
-
-    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
-        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
-
-    def forward(self, x):
-        return super().forward(x).type_as(x)
-
-
-class WanCrossAttention(nn.Module):
-    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-        self.q = ReplicatedLinear(dim, dim)
-        self.k = ReplicatedLinear(dim, dim)
-        self.v = ReplicatedLinear(dim, dim)
-        self.o = ReplicatedLinear(dim, dim)
-        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-
-    def _kv(self, context, b, n, d):
-        k, _ = self.k(context)
-        k = self.norm_k(k).view(b, -1, n, d)
-        v, _ = self.v(context)
-        v = v.view(b, -1, n, d)
-        return k, v
-
-    def forward(self, x, context, context_lens, crossattn_cache=None):
-        b, n, d = x.size(0), self.num_heads, self.head_dim
-        q, _ = self.q(x)
-        q = self.norm_q(q).view(b, -1, n, d)
-
-        if crossattn_cache is not None:
-            if not crossattn_cache["is_init"]:
-                crossattn_cache["is_init"] = True
-                k, v = self._kv(context, b, n, d)
-                crossattn_cache["k"] = k
-                crossattn_cache["v"] = v
-            else:
-                k = crossattn_cache["k"]
-                v = crossattn_cache["v"]
-        else:
-            k, v = self._kv(context, b, n, d)
-
-        x = attention(q, k, v)
-        x = x.flatten(2)
-        out, _ = self.o(x)
-        return out
-
-
-def block_relativistic_rope(x, grid_sizes, freqs, start_frame=0, relative_frame_indices=None):
-    """
-    Apply Block-Relativistic RoPE to input tensor.
-    Adapted from Infinity-RoPE (https://arxiv.org/abs/2511.20649).
-
-    Args:
-        x: Input tensor [B, L, num_heads, head_dim]
-        grid_sizes: Tensor [B, 3] containing (F, H, W)
-        freqs: RoPE frequencies
-        start_frame: Starting frame index for sequential RoPE
-        relative_frame_indices: Optional tensor [F] specifying explicit frame indices
-                               for Block-Relativistic RoPE. Overrides start_frame if provided.
-    """
-    n, c = x.size(2), x.size(3) // 2
-    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
-
-    output = []
-    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
-        seq_len = f * h * w
-        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
-            seq_len, n, -1, 2))
-
-        if relative_frame_indices is not None:
-            frame_indices = relative_frame_indices.long()
-            freqs_temporal = freqs[0][frame_indices].view(f, 1, 1, -1).expand(f, h, w, -1)
-        else:
-            freqs_temporal = freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1)
-
-        freqs_i = torch.cat([
-            freqs_temporal,
-            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
-            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
-        ], dim=-1).reshape(seq_len, 1, -1)
-
-        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
-        x_i = torch.cat([x_i, x[i, seq_len:]])
-        output.append(x_i)
-
-    return torch.stack(output).type_as(x)
-
-
-class CausalWanSelfAttention(nn.Module):
-    """Self-attention with KV cache and Block-Relativistic RoPE for causal inference."""
-
-    def __init__(self, dim, num_heads, local_attn_size=6, sink_size=1,
-                 qk_norm=True, eps=1e-6):
-        assert dim % num_heads == 0
-        super().__init__()
-        self.dim = dim
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-        self.local_attn_size = local_attn_size
-        self.sink_size = sink_size
-        self.qk_norm = qk_norm
-        self.eps = eps
-        self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
-
-        self.q = ReplicatedLinear(dim, dim)
-        self.k = ReplicatedLinear(dim, dim)
-        self.v = ReplicatedLinear(dim, dim)
-        self.o = ReplicatedLinear(dim, dim)
-        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-
-    def forward(self, x, seq_lens, grid_sizes, freqs, kv_cache,
-                current_start=0, cache_start=None, sink_recache_after_switch=False):
-        """
-        Args:
-            x: Shape [B, L, C]
-            seq_lens: Shape [B]
-            grid_sizes: Shape [B, 3] containing (F, H, W)
-            freqs: RoPE frequencies [1024, head_dim / 2]
-            kv_cache: Dict with 'k', 'v', 'global_end_index', 'local_end_index'
-            current_start: Current position in the global token sequence
-        """
-        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
-        if cache_start is None:
-            cache_start = current_start
-
-        q, _ = self.q(x)
-        q = self.norm_q(q).view(b, s, n, d)
-        k, _ = self.k(x)
-        k = self.norm_k(k).view(b, s, n, d)
-        v, _ = self.v(x)
-        v = v.view(b, s, n, d)
-
-        frame_seqlen = math.prod(grid_sizes[0][1:]).item()
-        num_new_frames = grid_sizes[0][0].item()
-        current_end = current_start + q.shape[1]
-        sink_tokens = self.sink_size * frame_seqlen
-        kv_cache_size = kv_cache["k"].shape[1]
-        num_new_tokens = q.shape[1]
-
-        cache_update_info = None
-        is_recompute = current_end <= kv_cache["global_end_index"].item() and current_start > 0
-
-        if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
-                num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
-            # === ROLLING MODE: cache full, evict oldest non-sink tokens ===
-            num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
-            num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
-            local_end_index = kv_cache["local_end_index"].item() + current_end - \
-                kv_cache["global_end_index"].item() - num_evicted_tokens
-            local_start_index = local_end_index - num_new_tokens
-
-            temp_k = kv_cache["k"].detach().clone()
-            temp_v = kv_cache["v"].detach().clone()
-            temp_k[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                temp_k[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-            temp_v[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                temp_v[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-
-            write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
-            roped_offset = max(0, write_start_index - local_start_index)
-            write_len = max(0, local_end_index - write_start_index)
-            if write_len > 0:
-                temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
-                temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
-
-            # Block-Relativistic RoPE: query uses window-relative indices
-            query_relative_indices = torch.arange(
-                self.local_attn_size - num_new_frames, self.local_attn_size, device=q.device)
-            roped_query = block_relativistic_rope(
-                q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)
-
-            # Block-Relativistic RoPE: cached K uses position-in-window indices
-            num_cache_frames = local_end_index // frame_seqlen
-            cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
-            cache_grid_sizes = grid_sizes.clone()
-            cache_grid_sizes[0, 0] = num_cache_frames
-            roped_temp_k = block_relativistic_rope(
-                temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
-                cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)
-
-            cache_update_info = {
-                "action": "roll_and_insert",
-                "sink_tokens": sink_tokens,
-                "num_rolled_tokens": num_rolled_tokens,
-                "num_evicted_tokens": num_evicted_tokens,
-                "local_start_index": local_start_index,
-                "local_end_index": local_end_index,
-                "write_start_index": write_start_index,
-                "write_end_index": local_end_index,
-                "new_k": k[:, roped_offset:roped_offset + write_len],
-                "new_v": v[:, roped_offset:roped_offset + write_len],
-                "current_end": current_end,
-                "is_recompute": is_recompute
-            }
-        else:
-            # === DIRECT INSERT MODE: cache not yet full ===
-            local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
-            local_start_index = local_end_index - num_new_tokens
-
-            temp_k = kv_cache["k"].detach().clone()
-            temp_v = kv_cache["v"].detach().clone()
-
-            write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
-            if sink_recache_after_switch:
-                write_start_index = local_start_index
-            roped_offset = max(0, write_start_index - local_start_index)
-            write_len = max(0, local_end_index - write_start_index)
-            if write_len > 0:
-                temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
-                temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
-
-            # RoPE with relative indices (growing sequentially before cache fills)
-            current_frame_in_window = local_start_index // frame_seqlen
-            query_relative_indices = torch.arange(
-                current_frame_in_window, current_frame_in_window + num_new_frames, device=q.device)
-            roped_query = block_relativistic_rope(
-                q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)
-
-            num_cache_frames = local_end_index // frame_seqlen
-            cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
-            cache_grid_sizes = grid_sizes.clone()
-            cache_grid_sizes[0, 0] = num_cache_frames
-            roped_temp_k = block_relativistic_rope(
-                temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
-                cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)
-
-            cache_update_info = {
-                "action": "direct_insert",
-                "local_start_index": local_start_index,
-                "local_end_index": local_end_index,
-                "write_start_index": write_start_index,
-                "write_end_index": local_end_index,
-                "new_k": k[:, roped_offset:roped_offset + write_len],
-                "new_v": v[:, roped_offset:roped_offset + write_len],
-                "current_end": current_end,
-                "is_recompute": is_recompute
-            }
-
-        # Attention: sink tokens + local window
-        if sink_tokens > 0:
-            local_budget = self.max_attention_size - sink_tokens
-            k_sink = roped_temp_k[:, :sink_tokens]
-            v_sink = temp_v[:, :sink_tokens]
-            if local_budget > 0:
-                local_start_for_window = max(sink_tokens, local_end_index - local_budget)
-                k_local = roped_temp_k[:, local_start_for_window:local_end_index]
-                v_local = temp_v[:, local_start_for_window:local_end_index]
-                k_cat = torch.cat([k_sink, k_local], dim=1)
-                v_cat = torch.cat([v_sink, v_local], dim=1)
-            else:
-                k_cat = k_sink
-                v_cat = v_sink
-            x = attention(roped_query, k_cat, v_cat)
-        else:
-            window_start = max(0, local_end_index - self.max_attention_size)
-            x = attention(
-                roped_query,
-                roped_temp_k[:, window_start:local_end_index],
-                temp_v[:, window_start:local_end_index])
-
-        x = x.flatten(2)
-        x, _ = self.o(x)
-        return x, (current_end, local_end_index, cache_update_info)
-
-
-class CausalPropeSelfAttention(nn.Module):
-    """PRoPE self-attention with optional KV cache for camera-controlled inference."""
-
-    def __init__(self, dim, attn_dim, num_heads, window_size=(-1, -1),
-                 local_attn_size=-1, sink_size=0, qk_norm=True, eps=1e-6):
-        assert dim % num_heads == 0
-        assert attn_dim % num_heads == 0
-        super().__init__()
-        self.dim = dim
-        self.attn_dim = attn_dim
-        self.num_heads = num_heads
-        self.head_dim = attn_dim // num_heads
-        self.local_attn_size = local_attn_size
-        self.sink_size = sink_size
-        self.qk_norm = qk_norm
-        self.eps = eps
-        self.window_size = window_size
-        self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
-
-        self.q_proj = ReplicatedLinear(dim, attn_dim)
-        self.k_proj = ReplicatedLinear(dim, attn_dim)
-        self.v_proj = ReplicatedLinear(dim, attn_dim)
-        self.out_proj = ReplicatedLinear(attn_dim, dim)
-
-        self.norm_q = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
-
-        nn.init.zeros_(self.out_proj.weight)
-        nn.init.zeros_(self.out_proj.bias)
-
-    def forward(self, x, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
-                kv_cache=None, current_start=0, cache_start=None,
-                sink_recache_after_switch=False, cache_update_policy="commit_detached"):
-        """
-        Args:
-            x: Shape [B, L, C]
-            cam_viewmats: Camera view matrices
-            cam_K: Camera intrinsics
-            kv_cache: Optional KV cache dict. When None, runs full attention over current chunk.
-        """
-        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
-        if cache_start is None:
-            cache_start = current_start
-
-        q, _ = self.q_proj(x)
-        q = self.norm_q(q).view(b, s, n, d)
-        k, _ = self.k_proj(x)
-        k = self.norm_k(k).view(b, s, n, d)
-        v, _ = self.v_proj(x)
-        v = v.view(b, s, n, d)
-
-        # Apply PRoPE (Positional Rotary Position Embedding from camera parameters)
-        q_t, k_t, v_t, apply_fn_o = prope_qkv(
-            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
-            viewmats=cam_viewmats, Ks=cam_K)
-        proped_q = q_t.transpose(1, 2)
-        proped_k = k_t.transpose(1, 2)
-        proped_v = v_t.transpose(1, 2)
-
-        if kv_cache is None:
-            # No cache: full attention over current chunk
-            x_out = attention(proped_q, proped_k, proped_v)
-        else:
-            # KV cache mode with rolling cache support
-            frame_seqlen = math.prod(grid_sizes[0][1:]).item()
-            num_new_tokens = s
-            current_end = current_start + num_new_tokens
-            sink_tokens = self.sink_size * frame_seqlen
-            kv_cache_size = kv_cache["k"].shape[1]
-            is_recompute = (current_end <= kv_cache["global_end_index"].item()) and (current_start > 0)
-
-            if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
-                    num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
-                # === ROLLING MODE ===
-                num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
-                num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
-                local_end_index = kv_cache["local_end_index"].item() + current_end - \
-                    kv_cache["global_end_index"].item() - num_evicted_tokens
-                local_start_index = local_end_index - num_new_tokens
-
-                if cache_update_policy != "none":
-                    with torch.no_grad():
-                        kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                        kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-
-                    write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
-                    roped_offset = max(0, write_start_index - local_start_index)
-                    write_len = max(0, local_end_index - write_start_index)
-                    if write_len > 0:
-                        with torch.no_grad():
-                            kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
-                            kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
-            else:
-                # === DIRECT INSERT MODE ===
-                local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
-                local_start_index = local_end_index - num_new_tokens
-
-                if cache_update_policy != "none":
-                    write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
-                    if sink_recache_after_switch:
-                        write_start_index = local_start_index
-                    roped_offset = max(0, write_start_index - local_start_index)
-                    write_len = max(0, local_end_index - write_start_index)
-                    if write_len > 0:
-                        with torch.no_grad():
-                            kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
-                            kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
-
-            # Attention: sink tokens + local window
-            if sink_tokens > 0:
-                local_budget = self.max_attention_size - sink_tokens
-                k_sink = kv_cache["k"][:, :sink_tokens].detach()
-                v_sink = kv_cache["v"][:, :sink_tokens].detach()
-                if local_budget > 0:
-                    local_start_for_window = max(sink_tokens, local_end_index - local_budget)
-                    k_local = kv_cache["k"][:, local_start_for_window:local_end_index].detach()
-                    v_local = kv_cache["v"][:, local_start_for_window:local_end_index].detach()
-                    k_cat = torch.cat([k_sink, k_local], dim=1)
-                    v_cat = torch.cat([v_sink, v_local], dim=1)
-                else:
-                    k_cat = k_sink
-                    v_cat = v_sink
-                x_out = attention(proped_q, k_cat, v_cat)
-            else:
-                window_start = max(0, local_end_index - self.max_attention_size)
-                x_out = attention(
-                    proped_q,
-                    kv_cache["k"][:, window_start:local_end_index].detach(),
-                    kv_cache["v"][:, window_start:local_end_index].detach())
-
-            if not is_recompute and cache_update_policy != "none":
-                kv_cache["global_end_index"].fill_(current_end)
-                kv_cache["local_end_index"].fill_(local_end_index)
-
-        # Apply inverse PRoPE
-        x = apply_fn_o(x_out.transpose(1, 2)).transpose(1, 2)
-        x = x.flatten(2)
-        x, _ = self.out_proj(x)
-        return x
-
-
-class CausalWanAttentionBlock(nn.Module):
-
-    def __init__(self, dim, ffn_dim, num_heads, local_attn_size=-1, sink_size=0,
-                 qk_norm=True, cross_attn_norm=False, eps=1e-6, **kwargs):
-        super().__init__()
-        self.dim = dim
-        self.ffn_dim = ffn_dim
-        self.num_heads = num_heads
-        self.local_attn_size = local_attn_size
-        self.qk_norm = qk_norm
-        self.cross_attn_norm = cross_attn_norm
-        self.eps = eps
-
-        self.add_control_adapter = kwargs.get('add_control_adapter', False)
-        self.cam_method = kwargs.get('cam_method')
-        self.attn_compress = kwargs.get('attn_compress', 1)
-        self.layer_idx = kwargs.get('layer_idx')
-        cam_self_attn_layers = kwargs.get('cam_self_attn_layers')
-
-        # layers
-        self.norm1 = WanLayerNorm(dim, eps)
-        self.self_attn = CausalWanSelfAttention(
-            dim, num_heads, local_attn_size, sink_size, qk_norm, eps)
-        self.norm3 = WanLayerNorm(
-            dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
-        self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
-        self.norm2 = WanLayerNorm(dim, eps)
-        # nn.Linear (not ReplicatedLinear) on purpose: the official checkpoint
-        # stores these as positional Sequential keys (ffn.0 / ffn.2) that the
-        # copy-only converter and the strict-load tests require verbatim, and
-        # ReplicatedLinear's (out, bias) tuple return cannot compose inside
-        # nn.Sequential without renaming the state-dict surface.
-        self.ffn = nn.Sequential(
-            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
-            nn.Linear(ffn_dim, dim))
-
-        # PRoPE self-attention branch for camera control
-        add_cam_attn = self.add_control_adapter and self.cam_method == 'prope'
-        if add_cam_attn and cam_self_attn_layers is not None:
-            add_cam_attn = self.layer_idx in cam_self_attn_layers
-        if add_cam_attn:
-            self.cam_self_attn = CausalPropeSelfAttention(
-                dim, dim // self.attn_compress, num_heads,
-                local_attn_size=local_attn_size, sink_size=sink_size,
-                qk_norm=qk_norm, eps=eps)
-
-        # modulation
-        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
-
-    def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens,
-                kv_cache, crossattn_cache=None, current_start=0, cache_start=None,
-                cam_viewmats=None, cam_K=None, sink_recache_after_switch=False,
-                cache_update_policy="commit_detached"):
-        num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
-        e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
-
-        # self-attention
-        attn_input = (self.norm1(x).unflatten(
-            dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(1, 2)
-        y, cache_update_info = self.self_attn(
-            attn_input, seq_lens, grid_sizes, freqs, kv_cache,
-            current_start, cache_start, sink_recache_after_switch)
-
-        # PRoPE camera attention (parallel branch)
-        if hasattr(self, 'cam_self_attn') and cam_viewmats is not None and cam_K is not None:
-            prope_kv_cache = None
-            if kv_cache is not None and "prope_k" in kv_cache:
-                prope_kv_cache = {
-                    "k": kv_cache["prope_k"],
-                    "v": kv_cache["prope_v"],
-                    "global_end_index": kv_cache["prope_global_end_index"],
-                    "local_end_index": kv_cache["prope_local_end_index"],
-                }
-            y = y + self.cam_self_attn(
-                attn_input, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
-                kv_cache=prope_kv_cache, current_start=current_start,
-                cache_start=cache_start, cache_update_policy=cache_update_policy)
-
-        x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2)
-
-        # cross-attention & FFN
-        x = x + self.cross_attn(self.norm3(x), context, context_lens,
-                                crossattn_cache=crossattn_cache)
-        y = self.ffn(
-            (self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
-             * (1 + e[4]) + e[3]).flatten(1, 2))
-        x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[5]).flatten(1, 2)
-
-        return x, cache_update_info
-
-
-class CausalHead(nn.Module):
-
-    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
-        super().__init__()
-        self.dim = dim
-        self.out_dim = out_dim
-        self.patch_size = patch_size
-        self.eps = eps
-
-        out_dim = math.prod(patch_size) * out_dim
-        self.norm = WanLayerNorm(dim, eps)
-        self.head = ReplicatedLinear(dim, out_dim)
-        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
-
-    def forward(self, x, e):
-        num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
-        e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
-        x, _ = self.head(
-            self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
-            * (1 + e[1]) + e[0])
-        return x
-
-
-class DreamXWorldARTransformer3DModel(BaseDiT):
-    """DreamX-World-5B autoregressive causal transformer."""
-
-    _fsdp_shard_conditions = DreamXWorldARConfig()._fsdp_shard_conditions
-    _compile_conditions = DreamXWorldARConfig()._compile_conditions
-    _supported_attention_backends = DreamXWorldARConfig()._supported_attention_backends
-    param_names_mapping = DreamXWorldARConfig().param_names_mapping
-    reverse_param_names_mapping = DreamXWorldARConfig().reverse_param_names_mapping
-    lora_param_names_mapping = DreamXWorldARConfig().lora_param_names_mapping
-    _no_split_modules = ["CausalWanAttentionBlock"]
-
-    def __init__(self, config: DreamXWorldARConfig, hf_config: dict[str, Any]) -> None:
-        super().__init__(config=config, hf_config=hf_config)
-
-        model_type = config.model_type
-        patch_size = config.patch_size
-        text_len = config.text_len
-        in_dim = config.in_channels
-        dim = config.hidden_size
-        ffn_dim = config.ffn_dim
-        freq_dim = config.freq_dim
-        text_dim = config.text_dim
-        out_dim = config.out_channels
-        num_heads = config.num_attention_heads
-        num_layers = config.num_layers
-        local_attn_size = config.local_attn_size
-        sink_size = config.sink_size
-        qk_norm = bool(config.qk_norm)
-        cross_attn_norm = config.cross_attn_norm
-        eps = config.eps
-        add_control_adapter = config.add_control_adapter
-        cam_method = config.cam_method
-        attn_compress = config.attn_compress
-        cam_self_attn_layers = config.cam_self_attn_layers
-
-        assert model_type in ['t2v', 'i2v', 'ti2v']
-        self.model_type = model_type
-        self.patch_size = patch_size
-        self.text_len = text_len
-        self.in_dim = in_dim
-        self.dim = dim
-        self.ffn_dim = ffn_dim
-        self.freq_dim = freq_dim
-        self.text_dim = text_dim
-        self.out_dim = out_dim
-        self.num_heads = num_heads
-        self.num_layers = num_layers
-        self.local_attn_size = local_attn_size
-        self.qk_norm = qk_norm
-        self.cross_attn_norm = cross_attn_norm
-        self.eps = eps
-
-        # embeddings — nn.Linear inside nn.Sequential on purpose: the official
-        # checkpoint keys are positional (text_embedding.0/.2, time_embedding.0/.2,
-        # time_projection.1) and must load verbatim (see ffn comment above).
-        self.patch_embedding = nn.Conv3d(
-            in_dim, dim, kernel_size=patch_size, stride=patch_size)
-        self.text_embedding = nn.Sequential(
-            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
-            nn.Linear(dim, dim))
-        self.time_embedding = nn.Sequential(
-            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
-        self.time_projection = nn.Sequential(
-            nn.SiLU(), nn.Linear(dim, dim * 6))
-
-        # transformer blocks
-        self.blocks = nn.ModuleList([
-            CausalWanAttentionBlock(
-                dim, ffn_dim, num_heads, local_attn_size, sink_size,
-                qk_norm, cross_attn_norm, eps,
-                add_control_adapter=add_control_adapter,
-                cam_method=cam_method,
-                attn_compress=attn_compress,
-                layer_idx=layer_idx,
-                cam_self_attn_layers=cam_self_attn_layers)
-            for layer_idx in range(num_layers)
-        ])
-        for layer_idx, block in enumerate(self.blocks):
-            block.self_attn.layer_idx = layer_idx
-            block.self_attn.num_layers = self.num_layers
-
-        # head
-        self.head = CausalHead(dim, out_dim, patch_size, eps)
-
-        # RoPE frequencies
-        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
-        d = dim // num_heads
-        self.freqs = torch.cat([
-            rope_params(1024, d - 4 * (d // 6)),
-            rope_params(1024, 2 * (d // 6)),
-            rope_params(1024, 2 * (d // 6))
-        ], dim=1)
-
-        self.num_attention_heads = num_heads
-        self.attention_head_dim = dim // num_heads
-        self.hidden_size = dim
-        self.in_channels = in_dim
-        self.out_channels = out_dim
-        self.num_channels_latents = out_dim
-        self.init_weights()
-        self.num_frame_per_block = config.arch_config.num_frames_per_block
-        self.__post_init__()
-
-    def forward(self, x=None, t=None, context=None, seq_len=None, y=None, y_camera=None,
-                kv_cache=None, crossattn_cache=None, current_start=0,
-                cache_start=0, cache_update_policy="commit_detached",
-                hidden_states=None, encoder_hidden_states=None, timestep=None, **kwargs):
-        """
-        Causal inference with KV caching.
-        See Algorithm 2 of CausVid (https://arxiv.org/abs/2412.07772).
-
-        Args:
-            x: List of input video tensors [C_in, F, H, W]
-            t: Timestep tensor [B, L]
-            context: List of text embeddings [L, C]
-            seq_len: Maximum sequence length for positional encoding
-            y: Optional conditional video inputs (I2V mode)
-            y_camera: Camera parameters dict {'viewmats': ..., 'K': ...}
-            kv_cache: List of KV cache dicts per transformer block
-            crossattn_cache: List of cross-attention cache dicts
-            current_start: Current position in global token sequence
-            cache_start: Cache start position
-            cache_update_policy: Cache update strategy ('commit_detached' or 'none')
-
-        Returns:
-            Stacked output tensors [B, C_out, F, H/8, W/8]
-        """
-        if x is None and hidden_states is not None:
-            x = [sample for sample in hidden_states]
-        if t is None and timestep is not None:
-            t = timestep
-        if context is None and encoder_hidden_states is not None:
-            if isinstance(encoder_hidden_states, torch.Tensor):
-                context = [sample for sample in encoder_hidden_states]
-            else:
-                context = encoder_hidden_states
-        if seq_len is None:
-            if torch.is_tensor(t):
-                seq_len = int(t.shape[1]) if t.dim() > 1 else int(t.numel())
-            elif x is not None:
-                sample = x[0]
-                seq_len = (sample.shape[1] // self.patch_size[0]) * (sample.shape[2] // self.patch_size[1]) * (sample.shape[3] // self.patch_size[2])
-        if x is None or t is None or context is None or seq_len is None:
-            raise ValueError("DreamXWorldARTransformer3DModel requires x/t/context/seq_len or FastVideo aliases")
-
-        device = self.patch_embedding.weight.device
-        if self.freqs.is_meta or self.freqs.device != device:
-            d = self.dim // self.num_heads
-            self.freqs = torch.cat([
-                rope_params(1024, d - 4 * (d // 6)),
-                rope_params(1024, 2 * (d // 6)),
-                rope_params(1024, 2 * (d // 6)),
-            ], dim=1).to(device)
-
-        if y is not None:
-            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)]
-
-        # patch embedding
-        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
-        grid_sizes = torch.stack(
-            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
-        x = [u.flatten(2).transpose(1, 2) for u in x]
-        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
-        assert seq_lens.max() <= seq_len
-        x = torch.cat(x)
-
-        # time embedding
-        e = self.time_embedding(
-            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
-        e0 = self.time_projection(e).unflatten(
-            1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
-
-        # text embedding
-        context_lens = None
-        context = self.text_embedding(
-            torch.stack([
-                torch.cat(
-                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
-                for u in context
-            ]))
-
-        # camera parameters
-        if y_camera is not None and isinstance(y_camera, dict):
-            cam_viewmats = y_camera['viewmats']
-            cam_K = y_camera['K']
-        else:
-            cam_viewmats = None
-            cam_K = None
-
-        block_kwargs = dict(
-            e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=self.freqs,
-            context=context, context_lens=context_lens,
-            cam_viewmats=cam_viewmats, cam_K=cam_K,
-            cache_update_policy=cache_update_policy,
-        )
-
-        cache_update_infos = []
-        for block_index, block in enumerate(self.blocks):
-            block_kwargs.update({
-                "kv_cache": kv_cache[block_index] if kv_cache is not None else None,
-                "crossattn_cache": crossattn_cache[block_index] if crossattn_cache is not None else None,
-                "current_start": current_start,
-                "cache_start": cache_start,
-            })
-            x, block_cache_update_info = block(x, **block_kwargs)
-            if kv_cache is not None:
-                cache_update_infos.append((block_index, block_cache_update_info))
-
-        # Apply deferred cache updates
-        if kv_cache is not None and cache_update_infos and cache_update_policy != "none":
-            self._apply_cache_updates(kv_cache, cache_update_infos)
-
-        # head & unpatchify
-        x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
-        x = self.unpatchify(x, grid_sizes)
-        return torch.stack(x)
-
-    def _apply_cache_updates(self, kv_cache, cache_update_infos):
-        """Apply deferred cache updates collected from all transformer blocks.
-
-        For Block-Relativistic RoPE, this stores un-roped K values in the cache.
-        RoPE is applied dynamically during attention based on each token's current
-        relative position in the sliding window.
-        """
-        with torch.no_grad():
-            for block_index, (current_end, local_end_index, update_info) in cache_update_infos:
-                if update_info is not None:
-                    cache = kv_cache[block_index]
-
-                    if update_info["action"] == "roll_and_insert":
-                        sink_tokens = update_info["sink_tokens"]
-                        num_rolled_tokens = update_info["num_rolled_tokens"]
-                        num_evicted_tokens = update_info["num_evicted_tokens"]
-                        write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
-                        write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
-                        new_k = update_info["new_k"].detach()
-                        new_v = update_info["new_v"].detach()
-
-                        cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                        cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-
-                        if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index):
-                            cache["k"][:, write_start_index:write_end_index] = new_k
-                            cache["v"][:, write_start_index:write_end_index] = new_v
-
-                    elif update_info["action"] == "direct_insert":
-                        write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
-                        write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
-                        new_k = update_info["new_k"].detach()
-                        new_v = update_info["new_v"].detach()
-
-                        if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index):
-                            cache["k"][:, write_start_index:write_end_index] = new_k
-                            cache["v"][:, write_start_index:write_end_index] = new_v
-
-                is_recompute = False if update_info is None else update_info.get("is_recompute", False)
-                if not is_recompute:
-                    kv_cache[block_index]["global_end_index"].fill_(current_end)
-                    kv_cache[block_index]["local_end_index"].fill_(local_end_index)
-
-    def unpatchify(self, x, grid_sizes):
-        """Reconstruct video tensors from patch embeddings."""
-        c = self.out_dim
-        out = []
-        for u, v in zip(x, grid_sizes.tolist(), strict=True):
-            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
-            u = torch.einsum('fhwpqrc->cfphqwr', u)
-            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)])
-            out.append(u)
-        return out
-
-    def init_weights(self):
-        """Initialize model parameters using Xavier initialization."""
-        for m in self.modules():
-            if isinstance(m, (nn.Linear, ReplicatedLinear)):
-                nn.init.xavier_uniform_(m.weight)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-
-        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
-        for m in self.text_embedding.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, std=.02)
-        for m in self.time_embedding.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, std=.02)
-
-        nn.init.zeros_(self.head.head.weight)
-
-
-EntryClass = DreamXWorldARTransformer3DModel

fastvideo/models/dits/hunyuanvideo15.py

@@ -699,8 +699,8 @@ class IndividualTokenRefinerBlock(nn.Module):
         if mask is None:
             attn_output = self.attn(q, k, v) 
         else:
-            from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder
-            attn_metadata = SDPAMetadataBuilder().build(
+            from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+            attn_metadata = FlashAttnMetadataBuilder().build(
                 current_timestep=0,
                 attn_mask=mask,
             )

fastvideo/models/dits/hyworld/hyworld.py

@@ -173,11 +173,8 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
         img_v_prope = img_v_prope.permute(0, 2, 1, 3)  # [batch, seqlen, num_heads, head_dim]
         # end hyworld
 
-        # The metadata only carries the text padding mask; the executing
-        # attention kernel is still whatever the layer's selector picked
-        # (flash-attn when installed), not necessarily torch SDPA.
-        from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder
-        attn_metadata = SDPAMetadataBuilder().build(
+        from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+        attn_metadata = FlashAttnMetadataBuilder().build(
             current_timestep=0,
             attn_mask=encoder_attention_mask,
         )
@@ -195,9 +192,14 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
             )
         
         # begin hyworld
-        # attention with prope (same text mask, so reuse the metadata)
+        # attention with prope
+        from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+        attn_metadata_prope = FlashAttnMetadataBuilder().build(
+            current_timestep=0,
+            attn_mask=encoder_attention_mask,
+        )       
         # NOTE: Do NOT pass freqs_cis to prope attention - HY-WorldPlay does not apply RoPE to prope
-        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
+        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata_prope):
             img_attn_prope, _ = self.attn(
                 img_q_prope,
                 img_k_prope,

fastvideo/models/dits/matrixgame2/action_module.py

@@ -16,8 +16,6 @@ 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(
@@ -102,7 +100,6 @@ 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,
@@ -136,6 +133,7 @@ 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]
@@ -208,22 +206,16 @@ 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[
-            : kv_cache["k"].shape[0]
-        ]
-        kv_cache["v"][:, local_start_index:local_end_index] = v[
-            : kv_cache["v"].shape[0]
-        ]
+        kv_cache["k"][:, local_start_index:local_end_index] = k[:1]
+        kv_cache["v"][:, local_start_index:local_end_index] = v[:1]
     else:
         kv_cache["k"][:, local_start_index:local_end_index] = k
         kv_cache["v"][:, local_start_index:local_end_index] = v
 
-    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,
-    )
+    # 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]
 
     if repeat_factor is not None:
         cached_k = cached_k.repeat(repeat_factor, 1, 1, 1)
@@ -269,7 +261,6 @@ class ActionModule(nn.Module):
         enable_mouse=True,
         enable_keyboard=True,
         local_attn_size=6,
-        sink_size=0,
         blocks: list | None = None,
     ):
         super().__init__()
@@ -283,7 +274,6 @@ 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
 
@@ -575,7 +565,6 @@ class ActionModule(nn.Module):
                     start_frame,
                     num_frame_per_block,
                     self.local_attn_size,
-                    self.sink_size,
                     use_k_for_num_tokens=False,
                 )
         else:
@@ -700,7 +689,6 @@ 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,
@@ -731,7 +719,6 @@ class ActionModule(nn.Module):
                         start_frame,
                         num_frame_per_block,
                         self.local_attn_size,
-                        self.sink_size,
                         use_k_for_num_tokens=True,
                     )
             else:

fastvideo/models/dits/matrixgame2/causal_model.py

@@ -46,7 +46,6 @@ 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__)
 
@@ -346,12 +345,9 @@ 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
 
-            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,
-            )
+            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]
 
             if relativistic:
                 window_len, query_lo, _ = relativistic_window_offsets(
@@ -422,7 +418,6 @@ 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)
@@ -479,7 +474,6 @@ 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[
@@ -653,8 +647,6 @@ 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(
@@ -670,12 +662,6 @@ 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"))
@@ -763,52 +749,6 @@ 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,
@@ -816,7 +756,6 @@ 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
@@ -843,10 +782,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
             else:
                 return (
                     (kv_idx < ends[q_idx])
-                    & (
-                        (kv_idx < sink_size * frame_seqlen)
-                        | (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
-                    )
+                    & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
                 ) | (q_idx == kv_idx)
 
         block_mask = create_block_mask(
@@ -860,9 +796,8 @@ class CausalMatrixGame2WanModel(BaseDiT):
         )
 
         if not dist.is_initialized() or dist.get_rank() == 0:
-            logger.info(
-                "cache a block wise causal mask with block size of %s frames",
-                num_frame_per_block,
+            print(
+                f" cache a block wise causal mask with block size of {num_frame_per_block} frames"
             )
 
         return block_mask
@@ -874,7 +809,6 @@ 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
@@ -900,7 +834,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
             else:
                 return (
                     (kv_idx < ends2[q_idx])
-                    & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))
+                    & (kv_idx >= (ends2[q_idx] - local_attn_size))
                 ) | (q_idx == kv_idx)
 
         block_mask2 = create_block_mask(
@@ -914,9 +848,8 @@ class CausalMatrixGame2WanModel(BaseDiT):
         )
 
         if not dist.is_initialized() or dist.get_rank() == 0:
-            logger.info(
-                "cache a block wise causal mask for keyboard with block size of %s frames",
-                num_frame_per_block,
+            print(
+                f" cache a block wise causal mask for keyboard with block size of {num_frame_per_block} frames"
             )
 
         return block_mask2
@@ -928,7 +861,6 @@ 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
@@ -954,7 +886,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
             else:
                 return (
                     (kv_idx < ends2[q_idx])
-                    & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))
+                    & (kv_idx >= (ends2[q_idx] - local_attn_size))
                 ) | (q_idx == kv_idx)
 
         block_mask2 = create_block_mask(
@@ -968,9 +900,8 @@ class CausalMatrixGame2WanModel(BaseDiT):
         )
 
         if not dist.is_initialized() or dist.get_rank() == 0:
-            logger.info(
-                "cache a block wise causal mask for action with block size of %s frames",
-                num_frame_per_block,
+            print(
+                f" cache a block wise causal mask for action with block size of {num_frame_per_block} frames"
             )
 
         return block_mask2
@@ -1132,7 +1063,6 @@ 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(
@@ -1141,7 +1071,6 @@ 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(
@@ -1150,7 +1079,6 @@ 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,
@@ -1158,7 +1086,6 @@ 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)
@@ -1312,7 +1239,6 @@ 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:
@@ -1322,7 +1248,6 @@ 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(
@@ -1331,7 +1256,6 @@ 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(
@@ -1340,7 +1264,6 @@ 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)

fastvideo/models/dits/matrixgame2/model.py

@@ -452,7 +452,7 @@ class MatrixGame2WanModel(BaseDiT):
             and len(encoder_hidden_states_image) > 0
         ):
             encoder_hidden_states_image = encoder_hidden_states_image[0]
-        elif not isinstance(encoder_hidden_states_image, torch.Tensor):
+        else:
             encoder_hidden_states_image = None
 
         batch_size, num_channels, num_frames, height, width = (

fastvideo/models/dits/matrixgame2/utils.py

@@ -4,18 +4,12 @@ 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 = {
@@ -40,256 +34,6 @@ 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 = {}

fastvideo/models/dits/wanvideo.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import copy
 import math
 from typing import Any
 
@@ -60,11 +59,6 @@ class WanTimeTextImageEmbedding(nn.Module):
         time_freq_dim: int,
         text_embed_dim: int,
         image_embed_dim: int | None = None,
-        *,
-        r_embedder: bool = False,
-        r_embedder_fusion: str = "additive",
-        r_embedder_gate_value: float = 0.25,
-        r_embedder_deltatime_type: str = "r",
     ):
         super().__init__()
 
@@ -83,57 +77,14 @@ class WanTimeTextImageEmbedding(nn.Module):
         if image_embed_dim is not None:
             self.image_embedder = WanImageEmbedding(image_embed_dim, dim)
 
-        # AnyFlow dual-timestep support. When r_embedder is False the forward
-        # path bypasses delta_embedder entirely and the output is byte-identical
-        # to the legacy single-timestep implementation.
-        self._r_embedder_enabled = bool(r_embedder)
-        self._r_embedder_fusion = r_embedder_fusion
-        self._r_embedder_deltatime_type = r_embedder_deltatime_type
-        if self._r_embedder_enabled:
-            if r_embedder_fusion not in ("additive", "gated"):
-                raise ValueError(
-                    "r_embedder_fusion must be one of {additive, gated}, "
-                    f"got {r_embedder_fusion!r}")
-            if r_embedder_deltatime_type not in ("r", "t-r"):
-                raise ValueError(
-                    "r_embedder_deltatime_type must be one of {r, t-r}, "
-                    f"got {r_embedder_deltatime_type!r}")
-            # Deep-copy preserves identical initialization with time_embedder,
-            # matching AnyFlow reference setup_flowmap_model() behavior.
-            self.delta_embedder = copy.deepcopy(self.time_embedder)
-            # Non-persistent buffer — gate is a hyperparameter, not learned.
-            self.register_buffer(
-                "_r_embedder_gate",
-                torch.tensor(float(r_embedder_gate_value)),
-                persistent=False,
-            )
-        else:
-            self.delta_embedder = None
-
     def forward(
         self,
         timestep: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         encoder_hidden_states_image: torch.Tensor | None = None,
         timestep_seq_len: int | None = None,
-        r_timestep: torch.Tensor | None = None,
     ):
         temb = self.time_embedder(timestep, timestep_seq_len)
-
-        if self._r_embedder_enabled and r_timestep is not None:
-            assert self.delta_embedder is not None
-            if self._r_embedder_deltatime_type == "r":
-                delta_input = r_timestep
-            else:
-                delta_input = timestep - r_timestep
-            delta_emb = self.delta_embedder(delta_input, timestep_seq_len)
-            gate = self._r_embedder_gate
-            if self._r_embedder_fusion == "gated":
-                temb = (1.0 - gate) * temb + gate * delta_emb
-            else:
-                # Additive (additional channel, no convex blend).
-                temb = temb + gate * delta_emb
-
         timestep_proj = self.time_modulation(temb)
 
         if self.text_embedder is not None:
@@ -644,10 +595,6 @@ class WanTransformer3DModel(BaseDiT):
             time_freq_dim=config.freq_dim,
             text_embed_dim=config.text_dim,
             image_embed_dim=config.image_dim,
-            r_embedder=config.r_embedder,
-            r_embedder_fusion=config.r_embedder_fusion,
-            r_embedder_gate_value=config.r_embedder_gate_value,
-            r_embedder_deltatime_type=config.r_embedder_deltatime_type,
         )
 
         # 3. Transformer blocks
@@ -689,7 +636,6 @@ class WanTransformer3DModel(BaseDiT):
                 encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                 | None = None,
                 guidance=None,
-                r_timestep: torch.Tensor | None = None,
                 **kwargs) -> torch.Tensor:
         orig_dtype = hidden_states.dtype
         if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
@@ -738,17 +684,8 @@ class WanTransformer3DModel(BaseDiT):
         else:
             ts_seq_len = None
 
-        # AnyFlow dual-timestep — match timestep's flattening so embedder
-        # sees aligned shapes.
-        if r_timestep is not None and r_timestep.dim() == 2:
-            r_timestep = r_timestep.flatten()
-
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep,
-            encoder_hidden_states,
-            encoder_hidden_states_image,
-            timestep_seq_len=ts_seq_len,
-            r_timestep=r_timestep)
+            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
         if ts_seq_len is not None:
             # batch_size, seq_len, 6, inner_dim
             timestep_proj = timestep_proj.unflatten(2, (6, -1))

fastvideo/models/registry.py

@@ -11,7 +11,7 @@ import tempfile
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Set
 from dataclasses import dataclass, field
-from functools import cache, lru_cache
+from functools import lru_cache
 from typing import NoReturn, TypeVar, cast
 
 import cloudpickle
@@ -32,8 +32,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "HYWorldTransformer3DModel":
     ("dits", "hyworld", "HYWorldTransformer3DModel"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
-    "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
-    "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
     "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
@@ -50,8 +48,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
 _IMAGE_TO_VIDEO_DIT_MODELS = {
     # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
-    "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
-    "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "MatrixGame2WanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"),
     "CausalMatrixGame2WanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"),
@@ -145,7 +141,7 @@ _LEGACY_FAST_VIDEO_MODELS = {
 MODELS_PATH = os.path.dirname(__file__)
 
 
-@cache
+@lru_cache(maxsize=None)
 def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
     discovered_models: dict[str, tuple[str, str, str]] = {}
     for root, dirs, files in os.walk(MODELS_PATH):
@@ -160,7 +156,7 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
 
             filepath = os.path.join(root, filename)
             try:
-                with open(filepath, encoding="utf-8") as f:
+                with open(filepath, "r", encoding="utf-8") as f:
                     source = f.read()
                 tree = ast.parse(source, filename=filename)
 

fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py

@@ -1,202 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Flow-map any-step Euler scheduler for AnyFlow.
-
-The model predicts the *average* velocity ``u_θ(x_t, t, r)`` from time
-``t`` back to time ``r``, so one Euler step is
-
-    x_r = x_t - ((t - r) / num_train_timesteps) * u_θ(x_t, t, r)
-
-regardless of how far apart ``t`` and ``r`` are. The scheduler also
-provides the AnyFlow training-time helpers ``apply_shift`` (flow-matching
-shift transform) and ``get_train_weight`` (per-timestep loss weight,
-including ``beta08``).
-
-Standalone — does not depend on diffusers' ConfigMixin/SchedulerMixin.
-"""
-
-from __future__ import annotations
-
-from typing import Literal
-
-import torch
-
-from fastvideo.models.schedulers.base import BaseScheduler
-
-
-WeightType = Literal["uniform", "gaussian", "beta08"]
-
-
-class FlowMapEulerDiscreteScheduler(BaseScheduler):
-    """Minimal flow-map scheduler.
-
-    Parameters
-    ----------
-    num_train_timesteps:
-        Discretization granularity for training. ``t`` is expressed in
-        absolute units in ``[0, num_train_timesteps]``.
-    shift:
-        Flow-matching shift parameter (Wan video default: ``5.0``). Set
-        to ``1.0`` for an identity shift.
-    """
-
-    order: int = 1
-
-    def __init__(
-        self,
-        *,
-        num_train_timesteps: int = 1000,
-        shift: float = 1.0,
-    ) -> None:
-        self.num_train_timesteps = int(num_train_timesteps)
-        self.shift = float(shift)
-        self.timesteps: torch.Tensor = torch.empty(0)
-        self.sigmas: torch.Tensor = torch.empty(0)
-        super().__init__()
-
-    # ------------------------------------------------------------------
-    # BaseScheduler abstract surface
-
-    def set_shift(self, shift: float) -> None:
-        self.shift = float(shift)
-
-    def scale_model_input(
-        self,
-        sample: torch.Tensor,
-        timestep: int | None = None,
-    ) -> torch.Tensor:
-        # Flow-matching has no per-step input scaling; pass through.
-        del timestep
-        return sample
-
-    # ------------------------------------------------------------------
-    # Public helpers used by the AnyFlow pretrain method.
-
-    def apply_shift(
-        self,
-        t: torch.Tensor,
-        *,
-        shift: float | None = None,
-    ) -> torch.Tensor:
-        """Apply the flow-matching shift: ``t' = s * t / (1 + (s - 1) * t)``.
-
-        Operates in the normalized ``[0, 1]`` domain — callers should pass
-        ``t / num_train_timesteps`` (or sample ``t`` directly from
-        ``[0, 1]``).
-        """
-        s = self.shift if shift is None else float(shift)
-        if s == 1.0:
-            return t
-        return s * t / (1.0 + (s - 1.0) * t)
-
-    def get_train_weight(
-        self,
-        t: torch.Tensor,
-        *,
-        weight_type: WeightType = "beta08",
-    ) -> torch.Tensor:
-        """Per-timestep training weight, renormalized so the total weight
-        mass equals ``num_train_timesteps`` (matching AnyFlow reference's
-        ``scheduling_flowmap_euler_discrete.py``).
-
-        ``beta08``: ``w(t) = t * sqrt(1 - t)`` (in normalized t-space).
-        """
-        # Auto-detect domain: if t was given in absolute units, normalize.
-        t_f = t.float()
-        max_val = t_f.max() if t_f.numel() > 0 else torch.tensor(0.0)
-        if max_val > 1.0 + 1e-6:
-            t_norm = t_f / self.num_train_timesteps
-        else:
-            t_norm = t_f
-        t_norm = t_norm.clamp(min=0.0, max=1.0)
-
-        if weight_type == "uniform":
-            w = torch.ones_like(t_norm)
-        elif weight_type == "gaussian":
-            w = torch.exp(-0.5 * ((t_norm - 0.5) / 0.2) ** 2)
-        elif weight_type == "beta08":
-            w = t_norm.pow(1.0) * (1.0 - t_norm).clamp_min(0.0).pow(0.5)
-        else:
-            raise ValueError(f"Unknown weight_type: {weight_type!r}")
-
-        denom = w.sum().clamp_min(1e-8)
-        return w * (float(self.num_train_timesteps) / denom)
-
-    # ------------------------------------------------------------------
-
-    def set_timesteps(
-        self,
-        *,
-        num_inference_steps: int,
-        device: torch.device | str = "cpu",
-        custom_timesteps: list[float] | torch.Tensor | None = None,
-    ) -> None:
-        """Build a descending timestep schedule ending at 0.
-
-        With ``num_inference_steps=N`` the schedule has ``N + 1`` entries
-        ``[T_max, ..., 0]`` so a rollout consumes ``N`` Euler steps.
-
-        ``custom_timesteps`` overrides the linspace+shift schedule with
-        a pinned list (in absolute train-timestep units), useful for the
-        AnyFlow paper's hand-tuned ``[999, 937, 833, 624, 0]`` schedule.
-        """
-        if num_inference_steps <= 0:
-            raise ValueError(
-                "num_inference_steps must be positive, "
-                f"got {num_inference_steps}")
-        device = torch.device(device)
-
-        if custom_timesteps is not None:
-            ts = torch.as_tensor(
-                custom_timesteps, dtype=torch.float32, device=device)
-            if ts.ndim != 1:
-                raise ValueError(
-                    "custom_timesteps must be 1-D, got shape "
-                    f"{tuple(ts.shape)}")
-            if not torch.all(ts[:-1] >= ts[1:]):
-                raise ValueError(
-                    "custom_timesteps must be descending (largest first)")
-        else:
-            ts_norm = torch.linspace(
-                1.0, 0.0, num_inference_steps + 1, device=device)
-            ts_norm = self.apply_shift(ts_norm)
-            ts = ts_norm * self.num_train_timesteps
-
-        self.timesteps = ts
-        self.sigmas = ts / self.num_train_timesteps
-
-    def step(
-        self,
-        model_output: torch.Tensor,
-        *,
-        sample: torch.Tensor,
-        timestep: torch.Tensor,
-        r_timestep: torch.Tensor,
-    ) -> torch.Tensor:
-        """One Euler step from ``t`` to ``r``.
-
-        ``model_output`` is the average-velocity prediction
-        ``u_θ(x_t, t, r)``. Both ``timestep`` and ``r_timestep`` are in
-        absolute train-timestep units (``[0, num_train_timesteps]``).
-        """
-        t = timestep.to(sample.device, dtype=sample.dtype)
-        r = r_timestep.to(sample.device, dtype=sample.dtype)
-        dt_norm = (t - r) / float(self.num_train_timesteps)
-        # Broadcast dt over channel/spatial dims.
-        view: list[int] = [-1] + [1] * (sample.ndim - 1)
-        return sample - dt_norm.view(*view) * model_output
-
-    def add_noise(
-        self,
-        original_samples: torch.Tensor,
-        noise: torch.Tensor,
-        timestep: torch.Tensor,
-    ) -> torch.Tensor:
-        """Linear flow-matching interpolation: ``x_t = (1 - σ) * x_0 + σ * ε``,
-        where ``σ = t / num_train_timesteps``.
-        """
-        sigma = (timestep.to(original_samples.device,
-                             dtype=original_samples.dtype)
-                 / float(self.num_train_timesteps))
-        view: list[int] = [-1] + [1] * (original_samples.ndim - 1)
-        sigma = sigma.view(*view)
-        return (1.0 - sigma) * original_samples + sigma * noise

fastvideo/models/vision_utils.py

@@ -1,9 +1,7 @@
 # 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
@@ -18,10 +16,6 @@ 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 = {
@@ -95,39 +89,6 @@ 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,
@@ -149,7 +110,7 @@ def load_image(
     """
     if isinstance(image, str):
         if image.startswith("http://") or image.startswith("https://"):
-            image = PIL.Image.open(io.BytesIO(_fetch_image_bytes(image)))
+            image = PIL.Image.open(requests.get(image, stream=True).raw)
         elif os.path.isfile(image):
             image = PIL.Image.open(image)
         else:

fastvideo/performance/hf_store.py

@@ -268,31 +268,16 @@ def load_records_for_model(
     model_id: str,
     gpu_type: str | None = None,
     *,
-    workload_id: str | None = None,
-    variant_id: str | None = None,
-    benchmark_version: str | None = None,
-    recipe_fingerprint: str | None = None,
-    hardware_profile_id: str | None = None,
-    software_profile_id: str | None = None,
     last_n: int | None = None,
     successful_only: bool = True,
     baseline_eligible_only: bool = False,
 ) -> list[dict[str, Any]]:
-    """Return records for a specific *model_id*, optionally filtered by cohort.
+    """Return records for a specific *model_id*, optionally filtered by GPU.
 
     Args:
         local_dir: Root directory previously populated by :func:`sync_from_hf`.
         model_id: Matches the ``model_id`` field inside each JSON record.
         gpu_type: When set, only records whose ``gpu_type`` matches are returned.
-        workload_id: When set, only records from the same workload are returned.
-        variant_id: When set, only records from the same workload variant are returned.
-        benchmark_version: When set, only records from the same benchmark version are returned.
-        recipe_fingerprint: When set, only records from the same benchmark
-            recipe are returned.
-        hardware_profile_id: When set, only records from the same hardware
-            cohort are returned.
-        software_profile_id: When set, only records from the same software
-            cohort are returned.
         last_n: When set, return only the most recent *n* records (after all
             other filters). Useful for sliding-window baseline calculations.
         successful_only: Passed through to :func:`load_records`.
@@ -314,18 +299,6 @@ def load_records_for_model(
     if gpu_type is not None:
         records = [r for r in records if r.get("gpu_type") == gpu_type]
 
-    identity_filters = {
-        "workload_id": workload_id,
-        "variant_id": variant_id,
-        "benchmark_version": benchmark_version,
-        "recipe_fingerprint": recipe_fingerprint,
-        "hardware_profile_id": hardware_profile_id,
-        "software_profile_id": software_profile_id,
-    }
-    for key, expected in identity_filters.items():
-        if expected is not None:
-            records = [r for r in records if str(r.get(key)) == str(expected)]
-
     if last_n is not None:
         records = records[-last_n:]
 

fastvideo/performance_dashboard/service.py

@@ -17,15 +17,6 @@ from fastvideo.performance.hf_store import is_baseline_eligible_record, safe_flo
 from fastvideo.performance.metric_policy import regression_delta, resolve_metric_policies
 
 Record = dict[str, Any]
-CohortKey = tuple[str, str, str, str, str, str, str, str]
-COMPARISON_COHORT_KEYS = (
-    "workload_id",
-    "variant_id",
-    "benchmark_version",
-    "recipe_fingerprint",
-    "hardware_profile_id",
-    "software_profile_id",
-)
 
 
 def parse_timestamp(value: Any) -> datetime | None:
@@ -86,56 +77,15 @@ 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: _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"),
-        _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")),
-    )
-
-
-def group_by_comparison_cohort(records: list[Record]) -> dict[CohortKey, list[Record]]:
-    groups: dict[CohortKey, list[Record]] = defaultdict(list)
+def group_by_model_gpu(records: list[Record]) -> dict[tuple[str, str], list[Record]]:
+    groups: dict[tuple[str, str], list[Record]] = defaultdict(list)
     for record in records:
-        groups[comparison_cohort_key(record)].append(record)
+        model_id = str(record.get("model_id") or "unknown")
+        gpu_type = str(record.get("gpu_type") or "unknown")
+        groups[(model_id, gpu_type)].append(record)
     return {key: sorted(value, key=record_sort_key) for key, value in groups.items()}
 
 
-def comparison_sort_key(record: Record) -> CohortKey:
-    return comparison_cohort_key(record)
-
-
-def latest_row_sort_key(row: Record) -> tuple[Any, ...]:
-    return (row["status"] != "fail", *comparison_sort_key(row))
-
-
-def group_identity(key: CohortKey) -> tuple[str, str]:
-    model_id = key[0]
-    gpu_type = key[1]
-    return model_id, gpu_type
-
-
 def baseline_value(records: list[Record], metric_key: str) -> float | None:
     values = [safe_float(record.get(metric_key)) for record in records]
     values = [value for value in values if value is not None]
@@ -149,8 +99,7 @@ def build_latest_summary(records: list[Record],
                          baseline_window: int = 5,
                          run_source: str | None = None) -> list[Record]:
     rows: list[Record] = []
-    for key, group in group_by_comparison_cohort(records).items():
-        model_id, gpu_type = group_identity(key)
+    for (model_id, gpu_type), group in group_by_model_gpu(records).items():
         latest_candidates = group
         if run_source:
             latest_candidates = [record for record in group if record_run_source(record) == run_source]
@@ -158,10 +107,9 @@ 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[:latest_index]
-            if record.get("success", True) and is_baseline_eligible_record(record)
+            record for record in group
+            if record is not latest and record.get("success", True) and is_baseline_eligible_record(record)
         ]
         baseline_records = baseline_pool[-baseline_window:]
         metric_policies = resolve_metric_policies(latest.get("regression_thresholds"))
@@ -208,7 +156,6 @@ def build_latest_summary(records: list[Record],
             "timestamp": latest.get("timestamp"),
             "commit_sha": latest.get("commit_sha"),
             **record_metadata(latest),
-            **record_comparison_metadata(latest),
             "success": success,
             "baseline_n": len(baseline_records),
             "worst_regression_pct": worst_regression,
@@ -219,13 +166,12 @@ def build_latest_summary(records: list[Record],
             "metrics": metrics,
         })
 
-    return sorted(rows, key=latest_row_sort_key)
+    return sorted(rows, key=lambda row: (row["status"] != "fail", row["model_id"], row["gpu_type"]))
 
 
 def build_trends(records: list[Record]) -> list[Record]:
     trends: list[Record] = []
-    for key, group in group_by_comparison_cohort(records).items():
-        model_id, gpu_type = group_identity(key)
+    for (model_id, gpu_type), group in group_by_model_gpu(records).items():
         points = []
         for record in group:
             metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
@@ -233,7 +179,6 @@ def build_trends(records: list[Record]) -> list[Record]:
                 "timestamp": record.get("timestamp"),
                 "commit_sha": record.get("commit_sha"),
                 **record_metadata(record),
-                **record_comparison_metadata(record),
                 "success": bool(record.get("success", True)),
                 "metrics": {
                     policy.key: safe_float(record.get(policy.key))
@@ -244,7 +189,6 @@ def build_trends(records: list[Record]) -> list[Record]:
         trends.append({
             "model_id": model_id,
             "gpu_type": gpu_type,
-            **record_comparison_metadata(group[-1]),
             "points": points,
         })
-    return sorted(trends, key=comparison_sort_key)
+    return sorted(trends, key=lambda trend: (trend["model_id"], trend["gpu_type"]))

fastvideo/pipelines/basic/dreamx_world/__init__.py

@@ -1,28 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo.configs.pipelines.dreamx_world import (
-    DreamXWorld5BARPipelineConfig,
-    DreamXWorld5BCamPipelineConfig,
-    make_dreamx_world_5b_ar_dit_config,
-    make_dreamx_world_5b_cam_dit_config,
-    make_dreamx_world_5b_cam_text_encoder_config,
-    make_dreamx_world_5b_cam_vae_config,
-)
-from fastvideo.pipelines.basic.dreamx_world.dreamx_world_ar_pipeline import DreamXWorldARPipeline
-from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import DreamXWorldPipeline
-from fastvideo.pipelines.basic.dreamx_world.stages import (
-    DREAMX_Y_CAMERA_KEY,
-    DreamXWorldCameraConditioningStage,
-)
-
-__all__ = [
-    "DREAMX_Y_CAMERA_KEY",
-    "DreamXWorld5BARPipelineConfig",
-    "DreamXWorld5BCamPipelineConfig",
-    "DreamXWorldCameraConditioningStage",
-    "DreamXWorldARPipeline",
-    "DreamXWorldPipeline",
-    "make_dreamx_world_5b_ar_dit_config",
-    "make_dreamx_world_5b_cam_dit_config",
-    "make_dreamx_world_5b_cam_text_encoder_config",
-    "make_dreamx_world_5b_cam_vae_config",
-]

fastvideo/pipelines/basic/dreamx_world/ar_denoising.py

@@ -1,219 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World autoregressive causal denoising stage."""
-
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-from tqdm.auto import tqdm
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.basic.dreamx_world.stages import DREAMX_Y_CAMERA_KEY
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.denoising import DenoisingStage
-
-
-class DreamXWorldARCausalDenoisingStage(DenoisingStage):
-    """Official DreamX AR-forcing denoising loop with KV cache."""
-
-    _AR_NOISE_SEED_OFFSET = 1_000_003
-
-    def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None:
-        super().__init__(transformer=transformer, scheduler=scheduler, pipeline=pipeline, vae=vae)
-        self.num_transformer_blocks = len(self.transformer.blocks)
-        self.num_frame_per_block = int(getattr(self.transformer, "num_frame_per_block", 3))
-        self.local_attn_size = int(getattr(self.transformer, "local_attn_size", 12))
-
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        assert batch.latents is not None, "latents must be prepared before DreamX AR denoising"
-        assert batch.prompt_embeds, "prompt embeds must be prepared before DreamX AR denoising"
-        latents = batch.latents
-        device = latents.device
-        target_dtype = torch.bfloat16
-        autocast_enabled = device.type == "cuda" and not fastvideo_args.disable_autocast
-
-        frame_seq_length = (latents.shape[-2] // self.transformer.patch_size[1]) * (latents.shape[-1] //
-                                                                                    self.transformer.patch_size[2])
-        timesteps = torch.tensor(
-            tuple(getattr(fastvideo_args.pipeline_config, "dmd_denoising_steps", (1000, 750, 500, 250))),
-            dtype=torch.long,
-        ).cpu()
-        if getattr(fastvideo_args.pipeline_config, "warp_denoising_step", True):
-            self.scheduler.set_timesteps(1000)
-            scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
-            timesteps = scheduler_timesteps[1000 - timesteps]
-        timesteps = timesteps.to(device)
-
-        if latents.shape[2] % self.num_frame_per_block != 0:
-            raise ValueError("DreamX AR latent frames must be divisible by num_frame_per_block")
-
-        y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY, batch.extra.get("y_camera"))
-        if isinstance(y_camera, dict):
-            y_camera = {
-                k: v.to(device=device, dtype=target_dtype) if torch.is_tensor(v) else v
-                for k, v in y_camera.items()
-            }
-
-        if batch.image_latent is not None and batch.image_latent.shape[1] == latents.shape[1]:
-            latents[:, :, :batch.image_latent.shape[2]] = batch.image_latent.to(device=device, dtype=latents.dtype)
-
-        kv_cache = self._initialize_kv_cache(latents.shape[0], target_dtype, device, frame_seq_length)
-        crossattn_cache = self._initialize_crossattn_cache(latents.shape[0], target_dtype, device)
-        prompt = batch.prompt_embeds[0]
-        if torch.is_tensor(prompt):
-            prompt = prompt.to(device=device, dtype=target_dtype)
-            context = [sample for sample in prompt]
-        else:
-            context = prompt
-
-        num_blocks = latents.shape[2] // self.num_frame_per_block
-        start = 0
-        first_frame_mask = torch.ones_like(latents)
-        first_frame_mask[:, :, 0] = 0
-        base_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
-        noise_generator = self._make_noise_generator(base_generator, device)
-
-        with tqdm(total=num_blocks * len(timesteps), desc="DreamX AR denoising", leave=False) as progress:
-            for _ in range(num_blocks):
-                current_num_frames = self.num_frame_per_block
-                block_latents = latents[:, :, start:start + current_num_frames]
-                noisy_input = block_latents.clone()
-                mask_block = first_frame_mask[:, :, start:start + current_num_frames]
-                camera_block = self._slice_camera(y_camera, start, current_num_frames)
-
-                for idx, current_timestep in enumerate(timesteps):
-                    timestep = torch.full(
-                        (latents.shape[0], current_num_frames * frame_seq_length),
-                        int(current_timestep.item()),
-                        device=device,
-                        dtype=torch.long,
-                    )
-                    if start == 0:
-                        timestep[:, :frame_seq_length] = 0
-                    with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
-                        denoised = self.transformer(
-                            hidden_states=block_latents.to(target_dtype),
-                            encoder_hidden_states=torch.stack(context).to(target_dtype),
-                            timestep=timestep,
-                            y_camera=camera_block,
-                            kv_cache=kv_cache,
-                            crossattn_cache=crossattn_cache,
-                            current_start=start * frame_seq_length,
-                        )
-                    denoised = denoised.to(latents.dtype)
-                    if idx < len(timesteps) - 1:
-                        next_timestep = torch.full((latents.shape[0], current_num_frames),
-                                                   int(timesteps[idx + 1].item()),
-                                                   device=device,
-                                                   dtype=torch.long)
-                        noise_kwargs = {"device": device, "dtype": denoised.dtype}
-                        if noise_generator is not None:
-                            noise_kwargs["generator"] = noise_generator
-                        noise = torch.randn(denoised.permute(0, 2, 1, 3, 4).shape, **noise_kwargs)
-                        block_btchw = self.scheduler.add_noise(
-                            denoised.permute(0, 2, 1, 3, 4).flatten(0, 1),
-                            noise.flatten(0, 1),
-                            next_timestep.flatten(),
-                        ).unflatten(0, (latents.shape[0], current_num_frames))
-                        block_latents = block_btchw.permute(0, 2, 1, 3, 4)
-                        block_latents = block_latents * mask_block + noisy_input * (1 - mask_block)
-                    else:
-                        block_latents = denoised * mask_block + noisy_input * (1 - mask_block)
-                    progress.update()
-
-                latents[:, :, start:start + current_num_frames] = block_latents
-                self._update_context_cache(block_latents, context, camera_block, kv_cache, crossattn_cache, start,
-                                           frame_seq_length, target_dtype, autocast_enabled,
-                                           float(getattr(fastvideo_args.pipeline_config, "context_noise", 0.1)))
-                start += current_num_frames
-
-        batch.latents = latents
-        return batch
-
-    def _make_noise_generator(self, generator: torch.Generator | None, device: torch.device) -> torch.Generator | None:
-        if generator is None:
-            return None
-        if getattr(generator, "device", None) == device:
-            return generator
-        seed = int(generator.initial_seed()) + self._AR_NOISE_SEED_OFFSET
-        return torch.Generator(device=device).manual_seed(seed)
-
-    @staticmethod
-    def _context_noise_timestep(context_noise: float) -> int:
-        if 0.0 < context_noise <= 1.0:
-            return int(context_noise * 1000)
-        return int(context_noise)
-
-    def _slice_camera(self, y_camera: Any, start: int, num_frames: int):
-        if not isinstance(y_camera, dict):
-            return y_camera
-        return {
-            "viewmats": y_camera["viewmats"][:, start:start + num_frames],
-            "K": y_camera["K"][:, start:start + num_frames],
-        }
-
-    def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device,
-                             frame_seq_length: int) -> list[dict[str, Any]]:
-        size = self.local_attn_size * frame_seq_length if self.local_attn_size != -1 else 18480
-        heads = self.transformer.num_attention_heads
-        head_dim = self.transformer.attention_head_dim
-        cam_self_attn = next(
-            (getattr(block, "cam_self_attn", None)
-             for block in self.transformer.blocks if getattr(block, "cam_self_attn", None) is not None),
-            None,
-        )
-
-        caches = []
-        for _ in range(self.num_transformer_blocks):
-            cache = {
-                "k": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
-                "v": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
-                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
-                "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
-            }
-            if cam_self_attn is not None:
-                cam_heads = int(cam_self_attn.num_heads)
-                cam_head_dim = int(cam_self_attn.head_dim)
-                cache.update({
-                    "prope_k":
-                    torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
-                    "prope_v":
-                    torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
-                    "prope_global_end_index":
-                    torch.tensor([0], dtype=torch.long, device=device),
-                    "prope_local_end_index":
-                    torch.tensor([0], dtype=torch.long, device=device),
-                })
-            caches.append(cache)
-        return caches
-
-    def _initialize_crossattn_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device):
-        heads = self.transformer.num_attention_heads
-        head_dim = self.transformer.attention_head_dim
-        return [{
-            "k": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
-            "v": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
-            "is_init": False,
-        } for _ in range(self.num_transformer_blocks)]
-
-    def _update_context_cache(self, block_latents: torch.Tensor, context: Any, camera_block: Any,
-                              kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]], start: int,
-                              frame_seq_length: int, target_dtype: torch.dtype, autocast_enabled: bool,
-                              context_noise: float) -> None:
-        timestep = torch.full(
-            (block_latents.shape[0], block_latents.shape[2] * frame_seq_length),
-            self._context_noise_timestep(context_noise),
-            device=block_latents.device,
-            dtype=torch.long,
-        )
-        with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
-            self.transformer(
-                hidden_states=block_latents.to(target_dtype),
-                encoder_hidden_states=torch.stack(context).to(target_dtype),
-                timestep=timestep,
-                y_camera=camera_block,
-                kv_cache=kv_cache,
-                crossattn_cache=crossattn_cache,
-                current_start=start * frame_seq_length,
-            )

fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py

@@ -1,228 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-import math
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-from scipy.interpolate import interp1d
-from scipy.spatial.transform import Rotation, Slerp
-
-_ACTION_TO_MOTION = {
-    "w": "forward",
-    "a": "left",
-    "d": "right",
-    "s": "backward",
-    "j": "left_rot",
-    "l": "right_rot",
-    "i": "up_rot",
-    "k": "down_rot",
-}
-_TRANSLATION_BASE_UNIT = 1.0
-_ROTATION_BASE_UNIT = 10.0
-_INTRINSIC_ROW = [0.8, 0.5, 0.5, 0.5]
-
-
-@dataclass
-class DreamXCamera:
-    fx: float
-    fy: float
-    cx: float
-    cy: float
-    w2c_mat: np.ndarray
-
-    @property
-    def c2w_mat(self) -> np.ndarray:
-        return np.linalg.inv(self.w2c_mat)
-
-    @classmethod
-    def from_pose_row(cls, row: list[float]) -> DreamXCamera:
-        w2c_mat = np.eye(4, dtype=np.float64)
-        w2c_mat[:3, :] = np.asarray(row[7:], dtype=np.float64).reshape(3, 4)
-        return cls(
-            fx=float(row[1]),
-            fy=float(row[2]),
-            cx=float(row[3]),
-            cy=float(row[4]),
-            w2c_mat=w2c_mat,
-        )
-
-
-def _translation_step(motion_type: str, current_pose: dict[str, np.ndarray], value: float, duration: int) -> np.ndarray:
-    if motion_type in ("forward", "backward"):
-        yaw = np.radians(current_pose["rotation"][1])
-        pitch = np.radians(current_pose["rotation"][0])
-        forward = np.array([-math.sin(yaw) * math.cos(pitch), math.sin(pitch), math.cos(yaw) * math.cos(pitch)])
-        direction = 1 if motion_type == "forward" else -1
-        return forward * value * direction / duration
-    if motion_type in ("left", "right"):
-        yaw = np.radians(current_pose["rotation"][1])
-        right = np.array([math.cos(yaw), 0.0, math.sin(yaw)])
-        direction = -1 if motion_type == "left" else 1
-        return right * value * direction / duration
-    return np.zeros(3)
-
-
-def _rotation_step(motion_type: str, value: float, duration: int) -> np.ndarray:
-    if not motion_type.endswith("rot"):
-        return np.zeros(3)
-    axis = motion_type.split("_")[0]
-    rotation = np.zeros(3)
-    if axis == "left":
-        rotation[1] = value
-    elif axis == "right":
-        rotation[1] = -value
-    elif axis == "up":
-        rotation[0] = -value
-    elif axis == "down":
-        rotation[0] = value
-    return rotation / duration
-
-
-def _euler_to_quaternion(angles: np.ndarray) -> list[float]:
-    pitch, yaw, roll = np.radians(angles)
-    cy = math.cos(yaw * 0.5)
-    sy = math.sin(yaw * 0.5)
-    cp = math.cos(pitch * 0.5)
-    sp = math.sin(pitch * 0.5)
-    cr = math.cos(roll * 0.5)
-    sr = math.sin(roll * 0.5)
-    return [
-        cy * cp * cr + sy * sp * sr,
-        cy * sp * cr + sy * cp * sr,
-        sy * cp * cr - cy * sp * sr,
-        cy * cp * sr - sy * sp * cr,
-    ]
-
-
-def _quaternion_to_rotation_matrix(quaternion: list[float]) -> np.ndarray:
-    qw, qx, qy, qz = quaternion
-    return np.array([
-        [1 - 2 * (qy**2 + qz**2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
-        [2 * (qx * qy + qw * qz), 1 - 2 * (qx**2 + qz**2), 2 * (qy * qz - qw * qx)],
-        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx**2 + qy**2)],
-    ])
-
-
-def _pose_rows_from_actions(action_seq: list[str], action_speed_list: list[float], duration: int) -> list[list[float]]:
-    if len(action_seq) != len(action_speed_list):
-        raise ValueError("action_seq and action_speed_list must have the same length")
-
-    positions: list[np.ndarray] = []
-    rotations: list[np.ndarray] = []
-    current_pose = {
-        "position": np.array([0.0, 0.0, 0.0]),
-        "rotation": np.array([0.0, 0.0, 0.0]),
-    }
-
-    for action_id, speed in zip(action_seq, action_speed_list, strict=True):
-        motion_types = [_ACTION_TO_MOTION[key] for key in list(action_id)]
-        translation_step = np.zeros(3)
-        rotation_step = np.zeros(3)
-        for motion_type in motion_types:
-            translation_step += _translation_step(motion_type, current_pose,
-                                                  float(speed) * _TRANSLATION_BASE_UNIT, duration)
-            rotation_step += _rotation_step(motion_type, float(speed) * _ROTATION_BASE_UNIT, duration)
-
-        segment_positions = []
-        segment_rotations = []
-        for index in range(1, duration + 1):
-            segment_positions.append(current_pose["position"] + translation_step * index)
-            segment_rotations.append(current_pose["rotation"] + rotation_step * index)
-        current_pose["position"] = segment_positions[-1].copy()
-        current_pose["rotation"] = segment_rotations[-1].copy()
-        positions.extend(segment_positions)
-        rotations.extend(segment_rotations)
-
-    rows: list[list[float]] = [[0.0] + _INTRINSIC_ROW + [0.0, 0.0] +
-                               [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]]
-    for index, (position, rotation) in enumerate(zip(positions, rotations, strict=False)):
-        rotation_matrix = _quaternion_to_rotation_matrix(_euler_to_quaternion(rotation))
-        translation = -rotation_matrix @ position
-        extrinsic = np.hstack([rotation_matrix, translation.reshape(3, 1)])
-        rows.append([float(index)] + _INTRINSIC_ROW + [0.0, 0.0] + extrinsic.flatten().tolist())
-    return rows
-
-
-def _interpolate_camera_poses(
-    cameras: list[DreamXCamera],
-    src_indices: np.ndarray,
-    tgt_indices: np.ndarray,
-) -> list[DreamXCamera]:
-    if len(cameras) <= 1:
-        return [cameras[0]] * len(tgt_indices) if cameras else []
-    src_rot_mat = np.array([camera.w2c_mat[:3, :3] for camera in cameras])
-    src_trans_vec = np.array([camera.w2c_mat[:3, 3] for camera in cameras])
-
-    dets = np.linalg.det(src_rot_mat)
-    flip_handedness = dets.size > 0 and np.median(dets) < 0.0
-    if flip_handedness:
-        flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype)
-        src_rot_mat = src_rot_mat @ flip_mat
-
-    trans = interp1d(src_indices, src_trans_vec, axis=0, kind="linear", bounds_error=False,
-                     fill_value="extrapolate")(tgt_indices)
-    quats = Rotation.from_matrix(src_rot_mat).as_quat().copy()
-    for index in range(1, len(quats)):
-        if np.dot(quats[index], quats[index - 1]) < 0:
-            quats[index] = -quats[index]
-    rot = Slerp(src_indices, Rotation.from_quat(quats))(tgt_indices).as_matrix()
-    if flip_handedness:
-        rot = rot @ flip_mat
-
-    ref = cameras[0]
-    result = []
-    for index in range(len(tgt_indices)):
-        w2c_mat = np.eye(4, dtype=np.float64)
-        w2c_mat[:3, :] = np.hstack([rot[index], trans[index].reshape(3, 1)])
-        result.append(DreamXCamera(ref.fx, ref.fy, ref.cx, ref.cy, w2c_mat))
-    return result
-
-
-def _relative_c2w_poses(cameras: list[DreamXCamera]) -> np.ndarray:
-    abs_w2cs = [camera.w2c_mat for camera in cameras]
-    abs_c2ws = [camera.c2w_mat for camera in cameras]
-    target_cam_c2w = np.eye(4, dtype=np.float64)
-    abs2rel = target_cam_c2w @ abs_w2cs[0]
-    poses = [target_cam_c2w] + [abs2rel @ c2w for c2w in abs_c2ws[1:]]
-    return np.asarray(poses, dtype=np.float32)
-
-
-def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
-    rotation_inv = transforms[..., :3, :3].transpose(-1, -2)
-    output = torch.zeros_like(transforms)
-    output[..., :3, :3] = rotation_inv
-    output[..., :3, 3] = -torch.einsum("...ij,...j->...i", rotation_inv, transforms[..., :3, 3])
-    output[..., 3, 3] = 1.0
-    return output
-
-
-def build_dreamx_camera_condition(
-    action_seq: list[str],
-    action_speed_list: list[float],
-    *,
-    num_frames: int,
-    height: int,
-    width: int,
-    dtype: torch.dtype = torch.float32,
-    device: torch.device | str = "cpu",
-) -> dict[str, torch.Tensor]:
-    del height, width  # DreamX-World-5B-Cam uses fixed normalized intrinsics.
-    duration = math.ceil(num_frames / len(action_seq))
-    rows = _pose_rows_from_actions(action_seq, action_speed_list, duration)[:num_frames]
-    cameras = [DreamXCamera.from_pose_row(row) for row in rows]
-
-    latent_frame_count = 1 + (len(cameras) - 1) // 4
-    src_indices = np.arange(len(cameras), dtype=np.float64)
-    tgt_indices = np.linspace(0, len(cameras) - 1, latent_frame_count)
-    cameras = _interpolate_camera_poses(cameras, src_indices, tgt_indices)
-
-    c2ws = torch.as_tensor(_relative_c2w_poses(cameras), dtype=dtype, device=device)
-    viewmats = _invert_se3(c2ws)
-
-    intrinsics = torch.zeros((latent_frame_count, 3, 3), dtype=dtype, device=device)
-    intrinsics[:, 0, 0] = 969.6969696969696 / (960.0 * 2)
-    intrinsics[:, 1, 1] = 969.6969696969696 / (540.0 * 2)
-    intrinsics[:, 2, 2] = 1.0
-    return {"viewmats": viewmats, "K": intrinsics}

fastvideo/pipelines/basic/dreamx_world/config.py

@@ -1,20 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Compatibility exports for DreamX-World pipeline configs."""
-
-from fastvideo.configs.pipelines.dreamx_world import (
-    DreamXWorld5BARPipelineConfig,
-    DreamXWorld5BCamPipelineConfig,
-    make_dreamx_world_5b_ar_dit_config,
-    make_dreamx_world_5b_cam_dit_config,
-    make_dreamx_world_5b_cam_text_encoder_config,
-    make_dreamx_world_5b_cam_vae_config,
-)
-
-__all__ = [
-    "DreamXWorld5BARPipelineConfig",
-    "DreamXWorld5BCamPipelineConfig",
-    "make_dreamx_world_5b_ar_dit_config",
-    "make_dreamx_world_5b_cam_dit_config",
-    "make_dreamx_world_5b_cam_text_encoder_config",
-    "make_dreamx_world_5b_cam_vae_config",
-]

fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py

@@ -1,67 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World-5B autoregressive pipeline entrypoint."""
-
-from fastvideo.api.sampling_param import SamplingParam
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
-from fastvideo.pipelines.basic.dreamx_world.ar_denoising import DreamXWorldARCausalDenoisingStage
-from fastvideo.pipelines.basic.dreamx_world.stages import (
-    DreamXWorldCameraConditioningStage,
-    DreamXWorldImageVAEEncodingStage,
-)
-from fastvideo.pipelines.stages import (
-    ConditioningStage,
-    DecodingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
-
-logger = init_logger(__name__)
-
-
-class DreamXWorldARPipeline(LoRAPipeline, ComposedPipelineBase):
-    """DreamX-World-5B autoregressive causal camera pipeline."""
-
-    _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
-    pipeline_config_cls = DreamXWorld5BARPipelineConfig
-    sampling_params_cls = SamplingParam
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
-        self.modules["scheduler"].set_timesteps(1000)
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
-        self.add_stage(stage_name="prompt_encoding_stage",
-                       stage=TextEncodingStage(
-                           text_encoders=[self.get_module("text_encoder")],
-                           tokenizers=[self.get_module("tokenizer")],
-                       ))
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-        self.add_stage(stage_name="timestep_preparation_stage",
-                       stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
-        self.add_stage(stage_name="latent_preparation_stage",
-                       stage=LatentPreparationStage(
-                           scheduler=self.get_module("scheduler"),
-                           transformer=self.get_module("transformer", None),
-                       ))
-        self.add_stage(stage_name="image_latent_preparation_stage",
-                       stage=DreamXWorldImageVAEEncodingStage(vae=self.get_module("vae")))
-        self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
-        self.add_stage(stage_name="denoising_stage",
-                       stage=DreamXWorldARCausalDenoisingStage(
-                           transformer=self.get_module("transformer"),
-                           scheduler=self.get_module("scheduler"),
-                           vae=self.get_module("vae"),
-                           pipeline=self,
-                       ))
-        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
-        logger.info("DreamXWorldARPipeline initialized with autoregressive causal denoising")
-
-
-EntryClass = DreamXWorldARPipeline

fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py

@@ -1,78 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World video pipeline entrypoint."""
-
-from fastvideo.api.sampling_param import SamplingParam
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
-from fastvideo.pipelines.basic.dreamx_world.stages import DreamXWorldCameraConditioningStage
-from fastvideo.pipelines.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
-
-logger = init_logger(__name__)
-
-
-class DreamXWorldPipeline(LoRAPipeline, ComposedPipelineBase):
-    """DreamX-World-5B-Cam pipeline with native FastVideo camera conditioning."""
-
-    _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
-    pipeline_config_cls = DreamXWorld5BCamPipelineConfig
-    sampling_params_cls = SamplingParam
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer", None),
-            ),
-        )
-
-        self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                transformer_2=self.get_module("transformer_2", None),
-                scheduler=self.get_module("scheduler"),
-                vae=self.get_module("vae"),
-                pipeline=self,
-            ),
-        )
-
-        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
-
-        logger.info("DreamXWorldPipeline initialized with native camera conditioning")
-
-
-EntryClass = DreamXWorldPipeline

fastvideo/pipelines/basic/dreamx_world/presets.py

@@ -1,58 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World model family pipeline presets."""
-
-from fastvideo.api.presets import InferencePreset, PresetStageSpec
-
-_NEGATIVE_PROMPT_CN = ("色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，"
-                       "静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，"
-                       "多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，"
-                       "形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，"
-                       "背景人很多，倒着走")
-
-_DENOISE_STAGE = PresetStageSpec(
-    name="denoise",
-    kind="denoising",
-    description="DreamX-World camera-conditioned denoising pass",
-    allowed_overrides=frozenset({
-        "num_inference_steps",
-        "guidance_scale",
-    }),
-)
-
-DREAMX_WORLD_5B_CAM = InferencePreset(
-    name="dreamx_world_5b_cam",
-    version=1,
-    model_family="dreamx_world",
-    description="DreamX-World 5B camera-control video generation",
-    workload_type="i2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "height": 480,
-        "width": 832,
-        "num_frames": 161,
-        "fps": 16,
-        "guidance_scale": 5.0,
-        "num_inference_steps": 30,
-        "negative_prompt": _NEGATIVE_PROMPT_CN,
-    },
-)
-
-DREAMX_WORLD_5B_AR = InferencePreset(
-    name="dreamx_world_5b_ar",
-    version=1,
-    model_family="dreamx_world",
-    description="DreamX-World 5B autoregressive camera-control generation",
-    workload_type="i2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "height": 704,
-        "width": 1280,
-        "num_frames": 1005,
-        "fps": 16,
-        "guidance_scale": 1.0,
-        "num_inference_steps": 4,
-        "negative_prompt": _NEGATIVE_PROMPT_CN,
-    },
-)
-
-ALL_PRESETS = (DREAMX_WORLD_5B_CAM, DREAMX_WORLD_5B_AR)

fastvideo/pipelines/basic/dreamx_world/stages.py

@@ -1,156 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World pipeline stages."""
-
-from __future__ import annotations
-
-import torch
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import (
-    build_dreamx_camera_condition, )
-
-DREAMX_Y_CAMERA_KEY = "dreamx_y_camera"
-
-logger = init_logger(__name__)
-
-
-class DreamXWorldCameraConditioningStage(PipelineStage):
-    """Build PRoPE camera conditioning for DreamX-World-5B-Cam."""
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        del fastvideo_args
-        if DREAMX_Y_CAMERA_KEY in batch.extra:
-            return batch
-
-        action_seq = batch.extra.get("dreamx_action_seq", batch.action_list)
-        action_speed_list = batch.extra.get("dreamx_action_speed_list", batch.action_speed_list)
-        if action_seq is None:
-            action_seq = ["w"]
-        if action_speed_list is None:
-            action_speed_list = [4]
-
-        if isinstance(action_seq, str):
-            action_seq = [action_seq]
-        if isinstance(action_speed_list, int | float):
-            action_speed_list = [action_speed_list]
-        if len(action_speed_list) == 1 and len(action_seq) > 1:
-            action_speed_list = list(action_speed_list) * len(action_seq)
-        action_speed_list = [float(speed) for speed in action_speed_list]
-
-        height = int(batch.height) if batch.height is not None else 704
-        width = int(batch.width) if batch.width is not None else 1280
-        num_frames = int(batch.num_frames)
-        dtype = batch.latents.dtype if torch.is_tensor(batch.latents) else torch.float32
-        device = batch.latents.device if torch.is_tensor(batch.latents) else "cpu"
-
-        y_camera = build_dreamx_camera_condition(
-            list(action_seq),
-            action_speed_list,
-            num_frames=num_frames,
-            height=height,
-            width=width,
-            dtype=dtype,
-            device=device,
-        )
-        batch.extra[DREAMX_Y_CAMERA_KEY] = {key: value.unsqueeze(0) for key, value in y_camera.items()}
-        return batch
-
-    def verify_output(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> VerificationResult:
-        del fastvideo_args
-        result = VerificationResult()
-        y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY)
-        result.add_check("dreamx_y_camera", y_camera, lambda value: isinstance(value, dict))
-        if isinstance(y_camera, dict):
-            result.add_check("dreamx_y_camera.viewmats", y_camera.get("viewmats"), torch.is_tensor)
-            result.add_check("dreamx_y_camera.K", y_camera.get("K"), torch.is_tensor)
-        return result
-
-
-class DreamXWorldImageVAEEncodingStage(PipelineStage):
-    """Encode the conditioning image into the first-frame latent.
-
-    Official AR-forcing flow (AMAP-ML/DreamX-World inference_ar_forcing.py):
-    the input image is resized, normalized to [-1, 1], VAE-encoded
-    deterministically, and written into frame 0 of the noise — the causal
-    denoiser then treats frame 0 as clean context. This stage produces
-    ``batch.image_latent`` ([B, C, 1, H_lat, W_lat]); the injection into
-    the latents happens in DreamXWorldARCausalDenoisingStage.
-    """
-
-    def __init__(self, vae) -> None:
-        super().__init__()
-        self.vae = vae
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        if batch.pil_image is None:
-            # No conditioning image: the causal denoiser falls back to
-            # running from pure noise (frame 0 uninitialized). Warn loudly —
-            # this pipeline is registered I2V and the official flow always
-            # forces from a frame.
-            logger.warning("DreamXWorldARPipeline called without an input image; "
-                           "first-frame context will be noise (T2V-style). Pass an "
-                           "image for the official AR-forcing behavior.")
-            return batch
-
-        from fastvideo.platforms import get_local_torch_device
-        from fastvideo.utils import PRECISION_TO_TYPE
-
-        device = get_local_torch_device()
-        image = batch.pil_image
-        if not isinstance(image, torch.Tensor):
-            import numpy as np
-            import PIL.Image
-            assert isinstance(image, PIL.Image.Image)
-            width = batch.width if isinstance(batch.width, int) else batch.width[0]
-            height = batch.height if isinstance(batch.height, int) else batch.height[0]
-            image = image.convert("RGB").resize((width, height), PIL.Image.Resampling.LANCZOS)
-            arr = torch.from_numpy(np.asarray(image)).float().permute(2, 0, 1) / 255.0
-            image = (arr - 0.5) / 0.5  # official Normalize([0.5], [0.5])
-            image = image.unsqueeze(0)  # [1, C, H, W]
-        if image.dim() == 4:
-            image = image.unsqueeze(2)  # [B, C, 1, H, W]
-        elif image.dim() == 5:
-            image = image[:, :, :1]
-        image = image.to(device=device, dtype=torch.float32)
-
-        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
-        vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
-        self.vae = self.vae.to(device)
-        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
-            if not vae_autocast_enabled:
-                image = image.to(vae_dtype)
-            encoder_output = self.vae.encode(image)
-
-        # Official encode_to_latent is deterministic ((mean - mu) / sigma per
-        # channel); the posterior mean + shift/scale is the FastVideo
-        # equivalent of that normalization.
-        latent = encoder_output.mean
-        if getattr(self.vae, "shift_factor", None) is not None:
-            shift = self.vae.shift_factor
-            latent = latent - (shift.to(latent.device, latent.dtype) if isinstance(shift, torch.Tensor) else shift)
-        scale = self.vae.scaling_factor
-        latent = latent * (scale.to(latent.device, latent.dtype) if isinstance(scale, torch.Tensor) else scale)
-
-        batch.image_latent = latent
-        return batch
-
-    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        return result

fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py

@@ -19,39 +19,6 @@ 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 = [
@@ -222,16 +189,28 @@ 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_upsampler_path(self.model_path, model_index))
+            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"))
         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(
-                self.model_path, model_index.get("fastvideo_refine_transformer_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(self.model_path,
-                                                                        model_index.get("fastvideo_refine_lora_path"))
+            fastvideo_args.ltx2_refine_lora_path = _resolve_refine_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):
@@ -247,11 +226,10 @@ 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(self.model_path,
-                                                                         model_index.get("fastvideo_refine_noise_path"))
+            fastvideo_args.ltx2_refine_noise_path = _resolve_refine_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(
-                self.model_path, model_index.get("fastvideo_refine_audio_noise_path")))
+                model_index.get("fastvideo_refine_audio_noise_path")))
 
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")

fastvideo/pipelines/stages/denoising.py

@@ -191,19 +191,6 @@ class DenoisingStage(PipelineStage):
             },
         )
 
-        dreamx_y_camera = batch.extra.get("dreamx_y_camera", batch.extra.get("y_camera"))
-        if isinstance(dreamx_y_camera, dict):
-            dreamx_y_camera = {
-                key: value.to(device=local_device, dtype=target_dtype) if torch.is_tensor(value) else value
-                for key, value in dreamx_y_camera.items()
-            }
-        dreamx_camera_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "y_camera": dreamx_y_camera,
-            },
-        )
-
         for key in ("flux2_txt_ids", "flux2_img_ids"):
             value = batch.extra.get(key)
             if torch.is_tensor(value):
@@ -255,11 +242,7 @@ class DenoisingStage(PipelineStage):
             # the image latent instead of appending along the channel dim
             assert batch.image_latent is None, "TI2V task should not have image latents"
             assert self.vae is not None, "VAE is not provided for TI2V task"
-            vae_device = next(self.vae.parameters()).device
-            self.vae = self.vae.to(local_device)
             z = self.vae.encode(batch.pil_image).mean.float()
-            if getattr(fastvideo_args, "vae_cpu_offload", False):
-                self.vae = self.vae.to(vae_device)
             if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
                 if isinstance(self.vae.shift_factor, torch.Tensor):
                     z -= self.vae.shift_factor.to(z.device, z.dtype)
@@ -512,7 +495,6 @@ class DenoisingStage(PipelineStage):
                             **pos_cond_kwargs,
                             **action_kwargs,
                             **camera_kwargs,
-                            **dreamx_camera_kwargs,
                             **timesteps_r_kwarg,
                             **flux2_id_kwargs,
                         )
@@ -555,7 +537,6 @@ class DenoisingStage(PipelineStage):
                                     **neg_cond_kwargs,
                                     **action_kwargs,
                                     **camera_kwargs,
-                                    **dreamx_camera_kwargs,
                                     **timesteps_r_kwarg,
                                     **flux2_id_kwargs,
                                 )

fastvideo/pipelines/stages/kandinsky5.py

@@ -51,11 +51,9 @@ class Kandinsky5LatentPreparationStage(PipelineStage):
             num_frames = num_frames // temporal_ratio * temporal_ratio + 1
             batch.num_frames = num_frames
 
-        required_divisor_h = spatial_ratio * patch_size[1]
-        required_divisor_w = spatial_ratio * patch_size[2]
-        if height % required_divisor_h != 0 or width % required_divisor_w != 0:
-            raise ValueError(f"Kandinsky5 height must be divisible by {required_divisor_h} and width by "
-                             f"{required_divisor_w}; "
+        required_divisor = spatial_ratio * patch_size[1]
+        if height % required_divisor != 0 or width % required_divisor != 0:
+            raise ValueError(f"Kandinsky5 height/width must be divisible by {required_divisor}; "
                              f"got height={height}, width={width}.")
 
         # NABLA sparse attention (Pro checkpoints) reshapes the post-patch grid
@@ -63,11 +61,9 @@ class Kandinsky5LatentPreparationStage(PipelineStage):
         # all the encoding work is done.
         arch_cfg = getattr(self.transformer, "config", None) or fastvideo_args.pipeline_config.dit_config.arch_config
         if getattr(arch_cfg, "attention_type", "regular") == "nabla":
-            nabla_divisor_h = required_divisor_h * 8
-            nabla_divisor_w = required_divisor_w * 8
-            if height % nabla_divisor_h != 0 or width % nabla_divisor_w != 0:
-                raise ValueError(f"Kandinsky5 NABLA checkpoints require height divisible by {nabla_divisor_h} and "
-                                 f"width divisible by {nabla_divisor_w}; "
+            nabla_divisor = required_divisor * 8
+            if height % nabla_divisor != 0 or width % nabla_divisor != 0:
+                raise ValueError(f"Kandinsky5 NABLA checkpoints require height/width divisible by {nabla_divisor}; "
                                  f"got height={height}, width={width}.")
 
         if isinstance(batch.prompt, list):
@@ -97,21 +93,14 @@ class Kandinsky5LatentPreparationStage(PipelineStage):
         if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
             raise ValueError(f"generator list length {len(batch.generator)} does not match batch size {batch_size}.")
 
-        visual_cond = getattr(self.transformer, "visual_cond", False)
-
         if batch.latents is None:
             latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
             if hasattr(self.scheduler, "init_noise_sigma"):
                 latents = latents * self.scheduler.init_noise_sigma
         else:
-            valid_shapes = [shape]
-            if visual_cond:
-                valid_shapes.append((*shape[:-1], num_channels * 2 + 1))
-            if tuple(batch.latents.shape) not in valid_shapes:
-                raise ValueError(f"Provided latents shape {list(batch.latents.shape)} does not match expected "
-                                 f"Kandinsky5 latent shape(s): {[list(s) for s in valid_shapes]}.")
             latents = batch.latents.to(device=device, dtype=dtype)
 
+        visual_cond = getattr(self.transformer, "visual_cond", False)
         if visual_cond and latents.shape[-1] == num_channels:
             cond = torch.zeros_like(latents)
             cond_mask = torch.zeros(
@@ -320,7 +309,7 @@ class Kandinsky5DenoisingStage(PipelineStage):
         with tqdm(total=batch.num_inference_steps, desc="Kandinsky5 Denoising") as progress_bar:
             for i, timestep in enumerate(batch.timesteps):
                 if hasattr(self, "interrupt") and self.interrupt:
-                    break
+                    continue
 
                 t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
                 attn_metadata = None
@@ -387,7 +376,7 @@ class Kandinsky5DenoisingStage(PipelineStage):
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
         result = VerificationResult()
         result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
-        result.add_check("prompt_embeds", batch.prompt_embeds, V.min_list_length(2))
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
         return result
 
 

fastvideo/pipelines/stages/validators.py

@@ -219,16 +219,6 @@ class StageValidators:
 
         return validator
 
-    @staticmethod
-    def min_list_length(min_length: int) -> Callable[[Any], bool]:
-        """Return a validator that checks if a list has at least min_length items."""
-
-        def validator(value: Any) -> bool:
-            return StageValidators.list_min_length(value, min_length)
-
-        validator.__name__ = f"list_min_length_{min_length}"
-        return validator
-
     @staticmethod
     def min_dims(min_dims: int) -> Callable[[Any], bool]:
         """Return a validator that checks if tensor has at least min_dims dimensions and no NaN values."""

fastvideo/registry.py

@@ -20,7 +20,6 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
     Cosmos25Config,
     Cosmos25_14BConfig,
 )
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.gen3c import Gen3CConfig
@@ -718,12 +717,6 @@ 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 (
@@ -943,40 +936,6 @@ def _register_configs() -> None:
         model_family="wan",
         default_preset="wan_2_2_ti2v_5b",
     )
-
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=DreamXWorld5BCamPipelineConfig,
-        workload_types=(WorkloadType.I2V, ),
-        hf_model_paths=[
-            "FastVideo/DreamX-World-5B-Cam-Diffusers",
-        ],
-        model_detectors=[
-            # Mutually exclusive with the AR detector below: Cam requires an
-            # explicit "cam" marker so hyphenated AR local paths (e.g.
-            # /ckpts/dreamx-world-5b-converted) don't first-match here —
-            # detector resolution is first-match in registration order.
-            lambda path:
-            ("dreamx-world" in path.lower() and "cam" in path.lower()) or "dreamxworldpipeline" in path.lower()
-        ],
-        model_family="dreamx_world",
-        default_preset="dreamx_world_5b_cam",
-    )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=DreamXWorld5BARPipelineConfig,
-        workload_types=(WorkloadType.I2V, ),
-        hf_model_paths=[
-            "FastVideo/DreamX-World-5B-Diffusers",
-        ],
-        model_detectors=[
-            lambda path:
-            ("dreamx-world-5b" in path.lower() and "cam" not in path.lower()) or "dreamxworldarpipeline" in path.lower(
-            )
-        ],
-        model_family="dreamx_world",
-        default_preset="dreamx_world_5b_ar",
-    )
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
@@ -1164,8 +1123,6 @@ def _register_presets() -> None:
     from fastvideo.api.presets import register_preset
     from fastvideo.pipelines.basic.cosmos.presets import (
         ALL_PRESETS as COSMOS_PRESETS, )
-    from fastvideo.pipelines.basic.dreamx_world.presets import (
-        ALL_PRESETS as DREAMX_WORLD_PRESETS, )
     from fastvideo.pipelines.basic.gamecraft.presets import (
         ALL_PRESETS as GAMECRAFT_PRESETS, )
     from fastvideo.pipelines.basic.gen3c.presets import (
@@ -1201,7 +1158,6 @@ def _register_presets() -> None:
 
     all_preset_groups = (
         COSMOS_PRESETS,
-        DREAMX_WORLD_PRESETS,
         FLUX2_PRESETS,
         GAMECRAFT_PRESETS,
         GEN3C_PRESETS,

fastvideo/tests/api/test_compat_translation.py

@@ -17,6 +17,46 @@ class TestLegacyTorchCompileKwargsTranslation:
     first-class :class:`CompileConfig` fields and anything unknown falls
     into ``extras``."""
 
+    def test_all_typed_keys_promoted(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {
+                "enable_torch_compile": True,
+                "torch_compile_kwargs": {
+                    "backend": "inductor",
+                    "fullgraph": True,
+                    "mode": "max-autotune-no-cudagraphs",
+                    "dynamic": False,
+                },
+            },
+        )
+        compile_config = config.engine.compile
+        assert compile_config.enabled is True
+        assert compile_config.backend == "inductor"
+        assert compile_config.fullgraph is True
+        assert compile_config.mode == "max-autotune-no-cudagraphs"
+        assert compile_config.dynamic is False
+        assert compile_config.extras == {}
+
+    def test_unknown_keys_land_in_extras(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {
+                "enable_torch_compile": True,
+                "torch_compile_kwargs": {
+                    "backend": "inductor",
+                    "options": {"triton.cudagraphs": False},
+                    "disable": False,
+                },
+            },
+        )
+        compile_config = config.engine.compile
+        assert compile_config.backend == "inductor"
+        assert compile_config.extras == {
+            "options": {"triton.cudagraphs": False},
+            "disable": False,
+        }
+
     def test_empty_kwargs_produces_empty_extras(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",
@@ -77,6 +117,13 @@ class TestLegacyLtx2VaeTilingTranslation:
     ``generator.pipeline.vae_tiling``; reverse direction emits the
     legacy name back to FastVideoArgs."""
 
+    def test_forward_routes_to_pipeline_vae_tiling(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {"ltx2_vae_tiling": False},
+        )
+        assert config.pipeline.vae_tiling is False
+
     def test_true_round_trips(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",

fastvideo/tests/attention/test_flash_attn_default_custom_op.py

@@ -31,29 +31,23 @@ def fa_default_impls():
     if not torch.cuda.is_available():
         pytest.skip("CUDA is required for the FA2/FA3 default custom-op tests")
 
-    # The registration + dispatcher live in `attention/utils/`, alongside the
-    # FP4 cute template and the masked/varlen wrappers. The backend
-    # (`attention/backends/flash_attn.py`) just imports `flash_attn_func_compilable`
-    # from there, so test references the utils module directly.
     try:
-        from fastvideo.attention.utils import flash_attn_default as fa_module
+        from fastvideo.attention.backends import flash_attn as fa_backend
     except ImportError as exc:
-        pytest.skip(f"flash_attn_default not importable: {exc}")
+        pytest.skip(f"flash_attn backend not importable: {exc}")
 
-    if fa_module.fa_version not in ("2", "3"):
+    if fa_backend.fa_version not in ("2", "3"):
         pytest.skip(
             f"FA2/FA3 default custom op only exists for fa_version in (2, 3); "
-            f"got {fa_module.fa_version!r}"
+            f"got {fa_backend.fa_version!r}"
         )
 
-    # compilable dispatcher, the original FA wrapper it falls back to, the
-    # raw custom op for opcheck, and the fa_version (FA2 has full register_
-    # autograd; FA3 keeps the carve-out so some tests gate on this).
+    # compilable dispatcher, the original FA wrapper it falls back to, and
+    # the raw custom op for opcheck.
     return (
-        fa_module.flash_attn_func_compilable,
-        fa_module._fa_default,
+        fa_backend.flash_attn_func_compilable,
+        fa_backend._fa_default,
         torch.ops.fastvideo._flash_attn_default_forward,
-        fa_module.fa_version,
     )
 
 
@@ -77,7 +71,7 @@ def test_default_compilable_inference_matches_original(fa_default_impls, dtype,
     """No-grad path routes through the custom op and is numerically identical."""
     if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
         pytest.skip("bfloat16 is not supported on this GPU")
-    compilable, original, _, _ = fa_default_impls
+    compilable, original, _ = fa_default_impls
 
     torch.manual_seed(0)
     q, k, v = _qkv(dtype, requires_grad=False)
@@ -97,7 +91,7 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau
     """
     if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
         pytest.skip("bfloat16 is not supported on this GPU")
-    compilable, original, _, _ = fa_default_impls
+    compilable, original, _ = fa_default_impls
 
     torch.manual_seed(0)
     q_ref, k_ref, v_ref = _qkv(dtype, requires_grad=True)
@@ -121,65 +115,7 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau
 @pytest.mark.parametrize("causal", [False, True])
 def test_default_forward_opcheck(fa_default_impls, causal):
     """Schema / fake-kernel consistency for the custom op (forward only)."""
-    _, _, op, _ = fa_default_impls
+    _, _, op = fa_default_impls
     torch.manual_seed(0)
     q, k, v = _qkv(torch.float16, requires_grad=False)
     torch.library.opcheck(op, (q, k, v, None, causal))
-
-
-# --------------------------------------------------------------------------- #
-# FA2-only: backward is registered on the custom op itself. Exercises the     #
-# `register_autograd` wiring directly (not via the dispatcher's carve-out).   #
-# Skipped on FA3 until Kuan-Hao's Modal FA3 setup PR lands and we mirror the  #
-# pattern there.                                                              #
-# --------------------------------------------------------------------------- #
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("causal", [False, True])
-def test_default_op_backward_through_registered_autograd(fa_default_impls, dtype, causal):
-    """FA2: gradients flow through ``torch.ops.fastvideo._flash_attn_default_forward``
-    itself (no dispatcher carve-out involved) and match the original
-    ``flash_attn_func``'s gradients."""
-    if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
-        pytest.skip("bfloat16 is not supported on this GPU")
-    _, original, op, fa_version = fa_default_impls
-    if fa_version != "2":
-        pytest.skip(f"register_autograd is only wired for FA2 right now; got {fa_version!r}")
-
-    torch.manual_seed(0)
-    q_ref, k_ref, v_ref = _qkv(dtype, requires_grad=True)
-    q_test, k_test, v_test = _clone(q_ref, k_ref, v_ref)
-
-    # Reference grads via the original autograd.Function.
-    out_ref = original(q_ref, k_ref, v_ref, softmax_scale=None, causal=causal)
-    # Custom-op grads via the registered backward — unpack (out, lse), discard lse.
-    out_test, _ = op(q_test, k_test, v_test, None, causal)
-
-    torch.testing.assert_close(out_test, out_ref,
-                               atol=0 if dtype == torch.float16 else 1e-3,
-                               rtol=0 if dtype == torch.float16 else 1e-3)
-
-    dout = torch.randn_like(out_ref)
-    dq_ref, dk_ref, dv_ref = torch.autograd.grad(
-        (out_ref * dout).sum(), (q_ref, k_ref, v_ref))
-    dq_test, dk_test, dv_test = torch.autograd.grad(
-        (out_test * dout).sum(), (q_test, k_test, v_test))
-
-    atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
-    torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol)
-    torch.testing.assert_close(dk_test, dk_ref, atol=atol, rtol=rtol)
-    torch.testing.assert_close(dv_test, dv_ref, atol=atol, rtol=rtol)
-
-
-@pytest.mark.parametrize("causal", [False, True])
-def test_default_op_opcheck_with_grad_inputs(fa_default_impls, causal):
-    """FA2: full ``opcheck`` including ``test_autograd_registration`` —
-    catches a missing/inconsistent backward at unit-test time, which was
-    exactly the gap that #1373's first revision shipped."""
-    _, _, op, fa_version = fa_default_impls
-    if fa_version != "2":
-        pytest.skip(f"autograd registration only wired for FA2; got {fa_version!r}")
-    torch.manual_seed(0)
-    q, k, v = _qkv(torch.float16, requires_grad=True)
-    torch.library.opcheck(op, (q, k, v, None, causal))

fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py

@@ -1,239 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regression guard for the masked/varlen custom ops.
-
-`flash_attn_no_pad_compilable` / `flash_attn_varlen_qk_no_pad_compilable` wrap
-the whole masked-attention functions in `torch.library.custom_op`s so dynamo
-sees one traceable node (the internal unpad/pad bookkeeping runs eager inside).
-
-On FA2 these ops register a real backward — `softmax_lse` is padded back to a
-statically-shaped `[batch, nheads, seqlen]` form on the way out and re-unpadded
-in backward, which calls FA2's `_flash_attn_varlen_backward`. So training also
-backprops through the op (no graph break on the training path either).
-
-On FA3/FA4 these ops are forward+fake only and the `*_compilable` dispatchers
-carve out to the autograd.Function for grad-enabled calls (PR #1373 pattern).
-Tests gating on FA2 are skipped on FA3/FA4.
-
-These tests pin:
-  - inference (no grad): output through the custom op is bit-identical to the
-    original function;
-  - training (requires_grad): gradients through the registered op match the
-    original autograd.Function;
-  - schema/fake-kernel consistency via torch.library.opcheck, both with and
-    without grad-requiring inputs (the latter exercises
-    `test_autograd_registration` — would have caught a missing/broken backward
-    at unit-test time).
-
-GPU assumptions: requires CUDA and the FA2 `flash_attn` varlen package.
-Skips on CPU and when flash_attn is unavailable.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-
-@pytest.fixture(scope="module")
-def no_pad_impls():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA is required for the masked/varlen custom-op tests")
-    try:
-        from fastvideo.attention.utils import flash_attn_no_pad as mod
-    except ImportError as exc:
-        pytest.skip(f"flash_attn_no_pad not importable (flash_attn missing?): {exc}")
-    return mod
-
-
-def _dtype_skip(dtype):
-    if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
-        pytest.skip("bfloat16 is not supported on this GPU")
-
-
-def _fa2_only(mod):
-    if mod._FA_VARLEN_VERSION != "2":
-        pytest.skip(
-            f"register_autograd is only wired for FA2; got "
-            f"_FA_VARLEN_VERSION={mod._FA_VARLEN_VERSION!r}"
-        )
-
-
-def _padding_mask(batch, seqlen, valid_lens, device):
-    mask = torch.zeros(batch, seqlen, dtype=torch.bool, device=device)
-    for i, n in enumerate(valid_lens):
-        mask[i, :n] = True
-    return mask
-
-
-# --------------------------------------------------------------------------- #
-# flash_attn_no_pad (masked self-attention)                                    #
-# --------------------------------------------------------------------------- #
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_no_pad_inference_matches_original(no_pad_impls, dtype):
-    """No-grad path through the custom op is bit-identical to the original."""
-    _dtype_skip(dtype)
-    mod = no_pad_impls
-    torch.manual_seed(0)
-    device = torch.device("cuda")
-    b, s, h, d = 2, 64, 4, 64
-    qkv = torch.randn(b, s, 3, h, d, device=device, dtype=dtype)
-    mask = _padding_mask(b, s, [64, 48], device)
-
-    with torch.inference_mode():
-        out_ref = mod.flash_attn_no_pad(qkv, mask, causal=False, dropout_p=0.0)
-        out_test = mod.flash_attn_no_pad_compilable(qkv, mask, causal=False, dropout_p=0.0)
-    torch.testing.assert_close(out_test, out_ref, atol=0, rtol=0)
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_no_pad_training_backward_through_registered_autograd(no_pad_impls, dtype):
-    """FA2: grads flow through the registered op and match the original."""
-    _dtype_skip(dtype)
-    mod = no_pad_impls
-    _fa2_only(mod)
-    torch.manual_seed(0)
-    device = torch.device("cuda")
-    b, s, h, d = 2, 64, 4, 64
-    mask = _padding_mask(b, s, [64, 48], device)
-
-    qkv_ref = torch.randn(b, s, 3, h, d, device=device, dtype=dtype, requires_grad=True)
-    qkv_test = qkv_ref.detach().clone().requires_grad_(True)
-
-    out_ref = mod.flash_attn_no_pad(qkv_ref, mask, causal=False, dropout_p=0.0)
-    # Go through the compilable wrapper (which on FA2 unconditionally routes
-    # to the registered op — no carve-out).
-    out_test = mod.flash_attn_no_pad_compilable(qkv_test, mask, causal=False, dropout_p=0.0)
-
-    dout = torch.randn_like(out_ref)
-    (dqkv_ref,) = torch.autograd.grad((out_ref * dout).sum(), (qkv_ref,))
-    (dqkv_test,) = torch.autograd.grad((out_test * dout).sum(), (qkv_test,))
-
-    atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
-    torch.testing.assert_close(dqkv_test, dqkv_ref, atol=atol, rtol=rtol)
-
-
-def test_no_pad_forward_opcheck(no_pad_impls):
-    """Schema/fake-kernel consistency for the custom op (forward only)."""
-    torch.manual_seed(0)
-    device = torch.device("cuda")
-    b, s, h, d = 2, 64, 4, 64
-    qkv = torch.randn(b, s, 3, h, d, device=device, dtype=torch.float16)
-    mask = _padding_mask(b, s, [64, 48], device)
-    torch.library.opcheck(
-        torch.ops.fastvideo._flash_attn_no_pad_forward,
-        (qkv, mask, False, 0.0, None, False),
-    )
-
-
-def test_no_pad_opcheck_with_grad_inputs(no_pad_impls):
-    """FA2: full opcheck including ``test_autograd_registration`` — catches a
-    missing/inconsistent backward at unit-test time."""
-    mod = no_pad_impls
-    _fa2_only(mod)
-    torch.manual_seed(0)
-    device = torch.device("cuda")
-    b, s, h, d = 2, 64, 4, 64
-    qkv = torch.randn(b, s, 3, h, d, device=device, dtype=torch.float16, requires_grad=True)
-    mask = _padding_mask(b, s, [64, 48], device)
-    torch.library.opcheck(
-        torch.ops.fastvideo._flash_attn_no_pad_forward,
-        (qkv, mask, False, 0.0, None, False),
-    )
-
-
-# --------------------------------------------------------------------------- #
-# flash_attn_varlen_qk_no_pad (cross-attn / unequal q-k seqlen)                #
-# --------------------------------------------------------------------------- #
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_varlen_qk_inference_matches_original(no_pad_impls, dtype):
-    _dtype_skip(dtype)
-    mod = no_pad_impls
-    # The varlen-qk custom op's real forward goes through FA's varlen func; off
-    # FA2 that path is the pre-existing FA3 `dropout_p` carve-out (out of scope
-    # here), so skip cleanly like the autograd tests until that's fixed.
-    _fa2_only(mod)
-    torch.manual_seed(1)
-    device = torch.device("cuda")
-    b, sq, sk, h, d = 2, 48, 64, 4, 64
-    q = torch.randn(b, sq, h, d, device=device, dtype=dtype)
-    k = torch.randn(b, sk, h, d, device=device, dtype=dtype)
-    v = torch.randn(b, sk, h, d, device=device, dtype=dtype)
-    qmask = _padding_mask(b, sq, [48, 40], device)
-    kmask = _padding_mask(b, sk, [64, 56], device)
-
-    with torch.inference_mode():
-        out_ref = mod.flash_attn_varlen_qk_no_pad(q, k, v, qmask, kmask, causal=False, dropout_p=0.0)
-        out_test = mod.flash_attn_varlen_qk_no_pad_compilable(q, k, v, qmask, kmask, causal=False, dropout_p=0.0)
-    torch.testing.assert_close(out_test, out_ref, atol=0, rtol=0)
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_varlen_qk_training_backward_through_registered_autograd(no_pad_impls, dtype):
-    """FA2: grads through q, k, v all flow via the registered op and match."""
-    _dtype_skip(dtype)
-    mod = no_pad_impls
-    _fa2_only(mod)
-    torch.manual_seed(1)
-    device = torch.device("cuda")
-    b, sq, sk, h, d = 2, 48, 64, 4, 64
-    qmask = _padding_mask(b, sq, [48, 40], device)
-    kmask = _padding_mask(b, sk, [64, 56], device)
-
-    q_ref = torch.randn(b, sq, h, d, device=device, dtype=dtype, requires_grad=True)
-    k_ref = torch.randn(b, sk, h, d, device=device, dtype=dtype, requires_grad=True)
-    v_ref = torch.randn(b, sk, h, d, device=device, dtype=dtype, requires_grad=True)
-    q_test = q_ref.detach().clone().requires_grad_(True)
-    k_test = k_ref.detach().clone().requires_grad_(True)
-    v_test = v_ref.detach().clone().requires_grad_(True)
-
-    out_ref = mod.flash_attn_varlen_qk_no_pad(q_ref, k_ref, v_ref, qmask, kmask, causal=False, dropout_p=0.0)
-    out_test = mod.flash_attn_varlen_qk_no_pad_compilable(q_test, k_test, v_test, qmask, kmask,
-                                                          causal=False, dropout_p=0.0)
-
-    dout = torch.randn_like(out_ref)
-    dq_ref, dk_ref, dv_ref = torch.autograd.grad((out_ref * dout).sum(), (q_ref, k_ref, v_ref))
-    dq_test, dk_test, dv_test = torch.autograd.grad((out_test * dout).sum(), (q_test, k_test, v_test))
-
-    atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
-    torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol)
-    torch.testing.assert_close(dk_test, dk_ref, atol=atol, rtol=rtol)
-    torch.testing.assert_close(dv_test, dv_ref, atol=atol, rtol=rtol)
-
-
-def test_varlen_qk_forward_opcheck(no_pad_impls):
-    mod = no_pad_impls
-    _fa2_only(mod)
-    torch.manual_seed(1)
-    device = torch.device("cuda")
-    b, sq, sk, h, d = 2, 48, 64, 4, 64
-    q = torch.randn(b, sq, h, d, device=device, dtype=torch.float16)
-    k = torch.randn(b, sk, h, d, device=device, dtype=torch.float16)
-    v = torch.randn(b, sk, h, d, device=device, dtype=torch.float16)
-    qmask = _padding_mask(b, sq, [48, 40], device)
-    kmask = _padding_mask(b, sk, [64, 56], device)
-    torch.library.opcheck(
-        torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward,
-        (q, k, v, qmask, kmask, False, 0.0, None, False),
-    )
-
-
-def test_varlen_qk_opcheck_with_grad_inputs(no_pad_impls):
-    """FA2: full opcheck with requires_grad inputs (autograd-registration check)."""
-    mod = no_pad_impls
-    _fa2_only(mod)
-    torch.manual_seed(1)
-    device = torch.device("cuda")
-    b, sq, sk, h, d = 2, 48, 64, 4, 64
-    q = torch.randn(b, sq, h, d, device=device, dtype=torch.float16, requires_grad=True)
-    k = torch.randn(b, sk, h, d, device=device, dtype=torch.float16, requires_grad=True)
-    v = torch.randn(b, sk, h, d, device=device, dtype=torch.float16, requires_grad=True)
-    qmask = _padding_mask(b, sq, [48, 40], device)
-    kmask = _padding_mask(b, sk, [64, 56], device)
-    torch.library.opcheck(
-        torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward,
-        (q, k, v, qmask, kmask, False, 0.0, None, False),
-    )

fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py

@@ -1,43 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""CPU-only regression test for the SDPA metadata mask contract.
-
-SDPAMetadata is cross-backend: HYWorld and HunyuanVideo15 build it while the
-layer's attention selector may pick FLASH_ATTN, and the shared convention
-for padding masks is the tokenizer-style 2D [batch, key_len]. The builder
-must store 2D masks unchanged; only the torch-sdpa impl lifts them to 4D
-internally.
-
-Regression guard for the CI break where the builder lifted 2D -> 4D and the
-FLASH_ATTN equal-length branch padded ``attn_mask.shape[1]`` (== 1 for 4D)
-into a garbage-length mask (CUDA device-side assert, SIGABRT).
-"""
-
-from __future__ import annotations
-
-import torch
-
-from fastvideo.attention.backends.sdpa import (SDPAMetadataBuilder, _normalize_attn_mask_for_sdpa)
-
-
-def test_builder_keeps_2d_padding_mask_2d() -> None:
-    """A tokenizer-style 2D int mask must be stored as-is (cross-backend contract)."""
-    mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.int64)
-    md = SDPAMetadataBuilder().build(current_timestep=0, attn_mask=mask)
-    assert md.attn_mask is not None
-    assert md.attn_mask.dim() == 2
-    assert md.attn_mask.shape == (2, 3)
-    assert torch.equal(md.attn_mask, mask)
-
-
-def test_normalize_lifts_2d_int_mask_to_4d_bool_for_torch_sdpa() -> None:
-    """The torch-sdpa consumer coerces int -> bool and lifts 2D -> [B,1,1,K]."""
-    batch, q_len, k_len, heads, head_dim = 2, 4, 3, 1, 8
-    query = torch.zeros(batch, heads, q_len, head_dim)
-    key = torch.zeros(batch, heads, k_len, head_dim)
-    mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.int64)
-
-    out = _normalize_attn_mask_for_sdpa(mask, query, key)
-    assert out is not None
-    assert out.dtype == torch.bool
-    assert out.shape == (batch, 1, 1, k_len)
-    assert torch.equal(out[:, 0, 0, :], mask != 0)